[mlir][linalg][bufferize] Limited support for scf.execute_region

[springerm]

Summary

Add support for analysis only.

Depends On D114040

[nicolasvasilache]

Please add a TODO that bufferization can be supported when needed.

[springerm]

address comments

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

bool isWritable(Operation *op, Value value) const {
     return false;
   }
 };
 
 } // namespace arith_ext
 
 namespace scf_ext {
 
+struct ExecuteRegionOpInterface
+    : public BufferizableOpInterface::ExternalModel<ExecuteRegionOpInterface,
+                                                    scf::ExecuteRegionOp> {
+  SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
+                                                OpResult opResult) const {
+    // ExecuteRegionOps do not have tensor OpOperands. The yielded value can be
+    // any SSA value that is in scope. To allow for use-def chain traversal
+    // through ExecuteRegionOps in the analysis, the corresponding yield value
+    // is considered to be aliasing with the result.
+    auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
+    size_t resultNum = std::distance(op->getOpResults().begin(),
+                                     llvm::find(op->getOpResults(), opResult));
+    assert(executeRegionOp.region().getBlocks().size() == 1 &&
+           "expected exactly 1 block");
+    auto yieldOp = dyn_cast<scf::YieldOp>(
+        executeRegionOp.region().front().getTerminator());
+    assert(yieldOp && "expected scf.yield terminator in scf.execute_region");
+    return {&yieldOp->getOpOperand(resultNum)};
+  }
+
+  bool mustBufferizeInPlace(Operation *op, OpResult opResult) const {
+    // ExecuteRegionOp results always bufferize in-place. Since they have no
+    // OpOperands, they are mostly ignored by the analysis once alias sets are
+    // set up.
+    return true;
+  }
+
+  LogicalResult bufferize(Operation *op, OpBuilder &b,
+                          BufferizationState &state) const {
+    // TODO: Add bufferization support when needed. scf.execute_region should be
+    // bufferized similar to scf.if.
+    bool hasTensorReturnType = any_of(
+        op->getResultTypes(), [](Type t) { return t.isa<TensorType>(); });
+    if (hasTensorReturnType)
+      return op->emitError(
+          "scf.execute_region with tensor result not supported");
+    return success();
+  }
+};
+
 struct IfOpInterface
     : public BufferizableOpInterface::ExternalModel<IfOpInterface, scf::IfOp> {
   SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
                                                 OpResult opResult) const {
     // IfOps do not have tensor OpOperands. The yielded value can be any SSA
     // value that is in scope. To allow for use-def chain traversal through
     // IfOps in the analysis, both corresponding yield values from the then/else
     // branches are considered to be aliasing with the result.
LogicalResult bufferize(Operation *op, OpBuilder &b,
     return success();
   }
 };
 
 } // namespace vector_ext
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   registry.addOpInterface<arith::ConstantOp, arith_ext::ConstantOpInterface>();
+  registry.addOpInterface<scf::ExecuteRegionOp,
+                          scf_ext::ExecuteRegionOpInterface>();
   registry.addOpInterface<scf::ForOp, scf_ext::ForOpInterface>();
   registry.addOpInterface<scf::IfOp, scf_ext::IfOpInterface>();
   registry.addOpInterface<scf::YieldOp, scf_ext::YieldOpInterface>();
   registry.addOpInterface<CallOp, std_ext::CallOpInterface>();
   registry.addOpInterface<ReturnOp, std_ext::ReturnOpInterface>();
   registry.addOpInterface<tensor::CastOp, tensor_ext::CastOpInterface>();
   registry.addOpInterface<tensor::DimOp, tensor_ext::DimOpInterface>();
   registry.addOpInterface<tensor::ExtractSliceOp,

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir

func @mini_test_case1() -> tensor<10x20xf32> {
   %t = linalg.init_tensor [10, 20] : 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}}
+  // expected-error @+1 {{scf.execute_region with tensor result not supported}}
   %r = scf.execute_region -> tensor<4xi32> {
     %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
     scf.yield %A: tensor<4xi32>
   }
   return %r: tensor<4xi32>
 }