[Mlir-commits] [mlir] 0b3e478 - [mlir][GPUToSPIRV] Use default ABI only when none of the arguments

[llvmlistbot at llvm.org]

Author: MaheshRavishankar
Date: 2020-05-14T21:48:51-07:00
New Revision: 0b3e478b10e233460073d24df527253ab5edac2a

URL: https://github.com/llvm/llvm-project/commit/0b3e478b10e233460073d24df527253ab5edac2a
DIFF: https://github.com/llvm/llvm-project/commit/0b3e478b10e233460073d24df527253ab5edac2a.diff

LOG: [mlir][GPUToSPIRV] Use default ABI only when none of the arguments
have abi attributes.

To ensure there is no conflict, use the default ABI only when none of
the arguments have the spv.interface_var_abi attribute. This also
implies that if one of the arguments has a spv.interface_var_abi
attribute, all of them should have it as well.

Differential Revision: https://reviews.llvm.org/D77232

Added: 
    

Modified: 
    mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp
    mlir/test/Conversion/GPUToSPIRV/simple.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp
index cae7239398f4..15bd2d923c16 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp
@@ -343,6 +343,26 @@ lowerAsEntryFunction(gpu::GPUFuncOp funcOp, SPIRVTypeConverter &typeConverter,
   return newFuncOp;
 }
 
+/// Populates `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 {
@@ -350,22 +370,21 @@ LogicalResult GPUFuncOpConversion::matchAndRewrite(
     return failure();
 
   SmallVector<spirv::InterfaceVarABIAttr, 4> argABI;
-  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 (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) {
+        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);

diff  --git a/mlir/test/Conversion/GPUToSPIRV/simple.mlir b/mlir/test/Conversion/GPUToSPIRV/simple.mlir
index f68823321568..0ecb83ebdcee 100644
--- a/mlir/test/Conversion/GPUToSPIRV/simple.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/simple.mlir
@@ -26,6 +26,30 @@ module attributes {gpu.container_module} {
 
 // -----
 
+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
+    }
+  }
+}
+
+// -----
+
 module attributes {gpu.container_module} {
   gpu.module @kernels {
     // expected-error @below {{failed to legalize operation 'gpu.func'}}
@@ -44,3 +68,37 @@ module attributes {gpu.container_module} {
     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
+    }
+  }
+}
+
+// -----
+
+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
+    }
+  }
+}