[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

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.td
index 2f7cedc..80f1715 100644 (file)
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.td
+++ b/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<

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h b/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h
index fbdc16a..0f5e40e 100644 (file)
--- a/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h
+++ b/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.

diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
index e78d7e8..311f272 100644 (file)
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
+++ b/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)

diff --git a/mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp b/mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp
index 866d414..645cf4e 100644 (file)
--- a/mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp
+++ b/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

diff --git a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
index bfe95c8..73a167c 100644 (file)
--- a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
+++ b/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());
 }
 

diff --git a/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp b/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp
index 107d961..b383c64 100644 (file)
--- a/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp
+++ b/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();
 }
 

diff --git a/mlir/test/Conversion/GPUToSPIRV/builtins.mlir b/mlir/test/Conversion/GPUToSPIRV/builtins.mlir
index 6414d29..7649687 100644 (file)
--- a/mlir/test/Conversion/GPUToSPIRV/builtins.mlir
+++ b/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

diff --git a/mlir/test/Conversion/GPUToSPIRV/entry-point.mlir b/mlir/test/Conversion/GPUToSPIRV/entry-point.mlir
index 8536b2f..99369d1 100644 (file)
--- a/mlir/test/Conversion/GPUToSPIRV/entry-point.mlir
+++ b/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 {

diff --git a/mlir/test/Conversion/GPUToSPIRV/gpu-to-spirv.mlir b/mlir/test/Conversion/GPUToSPIRV/gpu-to-spirv.mlir
index a823829..7bf6f84 100644 (file)
--- a/mlir/test/Conversion/GPUToSPIRV/gpu-to-spirv.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/gpu-to-spirv.mlir
@@ -6,9 +6,9 @@ module attributes {gpu.container_module} {
     // CHECK-LABEL: spirv.func @basic_module_structure
     // CHECK-SAME: {{%.*}}: f32 {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0), StorageBuffer>}
     // CHECK-SAME: {{%.*}}: !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>}
-    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]> : vector<3xi32>>
+    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
     gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, #spirv.storage_class<StorageBuffer>>) 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]>} {
       // CHECK: spirv.Return
       gpu.return
     }
