[mlir][spirv] Allow controlling subgroup size

This commit extends the `ResourceLimitsAttr` to support specifying
a minimal and maximal subgroup size, and extends `EntryPointABIAttr`
to support specifying the requested subgroup size. This is possible
now in Vulkan with the VK_EXT_subgroup_size_control extension.
For OpenCL it's possible to use the `SubgroupSize` execution mode
directly.

Reviewed By: ThomasRaoux

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

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.td

@@ -30,9 +30,13 @@ class SPIRV_Attr<string attrName, string attrMnemonic>
 
 // For entry functions, this attribute specifies information related to entry
 // points in the generated SPIR-V module:
-// 1) WorkGroup Size.
+// 1) [optional] Requested workgroup size.
+// 2) [optional] Requested subgroup size.
 def SPIRV_EntryPointABIAttr : SPIRV_Attr<"EntryPointABI", "entry_point_abi"> {
-  let parameters = (ins OptionalParameter<"DenseIntElementsAttr">:$local_size);
+  let parameters = (ins
+    OptionalParameter<"DenseI32ArrayAttr">:$workgroup_size,
+    OptionalParameter<"llvm::Optional<int>">:$subgroup_size
+  );
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
@@ -111,6 +115,11 @@ def SPIRV_ResourceLimitsAttr : SPIRV_Attr<"ResourceLimits", "resource_limits"> {
     // The default number of invocations in each subgroup.
     DefaultValuedParameter<"int", "32">:$subgroup_size,
 
+    // The minimum supported size if the subgroup size is controllable.
+    OptionalParameter<"mlir::Optional<int>">:$min_subgroup_size,
+    // The maximum supported size if the subgroup size is controllable.
+    OptionalParameter<"mlir::Optional<int>">:$max_subgroup_size,
+
     // The configurations of cooperative matrix operations
     // supported. Default is an empty list.
     DefaultValuedParameter<

mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h

@@ -86,8 +86,9 @@ bool needsInterfaceVarABIAttrs(TargetEnvAttr targetAttr);
 StringRef getEntryPointABIAttrName();
 
 /// Gets the EntryPointABIAttr given its fields.
-EntryPointABIAttr getEntryPointABIAttr(ArrayRef<int32_t> localSize,
-                                       MLIRContext *context);
+EntryPointABIAttr getEntryPointABIAttr(MLIRContext *context,
+                                       ArrayRef<int32_t> workgroupSize = {},
+                                       llvm::Optional<int> subgroupSize = {});
 
 /// Queries the entry point ABI on the nearest function-like op containing the
 /// given `op`. Returns null attribute if not found.
@@ -96,7 +97,7 @@ EntryPointABIAttr lookupEntryPointABI(Operation *op);
 /// Queries the local workgroup size from entry point ABI on the nearest
 /// function-like op containing the given `op`. Returns null attribute if not
 /// found.
-DenseIntElementsAttr lookupLocalWorkGroupSize(Operation *op);
+DenseI32ArrayAttr lookupLocalWorkGroupSize(Operation *op);
 
 /// Returns a default resource limits attribute that uses numbers from
 /// "Table 46. Required Limits" of the Vulkan spec.

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

@@ -171,12 +171,12 @@ SingleDimLaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
 LogicalResult WorkGroupSizeConversion::matchAndRewrite(
     gpu::BlockDimOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
-  auto workGroupSizeAttr = spirv::lookupLocalWorkGroupSize(op);
+  DenseI32ArrayAttr workGroupSizeAttr = spirv::lookupLocalWorkGroupSize(op);
   if (!workGroupSizeAttr)
     return failure();
 
-  auto val = workGroupSizeAttr
-                 .getValues<int32_t>()[static_cast<int32_t>(op.getDimension())];
+  int val =
+      workGroupSizeAttr.asArrayRef()[static_cast<int32_t>(op.getDimension())];
   auto convertedType =
       getTypeConverter()->convertType(op.getResult().getType());
   if (!convertedType)

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

@@ -119,14 +119,14 @@ LogicalResult SingleWorkgroupReduction::matchAndRewrite(
 
   // Query the shader interface for local workgroup size to make sure the
   // invocation configuration fits with the input memref's shape.
-  DenseIntElementsAttr localSize = spirv::lookupLocalWorkGroupSize(genericOp);
-  if (!localSize)
+  DenseI32ArrayAttr workgroupSize = spirv::lookupLocalWorkGroupSize(genericOp);
+  if (!workgroupSize)
     return failure();
 
-  if ((*localSize.begin()).getSExtValue() != originalInputType.getDimSize(0))
+  if (workgroupSize.asArrayRef()[0] != originalInputType.getDimSize(0))
     return failure();
-  if (llvm::any_of(llvm::drop_begin(localSize.getValues<APInt>(), 1),
-                   [](const APInt &size) { return !size.isOneValue(); }))
+  if (llvm::any_of(workgroupSize.asArrayRef().drop_front(),
+                   [](int size) { return size != 1; }))
     return failure();
 
   // TODO: Query the target environment to make sure the current

mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp

@@ -120,15 +120,16 @@ bool spirv::needsInterfaceVarABIAttrs(spirv::TargetEnvAttr targetAttr) {
 StringRef spirv::getEntryPointABIAttrName() { return "spirv.entry_point_abi"; }
 
 spirv::EntryPointABIAttr
-spirv::getEntryPointABIAttr(ArrayRef<int32_t> localSize, MLIRContext *context) {
-  if (localSize.empty())
-    return spirv::EntryPointABIAttr::get(context, nullptr);
-
-  assert(localSize.size() == 3);
-  return spirv::EntryPointABIAttr::get(
-      context, DenseElementsAttr::get<int32_t>(
-                   VectorType::get(3, IntegerType::get(context, 32)), localSize)
-                   .cast<DenseIntElementsAttr>());
+spirv::getEntryPointABIAttr(MLIRContext *context,
+                            ArrayRef<int32_t> workgroupSize,
+                            llvm::Optional<int> subgroupSize) {
+  DenseI32ArrayAttr workgroupSizeAttr;
+  if (!workgroupSize.empty()) {
+    assert(workgroupSize.size() == 3);
+    workgroupSizeAttr = DenseI32ArrayAttr::get(context, workgroupSize);
+  }
+  return spirv::EntryPointABIAttr::get(context, workgroupSizeAttr,
+                                       /*subgroupSize=*/llvm::None);
 }
 
 spirv::EntryPointABIAttr spirv::lookupEntryPointABI(Operation *op) {
@@ -144,9 +145,9 @@ spirv::EntryPointABIAttr spirv::lookupEntryPointABI(Operation *op) {
   return {};
 }
 
-DenseIntElementsAttr spirv::lookupLocalWorkGroupSize(Operation *op) {
+DenseI32ArrayAttr spirv::lookupLocalWorkGroupSize(Operation *op) {
   if (auto entryPoint = spirv::lookupEntryPointABI(op))
-    return entryPoint.getLocalSize();
+    return entryPoint.getWorkgroupSize();
 
   return {};
 }
@@ -162,6 +163,8 @@ spirv::getDefaultResourceLimits(MLIRContext *context) {
       /*max_compute_workgroup_invocations=*/128,
       /*max_compute_workgroup_size=*/b.getI32ArrayAttr({128, 128, 64}),
       /*subgroup_size=*/32,
+      /*min_subgroup_size=*/llvm::None,
+      /*max_subgroup_size=*/llvm::None,
       /*cooperative_matrix_properties_nv=*/ArrayAttr());
 }
 

mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp

@@ -13,10 +13,14 @@
 
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
 
+#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
+#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/Dialect/SPIRV/Utils/LayoutUtils.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/SetVector.h"
 
@@ -131,9 +135,10 @@ static LogicalResult lowerEntryPointABIAttr(spirv::FuncOp funcOp,
     return failure();
   }
 
-  spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnv(funcOp);
+  spirv::TargetEnvAttr targetEnvAttr = spirv::lookupTargetEnv(funcOp);
+  spirv::TargetEnv targetEnv(targetEnvAttr);
   FailureOr<spirv::ExecutionModel> executionModel =
-      spirv::getExecutionModel(targetEnv);
+      spirv::getExecutionModel(targetEnvAttr);
   if (failed(executionModel))
     return funcOp.emitRemark("lower entry point failure: could not select "
                              "execution model based on 'spirv.target_env'");
@@ -142,14 +147,36 @@ static LogicalResult lowerEntryPointABIAttr(spirv::FuncOp funcOp,
                                       funcOp, interfaceVars);
 
   // Specifies the spirv.ExecutionModeOp.
-  auto localSizeAttr = entryPointAttr.getLocalSize();
-  if (localSizeAttr) {
-    auto values = localSizeAttr.getValues<int32_t>();
-    SmallVector<int32_t, 3> localSize(values);
-    builder.create<spirv::ExecutionModeOp>(
-        funcOp.getLoc(), funcOp, spirv::ExecutionMode::LocalSize, localSize);
-    funcOp->removeAttr(entryPointAttrName);
+  if (DenseI32ArrayAttr workgroupSizeAttr = entryPointAttr.getWorkgroupSize()) {
+    Optional<ArrayRef<spirv::Capability>> caps =
+        spirv::getCapabilities(spirv::ExecutionMode::LocalSize);
+    if (!caps || targetEnv.allows(*caps)) {
+      builder.create<spirv::ExecutionModeOp>(funcOp.getLoc(), funcOp,
+                                             spirv::ExecutionMode::LocalSize,
+                                             workgroupSizeAttr.asArrayRef());
+      // Erase workgroup size.
+      entryPointAttr = spirv::EntryPointABIAttr::get(
+          entryPointAttr.getContext(), DenseI32ArrayAttr(),
+          entryPointAttr.getSubgroupSize());
+    }
   }
+  if (Optional<int> subgroupSize = entryPointAttr.getSubgroupSize()) {
+    Optional<ArrayRef<spirv::Capability>> caps =
+        spirv::getCapabilities(spirv::ExecutionMode::SubgroupSize);
+    if (!caps || targetEnv.allows(*caps)) {
+      builder.create<spirv::ExecutionModeOp>(funcOp.getLoc(), funcOp,
+                                             spirv::ExecutionMode::SubgroupSize,
+                                             *subgroupSize);
+      // Erase subgroup size.
+      entryPointAttr = spirv::EntryPointABIAttr::get(
+          entryPointAttr.getContext(), entryPointAttr.getWorkgroupSize(),
+          llvm::None);
+    }
+  }
+  if (entryPointAttr.getWorkgroupSize() || entryPointAttr.getSubgroupSize())
+    funcOp->setAttr(entryPointAttrName, entryPointAttr);
+  else
+    funcOp->removeAttr(entryPointAttrName);
   return success();
 }
 

mlir/test/Conversion/GPUToSPIRV/builtins.mlir

@@ -12,7 +12,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
   gpu.module @kernels {
     gpu.func @builtin_workgroup_id_x() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[WORKGROUPID]]
       // CHECK-NEXT: [[VEC:%.*]] = spirv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spirv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
@@ -38,7 +38,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
   gpu.module @kernels {
     gpu.func @builtin_workgroup_id_y() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[WORKGROUPID]]
       // CHECK-NEXT: [[VEC:%.*]] = spirv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spirv.CompositeExtract [[VEC]]{{\[}}1 : i32{{\]}}
@@ -62,7 +62,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
   gpu.module @kernels {
     gpu.func @builtin_workgroup_id_z() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[WORKGROUPID]]
       // CHECK-NEXT: [[VEC:%.*]] = spirv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spirv.CompositeExtract [[VEC]]{{\[}}2 : i32{{\]}}
@@ -85,7 +85,7 @@ module attributes {gpu.container_module} {
   // CHECK-LABEL:  spirv.module @{{.*}} Logical GLSL450
   gpu.module @kernels {
     gpu.func @builtin_workgroup_size_x() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 1, 1]>} {
       // The constant value is obtained from the spirv.entry_point_abi.
       // Note that this ignores the workgroup size specification in gpu.launch.
       // We may want to define gpu.workgroup_size and convert it to the entry
@@ -110,7 +110,7 @@ module attributes {gpu.container_module} {
   // CHECK-LABEL:  spirv.module @{{.*}} Logical GLSL450
   gpu.module @kernels {
     gpu.func @builtin_workgroup_size_y() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       // The constant value is obtained from the spirv.entry_point_abi.
       // CHECK: spirv.Constant 4 : i32
       %0 = gpu.block_dim y
@@ -132,7 +132,7 @@ module attributes {gpu.container_module} {
   // CHECK-LABEL:  spirv.module @{{.*}} Logical GLSL450
   gpu.module @kernels {
     gpu.func @builtin_workgroup_size_z() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       // The constant value is obtained from the spirv.entry_point_abi.
       // CHECK: spirv.Constant 1 : i32
       %0 = gpu.block_dim z
@@ -155,7 +155,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[LOCALINVOCATIONID:@.*]] built_in("LocalInvocationId")
   gpu.module @kernels {
     gpu.func @builtin_local_id_x() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[LOCALINVOCATIONID]]
       // CHECK-NEXT: [[VEC:%.*]] = spirv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spirv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
@@ -179,7 +179,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[NUMWORKGROUPS:@.*]] built_in("NumWorkgroups")
   gpu.module @kernels {
     gpu.func @builtin_num_workgroups_x() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[NUMWORKGROUPS]]
       // CHECK-NEXT: [[VEC:%.*]] = spirv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spirv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
@@ -196,7 +196,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[SUBGROUPID:@.*]] built_in("SubgroupId")
   gpu.module @kernels {
     gpu.func @builtin_subgroup_id() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[SUBGROUPID]]
       // CHECK-NEXT: {{%.*}} = spirv.Load "Input" [[ADDRESS]]
       %0 = gpu.subgroup_id : index
@@ -212,7 +212,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[NUMSUBGROUPS:@.*]] built_in("NumSubgroups")
   gpu.module @kernels {
     gpu.func @builtin_num_subgroups() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[NUMSUBGROUPS]]
       // CHECK-NEXT: {{%.*}} = spirv.Load "Input" [[ADDRESS]]
       %0 = gpu.num_subgroups : index
@@ -307,7 +307,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[GLOBALINVOCATIONID:@.*]] built_in("GlobalInvocationId")
   gpu.module @kernels {
     gpu.func @builtin_global_id_x() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[GLOBALINVOCATIONID]]
       // CHECK-NEXT: [[VEC:%.*]] = spirv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spirv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
@@ -331,7 +331,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[GLOBALINVOCATIONID:@.*]] built_in("GlobalInvocationId")
   gpu.module @kernels {
     gpu.func @builtin_global_id_y() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[GLOBALINVOCATIONID]]
       // CHECK-NEXT: [[VEC:%.*]] = spirv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spirv.CompositeExtract [[VEC]]{{\[}}1 : i32{{\]}}
@@ -355,7 +355,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[GLOBALINVOCATIONID:@.*]] built_in("GlobalInvocationId")
   gpu.module @kernels {
     gpu.func @builtin_global_id_z() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[GLOBALINVOCATIONID]]
       // CHECK-NEXT: [[VEC:%.*]] = spirv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spirv.CompositeExtract [[VEC]]{{\[}}2 : i32{{\]}}
@@ -373,7 +373,7 @@ module attributes {gpu.container_module} {
   // CHECK: spirv.GlobalVariable [[SUBGROUPSIZE:@.*]] built_in("SubgroupSize")
   gpu.module @kernels {
     gpu.func @builtin_subgroup_size() kernel
-      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
       // CHECK: [[ADDRESS:%.*]] = spirv.mlir.addressof [[SUBGROUPSIZE]]
       // CHECK-NEXT: {{%.*}} = spirv.Load "Input" [[ADDRESS]]
       %0 = gpu.subgroup_size : index

mlir/test/Conversion/GPUToSPIRV/entry-point.mlir

@@ -2,10 +2,10 @@
 // RUN: mlir-opt -test-spirv-entry-point-abi="workgroup-size=32" %s | FileCheck %s -check-prefix=WG32
 
 //      DEFAULT: gpu.func @foo()
-// DEFAULT-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<1> : vector<3xi32>>
+// DEFAULT-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>
 
 //      WG32: gpu.func @foo()
-// WG32-SAME:  spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 1, 1]> : vector<3xi32>>
+// WG32-SAME:  spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 1, 1]>
 
 gpu.module @kernels {
   gpu.func @foo() kernel {