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

[Aayush Shrivastava]

https://github.com/iamaayushrivastava created https://github.com/llvm/llvm-project/pull/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`.

>From 50df936c86ffc4dc1e8ae1a254775b61a7c9c56c Mon Sep 17 00:00:00 2001
From: iamaayushrivastava <iamaayushrivastava at gmail.com>
Date: Tue, 14 Jul 2026 15:35:31 +0530
Subject: [PATCH] [mlir] Fix mem2reg crash on scalable vector store/load with
 matching type

---
 mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp | 13 ++++++++
 mlir/test/Dialect/LLVMIR/mem2reg.mlir         | 32 +++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index 66d6a592a93df..1186e64b8934d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -256,6 +256,12 @@ 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 types that may not have one, such as
+  // scalable vectors.
+  if (srcType == targetType)
+    return srcValue;
+
   uint64_t srcTypeSize = dataLayout.getTypeSizeInBits(srcType);
   uint64_t targetTypeSize = dataLayout.getTypeSizeInBits(targetType);
   if (srcTypeSize == targetTypeSize)
@@ -291,6 +297,13 @@ 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 types that may not have one, such as
+  // scalable vectors.
+  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..07ad0b86043e9 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg.mlir
@@ -1035,6 +1035,38 @@ llvm.func @scalable_llvm_vector() -> i16 {
 
 // -----
 
+// Storing into a scalable vector slot with a type matching the slot's
+// element type must not query the type's bit size, as this is not possible
+// for scalable vectors. See #209065.
+
+// CHECK-LABEL: @scalable_vector_matching_store
+llvm.func @scalable_vector_matching_store(%arg: 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
+  // CHECK-NOT: llvm.store
+  llvm.store %arg, %1 : vector<[4]xi1>, !llvm.ptr
+  llvm.return
+}
+
+// -----
+
+// Loading back a scalable vector value that was stored with a matching type
+// must also avoid querying the type's bit size. See #209065.
+
+// CHECK-LABEL: @scalable_vector_matching_store_load
+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) {