[Mlir-commits] [mlir] [MLIR][XeVM] Fix mask lifetime in HandleVectorExtractPattern (PR #223879)

[Sang Ik Lee]

https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/223879

>From fc40ab5d394a0a6a89c455333ee85ecdd0d2463f Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Tue, 15 Sep 2026 23:50:38 +0000
Subject: [PATCH] [MLIR][XeVM] Fix mask lifetime in HandleVectorExtractPattern

When a contiguous-slice shuffle is hoisted through a bitcast that changes
the element count, the slice is rewritten against the bitcast source by
rescaling the mask into a local SmallVector and rebinding the `mask`
ArrayRef to it. That SmallVector was declared inside the
`if (maskScale != 1)` block, so it was destroyed before the two uses of
`mask` that follow the block - a stack-use-after-scope.

It is benign on uninstrumented builds, which is why precommit did not
catch it, but the sanitizer bots trap it on
Conversion/XeVMToLLVM/mxfp_scale_binary_hoist.mlir, the first test whose
chain drives the length-changing bitcast path:

  ERROR: AddressSanitizer: stack-use-after-scope
  READ of size 4
    #23 mlir::detail::DenseArrayAttrImpl<int>::get(MLIRContext*,
                                                  ArrayRef<int>)
    #24 mlir::LLVM::ShuffleVectorOp::build(..., ArrayRef<int>)
    #26 HandleVectorExtractPattern::matchAndRewrite
          XeVMToLLVM.cpp:1866:27
  Address is located in stack of thread T0 at offset 304 in frame
    #0 HandleVectorExtractPattern::matchAndRewrite XeVMToLLVM.cpp:1801
    [288, 352) 'newMask' (line 1856) <== Memory access at offset 304 is
                                         inside this variable
  SUMMARY: AddressSanitizer: stack-use-after-scope

ASan reports the above, MSan reports a read of uninitialized memory and
HWASan a tag mismatch. In each case mlir-opt dies partway through
--split-input-file, so the surface symptom is the CHECK-LABEL of the
second module in that file not being found:

  https://lab.llvm.org/buildbot/#/builders/52/builds/20099
  https://lab.llvm.org/buildbot/#/builders/164/builds/25415
  https://lab.llvm.org/buildbot/#/builders/169/builds/26571
  https://lab.llvm.org/buildbot/#/builders/55/builds/33165

Hoist the declaration so it outlives the uses, and add a lit test for the
length-changing bitcast path, which previously had no direct coverage -
the existing tests only bitcast between equal element counts.
---
 mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp  |  6 +++++-
 .../XeVMToLLVM/legalize_large_vector.mlir      | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index 8e4410729caa8..d9cb0e2eed3e7 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1846,6 +1846,11 @@ class HandleVectorExtractPattern
           return failure();
         }
         auto maskScale = srcResSize / srcInputSize;
+        // Storage for the rescaled mask. `mask` is an ArrayRef that gets
+        // rebound to this buffer below, so it has to outlive the `if` that
+        // fills it - otherwise the uses after the `if` read a destroyed
+        // SmallVector.
+        SmallVector<int32_t> newMask;
         if (maskScale != 1) {
           // The slice has to start at, and cover, whole source elements to be
           // expressible in terms of the bitcast source.
@@ -1853,7 +1858,6 @@ class HandleVectorExtractPattern
             return failure();
           }
           // Create a new mask that maps to the source vector
-          SmallVector<int32_t> newMask;
           int32_t newMaskSize = maskSize / maskScale;
           int32_t maskStart = mask[0] / maskScale;
           for (int32_t i = 0; i < newMaskSize; ++i) {
diff --git a/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir b/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir
index 21413e3e53a5f..9aa4892c65aa8 100644
--- a/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir
+++ b/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir
@@ -177,3 +177,21 @@ module @test_nary_no_hoist_through_call {
     llvm.return %1 : vector<16xf32>
   }
 }
+
+// -----
+
+// A bitcast that widens the element count (8xi32 -> 16xbf16) feeding a
+// contiguous slice: the slice can only be expressed against the bitcast
+// source by rescaling the mask, which is the one path that rebuilds the mask
+// into a local buffer. Guards the lifetime of that buffer - reading it after
+// scope exit is a stack-use-after-scope that only sanitizer builds observe.
+module @test_bitcast_mask_rescale {
+  // CHECK-LABEL: llvm.func @test_bitcast_mask_rescale
+  // CHECK:         %[[S:.*]] = llvm.shufflevector %arg0, %arg0 [0, 1, 2, 3] : vector<8xi32>
+  // CHECK:         llvm.bitcast %[[S]] : vector<4xi32> to vector<8xbf16>
+  llvm.func @test_bitcast_mask_rescale(%a: vector<8xi32>) -> vector<8xbf16> {
+    %0 = llvm.bitcast %a : vector<8xi32> to vector<16xbf16>
+    %1 = llvm.shufflevector %0, %0 [0, 1, 2, 3, 4, 5, 6, 7] : vector<16xbf16>
+    llvm.return %1 : vector<8xbf16>
+  }
+}