[Mlir-commits] [mlir] af9dce1 - [mlir][Linalg] Allow operand-less scf::ExecuteRegionOp to encapsulate scf::YieldOp

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2021-10-05T11:34:53Z
New Revision: af9dce18bfdb09ab1dd8b4c55979e67799a0b322

URL: https://github.com/llvm/llvm-project/commit/af9dce18bfdb09ab1dd8b4c55979e67799a0b322
DIFF: https://github.com/llvm/llvm-project/commit/af9dce18bfdb09ab1dd8b4c55979e67799a0b322.diff

LOG: [mlir][Linalg] Allow operand-less scf::ExecuteRegionOp to encapsulate scf::YieldOp

These are considered noops.
Buferization will still fail on scf.execute_region which yield values.
This is used to make comprehensive bufferization interoperate better with external clients.

Differential Revision: https://reviews.llvm.org/D111130

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
    mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir
    mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
index 6b5f6249246f5..aba7d108e38a6 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
@@ -2110,8 +2110,16 @@ static LogicalResult bufferize(OpBuilder &b, scf::YieldOp yieldOp,
   // Cannot create IR past a yieldOp.
   b.setInsertionPoint(yieldOp);
 
+  if (auto execOp = dyn_cast<scf::ExecuteRegionOp>(yieldOp->getParentOp())) {
+    if (execOp->getNumResults() != 0)
+      return execOp->emitError(
+          "expected result-less scf.execute_region containing op");
+    return success();
+  }
+
   scf::ForOp forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
-  assert(forOp && "only support scf::ForOp parent for scf::YieldOp");
+  if (!forOp)
+    return yieldOp->emitError("expected scf::ForOp parent for scf::YieldOp");
   for (OpOperand &operand : yieldOp->getOpOperands()) {
     auto tensorType = operand.get().getType().dyn_cast<TensorType>();
     if (!tensorType)

diff  --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir
index 60533c677911a..57d3f2215f2f1 100644
--- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir
+++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir
@@ -140,3 +140,27 @@ func @mini_test_case1() -> tensor<10x20xf32> {
   %r = linalg.fill(%f0, %t) : f32, tensor<10x20xf32> -> tensor<10x20xf32>
   return %r : tensor<10x20xf32>
 }
+
+// -----
+
+func @main() -> tensor<4xi32> {
+  // expected-error @+1 {{unsupported op with tensors}}
+  %r = scf.execute_region -> tensor<4xi32> {
+    %A = constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+    scf.yield %A: tensor<4xi32>
+  }
+  return %r: tensor<4xi32>
+}
+
+// -----
+
+func @main() -> i32 {
+  %c0 = constant 0: index
+  // expected-error @+1 {{expected result-less scf.execute_region containing op}}
+  %r = scf.execute_region -> i32 {
+    %A = constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+    %e = tensor.extract %A[%c0]: tensor<4xi32>
+    scf.yield %e: i32
+  }
+  return %r: i32
+}

diff  --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
index d06c5162c1ceb..ac9735dedec3c 100644
--- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
+++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
@@ -374,6 +374,29 @@ func @main() {
 
 //      CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
+//      CHECK: memref.global "private" constant @__constant_4xi32 : memref<4xi32> = dense<[1, 2, 3, 4]>
+//      CHECK: func private @some_external_func_within_scf_execute(memref<4xi32, #[[$DYN_1D_MAP]]>)
+func private @some_external_func_within_scf_execute(tensor<4xi32>)
+
+//      CHECK: func @main()
+func @main() {
+//      CHECK:   %[[A:.*]] = memref.get_global @__constant_4xi32 : memref<4xi32>
+  %A = constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+
+//      CHECK:   %[[B:.*]] = memref.cast %[[A]] : memref<4xi32> to memref<4xi32, #[[$DYN_1D_MAP]]>
+//      CHECK:   call @some_external_func_within_scf_execute(%[[B]]) : (memref<4xi32, #[[$DYN_1D_MAP]]>) -> ()
+  scf.execute_region {
+    call @some_external_func_within_scf_execute(%A) : (tensor<4xi32>) -> ()
+    scf.yield
+  }
+
+  return
+}
+
+// -----
+
+//      CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
 //      CHECK:  func private @some_external_func(memref<?xf32, #[[$DYN_1D_MAP]]>)
 func private @some_external_func(tensor<?xf32>)