Remove rotation semantic in gpu.shufflw up/down

There is no such semantic in SPIRV OpGroupNonUniformShuffleUp and
OpGroupNonUniformShuffleDown. In addition, there is no such semantic in
NVVM shfl intrinsics.

Refer to NVVM IR spec
https://docs.nvidia.com/cuda/archive/12.2.1/nvvm-ir-spec/index.html#data-movement

"If the computed source lane index j is in range, the returned i32 value
will be the value of %a from lane j; otherwise, it will be the the value
of %a from the current thread. If the thread corresponding to lane j is
inactive, then the returned i32 value is undefined."

Author: Hsiangkai

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

@@ -1332,7 +1332,8 @@ def GPU_ShuffleOp : GPU_Op<
     %3, %4 = gpu.shuffle down %0, %cst1, %width : f32
     ```
 
-    For lane `k`, returns the value from lane `(k + 1) % width`.
+    For lane `k`, returns the value from lane `(k + cst1)`. The resulting value
+    is undefined if the lane is out of bounds in the subgroup.
 
     `up` example:
 
@@ -1341,7 +1342,8 @@ def GPU_ShuffleOp : GPU_Op<
     %5, %6 = gpu.shuffle up %0, %cst1, %width : f32
     ```
 
-    For lane `k`, returns the value from lane `(k - 1) % width`.
+    For lane `k`, returns the value from lane `(k - cst1)`. The resulting value
+    is undefined if the lane is out of bounds in the subgroup.
 
     `idx` example:
 

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

@@ -430,12 +430,10 @@ LogicalResult GPUShuffleConversion::matchAndRewrite(
   unsigned subgroupSize =
       targetEnv.getAttr().getResourceLimits().getSubgroupSize();
   IntegerAttr widthAttr;
-  // The width argument specifies the number of lanes that participate in the
-  // shuffle. The width value should not exceed the subgroup limit.
   if (!matchPattern(shuffleOp.getWidth(), m_Constant(&widthAttr)) ||
-      widthAttr.getValue().getZExtValue() > subgroupSize)
+      widthAttr.getValue().getZExtValue() != subgroupSize)
     return rewriter.notifyMatchFailure(
-        shuffleOp, "shuffle width is larger than target subgroup size");
+        shuffleOp, "shuffle width and target subgroup size mismatch");
 
   Location loc = shuffleOp.getLoc();
   Value trueVal = spirv::ConstantOp::getOne(rewriter.getI1Type(),
@@ -453,19 +451,14 @@ LogicalResult GPUShuffleConversion::matchAndRewrite(
         loc, scope, adaptor.getValue(), adaptor.getOffset());
     break;
   case gpu::ShuffleMode::DOWN:
-    result = rewriter.create<spirv::GroupNonUniformRotateKHROp>(
-        loc, scope, adaptor.getValue(), adaptor.getOffset(),
-        shuffleOp.getWidth());
+    result = rewriter.create<spirv::GroupNonUniformShuffleDownOp>(
+        loc, scope, adaptor.getValue(), adaptor.getOffset());
     break;
-  case gpu::ShuffleMode::UP: {
-    Value offsetForShuffleDown = rewriter.create<arith::SubIOp>(
-        loc, shuffleOp.getWidth(), adaptor.getOffset());
-    result = rewriter.create<spirv::GroupNonUniformRotateKHROp>(
-        loc, scope, adaptor.getValue(), offsetForShuffleDown,
-        shuffleOp.getWidth());
+  case gpu::ShuffleMode::UP:
+    result = rewriter.create<spirv::GroupNonUniformShuffleUpOp>(
+        loc, scope, adaptor.getValue(), adaptor.getOffset());
     break;
   }
-  }
 
   rewriter.replaceOp(shuffleOp, {result, trueVal});
   return success();

mlir/test/Conversion/GPUToSPIRV/shuffle.mlir

@@ -26,6 +26,29 @@ gpu.module @kernels {
 
 // -----
 
+module attributes {
+  gpu.container_module,
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, GroupNonUniformShuffle], []>, #spirv.resource_limits<subgroup_size = 32>>
+} {
+
+gpu.module @kernels {
+  gpu.func @shuffle_xor() kernel
+    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
+
+    // Cannot convert due to shuffle width and target subgroup size mismatch
+    // expected-error @+1 {{failed to legalize operation 'gpu.shuffle'}}
+    %result, %valid = gpu.shuffle xor %val, %mask, %width : f32
+    gpu.return
+  }
+}
+
+}
+
+// -----
+
 module attributes {
   gpu.container_module,
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, GroupNonUniformShuffle], []>, #spirv.resource_limits<subgroup_size = 16>>
@@ -54,7 +77,7 @@ gpu.module @kernels {
 
 module attributes {
   gpu.container_module,
-  spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, GroupNonUniformShuffle, GroupNonUniformRotateKHR], []>,
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, GroupNonUniformShuffle, GroupNonUniformShuffleRelative], []>,
     #spirv.resource_limits<subgroup_size = 16>>
 } {
 
@@ -70,7 +93,7 @@ gpu.module @kernels {
     // CHECK: %[[WIDTH:.+]] = spirv.Constant 16 : i32
     // CHECK: %[[VAL:.+]] = spirv.Constant 4.200000e+01 : f32
     // CHECK: %{{.+}} = spirv.Constant true
-    // CHECK: %{{.+}} = spirv.GroupNonUniformRotateKHR <Subgroup> %[[VAL]], %[[OFFSET]], cluster_size(%[[WIDTH]]) : f32, i32, i32 -> f32
+    // CHECK: %{{.+}} = spirv.GroupNonUniformShuffleDown <Subgroup> %[[VAL]], %[[OFFSET]] : f32, i32
     %result, %valid = gpu.shuffle down %val, %offset, %width : f32
     gpu.return
   }
@@ -82,7 +105,7 @@ gpu.module @kernels {
 
 module attributes {
   gpu.container_module,
-  spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, GroupNonUniformShuffle, GroupNonUniformRotateKHR], []>,
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, GroupNonUniformShuffle, GroupNonUniformShuffleRelative], []>,
     #spirv.resource_limits<subgroup_size = 16>>
 } {
 
@@ -98,8 +121,7 @@ gpu.module @kernels {
     // CHECK: %[[WIDTH:.+]] = spirv.Constant 16 : i32
     // CHECK: %[[VAL:.+]] = spirv.Constant 4.200000e+01 : f32
     // CHECK: %{{.+}} = spirv.Constant true
-    // CHECK: %[[DOWN_OFFSET:.+]] = spirv.Constant 12 : i32
-    // CHECK: %{{.+}} = spirv.GroupNonUniformRotateKHR <Subgroup> %[[VAL]], %[[DOWN_OFFSET]], cluster_size(%[[WIDTH]]) : f32, i32, i32 -> f32
+    // CHECK: %{{.+}} = spirv.GroupNonUniformShuffleUp <Subgroup> %[[VAL]], %[[OFFSET]] : f32, i32
     %result, %valid = gpu.shuffle up %val, %offset, %width : f32
     gpu.return
   }