[mlir][spirv] Add instruction OpGroupNonUniformRotateKHR by Hsiangkai · Pull Request #133428 · llvm/llvm-project (original) (raw)
@llvm/pr-subscribers-mlir-spirv
@llvm/pr-subscribers-mlir
Author: Hsiangkai Wang (Hsiangkai)
Changes
Add an instruction under the extension SPV_KHR_subgroup_rotate.
The specification for the extension is here:
https://github.khronos.org/SPIRV-Registry/extensions/KHR/SPV_KHR_subgroup_rotate.html
Full diff: https://github.com/llvm/llvm-project/pull/133428.diff
3 Files Affected:
- (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td (+3-1)
- (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVNonUniformOps.td (+75)
- (modified) mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir (+23)
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td index d5359da2a590e..cd5d201c3d5da 100644 --- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td +++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td @@ -4489,6 +4489,7 @@ def SPIRV_OC_OpGroupNonUniformBitwiseXor : I32EnumAttrCase<"OpGroupNonUnifo def SPIRV_OC_OpGroupNonUniformLogicalAnd : I32EnumAttrCase<"OpGroupNonUniformLogicalAnd", 362>; def SPIRV_OC_OpGroupNonUniformLogicalOr : I32EnumAttrCase<"OpGroupNonUniformLogicalOr", 363>; def SPIRV_OC_OpGroupNonUniformLogicalXor : I32EnumAttrCase<"OpGroupNonUniformLogicalXor", 364>; +def SPIRV_OC_OpGroupNonUniformRotateKHR : I32EnumAttrCase<"OpGroupNonUniformRotateKHR", 4431>; def SPIRV_OC_OpSubgroupBallotKHR : I32EnumAttrCase<"OpSubgroupBallotKHR", 4421>; def SPIRV_OC_OpSDot : I32EnumAttrCase<"OpSDot", 4450>; def SPIRV_OC_OpUDot : I32EnumAttrCase<"OpUDot", 4451>; @@ -4598,7 +4599,8 @@ def SPIRV_OpcodeAttr : SPIRV_OC_OpGroupNonUniformFMax, SPIRV_OC_OpGroupNonUniformBitwiseAnd, SPIRV_OC_OpGroupNonUniformBitwiseOr, SPIRV_OC_OpGroupNonUniformBitwiseXor, SPIRV_OC_OpGroupNonUniformLogicalAnd, SPIRV_OC_OpGroupNonUniformLogicalOr,
SPIRV_OC_OpGroupNonUniformLogicalXor, SPIRV_OC_OpSubgroupBallotKHR,
SPIRV_OC_OpGroupNonUniformLogicalXor, SPIRV_OC_OpGroupNonUniformRotateKHR,SPIRV_OC_OpSubgroupBallotKHR, SPIRV_OC_OpSDot, SPIRV_OC_OpUDot, SPIRV_OC_OpSUDot, SPIRV_OC_OpSDotAccSat, SPIRV_OC_OpUDotAccSat, SPIRV_OC_OpSUDotAccSat, SPIRV_OC_OpTypeCooperativeMatrixKHR, SPIRV_OC_OpCooperativeMatrixLoadKHR,
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVNonUniformOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVNonUniformOps.td index 98e435c18d3d7..f195adfc0e73d 100644 --- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVNonUniformOps.td +++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVNonUniformOps.td @@ -1361,4 +1361,79 @@ def SPIRV_GroupNonUniformBallotBitCountOp : SPIRV_Op<"GroupNonUniformBallotBitCo
// -----
+def SPIRV_GroupNonUniformRotateKHR : SPIRV_Op<"GroupNonUniformRotateKHR", []> {
- let summary = [{
- Rotate values across invocations within a subgroup.
- }];
- let description = [{
- Return the Value of the invocation whose id within the group is calculated
- as follows:
- LocalId = SubgroupLocalInvocationId if Execution is Subgroup or
LocalInvocationId if Execution is Workgroup- RotationGroupSize = ClusterSize when ClusterSize is present, otherwise
- RotationGroupSize = SubgroupMaxSize if the Kernel capability is declared
and SubgroupSize if not.- Invocation ID = ( (LocalId + Delta) & (RotationGroupSize - 1) ) +
(LocalId & ~(RotationGroupSize - 1))- Result Type must be a scalar or vector of floating-point type, integer
- type, or Boolean type.
- Execution is a Scope. It must be either Workgroup or Subgroup.
- The type of Value must be the same as Result Type.
- Delta must be a scalar of integer type, whose Signedness operand is 0.
- Delta must be dynamically uniform within Execution.
- Delta is treated as unsigned and the resulting value is undefined if the
- selected lane is inactive.
- ClusterSize is the size of cluster to use. ClusterSize must be a scalar of
- integer type, whose Signedness operand is 0. ClusterSize must come from a
- constant instruction. Behavior is undefined unless ClusterSize is at least
- 1 and a power of 2. If ClusterSize is greater than the declared
- SubGroupSize, executing this instruction results in undefined behavior.
Example:
- ```mlir
- %four = spirv.Constant 4 : i32
- %0 = spirv.GroupNonUniformRotateKHR , %value, %delta : f32, i32 -> f32
- %1 = spirv.GroupNonUniformRotateKHR , %value, %delta,
clustersize(%four) : f32, i32, i32 -> f32- ```
- }];
- let availability = [
- MinVersion,
- MaxVersion,
- Extension<[]>,
- Capability<[SPIRV_C_GroupNonUniformRotateKHR]>
- ];
- let arguments = (ins
- SPIRV_ScopeAttr:$execution_scope,
- SPIRV_Type:$value,
- SPIRV_Integer:$delta,
- Optional:$cluster_size
- );
- let results = (outs
- SPIRV_Type:$result
- );
- let hasVerifier = 0;
- let assemblyFormat = [{
- executionscope‘,‘execution_scope
,executionscope‘,‘value,delta(‘,‘‘clustersize‘‘(‘delta (,cluster_size(delta(‘,‘‘clustersize‘‘(‘cluster_size^))? attr-dict:type($value),type($delta) (,type($cluster_size)^)?->type(results) - }]; +}
- +// -----
- #endif // MLIR_DIALECT_SPIRV_IR_NON_UNIFORM_OPS
diff --git a/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir b/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir index 60ae1584d29fb..60b99d51363e9 100644 --- a/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir +++ b/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir @@ -604,3 +604,26 @@ func.func @group_non_uniform_logical_xor(%val: i32) -> i32 { %0 = spirv.GroupNonUniformLogicalXor %val : i32 -> i32 return %0: i32 } + +// ----- + +//===----------------------------------------------------------------------===// +// spirv.GroupNonUniformRotateKHR +//===----------------------------------------------------------------------===// + +// CHECK-LABEL: @group_non_uniform_rotate_khr +func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32) -> f32 {
- // CHECK: %{{.+}} = spirv.GroupNonUniformRotateKHR , %{{.+}} : f32, i32 -> f32
- %0 = spirv.GroupNonUniformRotateKHR , %val, %delta : f32, i32 -> f32
- return %0: f32 +}
- +// -----
- +// CHECK-LABEL: @group_non_uniform_rotate_khr +func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32) -> f32 {
- // CHECK: %{{.+}} = spirv.GroupNonUniformRotateKHR , %{{.+}} : f32, i32, i32 -> f32
- %four = spirv.Constant 4 : i32
- %0 = spirv.GroupNonUniformRotateKHR , %val, %delta, cluster_size(%four) : f32, i32, i32 -> f32
- return %0: f32 +}