[Mlir-commits] [mlir] [mlir][gpu][spirv] Add patterns for gpu.shuffle up/down (PR #139105)

[Hsiangkai Wang]

================
@@ -430,10 +430,12 @@ 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 and target subgroup size mismatch");
+        shuffleOp, "shuffle width is larger than target subgroup size");
----------------
Hsiangkai wrote:

After revisiting spirv and nvvm specification, we should remove the rotation semantic in gpu.shuffle up/down.
nvvm:
https://docs.nvidia.com/cuda/archive/12.2.1/nvvm-ir-spec/index.html#data-movement
spirv:
https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpGroupNonUniformShuffleUp
https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpGroupNonUniformShuffleDown

In addition, we need another gpu.rotate operation for rotation semantics. I will have another patch for rotation operation.

https://github.com/llvm/llvm-project/pull/139105