diff options
| -rw-r--r-- | mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp | 11 | ||||
| -rw-r--r-- | mlir/test/Dialect/SPIRV/Transforms/inlining.mlir | 45 |
2 files changed, 56 insertions, 0 deletions
diff --git a/mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp b/mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp index 1460cf091eb..c99e7ca8b20 100644 --- a/mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp +++ b/mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp @@ -56,6 +56,17 @@ struct SPIRVInlinerInterface : public DialectInlinerInterface { /// Returns true if the given region 'src' can be inlined into the region /// 'dest' that is attached to an operation registered to the current dialect. + bool isLegalToInline(Region *dest, Region *src, + BlockAndValueMapping &) const final { + // Return true here when inlining into spv.selection and spv.loop + // operations. + auto op = dest->getParentOp(); + return isa<spirv::SelectionOp>(op) || isa<spirv::LoopOp>(op); + } + + /// Returns true if the given operation 'op', that is registered to this + /// dialect, can be inlined into the region 'dest' that is attached to an + /// operation registered to the current dialect. bool isLegalToInline(Operation *op, Region *dest, BlockAndValueMapping &) const final { // TODO(antiagainst): Enable inlining structured control flows with return. diff --git a/mlir/test/Dialect/SPIRV/Transforms/inlining.mlir b/mlir/test/Dialect/SPIRV/Transforms/inlining.mlir index 9837d7babb7..75360194d7b 100644 --- a/mlir/test/Dialect/SPIRV/Transforms/inlining.mlir +++ b/mlir/test/Dialect/SPIRV/Transforms/inlining.mlir @@ -180,3 +180,48 @@ spv.module "Logical" "GLSL450" { spv.Return } } + +// ----- + +spv.module "Logical" "GLSL450" { + spv.globalVariable @arg_0 bind(0, 0) : !spv.ptr<!spv.struct<i32 [0]>, StorageBuffer> + spv.globalVariable @arg_1 bind(0, 1) : !spv.ptr<!spv.struct<i32 [0]>, StorageBuffer> + // CHECK: func @inline_into_selection_region + func @inline_into_selection_region() { + %1 = spv.constant 0 : i32 + // CHECK-DAG: [[ADDRESS_ARG0:%.*]] = spv._address_of @arg_0 + // CHECK-DAG: [[ADDRESS_ARG1:%.*]] = spv._address_of @arg_1 + // CHECK-DAG: [[LOADPTR:%.*]] = spv.AccessChain [[ADDRESS_ARG0]] + // CHECK: [[VAL:%.*]] = spv.Load "StorageBuffer" [[LOADPTR]] + %2 = spv._address_of @arg_0 : !spv.ptr<!spv.struct<i32 [0]>, StorageBuffer> + %3 = spv._address_of @arg_1 : !spv.ptr<!spv.struct<i32 [0]>, StorageBuffer> + %4 = spv.AccessChain %2[%1] : !spv.ptr<!spv.struct<i32 [0]>, StorageBuffer> + %5 = spv.Load "StorageBuffer" %4 : i32 + %6 = spv.SGreaterThan %5, %1 : i32 + // CHECK: spv.selection + spv.selection { + spv.BranchConditional %6, ^bb1, ^bb2 + ^bb1: // pred: ^bb0 + // CHECK: [[STOREPTR:%.*]] = spv.AccessChain [[ADDRESS_ARG1]] + %7 = spv.AccessChain %3[%1] : !spv.ptr<!spv.struct<i32 [0]>, StorageBuffer> + // CHECK-NOT: spv.FunctionCall + // CHECK: spv.AtomicIAdd "Device" "AcquireRelease" [[STOREPTR]], [[VAL]] + // CHECK: spv.Branch + spv.FunctionCall @atomic_add(%5, %7) : (i32, !spv.ptr<i32, StorageBuffer>) -> () + spv.Branch ^bb2 + ^bb2 : // 2 preds: ^bb0, ^bb1 + spv._merge + } + // CHECK: spv.Return + spv.Return + } + func @atomic_add(%arg0: i32, %arg1: !spv.ptr<i32, StorageBuffer>) { + %0 = spv.AtomicIAdd "Device" "AcquireRelease" %arg1, %arg0 : !spv.ptr<i32, StorageBuffer> + spv.Return + } + spv.EntryPoint "GLCompute" @inline_into_selection_region + spv.ExecutionMode @inline_into_selection_region "LocalSize", 32, 1, 1 +} attributes {capabilities = ["Shader"], extensions = ["SPV_KHR_storage_buffer_storage_class"]} + +// TODO: Add tests for inlining structured control flow into +// structured control flow.
\ No newline at end of file |
