diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp @@ -2110,8 +2110,17 @@ // Cannot create IR past a yieldOp. b.setInsertionPoint(yieldOp); + auto execOp = dyn_cast(yieldOp->getParentOp()); + if (execOp) { + if (execOp->getNumResults() != 0) + return execOp->emitError( + "expected result-less scf.execute_region containing op"); + return success(); + } + scf::ForOp forOp = dyn_cast(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(); if (!tensorType) diff --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir --- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir +++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir @@ -140,3 +140,27 @@ %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 --- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir +++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir @@ -374,6 +374,29 @@ // 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) func private @some_external_func(tensor)