@@ -35,14 +35,14 @@ module attributes {gpu.container_module} {
     // CHECK-SAME: spirv.interface_var_abi = #spirv.interface_var_abi<(1, 2), StorageBuffer>
     // CHECK-SAME: !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>
     // CHECK-SAME: spirv.interface_var_abi = #spirv.interface_var_abi<(3, 0)>
-    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]> : vector<3xi32>>
+    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
     gpu.func @basic_module_structure_preset_ABI(
       %arg0 : f32
         {spirv.interface_var_abi = #spirv.interface_var_abi<(1, 2), StorageBuffer>},
       %arg1 : memref<12xf32, #spirv.storage_class<StorageBuffer>>
         {spirv.interface_var_abi = #spirv.interface_var_abi<(3, 0)>}) kernel
       attributes
-        {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]>: vector<3xi32>>} {
+        {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       // CHECK: spirv.Return
       gpu.return
     }
@@ -82,7 +82,7 @@ module attributes {gpu.container_module} {
         {spirv.interface_var_abi = #spirv.interface_var_abi<(1, 2), StorageBuffer>},
       %arg1 : memref<12xf32, #spirv.storage_class<StorageBuffer>>) kernel
     attributes
-      {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]>: vector<3xi32>>} {
+      {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       gpu.return
     }
   }
@@ -99,7 +99,7 @@ module attributes {gpu.container_module} {
       %arg1 : memref<12xf32, #spirv.storage_class<StorageBuffer>>
         {spirv.interface_var_abi = #spirv.interface_var_abi<(3, 0)>}) kernel
     attributes
-      {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]>: vector<3xi32>>} {
+      {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       gpu.return
     }
   }
@@ -111,7 +111,7 @@ module attributes {gpu.container_module} {
   gpu.module @kernels {
     // CHECK-LABEL: spirv.func @barrier
     gpu.func @barrier(%arg0 : f32, %arg1 : memref<12xf32, #spirv.storage_class<StorageBuffer>>) 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]>} {
       // CHECK: spirv.ControlBarrier <Workgroup>, <Workgroup>, <AcquireRelease|WorkgroupMemory>
       gpu.barrier
       gpu.return

diff --git a/mlir/test/Conversion/GPUToSPIRV/load-store.mlir b/mlir/test/Conversion/GPUToSPIRV/load-store.mlir
index 07fae0c..fa12da8 100644 (file)
--- a/mlir/test/Conversion/GPUToSPIRV/load-store.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/load-store.mlir
@@ -36,7 +36,7 @@ module attributes {
     // CHECK-SAME: %[[ARG5:.*]]: i32 {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 5), StorageBuffer>}
     // CHECK-SAME: %[[ARG6:.*]]: i32 {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 6), StorageBuffer>}
     gpu.func @load_store_kernel(%arg0: memref<12x4xf32, #spirv.storage_class<StorageBuffer>>, %arg1: memref<12x4xf32, #spirv.storage_class<StorageBuffer>>, %arg2: memref<12x4xf32, #spirv.storage_class<StorageBuffer>>, %arg3: index, %arg4: index, %arg5: index, %arg6: index) 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: %[[ADDRESSWORKGROUPID:.*]] = spirv.mlir.addressof @[[$WORKGROUPIDVAR]]
       // CHECK: %[[WORKGROUPID:.*]] = spirv.Load "Input" %[[ADDRESSWORKGROUPID]]
       // CHECK: %[[WORKGROUPIDX:.*]] = spirv.CompositeExtract %[[WORKGROUPID]]{{\[}}0 : i32{{\]}}

diff --git a/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir b/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir
index fa554f9..be2fcda 100644 (file)
--- a/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir
@@ -11,9 +11,9 @@ module attributes {
     //   CHECK-NOT:     spirv.interface_var_abi
     //  CHECK-SAME:     {{%.*}}: !spirv.ptr<!spirv.array<12 x f32>, CrossWorkgroup>
     //   CHECK-NOT:     spirv.interface_var_abi
-    //  CHECK-SAME:     spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]> : vector<3xi32>>
+    //  CHECK-SAME:     spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
     gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, #spirv.storage_class<CrossWorkgroup>>) 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]>} {
       gpu.return
     }
   }
@@ -44,9 +44,9 @@ module attributes {
     //   CHECK-NOT:     spirv.interface_var_abi
     //  CHECK-SAME:     {{%.*}}: !spirv.ptr<!spirv.array<12 x f32>, CrossWorkgroup>
     //   CHECK-NOT:     spirv.interface_var_abi
-    //  CHECK-SAME:     spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]> : vector<3xi32>>
+    //  CHECK-SAME:     spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
     gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, #spirv.storage_class<CrossWorkgroup>>) 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]>} {
       gpu.return
     }
   }

diff --git a/mlir/test/Conversion/GPUToSPIRV/shuffle.mlir b/mlir/test/Conversion/GPUToSPIRV/shuffle.mlir
index 2185884..d3d8ec0 100644 (file)
--- a/mlir/test/Conversion/GPUToSPIRV/shuffle.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/shuffle.mlir
@@ -8,7 +8,7 @@ module attributes {
 gpu.module @kernels {
   // CHECK-LABEL:  spirv.func @shuffle_xor()
   gpu.func @shuffle_xor() 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]>} {
     %mask = arith.constant 8 : i32
     %width = arith.constant 16 : i32
     %val = arith.constant 42.0 : f32
@@ -33,7 +33,7 @@ module attributes {
 
 gpu.module @kernels {
   gpu.func @shuffle_xor() 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]>} {
     %mask = arith.constant 8 : i32
     %width = arith.constant 16 : i32
     %val = arith.constant 42.0 : f32
@@ -57,7 +57,7 @@ module attributes {
 gpu.module @kernels {
   // CHECK-LABEL:  spirv.func @shuffle_idx()
   gpu.func @shuffle_idx() 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]>} {
     %mask = arith.constant 8 : i32
     %width = arith.constant 16 : i32
     %val = arith.constant 42.0 : f32

diff --git a/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv.mlir b/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv.mlir
index 9f2a27c..0c4b056 100644 (file)
--- a/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv.mlir
@@ -7,9 +7,9 @@ module attributes {
     // CHECK:       spirv.module @{{.*}} Logical GLSL450 {
     // CHECK-LABEL: spirv.func @gpu_wmma_load_op
     // CHECK-SAME: {{%.*}}: !spirv.ptr<!spirv.struct<(!spirv.array<512 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}
-    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]> : vector<3xi32>>
+    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
     gpu.func @gpu_wmma_load_op(%arg0 : memref<32x32xf16, #spirv.storage_class<StorageBuffer>>) 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]>} {
       %i = arith.constant 16 : index
       %j = arith.constant 16 : index
       // CHECK: {{%.*}} = spirv.NV.CooperativeMatrixLoad {{%.*}}, {{%.*}}, {{%.*}} :  !spirv.ptr<f32, StorageBuffer> as !spirv.coopmatrix<16x16xf16, Subgroup>
@@ -30,9 +30,9 @@ module attributes {
     // CHECK-LABEL: spirv.func @gpu_wmma_store_op
     // CHECK-SAME: {{%.*}}: !spirv.ptr<!spirv.struct<(!spirv.array<512 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}
     // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>})
-    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]> : vector<3xi32>>
+    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
     gpu.func @gpu_wmma_store_op(%arg0 : memref<32x32xf16, #spirv.storage_class<StorageBuffer>>, %arg1 : !gpu.mma_matrix<16x16xf16, "COp">) 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]>} {
       %i = arith.constant 16 : index
       %j = arith.constant 16 : index
       //  CHECK: spirv.NV.CooperativeMatrixStore {{%.*}}, {{%.*}}, {{%.*}}, {{%.*}} : !spirv.ptr<f32, StorageBuffer>, !spirv.coopmatrix<16x16xf16, Subgroup>
@@ -54,9 +54,9 @@ module attributes {
     // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}
     // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>}
     // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 2)>})
-    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 4, 1]> : vector<3xi32>>
+    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
     gpu.func @gpu_wmma_mma_op(%A : !gpu.mma_matrix<16x16xf16, "AOp">, %B : !gpu.mma_matrix<16x16xf16, "BOp">, %C : !gpu.mma_matrix<16x16xf16, "COp">) 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]>} {
       // CHECK: {{%.*}} = spirv.NV.CooperativeMatrixMulAdd {{%.*}}, {{%.*}}, {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup>, !spirv.coopmatrix<16x16xf16, Subgroup> -> !spirv.coopmatrix<16x16xf16, Subgroup>
       %D = gpu.subgroup_mma_compute %A, %B, %C : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
       // CHECK: spirv.Return
@@ -74,7 +74,7 @@ module attributes {
     // CHECK:       spirv.module @{{.*}} Logical GLSL450 {
     // CHECK-LABEL: spirv.func @gpu_wmma_constant_op
     gpu.func @gpu_wmma_constant_op() 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]>} {
       // CHECK: {{%.*}} = spirv.Constant
       %cst = arith.constant 1.0 : f16
       // CHECK: {{%.*}} = spirv.CompositeConstruct {{%.*}} : (f16) -> !spirv.coopmatrix<16x16xf16, Subgroup>
@@ -96,7 +96,7 @@ module attributes {
     // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}
     // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>})
     gpu.func @gpu_wmma_elementwise_op(%A : !gpu.mma_matrix<16x16xf16, "COp">, %B : !gpu.mma_matrix<16x16xf16, "COp">) 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]>} {
       // CHECK:  {{%.*}} = spirv.FAdd {{%.*}}, {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup>
       %C = gpu.subgroup_mma_elementwise addf %A, %B : (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">
       // CHECK:  {{%.*}} = spirv.FNegate {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup>

diff --git a/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir b/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir
index 17e8f45..fb9fff1 100644 (file)
--- a/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir
+++ b/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir
@@ -45,7 +45,7 @@ module attributes {
 // CHECK:        spirv.Return
 
 func.func @single_workgroup_reduction(%input: memref<16xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<1xi32, #spirv.storage_class<StorageBuffer>>) attributes {
-  spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>
+  spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>
 } {
   linalg.generic #single_workgroup_reduction_trait
       ins(%input : memref<16xi32, #spirv.storage_class<StorageBuffer>>)
@@ -104,7 +104,7 @@ module attributes {
     #spirv.vce<v1.3, [Shader, GroupNonUniformArithmetic], []>, #spirv.resource_limits<>>
 } {
 func.func @single_workgroup_reduction(%input: memref<16xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<1xi32, #spirv.storage_class<StorageBuffer>>) attributes {
-  spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 1, 1]>: vector<3xi32>>
+  spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 1, 1]>
 } {
   // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
   linalg.generic #single_workgroup_reduction_trait
@@ -135,7 +135,7 @@ module attributes {
     #spirv.vce<v1.3, [Shader, GroupNonUniformArithmetic], []>, #spirv.resource_limits<>>
 } {
 func.func @single_workgroup_reduction(%input: memref<16x8xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<16xi32, #spirv.storage_class<StorageBuffer>>) attributes {
-  spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 8, 1]>: vector<3xi32>>
+  spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 8, 1]>
 } {
   // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
   linalg.generic #single_workgroup_reduction_trait

diff --git a/mlir/test/Conversion/SPIRVToLLVM/lower-host-to-llvm-calls.mlir b/mlir/test/Conversion/SPIRVToLLVM/lower-host-to-llvm-calls.mlir
index 2c48194..f46b23c 100644 (file)
--- a/mlir/test/Conversion/SPIRVToLLVM/lower-host-to-llvm-calls.mlir
+++ b/mlir/test/Conversion/SPIRVToLLVM/lower-host-to-llvm-calls.mlir
@@ -32,7 +32,7 @@ module attributes {gpu.container_module, spirv.target_env = #spirv.target_env<#s
   }
 
   gpu.module @foo {
-    gpu.func @bar(%arg0: memref<6xi32>) kernel attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<1> : vector<3xi32>>} {
+    gpu.func @bar(%arg0: memref<6xi32>) kernel attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
       gpu.return
     }
   }

diff --git a/mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir b/mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir
index ed84746..82a7601 100644 (file)
--- a/mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir
@@ -34,16 +34,16 @@ func.func @spv_entry_point() attributes {
 // -----
 
 func.func @spv_entry_point() attributes {
-  // expected-error @+2 {{failed to parse SPIRV_EntryPointABIAttr parameter 'local_size' which is to be a `DenseIntElementsAttr`}}
-  // expected-error @+1 {{invalid kind of attribute specified}}
-  spirv.entry_point_abi = #spirv.entry_point_abi<local_size = 64>
+  // expected-error @+2 {{failed to parse SPIRV_EntryPointABIAttr parameter 'workgroup_size' which is to be a `DenseI32ArrayAttr`}}
+  // expected-error @+1 {{expected '['}}
+  spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = 64>
 } { return }
 
 // -----
 
 func.func @spv_entry_point() attributes {
-  // CHECK: {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[64, 1, 1]> : vector<3xi32>>}
-  spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[64, 1, 1]>: vector<3xi32>>
+  // CHECK: {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [64, 1, 1]>}
+  spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [64, 1, 1]>
 } { return }
 
 // -----
@@ -102,6 +102,26 @@ func.func @interface_var(
 // -----
 
 //===----------------------------------------------------------------------===//
+// spirv.resource_limits
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @resource_limits_all_default()
+func.func @resource_limits_all_default() attributes {
+  // CHECK-SAME: #spirv.resource_limits<>
+  limits = #spirv.resource_limits<>
+} { return }
+
+// -----
+
+// CHECK-LABEL: func @resource_limits_min_max_subgroup_size()
+func.func @resource_limits_min_max_subgroup_size() attributes {
+  // CHECK-SAME: #spirv.resource_limits<min_subgroup_size = 32, max_subgroup_size = 64>
+  limits = #spirv.resource_limits<min_subgroup_size = 32, max_subgroup_size=64>
+} { return }
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // spirv.target_env
 //===----------------------------------------------------------------------===//
 

diff --git a/mlir/test/Dialect/SPIRV/Linking/ModuleCombiner/deduplication.mlir b/mlir/test/Dialect/SPIRV/Linking/ModuleCombiner/deduplication.mlir
index 19169a2..1e06051 100644 (file)
--- a/mlir/test/Dialect/SPIRV/Linking/ModuleCombiner/deduplication.mlir
+++ b/mlir/test/Dialect/SPIRV/Linking/ModuleCombiner/deduplication.mlir
@@ -259,14 +259,14 @@ spirv.module Logical GLSL450 {
   spirv.func @kernel(
     %arg0: f32,
     %arg1: !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32>)>, CrossWorkgroup>) "None"
-  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]>} {
     spirv.Return
   }
 
   spirv.func @kernel_different_attr(
     %arg0: f32,
     %arg1: !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32>)>, CrossWorkgroup>) "None"
-  attributes {spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[64, 1, 1]> : vector<3xi32>>} {
+  attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [64, 1, 1]>} {
     spirv.Return
   }
 }

diff --git a/mlir/test/Dialect/SPIRV/Transforms/abi-interface-opencl.mlir b/mlir/test/Dialect/SPIRV/Transforms/abi-interface-opencl.mlir
index 07cc1c8..92efb0a 100644 (file)
--- a/mlir/test/Dialect/SPIRV/Transforms/abi-interface-opencl.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/abi-interface-opencl.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -spirv-lower-abi-attrs -verify-diagnostics %s -o - | FileCheck %s
+// RUN: mlir-opt -split-input-file -spirv-lower-abi-attrs %s | FileCheck %s
 
 module attributes {
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Kernel, Addresses], []>, #spirv.resource_limits<>>
@@ -6,12 +6,34 @@ module attributes {
   spirv.module Physical64 OpenCL {
     // CHECK-LABEL: spirv.module
     //       CHECK:   spirv.func [[FN:@.*]]({{%.*}}: f32, {{%.*}}: !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32>)>, CrossWorkgroup>
+    // We cannot generate SubgroupSize execution mode without necessary capability -- leave it alone.
+    //  CHECK-SAME:      #spirv.entry_point_abi<subgroup_size = 64>
     //       CHECK:   spirv.EntryPoint "Kernel" [[FN]]
     //       CHECK:   spirv.ExecutionMode [[FN]] "LocalSize", 32, 1, 1
     spirv.func @kernel(
       %arg0: f32,
       %arg1: !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32>)>, CrossWorkgroup>) "None"
-    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], subgroup_size = 64>} {
+      spirv.Return
+    }
+  }
+}
+
+// -----
+
+module attributes {
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Kernel, SubgroupDispatch], []>, #spirv.resource_limits<>>
+} {
+  spirv.module Physical64 OpenCL {
+    // CHECK-LABEL: spirv.module
+    //       CHECK:   spirv.func [[FN:@.*]]({{%.*}}: f32, {{%.*}}: !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32>)>, CrossWorkgroup>
+    //       CHECK:   spirv.EntryPoint "Kernel" [[FN]]
+    //       CHECK:   spirv.ExecutionMode [[FN]] "LocalSize", 32, 1, 1
+    //       CHECK:   spirv.ExecutionMode [[FN]] "SubgroupSize", 64
+    spirv.func @kernel(
+      %arg0: f32,
+      %arg1: !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32>)>, CrossWorkgroup>) "None"
+    attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 1, 1], subgroup_size = 64>} {
       spirv.Return
     }
   }

diff --git a/mlir/test/Dialect/SPIRV/Transforms/abi-interface.mlir b/mlir/test/Dialect/SPIRV/Transforms/abi-interface.mlir
index 8766177..4795a13 100644 (file)
--- a/mlir/test/Dialect/SPIRV/Transforms/abi-interface.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/abi-interface.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -spirv-lower-abi-attrs -verify-diagnostics %s -o - | FileCheck %s
+// RUN: mlir-opt -split-input-file -spirv-lower-abi-attrs %s | FileCheck %s
 
 module attributes {
   spirv.target_env = #spirv.target_env<
@@ -7,15 +7,17 @@ module attributes {
 
 // CHECK-LABEL: spirv.module
 spirv.module Logical GLSL450 {
-  // CHECK-DAG:    spirv.GlobalVariable [[VAR0:@.*]] bind(0, 0) : !spirv.ptr<!spirv.struct<(f32 [0])>, StorageBuffer>
-  // CHECK-DAG:    spirv.GlobalVariable [[VAR1:@.*]] bind(0, 1) : !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>
-  // CHECK:    spirv.func [[FN:@.*]]()
+  //  CHECK-DAG:    spirv.GlobalVariable [[VAR0:@.*]] bind(0, 0) : !spirv.ptr<!spirv.struct<(f32 [0])>, StorageBuffer>
+  //  CHECK-DAG:    spirv.GlobalVariable [[VAR1:@.*]] bind(0, 1) : !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>
+  //      CHECK:    spirv.func [[FN:@.*]]()
+  // We cannot generate SubgroupSize execution mode for Shader capability -- leave it alone.
+  // CHECK-SAME:      #spirv.entry_point_abi<subgroup_size = 64>
   spirv.func @kernel(
     %arg0: f32
            {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0), StorageBuffer>},
     %arg1: !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32>)>, StorageBuffer>
            {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>}) "None"
-  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], subgroup_size = 64>} {
     // CHECK: [[ARG1:%.*]] = spirv.mlir.addressof [[VAR1]]
     // CHECK: [[ADDRESSARG0:%.*]] = spirv.mlir.addressof [[VAR0]]
     // CHECK: [[CONST0:%.*]] = spirv.Constant 0 : i32

diff --git a/mlir/test/Dialect/SPIRV/Transforms/abi-load-store.mlir b/mlir/test/Dialect/SPIRV/Transforms/abi-load-store.mlir
index b7368b7..6a5edc7 100644 (file)
--- a/mlir/test/Dialect/SPIRV/Transforms/abi-load-store.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/abi-load-store.mlir
@@ -38,7 +38,7 @@ spirv.module Logical GLSL450 {
     {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 5), StorageBuffer>},
     %arg6: i32
     {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 6), StorageBuffer>}) "None"
-  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]>} {
     // CHECK: [[ADDRESSARG6:%.*]] = spirv.mlir.addressof [[VAR6]]
     // CHECK: [[CONST6:%.*]] = spirv.Constant 0 : i32
     // CHECK: [[ARG6PTR:%.*]] = spirv.AccessChain [[ADDRESSARG6]]{{\[}}[[CONST6]]

diff --git a/mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp b/mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp
index 5fba9a3..129ba72 100644 (file)
--- a/mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp
+++ b/mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp
@@ -59,7 +59,7 @@ void TestSpirvEntryPointABIPass::runOnOperation() {
                                              workgroupSize.end());
     workgroupSizeVec.resize(3, 1);
     gpuFunc->setAttr(attrName,
-                     spirv::getEntryPointABIAttr(workgroupSizeVec, context));
+                     spirv::getEntryPointABIAttr(context, workgroupSizeVec));
   }
 }
 

diff --git a/mlir/test/mlir-spirv-cpu-runner/double.mlir b/mlir/test/mlir-spirv-cpu-runner/double.mlir
index 577aff2..b9a3f0d 100644 (file)
--- a/mlir/test/mlir-spirv-cpu-runner/double.mlir
+++ b/mlir/test/mlir-spirv-cpu-runner/double.mlir
@@ -11,7 +11,7 @@ module attributes {
 } {
   gpu.module @kernels {
     gpu.func @double(%arg0 : memref<6xi32>, %arg1 : memref<6xi32>)
-      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[1, 1, 1]>: vector<3xi32>>} {
+      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
       %factor = arith.constant 2 : i32
 
       %i0 = arith.constant 0 : index

diff --git a/mlir/test/mlir-spirv-cpu-runner/simple_add.mlir b/mlir/test/mlir-spirv-cpu-runner/simple_add.mlir
index 0e222e3..7b8d964 100644 (file)
--- a/mlir/test/mlir-spirv-cpu-runner/simple_add.mlir
+++ b/mlir/test/mlir-spirv-cpu-runner/simple_add.mlir
@@ -11,7 +11,7 @@ module attributes {
 } {
   gpu.module @kernels {
     gpu.func @sum(%arg0 : memref<3xf32>, %arg1 : memref<3x3xf32>, %arg2 :  memref<3x3x3xf32>)
-      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[1, 1, 1]>: vector<3xi32>>} {
+      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
       %i0 = arith.constant 0 : index
       %i1 = arith.constant 1 : index
       %i2 = arith.constant 2 : index

diff --git a/mlir/test/mlir-vulkan-runner/addf.mlir b/mlir/test/mlir-vulkan-runner/addf.mlir
index 7d8a580..407325a 100644 (file)
--- a/mlir/test/mlir-vulkan-runner/addf.mlir
+++ b/mlir/test/mlir-vulkan-runner/addf.mlir
@@ -8,7 +8,7 @@ module attributes {
 } {
   gpu.module @kernels {
     gpu.func @kernel_add(%arg0 : memref<8xf32>, %arg1 : memref<8xf32>, %arg2 : memref<8xf32>)
-      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[1, 1, 1]>: vector<3xi32>>} {
+      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
       %0 = gpu.block_id x
       %1 = memref.load %arg0[%0] : memref<8xf32>
       %2 = memref.load %arg1[%0] : memref<8xf32>

diff --git a/mlir/test/mlir-vulkan-runner/addi.mlir b/mlir/test/mlir-vulkan-runner/addi.mlir
index 3dfbc4b..5490924 100644 (file)
--- a/mlir/test/mlir-vulkan-runner/addi.mlir
+++ b/mlir/test/mlir-vulkan-runner/addi.mlir
@@ -8,7 +8,7 @@ module attributes {
 } {
   gpu.module @kernels {
     gpu.func @kernel_addi(%arg0 : memref<8xi32>, %arg1 : memref<8x8xi32>, %arg2 : memref<8x8x8xi32>)
-      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[1, 1, 1]>: vector<3xi32>>} {
+      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
       %x = gpu.block_id x
       %y = gpu.block_id y
       %z = gpu.block_id z

diff --git a/mlir/test/mlir-vulkan-runner/addi8.mlir b/mlir/test/mlir-vulkan-runner/addi8.mlir
index 7b5bf38..13bdad6 100644 (file)
--- a/mlir/test/mlir-vulkan-runner/addi8.mlir
+++ b/mlir/test/mlir-vulkan-runner/addi8.mlir
@@ -8,7 +8,7 @@ module attributes {
 } {
   gpu.module @kernels {
     gpu.func @kernel_addi(%arg0 : memref<8xi8>, %arg1 : memref<8x8xi8>, %arg2 : memref<8x8x8xi32>)
-      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[1, 1, 1]>: vector<3xi32>>} {
+      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
       %x = gpu.block_id x
       %y = gpu.block_id y
       %z = gpu.block_id z

diff --git a/mlir/test/mlir-vulkan-runner/mulf.mlir b/mlir/test/mlir-vulkan-runner/mulf.mlir
index 41682a3..b87c006 100644 (file)
--- a/mlir/test/mlir-vulkan-runner/mulf.mlir
+++ b/mlir/test/mlir-vulkan-runner/mulf.mlir
@@ -8,7 +8,7 @@ module attributes {
 } {
   gpu.module @kernels {
     gpu.func @kernel_mul(%arg0 : memref<4x4xf32>, %arg1 : memref<4x4xf32>, %arg2 : memref<4x4xf32>)
-      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[1, 1, 1]>: vector<3xi32>>} {
+      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
       %x = gpu.block_id x
       %y = gpu.block_id y
       %1 = memref.load %arg0[%x, %y] : memref<4x4xf32>

diff --git a/mlir/test/mlir-vulkan-runner/subf.mlir b/mlir/test/mlir-vulkan-runner/subf.mlir
index 22d9cc1..28facaa 100644 (file)
--- a/mlir/test/mlir-vulkan-runner/subf.mlir
+++ b/mlir/test/mlir-vulkan-runner/subf.mlir
@@ -9,7 +9,7 @@ module attributes {
 } {
   gpu.module @kernels {
     gpu.func @kernel_sub(%arg0 : memref<8x4x4xf32>, %arg1 : memref<4x4xf32>, %arg2 : memref<8x4x4xf32>)
-      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[1, 1, 1]>: vector<3xi32>>} {
+      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
       %x = gpu.block_id x
       %y = gpu.block_id y
       %z = gpu.block_id z

diff --git a/mlir/test/mlir-vulkan-runner/time.mlir b/mlir/test/mlir-vulkan-runner/time.mlir
index 9e2c762..b814e7c 100644 (file)
--- a/mlir/test/mlir-vulkan-runner/time.mlir
+++ b/mlir/test/mlir-vulkan-runner/time.mlir
@@ -11,7 +11,7 @@ module attributes {
 } {
   gpu.module @kernels {
     gpu.func @kernel_add(%arg0 : memref<16384xf32>, %arg1 : memref<16384xf32>, %arg2 : memref<16384xf32>)
-      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[128, 1, 1]>: vector<3xi32>>} {
+      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [128, 1, 1]>} {
       %bid = gpu.block_id x
       %tid = gpu.thread_id x
       %cst = arith.constant 128 : index