Add missing virtual inliner interface method in SPIR-V dialect.

The inline interface uses two methods to check legality of inling:
1) Can a region be inlined into another.
2) Can an operation be inlined into another.
Setting the former to true, allows the inliner to use the second for
legality checks. Add this method to the SPIR-V dialect inlining
interface.

PiperOrigin-RevId: 286041734

diff --git a/mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp b/mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp
index 1460cf0..c99e7ca 100644 (file)
--- 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 9837d7b..7536019 100644 (file)
--- 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