Skip to content

Commit

Permalink
[mlir][Linalg] Allow operand-less scf::ExecuteRegionOp to encapsulate…
Browse files Browse the repository at this point in the history
… 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
  • Loading branch information
nicolasvasilache committed Oct 5, 2021
1 parent 1549be3 commit af9dce1
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 1 deletion.
10 changes: 9 additions & 1 deletion mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
Expand Up @@ -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)
Expand Down
Expand Up @@ -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
}
23 changes: 23 additions & 0 deletions mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
Expand Up @@ -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>)

Expand Down

0 comments on commit af9dce1

Please sign in to comment.