[Mlir-commits] [mlir] c34d2c3 - [mlir] Fix mem2reg crash on scalable vector store/load with matching type (#209426)

[llvmlistbot at llvm.org]

Author: Aayush Shrivastava
Date: 2026-07-20T08:52:49+02:00
New Revision: c34d2c30e86358bfe308558b5e0719e9f5305153

URL: https://github.com/llvm/llvm-project/commit/c34d2c30e86358bfe308558b5e0719e9f5305153
DIFF: https://github.com/llvm/llvm-project/commit/c34d2c30e86358bfe308558b5e0719e9f5305153.diff

LOG: [mlir] Fix mem2reg crash on scalable vector store/load with matching type (#209426)

Fixes #209065

`createInsertAndCast/createExtractAndCast` in `LLVMMemorySlot.cpp`
queried the bit size of the source/target types before checking whether
they were already identical. For scalable vector types (e.g.
`vector<[4]xi1>`), this size query implicitly converts a scalable
TypeSize to a scalar, which aborts. `mem2reg` hits this path for any
store or load whose value type exactly matches the slot's element type,
since `getStored/removeBlockingUses` run even when no cast is needed.

This fix adds an early return when the types are already equal, skipping
the size query, consistent with the existing short-circuit in
`areConversionCompatible/castSameSizedTypes`.

Added: 
    

Modified: 
    mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
    mlir/test/Dialect/LLVMIR/mem2reg.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index 66d6a592a93df..2416e01a67955 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -256,6 +256,11 @@ static Value createExtractAndCast(OpBuilder &builder, Location loc,
                                  /*narrowingConversion=*/true) &&
          "expected that the compatibility was checked before");
 
+  // Nothing has to be done if the types are already the same. This also
+  // avoids querying the bit size of scalable vector types below.
+  if (srcType == targetType)
+    return srcValue;
+
   uint64_t srcTypeSize = dataLayout.getTypeSizeInBits(srcType);
   uint64_t targetTypeSize = dataLayout.getTypeSizeInBits(targetType);
   if (srcTypeSize == targetTypeSize)
@@ -291,6 +296,12 @@ static Value createInsertAndCast(OpBuilder &builder, Location loc,
                                  srcValue.getType(),
                                  /*narrowingConversion=*/false) &&
          "expected that the compatibility was checked before");
+
+  // Nothing has to be done if the types are already the same. This also
+  // avoids querying the bit size of scalable vector types below.
+  if (srcValue.getType() == reachingDef.getType())
+    return srcValue;
+
   uint64_t valueTypeSize = dataLayout.getTypeSizeInBits(srcValue.getType());
   uint64_t slotTypeSize = dataLayout.getTypeSizeInBits(reachingDef.getType());
   if (slotTypeSize == valueTypeSize)

diff  --git a/mlir/test/Dialect/LLVMIR/mem2reg.mlir b/mlir/test/Dialect/LLVMIR/mem2reg.mlir
index cabbd1c1bc013..5f1029625b9b4 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg.mlir
@@ -1035,6 +1035,20 @@ llvm.func @scalable_llvm_vector() -> i16 {
 
 // -----
 
+// CHECK-LABEL: @scalable_vector_matching_store_load
+// CHECK-SAME: %[[ARG:.+]]: vector<[4]xi1>
+llvm.func @scalable_vector_matching_store_load(%arg: vector<[4]xi1>) -> vector<[4]xi1> {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK-NOT: llvm.alloca
+  %1 = llvm.alloca %0 x vector<[4]xi1> : (i32) -> !llvm.ptr
+  llvm.store %arg, %1 : vector<[4]xi1>, !llvm.ptr
+  %2 = llvm.load %1 : !llvm.ptr -> vector<[4]xi1>
+  // CHECK: llvm.return %[[ARG]] : vector<[4]xi1>
+  llvm.return %2 : vector<[4]xi1>
+}
+
+// -----
+
 // CHECK-LABEL: @smaller_store_forwarding
 // CHECK-SAME: %[[ARG:.+]]: i16
 llvm.func @smaller_store_forwarding(%arg : i16) {