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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Sang Ik Lee (silee2)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/220360.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp (+8-2) 
- (modified) mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir (+28) 


``````````diff
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
+  }
+}

``````````

</details>


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