[Mlir-commits] [mlir] [MLIR][XeVM] Legalize large vectors when shuffle second operand is poison (PR #220360)

[Sang Ik Lee]

https://github.com/silee2 created https://github.com/llvm/llvm-project/pull/220360

The SPIR-V backend now rejects OpTypeVector with a component count outside {2, 3, 4, 8, 16} with a fatal error, instead of silently emitting invalid SPIR-V that IGC happened to accept (36135544a13c). This exposed a gap in the large-vector legalization performed by convert-xevm-to-llvm.

HandleVectorExtractPattern exists to split vectors wider than 16 elements so that they never reach the backend, but its isExtractingContiguousSlice guard required both shufflevector source operands to be the same value. XeGPU lowering can produce a contiguous slice extraction whose second operand is poison instead of a copy of the first one:

  %0 = llvm.mlir.poison : vector<32xf16>
  %1 = llvm.load %p : !llvm.ptr -> vector<32xi16>
  %2 = llvm.bitcast %1 : vector<32xi16> to vector<32xf16>
  %3 = llvm.shufflevector %2, %0 [0, 1, 2, 3, 4, 5, 6, 7] : vector<32xf16>

For those the pattern bailed out, the vector<32x...> chain survived to the SPIR-V backend, and compilation aborted with:

  LLVM ERROR: OpTypeVector with 32 components requires the following
  SPIR-V extension: SPV_EXT_long_vector

The check was stricter than necessary. The mask bounds checks already ensure that every index addresses the first operand, so the second one is never read; accept it when it is poison or undef.

Also bounds-check the first mask element, which was previously unchecked and allowed a single-element mask referring to the second operand to slip through.

>From 0a3465176c5e4e7d41649c2a6a738168b4aa1a87 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Tue, 1 Sep 2026 19:53:01 +0000
Subject: [PATCH] [MLIR][XeVM] Legalize large vectors when shuffle second
 operand is poison

The SPIR-V backend now rejects OpTypeVector with a component count outside
{2, 3, 4, 8, 16} with a fatal error, instead of silently emitting invalid
SPIR-V that IGC happened to accept (36135544a13c). This exposed a gap in the
large-vector legalization performed by convert-xevm-to-llvm.

HandleVectorExtractPattern exists to split vectors wider than 16 elements so
that they never reach the backend, but its isExtractingContiguousSlice guard
required both shufflevector source operands to be the same value. XeGPU
lowering can produce a contiguous slice extraction whose second operand is
poison instead of a copy of the first one:

  %0 = llvm.mlir.poison : vector<32xf16>
  %1 = llvm.load %p : !llvm.ptr -> vector<32xi16>
  %2 = llvm.bitcast %1 : vector<32xi16> to vector<32xf16>
  %3 = llvm.shufflevector %2, %0 [0, 1, 2, 3, 4, 5, 6, 7] : vector<32xf16>

For those the pattern bailed out, the vector<32x...> chain survived to the
SPIR-V backend, and compilation aborted with:

  LLVM ERROR: OpTypeVector with 32 components requires the following
  SPIR-V extension: SPV_EXT_long_vector

The check was stricter than necessary. The mask bounds checks already ensure
that every index addresses the first operand, so the second one is never
read; accept it when it is poison or undef.

Also bounds-check the first mask element, which was previously unchecked and
allowed a single-element mask referring to the second operand to slip
through.
---
 mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 10 +++++--
 .../XeVMToLLVM/legalize_large_vector.mlir     | 28 +++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index ff4c69b23f8ed..d2594cafdc2ff 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1722,12 +1722,16 @@ class AllocaToGlobalPattern : public OpConversionPattern<LLVM::AllocaOp> {
 
 // Checks if shufflevector is used as a way to extract a contiguous slice
 // from a vector.
-// - source vector V1 and V2 are the same vector.
+// - source vector V2 is either the same as V1, or a poison/undef value. In
+//   both cases the mask can only meaningfully address elements of V1, which
+//   the mask checks below enforce.
 // - mask size is not greater than the source vector size
 // - mask values represent a sequence of consecutive increasing numbers
 //   that stay in bounds of the source vector when used for indexing.
 static bool isExtractingContiguousSlice(LLVM::ShuffleVectorOp op) {
-  if (op.getV1() != op.getV2())
+  if (op.getV1() != op.getV2() &&
+      !isa_and_present<LLVM::PoisonOp, LLVM::UndefOp>(
+          op.getV2().getDefiningOp()))
     return false;
   auto maskAttr = op.getMask();
   int64_t maskSize = static_cast<int64_t>(maskAttr.size());
@@ -1735,6 +1739,8 @@ static bool isExtractingContiguousSlice(LLVM::ShuffleVectorOp op) {
   if (maskSize > sourceSize)
     return false;
   int64_t firstIndex = maskAttr[0];
+  if (firstIndex < 0 || firstIndex >= sourceSize)
+    return false;
   for (int64_t i = 1; i < maskSize; ++i) {
     int64_t index = maskAttr[i];
     if (index != firstIndex + i)
diff --git a/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir b/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir
index b9721efd6e85e..9e2ee4c20679b 100644
--- a/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir
+++ b/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir
@@ -93,3 +93,31 @@ module @test_non_private_addrspace {
     llvm.return
   }
 }
+
+// -----
+
+// The second shuffle operand is poison rather than a copy of the first one.
+// The mask only addresses elements of the first operand, so the slice is still
+// extractable and the oversized load must be split.
+module @test_poison_second_operand {
+  // CHECK-LABEL: llvm.func @test_poison_second_operand
+  // CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr, %[[ARG1:.*]]: !llvm.ptr, %[[ARG2:.*]]: !llvm.ptr
+  llvm.func @test_poison_second_operand(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
+    // CHECK-NOT: vector<32x
+    // CHECK: %[[LOAD0:.*]] = llvm.load %[[ARG0]] : !llvm.ptr -> vector<8xi16>
+    // CHECK: %[[BITCAST0:.*]] = llvm.bitcast %[[LOAD0]] : vector<8xi16> to vector<8xf16>
+    // CHECK: %[[GEP0:.*]] = llvm.getelementptr %[[ARG0]][8] : (!llvm.ptr) -> !llvm.ptr, i16
+    // CHECK: %[[LOAD1:.*]] = llvm.load %[[GEP0]] : !llvm.ptr -> vector<8xi16>
+    // CHECK: %[[BITCAST1:.*]] = llvm.bitcast %[[LOAD1]] : vector<8xi16> to vector<8xf16>
+    // CHECK: llvm.store %[[BITCAST0]], %[[ARG1]] : vector<8xf16>, !llvm.ptr
+    // CHECK: llvm.store %[[BITCAST1]], %[[ARG2]] : vector<8xf16>, !llvm.ptr
+    %poison = llvm.mlir.poison : vector<32xf16>
+    %0 = llvm.load %arg0 : !llvm.ptr -> vector<32xi16>
+    %1 = llvm.bitcast %0 : vector<32xi16> to vector<32xf16>
+    %2 = llvm.shufflevector %1, %poison [0, 1, 2, 3, 4, 5, 6, 7] : vector<32xf16>
+    %3 = llvm.shufflevector %1, %poison [8, 9, 10, 11, 12, 13, 14, 15] : vector<32xf16>
+    llvm.store %2, %arg1 : vector<8xf16>, !llvm.ptr
+    llvm.store %3, %arg2 : vector<8xf16>, !llvm.ptr
+    llvm.return
+  }
+}