[Mlir-commits] [mlir] [mlir][SPIR-V] Preserve poison mask in scalarized vector.shuffle lowering (PR #214686)

[Arseniy Obolenskiy]

https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/214686

>From a8cd24852ceca91bd8c93cb71be19dd24a4b216c Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Fri, 7 Aug 2026 11:52:48 +0200
Subject: [PATCH 1/2] [mlir][SPIR-V] Preserve poison mask in scalarized
 vector.shuffle lowering

---
 .../VectorToSPIRV/VectorToSPIRV.cpp           |  6 ++++-
 .../VectorToSPIRV/vector-to-spirv.mlir        | 25 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
index 78693e924c4d9..62fe707b921bd 100644
--- a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
+++ b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
@@ -589,8 +589,12 @@ struct VectorShuffleOpConvert final
     // When at least one of the operands or the result becomes a scalar after
     // type conversion for SPIR-V, extract all the required elements and
     // construct the result vector.
-    auto getElementAtIdx = [&rewriter, loc = shuffleOp.getLoc()](
+    Type scalarType = oldV1Type.getElementType();
+    auto getElementAtIdx = [&rewriter, loc = shuffleOp.getLoc(), scalarType](
                                Value scalarOrVec, int32_t idx) -> Value {
+      if (idx == vector::ShuffleOp::kPoisonIndex)
+        return spirv::UndefOp::create(rewriter, loc, scalarType);
+
       if (auto vecTy = dyn_cast<VectorType>(scalarOrVec.getType()))
         return spirv::CompositeExtractOp::create(rewriter, loc, scalarOrVec,
                                                  idx);
diff --git a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
index f904dd9d35c37..b31e6842ef739 100644
--- a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
+++ b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
@@ -604,6 +604,31 @@ func.func @shuffle(%v0 : vector<4xi32>, %v1: vector<4xi32>) -> vector<4xi32> {
 
 // -----
 
+// CHECK-LABEL:  func @shuffle_poison_scalar_operand
+//  CHECK-SAME:  %[[ARG0:.+]]: vector<4xi32>, %[[ARG1:.+]]: vector<4xi32>
+//       CHECK:    %[[UNDEF:.+]] = spirv.Undef : i32
+//       CHECK:    %[[RES:.+]]  = builtin.unrealized_conversion_cast %[[UNDEF]] : i32 to vector<1xi32>
+//       CHECK:    return %[[RES]] : vector<1xi32>
+func.func @shuffle_poison_scalar_operand(%v0 : vector<4xi32>, %v1: vector<4xi32>) -> vector<1xi32> {
+  %shuffle = vector.shuffle %v0, %v1 [-1] : vector<4xi32>, vector<4xi32>
+  return %shuffle : vector<1xi32>
+}
+
+// -----
+
+// CHECK-LABEL:  func @shuffle_poison_scalar_result
+//  CHECK-SAME:  %[[ARG0:.+]]: vector<1xi32>, %[[ARG1:.+]]: vector<1xi32>
+//   CHECK-DAG:    %[[V1:.+]] = builtin.unrealized_conversion_cast %[[ARG1]] : vector<1xi32> to i32
+//   CHECK-DAG:    %[[UNDEF:.+]] = spirv.Undef : i32
+//       CHECK:    %[[RES:.+]] = spirv.CompositeConstruct %[[UNDEF]], %[[V1]] : (i32, i32) -> vector<2xi32>
+//       CHECK:    return %[[RES]]
+func.func @shuffle_poison_scalar_result(%v0 : vector<1xi32>, %v1: vector<1xi32>) -> vector<2xi32> {
+  %shuffle = vector.shuffle %v0, %v1 [-1, 1] : vector<1xi32>, vector<1xi32>
+  return %shuffle : vector<2xi32>
+}
+
+// -----
+
 // CHECK-LABEL: func @interleave
 //  CHECK-SAME: (%[[ARG0:.+]]: vector<2xf32>, %[[ARG1:.+]]: vector<2xf32>)
 //       CHECK: %[[SHUFFLE:.*]] = spirv.VectorShuffle [0 : i32, 2 : i32, 1 : i32, 3 : i32] %[[ARG0]], %[[ARG1]] : vector<2xf32>, vector<2xf32> -> vector<4xf32>

>From c35b46ea4c5add9c21fd25434d5cfb51c947eeb8 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Fri, 21 Aug 2026 17:51:26 +0200
Subject: [PATCH 2/2] address coment

---
 mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
index 62fe707b921bd..15121ba36d21b 100644
--- a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
+++ b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
@@ -589,8 +589,8 @@ struct VectorShuffleOpConvert final
     // When at least one of the operands or the result becomes a scalar after
     // type conversion for SPIR-V, extract all the required elements and
     // construct the result vector.
-    Type scalarType = oldV1Type.getElementType();
-    auto getElementAtIdx = [&rewriter, loc = shuffleOp.getLoc(), scalarType](
+    auto getElementAtIdx = [&rewriter, loc = shuffleOp.getLoc(),
+                            scalarType = oldV1Type.getElementType()](
                                Value scalarOrVec, int32_t idx) -> Value {
       if (idx == vector::ShuffleOp::kPoisonIndex)
         return spirv::UndefOp::create(rewriter, loc, scalarType);