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

Aayush Shrivastava llvmlistbot at llvm.org
Wed Jul 15 01:26:14 PDT 2026


https://github.com/iamaayushrivastava updated https://github.com/llvm/llvm-project/pull/209426

>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 1/3] [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) {

>From 4710e803e7180e7d4320cccf0b2eefe835385a5e Mon Sep 17 00:00:00 2001
From: iamaayushrivastava <iamaayushrivastava at gmail.com>
Date: Tue, 14 Jul 2026 21:03:31 +0530
Subject: [PATCH 2/3] [mlir] Address review comments on mem2reg scalable vector
 fix

---
 mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp |  6 ++----
 mlir/test/Dialect/LLVMIR/mem2reg.mlir         | 19 -------------------
 2 files changed, 2 insertions(+), 23 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index 1186e64b8934d..2416e01a67955 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -257,8 +257,7 @@ static Value createExtractAndCast(OpBuilder &builder, Location loc,
          "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.
+  // avoids querying the bit size of scalable vector types below.
   if (srcType == targetType)
     return srcValue;
 
@@ -299,8 +298,7 @@ static Value createInsertAndCast(OpBuilder &builder, Location loc,
          "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.
+  // avoids querying the bit size of scalable vector types below.
   if (srcValue.getType() == reachingDef.getType())
     return srcValue;
 
diff --git a/mlir/test/Dialect/LLVMIR/mem2reg.mlir b/mlir/test/Dialect/LLVMIR/mem2reg.mlir
index 07ad0b86043e9..bf7da45bfa97e 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg.mlir
@@ -1035,25 +1035,6 @@ 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

>From 37654ffd85b1c4ad02f24336b57779e2d2590cc0 Mon Sep 17 00:00:00 2001
From: iamaayushrivastava <iamaayushrivastava at gmail.com>
Date: Wed, 15 Jul 2026 13:55:47 +0530
Subject: [PATCH 3/3] [mlir] Bind test function argument via CHECK-SAME in
 mem2reg scalable vector test

---
 mlir/test/Dialect/LLVMIR/mem2reg.mlir | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/test/Dialect/LLVMIR/mem2reg.mlir b/mlir/test/Dialect/LLVMIR/mem2reg.mlir
index bf7da45bfa97e..5f1029625b9b4 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg.mlir
@@ -1036,13 +1036,14 @@ 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>
+  // CHECK: llvm.return %[[ARG]] : vector<[4]xi1>
   llvm.return %2 : vector<[4]xi1>
 }
 



More information about the Mlir-commits mailing list