[mlir][GPUToSPIRV] Make spv.interface_var_abi attribute on arguments either be unspecified on all arguments to use default ABI, or be present on all arguments.

[mravishankar]

Adding a fallback path to get a "default" ABI to not make this change
breaking.

[mehdi_amini]

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

347

The revision title says that it makes "spv.interface_var_abi attribute a required attribute", but this function seems to make it optional?

[mravishankar]

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

347

This is more a escape hatch and is legacy. When I initially set this up it was convenient, but this is very limited in its usage. I would actually prefer to not have this, and any realistic user of the SPIR-V lowering will not be able to use this on many devices. So I can change the comment to make this clearer.

[antiagainst]

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

346

Generates ... Returns false if any argument already has ABI specified.

356

Let's put a comment here to explain why this if statement like:

We need to wrap scalar values in a struct to satisfy Vulkan's interface variable requirements. We choose StorageBuffer for holding the struct here.

Otherwise it might be a bit mysterious.

mlir/test/Conversion/GPUToSPIRV/if.mlir

19 ↗

(On Diff #254311)

I guess this is changes made when this commit is truly breaking. But now we have a default behavior; we might want to revert these to simplify the tests.

mlir/test/Conversion/GPUToSPIRV/simple.mlir

32

Instead of checking the default behavior, test the explicit behavior? Also, please make sure the assignment is not sequential as the default behavior.

mlir/test/mlir-vulkan-runner/addf.mlir

13 ↗

(On Diff #254311)

I think we can omit the changes here; the in-tree Vulkan runner follows the sequential scheme.

[mravishankar]

Addressing comments and changing description to better reflect the
change post modifications.

[antiagainst]

mlir/test/Conversion/GPUToSPIRV/if.mlir

40 ↗

(On Diff #255381)

Let's revert changes to these tests. They should pass as-is. Applies to following unaffected tests.

mlir/test/Conversion/GPUToSPIRV/simple.mlir

38

I think we can just test the gpu.func itself without the wrapping gpu.module and the corresponding launch op? This way it's clearer regarding the intended use case of attaching these attributes: it gives us a way to use and control the CodeGen side as a stand-alone component. :)

[mravishankar]

Rebasing and addressing comments

[mravishankar]

Adressing comments

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

346

Made it populates.

mlir/test/Conversion/GPUToSPIRV/if.mlir

40 ↗

(On Diff #255381)

This is not really a change, the args are split, thats all. Will remove it anyway.

mlir/test/Conversion/GPUToSPIRV/simple.mlir

38

Apparently the gpu.module is needed (gpu.func must have this as parent). Removed the funcOp that invokes the kernel though.

[antiagainst]

It would be nice to have a shorter oneliner for the first paragraph of the commit message and put the details in the second paragraph.

mlir/test/Conversion/GPUToSPIRV/simple.mlir

38

Oh right. A recent change made gpu.func must appear in gpu.module

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
                               newFuncOp.end());
   rewriter.applySignatureConversion(&newFuncOp.getBody(), signatureConverter);
   rewriter.eraseOp(funcOp);
 
   spirv::setABIAttrs(newFuncOp, entryPointInfo, argABIInfo);
   return newFuncOp;
 }
 
+/// Updates `argABI` with spv.interface_var_abi attributes for lowering gpu.func
+/// to spv.func if no arguments have the attributes set already. Returns failure
+/// if any argument has the ABI attribute set already.
+static LogicalResult
+getDefaultABIAttrs(MLIRContext *context, gpu::GPUFuncOp funcOp,
+                   SmallVectorImpl<spirv::InterfaceVarABIAttr> &argABI) {
+  for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
+    if (funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
+            argIndex, spirv::getInterfaceVarABIAttrName()))
+      return failure();
+    // Vulkan's interface variable requirements needs scalars to be wrapped in a
+    // struct. The struct held in storage buffer.
+    Optional<spirv::StorageClass> sc;
+    if (funcOp.getArgument(argIndex).getType().isIntOrIndexOrFloat())
+      sc = spirv::StorageClass::StorageBuffer;
+    argABI.push_back(spirv::getInterfaceVarABIAttr(0, argIndex, sc, context));
+  }
+  return success();
+}
+
 LogicalResult GPUFuncOpConversion::matchAndRewrite(
     gpu::GPUFuncOp funcOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   if (!gpu::GPUDialect::isKernel(funcOp))
     return failure();
 
   SmallVector<spirv::InterfaceVarABIAttr, 4> argABI;
+  if (failed(getDefaultABIAttrs(rewriter.getContext(), funcOp, argABI))) {
+    argABI.clear();
     for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
       // If the ABI is already specified, use it.
       auto abiAttr = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
           argIndex, spirv::getInterfaceVarABIAttrName());
-    if (abiAttr) {
+      if (!abiAttr) {
+        funcOp.emitRemark(
+            "match failure: missing 'spv.interface_var_abi' attribute at "
+            "argument ")
+            << argIndex;
+        return failure();
+      }
       argABI.push_back(abiAttr);
-      continue;
     }
-    // todo(ravishankarm): Use the "default ABI". Remove this in a follow up
-    // CL. Staging this to make this easy to revert in case of breakages out of
-    // tree.
-    Optional<spirv::StorageClass> sc;
-    if (funcOp.getArgument(argIndex).getType().isIntOrIndexOrFloat())
-      sc = spirv::StorageClass::StorageBuffer;
-    argABI.push_back(
-        spirv::getInterfaceVarABIAttr(0, argIndex, sc, rewriter.getContext()));
   }
 
   auto entryPointAttr = spirv::lookupEntryPointABI(funcOp);
   if (!entryPointAttr) {
     funcOp.emitRemark("match failure: missing 'spv.entry_point_abi' attribute");
     return failure();
   }
   spirv::FuncOp newFuncOp = lowerAsEntryFunction(

mlir/test/Conversion/GPUToSPIRV/simple.mlir

func @main() {
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   gpu.module @kernels {
+    // CHECK:       spv.module Logical GLSL450 {
+    // CHECK-LABEL: spv.func @basic_module_structure_preset_ABI
+    // CHECK-SAME: {{%[a-zA-Z0-9_]*}}: f32
+    // CHECK-SAME: spv.interface_var_abi = #spv.interface_var_abi<(1, 2), StorageBuffer>
+    // CHECK-SAME: !spv.ptr<!spv.struct<!spv.array<12 x f32, stride=4> [0]>, StorageBuffer>
+    // CHECK-SAME: spv.interface_var_abi = #spv.interface_var_abi<(3, 0)>
+    // CHECK-SAME: spv.entry_point_abi = {local_size = dense<[32, 4, 1]> : vector<3xi32>}
+    gpu.func @basic_module_structure_preset_ABI(
+      %arg0 : f32
+        {spv.interface_var_abi = #spv.interface_var_abi<(1, 2), StorageBuffer>},
+      %arg1 : memref<12xf32>
+        {spv.interface_var_abi = #spv.interface_var_abi<(3, 0)>}) kernel
+      attributes
+        {spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
+      // CHECK: spv.Return
+      gpu.return
+    }
+  }
+
+  func @main() {
+    %0 = "op"() : () -> (f32)
+    %1 = "op"() : () -> (memref<12xf32>)
+    %cst = constant 1 : index
+    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1) { kernel = @kernels::@basic_module_structure_preset_ABI }
+        : (index, index, index, index, index, index, f32, memref<12xf32>) -> ()
+    return
+  }
+}
+
+// -----
+
+module attributes {gpu.container_module} {
+  gpu.module @kernels {
     // expected-error @below {{failed to legalize operation 'gpu.func'}}
     // expected-remark @below {{match failure: missing 'spv.entry_point_abi' attribute}}
     gpu.func @missing_entry_point_abi(%arg0 : f32, %arg1 : memref<12xf32>) kernel {
       gpu.return
     }
   }
 
   func @main() {
     %0 = "op"() : () -> (f32)
     %1 = "op"() : () -> (memref<12xf32>)
     %cst = constant 1 : index
     "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1) { kernel = @kernels::@missing_entry_point_abi }
         : (index, index, index, index, index, index, f32, memref<12xf32>) -> ()
     return
   }
 }
+
+// -----
+
+module attributes {gpu.container_module} {
+  gpu.module @kernels {
+    // expected-error @below {{failed to legalize operation 'gpu.func'}}
+    // expected-remark @below {{match failure: missing 'spv.interface_var_abi' attribute at argument 1}}
+    gpu.func @missing_entry_point_abi(
+      %arg0 : f32
+        {spv.interface_var_abi = #spv.interface_var_abi<(1, 2), StorageBuffer>},
+      %arg1 : memref<12xf32>) kernel
+    attributes
+      {spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
+      gpu.return
+    }
+  }
+
+  func @main() {
+    %0 = "op"() : () -> (f32)
+    %1 = "op"() : () -> (memref<12xf32>)
+    %cst = constant 1 : index
+    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1)
+      { kernel = @kernels::@missing_entry_point_abi }
+        : (index, index, index, index, index, index, f32, memref<12xf32>) -> ()
+    return
+  }
+}
+
+// -----
+
+module attributes {gpu.container_module} {
+  gpu.module @kernels {
+    // expected-error @below {{failed to legalize operation 'gpu.func'}}
+    // expected-remark @below {{match failure: missing 'spv.interface_var_abi' attribute at argument 0}}
+    gpu.func @missing_entry_point_abi(
+      %arg0 : f32,
+      %arg1 : memref<12xf32>
+        {spv.interface_var_abi = #spv.interface_var_abi<(3, 0)>}) kernel
+    attributes
+      {spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
+      gpu.return
+    }
+  }
+
+  func @main() {
+    %0 = "op"() : () -> (f32)
+    %1 = "op"() : () -> (memref<12xf32>)
+    %cst = constant 1 : index
+    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1)
+      { kernel = @kernels::@missing_entry_point_abi}
+        : (index, index, index, index, index, index, f32, memref<12xf32>) -> ()
+    return
+  }
+}