[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