diff --git a/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp --- a/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp +++ b/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp @@ -343,6 +343,23 @@ return newFuncOp; } +/// Generate default spv.interface_var_abi attributes for lowering gpu.func to +/// spv.func if no arguments have the attributes set already. +static LogicalResult +getDefaultABIAttrs(MLIRContext *context, gpu::GPUFuncOp funcOp, + SmallVectorImpl &argABI) { + for (auto argIndex : llvm::seq(0, funcOp.getNumArguments())) { + if (funcOp.getArgAttrOfType( + argIndex, spirv::getInterfaceVarABIAttrName())) + return failure(); + Optional 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 operands, ConversionPatternRewriter &rewriter) const { @@ -350,22 +367,21 @@ return failure(); SmallVector argABI; - for (auto argIndex : llvm::seq(0, funcOp.getNumArguments())) { - // If the ABI is already specified, use it. - auto abiAttr = funcOp.getArgAttrOfType( - argIndex, spirv::getInterfaceVarABIAttrName()); - if (abiAttr) { + if (failed(getDefaultABIAttrs(rewriter.getContext(), funcOp, argABI))) { + argABI.clear(); + for (auto argIndex : llvm::seq(0, funcOp.getNumArguments())) { + // If the ABI is already specified, use it. + auto abiAttr = funcOp.getArgAttrOfType( + 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 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/if.mlir b/mlir/test/Conversion/GPUToSPIRV/if.mlir --- a/mlir/test/Conversion/GPUToSPIRV/if.mlir +++ b/mlir/test/Conversion/GPUToSPIRV/if.mlir @@ -15,7 +15,14 @@ gpu.module @kernels { // CHECK-LABEL: @kernel_simple_selection - gpu.func @kernel_simple_selection(%arg2 : memref<10xf32>, %arg3 : i1) + gpu.func @kernel_simple_selection( + %arg2 : memref<10xf32> + {spv.interface_var_abi = {binding = 0 : i32, + descriptor_set = 0 : i32}}, + %arg3 : i1 + {spv.interface_var_abi = {binding = 1 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}) attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} { %value = constant 0.0 : f32 %i = constant 0 : index @@ -36,7 +43,21 @@ } // CHECK-LABEL: @kernel_nested_selection - gpu.func @kernel_nested_selection(%arg3 : memref<10xf32>, %arg4 : memref<10xf32>, %arg5 : i1, %arg6 : i1) + gpu.func @kernel_nested_selection( + %arg3 : memref<10xf32> + {spv.interface_var_abi = {binding = 0 : i32, + descriptor_set = 0 : i32}}, + %arg4 : memref<10xf32> + {spv.interface_var_abi = {binding = 1 : i32, + descriptor_set = 0 : i32}}, + %arg5 : i1 + {spv.interface_var_abi = {binding = 2 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg6 : i1 + {spv.interface_var_abi = {binding = 3 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}) attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} { %i = constant 0 : index %j = constant 9 : index diff --git a/mlir/test/Conversion/GPUToSPIRV/load-store.mlir b/mlir/test/Conversion/GPUToSPIRV/load-store.mlir --- a/mlir/test/Conversion/GPUToSPIRV/load-store.mlir +++ b/mlir/test/Conversion/GPUToSPIRV/load-store.mlir @@ -34,7 +34,32 @@ // CHECK-SAME: [[ARG4:%.*]]: i32 {spv.interface_var_abi = {binding = 4 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}} // CHECK-SAME: [[ARG5:%.*]]: i32 {spv.interface_var_abi = {binding = 5 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}} // CHECK-SAME: [[ARG6:%.*]]: i32 {spv.interface_var_abi = {binding = 6 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}} - gpu.func @load_store_kernel(%arg0: memref<12x4xf32>, %arg1: memref<12x4xf32>, %arg2: memref<12x4xf32>, %arg3: index, %arg4: index, %arg5: index, %arg6: index) + gpu.func @load_store_kernel( + %arg0: memref<12x4xf32> + {spv.interface_var_abi = {binding = 0 : i32, + descriptor_set = 0 : i32}}, + %arg1: memref<12x4xf32> + {spv.interface_var_abi = {binding = 1 : i32, + descriptor_set = 0 : i32}}, + %arg2: memref<12x4xf32> + {spv.interface_var_abi = {binding = 2 : i32, + descriptor_set = 0 : i32}}, + %arg3: index + {spv.interface_var_abi = {binding = 3 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg4: index + {spv.interface_var_abi = {binding = 4 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg5: index + {spv.interface_var_abi = {binding = 5 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg6: index + {spv.interface_var_abi = {binding = 6 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}) attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} { // CHECK: [[ADDRESSWORKGROUPID:%.*]] = spv._address_of [[WORKGROUPIDVAR]] // CHECK: [[WORKGROUPID:%.*]] = spv.Load "Input" [[ADDRESSWORKGROUPID]] diff --git a/mlir/test/Conversion/GPUToSPIRV/loop.mlir b/mlir/test/Conversion/GPUToSPIRV/loop.mlir --- a/mlir/test/Conversion/GPUToSPIRV/loop.mlir +++ b/mlir/test/Conversion/GPUToSPIRV/loop.mlir @@ -14,7 +14,13 @@ } gpu.module @kernels { - gpu.func @loop_kernel(%arg2 : memref<10xf32>, %arg3 : memref<10xf32>) + gpu.func @loop_kernel( + %arg2 : memref<10xf32> + {spv.interface_var_abi = {binding = 0 : i32, + descriptor_set = 0 : i32}}, + %arg3 : memref<10xf32> + {spv.interface_var_abi = {binding = 1 : i32, + descriptor_set = 0 : i32}}) attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} { // CHECK: [[LB:%.*]] = spv.constant 4 : i32 %lb = constant 4 : index diff --git a/mlir/test/Conversion/GPUToSPIRV/simple.mlir b/mlir/test/Conversion/GPUToSPIRV/simple.mlir --- a/mlir/test/Conversion/GPUToSPIRV/simple.mlir +++ b/mlir/test/Conversion/GPUToSPIRV/simple.mlir @@ -7,7 +7,14 @@ // CHECK-SAME: {{%.*}}: f32 {spv.interface_var_abi = {binding = 0 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}} // CHECK-SAME: {{%.*}}: !spv.ptr [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 1 : i32, descriptor_set = 0 : i32{{[}][}]}} // CHECK-SAME: spv.entry_point_abi = {local_size = dense<[32, 4, 1]> : vector<3xi32>} - gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32>) + gpu.func @basic_module_structure( + %arg0 : f32 + {spv.interface_var_abi = {binding = 0 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg1 : memref<12xf32> + {spv.interface_var_abi = {binding = 1 : i32, + descriptor_set = 0 : i32}}) attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} { // CHECK: spv.Return gpu.return @@ -26,11 +33,72 @@ // ----- +module attributes {gpu.container_module} { + gpu.module @kernels { + // CHECK: spv.module Logical GLSL450 { + // CHECK-LABEL: spv.func @basic_module_structure_default_ABI + // CHECK-SAME: {{%.*}}: f32 {spv.interface_var_abi = {binding = 0 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}} + // CHECK-SAME: {{%.*}}: !spv.ptr [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 1 : i32, descriptor_set = 0 : i32{{[}][}]}} + // CHECK-SAME: spv.entry_point_abi = {local_size = dense<[32, 4, 1]> : vector<3xi32>} + gpu.func @basic_module_structure_default_ABI(%arg0 : f32, %arg1 : memref<12xf32>) + attributes {gpu.kernel, 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 = "basic_module_structure_default_ABI", kernel_module = @kernels } + : (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>) attributes {gpu.kernel} { + gpu.func @missing_entry_point_abi( + %arg0 : f32 + {spv.interface_var_abi = {binding = 0 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg1 : memref<12xf32> + {spv.interface_var_abi = {binding = 1 : i32, + descriptor_set = 0 : i32}}) + attributes {gpu.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 = "missing_entry_point_abi", kernel_module = @kernels } + : (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 = {binding = 1 : i32, + descriptor_set = 0 : i32}}) + attributes {gpu.kernel} { gpu.return } } diff --git a/mlir/test/mlir-vulkan-runner/addf.mlir b/mlir/test/mlir-vulkan-runner/addf.mlir --- a/mlir/test/mlir-vulkan-runner/addf.mlir +++ b/mlir/test/mlir-vulkan-runner/addf.mlir @@ -9,7 +9,18 @@ max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}> } { gpu.module @kernels { - gpu.func @kernel_add(%arg0 : memref<8xf32>, %arg1 : memref<8xf32>, %arg2 : memref<8xf32>) + gpu.func @kernel_add(%arg0 : memref<8xf32> + {spv.interface_var_abi = {binding = 0 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg1 : memref<8xf32> + {spv.interface_var_abi = {binding = 1 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg2 : memref<8xf32> + {spv.interface_var_abi = {binding = 2 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}) attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[1, 1, 1]>: vector<3xi32>}} { %0 = "gpu.block_id"() {dimension = "x"} : () -> index %1 = load %arg0[%0] : memref<8xf32> @@ -48,4 +59,3 @@ func @fillResource1DFloat(%0 : memref, %1 : f32) func @print_memref_f32(%ptr : memref<*xf32>) } - diff --git a/mlir/test/mlir-vulkan-runner/mulf.mlir b/mlir/test/mlir-vulkan-runner/mulf.mlir --- a/mlir/test/mlir-vulkan-runner/mulf.mlir +++ b/mlir/test/mlir-vulkan-runner/mulf.mlir @@ -9,7 +9,18 @@ max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}> } { gpu.module @kernels { - gpu.func @kernel_mul(%arg0 : memref<4x4xf32>, %arg1 : memref<4x4xf32>, %arg2 : memref<4x4xf32>) + gpu.func @kernel_mul(%arg0 : memref<4x4xf32> + {spv.interface_var_abi = {binding = 0 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg1 : memref<4x4xf32> + {spv.interface_var_abi = {binding = 1 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg2 : memref<4x4xf32> + {spv.interface_var_abi = {binding = 2 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}) attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[1, 1, 1]>: vector<3xi32>}} { %x = "gpu.block_id"() {dimension = "x"} : () -> index %y = "gpu.block_id"() {dimension = "y"} : () -> index @@ -49,4 +60,3 @@ func @fillResource2DFloat(%0 : memref, %1 : f32) func @print_memref_f32(%ptr : memref<*xf32>) } - diff --git a/mlir/test/mlir-vulkan-runner/time.mlir b/mlir/test/mlir-vulkan-runner/time.mlir --- a/mlir/test/mlir-vulkan-runner/time.mlir +++ b/mlir/test/mlir-vulkan-runner/time.mlir @@ -12,7 +12,18 @@ max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}> } { gpu.module @kernels { - gpu.func @kernel_add(%arg0 : memref<16384xf32>, %arg1 : memref<16384xf32>, %arg2 : memref<16384xf32>) + gpu.func @kernel_add(%arg0 : memref<16384xf32> + {spv.interface_var_abi = {binding = 0 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg1 : memref<16384xf32> + {spv.interface_var_abi = {binding = 1 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}, + %arg2 : memref<16384xf32> + {spv.interface_var_abi = {binding = 2 : i32, + descriptor_set = 0 : i32, + storage_class = 12 : i32}}) attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[128, 1, 1]>: vector<3xi32>}} { %bid = "gpu.block_id"() {dimension = "x"} : () -> index %tid = "gpu.thread_id"() {dimension = "x"} : () -> index @@ -54,4 +65,3 @@ func @fillResource1DFloat(%0 : memref, %1 : f32) func @print_memref_f32(%ptr : memref<*xf32>) } -