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

[Hsiangkai Wang]

================
@@ -450,8 +452,19 @@ LogicalResult GPUShuffleConversion::matchAndRewrite(
     result = rewriter.create<spirv::GroupNonUniformShuffleOp>(
         loc, scope, adaptor.getValue(), adaptor.getOffset());
     break;
-  default:
-    return rewriter.notifyMatchFailure(shuffleOp, "unimplemented shuffle mode");
+  case gpu::ShuffleMode::DOWN:
+    result = rewriter.create<spirv::GroupNonUniformRotateKHROp>(
+        loc, scope, adaptor.getValue(), adaptor.getOffset(),
+        shuffleOp.getWidth());
+    break;
+  case gpu::ShuffleMode::UP: {
+    Value offsetForShuffleDown = rewriter.create<arith::SubIOp>(
+        loc, shuffleOp.getWidth(), adaptor.getOffset());
----------------
Hsiangkai wrote:

> Is this correct? Shouldn't we use subgroup size instead of width here? Width is just the cluster size which can be smaller than the subgroup size?

>From the MLIR document, `width` is the number of lanes participate in the shuffle. Take subgroup size 16 and width 4 as an example, shuffle up with offset 1 should be 3, 0, 1, 2, 4, 5, 6, ..., 14, 15. It's equivalent to spirv rotation with Delta 3 (i.e., 4 minus 1) and ClusterSize 4, isn't it? If we use subgroup size to minus offset, we will get a spirv rotation with Delta 15 and ClusterSize 4. It's a bit weird to have Delta larger than ClusterSize.

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