[Mlir-commits] [mlir] [MLIR][GPUToLLVMSPV] Relax the width check in gpu.shuffle lowering (PR #183445)

[Jianhui Li]

https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/183445

>From 087753f1566fe6c38087abc98c0291d48091bef3 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 26 Feb 2026 04:03:05 +0000
Subject: [PATCH 1/3] relax the width check in the shuffle lowering

---
 .../Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp  |  2 +-
 .../GPUToLLVMSPV/gpu-to-llvm-spv.mlir         | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
index 8c92b7e2b718b..87dca361e81e1 100644
--- a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
+++ b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
@@ -345,7 +345,7 @@ struct GPUShuffleConversion final : ConvertOpToLLVMPattern<gpu::ShuffleOp> {
   LogicalResult
   matchAndRewrite(gpu::ShuffleOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
-    if (!hasValidWidth(op))
+    if (getSubgroupSize(op) && !hasValidWidth(op))
       return rewriter.notifyMatchFailure(
           op, "shuffle width and subgroup size mismatch");
 
diff --git a/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir b/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir
index e56f3117e3b3c..ef5aa023bce51 100644
--- a/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir
+++ b/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir
@@ -344,6 +344,25 @@ gpu.module @shuffles {
 
 // -----
 
+// Check `gpu.shuffle` conversion with no explicit subgroup size.
+
+// CHECK: llvm.func spir_funccc @_Z17sub_group_shuffleij(i32, i32) -> i32 attributes {convergent, no_unwind, will_return}
+// CHECK-LABEL: llvm.func @gpu_shuffles(
+// CHECK-SAME: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
+// CHECK: %[[C16:.*]] = arith.constant 16 : i32
+// CHECK: %[[SHUF:.*]] = llvm.call spir_funccc @_Z17sub_group_shuffleij(%[[ARG0]], %[[ARG1]]) {convergent, no_unwind, will_return} : (i32, i32) -> i32
+// CHECK: %[[TRUE:.*]] = llvm.mlir.constant(true) : i1
+// CHECK: llvm.return
+gpu.module @shuffles_without_intel_reqd_sub_group_size_attribute {
+  llvm.func @gpu_shuffles(%val: i32, %id: i32) {
+    %width = arith.constant 16 : i32
+    %shuffleResult, %valid = gpu.shuffle idx %val, %id, %width : i32
+    llvm.return
+  }
+}
+
+// -----
+
 // Cannot convert due to shuffle width and target subgroup size mismatch
 
 gpu.module @shuffles_mismatch {

>From 11394c650182bba6442c5015133e24f2aed6e393 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 5 Mar 2026 20:06:28 +0000
Subject: [PATCH 2/3] address feedback

---
 mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
index 87dca361e81e1..96e3ee9e1073e 100644
--- a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
+++ b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
@@ -304,11 +304,11 @@ struct GPUShuffleConversion final : ConvertOpToLLVMPattern<gpu::ShuffleOp> {
     return parentFunc.getIntelReqdSubGroupSize();
   }
 
-  static bool hasValidWidth(gpu::ShuffleOp op) {
+  static bool hasValidWidth(gpu::ShuffleOp op, int subgroupSize) {
     llvm::APInt val;
     Value width = op.getWidth();
     return matchPattern(width, m_ConstantInt(&val)) &&
-           val == getSubgroupSize(op);
+           val == subgroupSize;
   }
 
   static Value bitcastOrExtBeforeShuffle(Value oldVal, Location loc,
@@ -345,7 +345,8 @@ struct GPUShuffleConversion final : ConvertOpToLLVMPattern<gpu::ShuffleOp> {
   LogicalResult
   matchAndRewrite(gpu::ShuffleOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
-    if (getSubgroupSize(op) && !hasValidWidth(op))
+    auto maybeSubgroupSize = getSubgroupSize(op);
+    if (maybeSubgroupSize && !hasValidWidth(op, maybeSubgroupSize.value()))
       return rewriter.notifyMatchFailure(
           op, "shuffle width and subgroup size mismatch");
 

>From 5f4a39a5d2918d7edd89c972a03f44def3e4cab6 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 6 Mar 2026 14:03:16 +0000
Subject: [PATCH 3/3] fix format

---
 mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
index 96e3ee9e1073e..6efd137f513e9 100644
--- a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
+++ b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
@@ -307,8 +307,7 @@ struct GPUShuffleConversion final : ConvertOpToLLVMPattern<gpu::ShuffleOp> {
   static bool hasValidWidth(gpu::ShuffleOp op, int subgroupSize) {
     llvm::APInt val;
     Value width = op.getWidth();
-    return matchPattern(width, m_ConstantInt(&val)) &&
-           val == subgroupSize;
+    return matchPattern(width, m_ConstantInt(&val)) && val == subgroupSize;
   }
 
   static Value bitcastOrExtBeforeShuffle(Value oldVal, Location loc,