[mlir][spirv] Add instruction OpGroupNonUniformRotateKHR

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

@@ -1361,7 +1361,8 @@ def SPIRV_GroupNonUniformBallotBitCountOp : SPIRV_Op<"GroupNonUniformBallotBitCo
 
 // -----
 
-def SPIRV_GroupNonUniformRotateKHR : SPIRV_Op<"GroupNonUniformRotateKHR", []> {
+def SPIRV_GroupNonUniformRotateKHROp : SPIRV_Op<"GroupNonUniformRotateKHR", [
+  Pure, AllTypesMatch<["value", "result"]>]> {
   let summary = [{
     Rotate values across invocations within a subgroup.
   }];
@@ -1424,11 +1425,9 @@ def SPIRV_GroupNonUniformRotateKHR : SPIRV_Op<"GroupNonUniformRotateKHR", []> {
   );
 
   let results = (outs
-    SPIRV_Type:$result
+    AnyTypeOf<[SPIRV_ScalarOrVectorOf<SPIRV_Float>, SPIRV_ScalarOrVectorOf<SPIRV_Integer>, SPIRV_ScalarOrVectorOf<SPIRV_Bool>]>:$result
   );
 
-  let hasVerifier = 0;
-
   let assemblyFormat = [{
     $execution_scope `,` $value `,` $delta (`,` `cluster_size` `(` $cluster_size^ `)`)? attr-dict `:` type($value) `,` type($delta) (`,` type($cluster_size)^)? `->` type(results)
   }];

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

@@ -304,6 +304,38 @@ LogicalResult GroupNonUniformLogicalXorOp::verify() {
   return verifyGroupNonUniformArithmeticOp<GroupNonUniformLogicalXorOp>(*this);
 }
 
+//===----------------------------------------------------------------------===//
+// spirv.GroupNonUniformRotateKHR
+//===----------------------------------------------------------------------===//
+
+LogicalResult GroupNonUniformRotateKHROp::verify() {
+  spirv::Scope scope = getExecutionScope();
+  if (scope != spirv::Scope::Workgroup && scope != spirv::Scope::Subgroup)
+    return emitOpError("execution scope must be 'Workgroup' or 'Subgroup'");
+
+  if (getDelta().getType().isSignedInteger())
+    return emitOpError("delta must be a singless/unsigned integer");
+
+  auto clusterSizeVal = getClusterSize();
+  if (clusterSizeVal) {
+    if (clusterSizeVal.getType().isSignedInteger())
+      return emitOpError("cluster size must be a singless/unsigned integer");
+
+    mlir::Operation *defOp = clusterSizeVal.getDefiningOp();
+    int32_t clusterSize = 0;
+
+    if (failed(extractValueFromConstOp(defOp, clusterSize)))
+      return emitOpError(
+          "cluster size operand must come from a constant op");
+
+    if (!llvm::isPowerOf2_32(clusterSize))
+      return emitOpError(
+          "cluster size operand must be a power of two");
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Group op verification
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir

@@ -627,3 +627,47 @@ func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32) -> f32 {
   %0 = spirv.GroupNonUniformRotateKHR <Workgroup>, %val, %delta, cluster_size(%four) : f32, i32, i32 -> f32
   return %0: f32
 }
+
+// -----
+
+func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32) -> f32 {
+  %four = spirv.Constant 4 : i32
+  // expected-error @+1 {{execution scope must be 'Workgroup' or 'Subgroup'}}
+  %0 = spirv.GroupNonUniformRotateKHR <Device>, %val, %delta, cluster_size(%four) : f32, i32, i32 -> f32
+  return %0: f32
+}
+
+// -----
+
+func.func @group_non_uniform_rotate_khr(%val: f32, %delta: si32) -> f32 {
+  %four = spirv.Constant 4 : i32
+  // expected-error @+1 {{delta must be a singless/unsigned integer}}
+  %0 = spirv.GroupNonUniformRotateKHR <Subgroup>, %val, %delta, cluster_size(%four) : f32, si32, i32 -> f32
+  return %0: f32
+}
+
+// -----
+
+func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32) -> f32 {
+  %four = spirv.Constant 4 : si32
+  // expected-error @+1 {{cluster size must be a singless/unsigned integer}}
+  %0 = spirv.GroupNonUniformRotateKHR <Subgroup>, %val, %delta, cluster_size(%four) : f32, i32, si32 -> f32
+  return %0: f32
+}
+
+// -----
+
+func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32, %four: i32) -> f32 {
+  // expected-error @+1 {{cluster size operand must come from a constant op}}
+  %0 = spirv.GroupNonUniformRotateKHR <Subgroup>, %val, %delta, cluster_size(%four) : f32, i32, i32 -> f32
+  return %0: f32
+}
+
+// -----
+
+func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32) -> f32 {
+  %five = spirv.Constant 5 : i32
+  // expected-error @+1 {{cluster size operand must be a power of two}}
+  %0 = spirv.GroupNonUniformRotateKHR <Subgroup>, %val, %delta, cluster_size(%five) : f32, i32, i32 -> f32
+  return %0: f32
+}