[Mlir-commits] [mlir] [MLIR] AllocaOp canonicalization for array allocation in LLVM dialect (PR #221508)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 7 19:53:46 PDT 2026
https://github.com/GabrielWu2004 updated https://github.com/llvm/llvm-project/pull/221508
>From cc30ec08f71a31fc019d649d38313913420bea89 Mon Sep 17 00:00:00 2001
From: GabrielWu2004 <gabrielwu2004 at gmail.com>
Date: Sat, 5 Sep 2026 22:31:48 -0400
Subject: [PATCH 1/3] add canonicalization for array alloca in MLIR LLVM
dialect
---
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td | 1 +
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp | 21 +++++++-
mlir/test/Dialect/LLVMIR/canonicalize.mlir | 59 +++++++++++++++++++++
3 files changed, 80 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index d16ea595d5f39..ceb418ebd2998 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -282,6 +282,7 @@ def LLVM_AllocaOp : LLVM_Op<"alloca",
];
let hasCustomAssemblyFormat = 1;
let hasVerifier = 1;
+ let hasCanonicalizeMethod = 1;
}
def LLVM_GEPOp : LLVM_Op<"getelementptr", [Pure,
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 87eab5ed9e6f3..e35785cc0cf02 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -397,7 +397,7 @@ OpFoldResult ICmpOp::fold(FoldAdaptor adaptor) {
}
//===----------------------------------------------------------------------===//
-// Printing, parsing and verification for LLVM::AllocaOp.
+// Printing, parsing, verification and canonicalization for LLVM::AllocaOp.
//===----------------------------------------------------------------------===//
void AllocaOp::print(OpAsmPrinter &p) {
@@ -472,6 +472,25 @@ LogicalResult AllocaOp::verify() {
return success();
}
+LogicalResult AllocaOp::canonicalize(AllocaOp op, PatternRewriter &rewriter) {
+ // Convert `alloca Ty, C` to the canonical `alloca [C x Ty], 1` form.
+ APInt numElements;
+ if (!matchPattern(op.getArraySize(), m_ConstantInt(&numElements)) ||
+ numElements.isOne() || numElements.getActiveBits() > 64)
+ return failure();
+
+ Type arrayType =
+ LLVMArrayType::get(op.getElemType(), numElements.getZExtValue());
+ Value one = ConstantOp::create(rewriter, op.getLoc(), rewriter.getI32Type(),
+ /*value=*/1);
+ auto newAlloca = AllocaOp::create(
+ rewriter, op.getLoc(), op.getType(), one, op.getAlignmentAttr(),
+ arrayType, op.getInalloca());
+ newAlloca->setDiscardableAttrs(op->getDiscardableAttrDictionary());
+ rewriter.replaceOp(op, newAlloca);
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// LLVM::BrOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
index 11e28d52355cd..42fef746fcce0 100644
--- a/mlir/test/Dialect/LLVMIR/canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
@@ -33,6 +33,65 @@ llvm.func @fold_icmp_alloca() -> i1 {
// -----
+// CHECK-LABEL: @canonicalize_constant_array_alloca
+llvm.func @canonicalize_constant_array_alloca() -> !llvm.ptr {
+ // CHECK-NEXT: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK-NEXT: %[[ALLOCA:.*]] = llvm.alloca inalloca %[[ONE]] x !llvm.array<4 x i32> {tag = "preserved", alignment = 16 : i64} : (i32) -> !llvm.ptr
+ %c4 = arith.constant 4 : i64
+ %alloca = llvm.alloca inalloca %c4 x i32 {alignment = 16 : i64, tag = "preserved"} : (i64) -> !llvm.ptr
+ // CHECK-NEXT: llvm.return %[[ALLOCA]] : !llvm.ptr
+ llvm.return %alloca : !llvm.ptr
+}
+
+// -----
+
+// CHECK-LABEL: @canonicalize_zero_array_alloca
+llvm.func @canonicalize_zero_array_alloca() -> !llvm.ptr {
+ // CHECK-NEXT: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK-NEXT: %[[ALLOCA:.*]] = llvm.alloca %[[ONE]] x !llvm.array<0 x i8> : (i32) -> !llvm.ptr
+ %c0 = llvm.mlir.constant(0 : i32) : i32
+ %alloca = llvm.alloca %c0 x i8 : (i32) -> !llvm.ptr
+ // CHECK-NEXT: llvm.return %[[ALLOCA]] : !llvm.ptr
+ llvm.return %alloca : !llvm.ptr
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_canonicalize_scalar_alloca
+llvm.func @do_not_canonicalize_scalar_alloca() -> !llvm.ptr {
+ // CHECK-NEXT: %[[ONE:.*]] = arith.constant 1 : i64
+ // CHECK-NEXT: %[[ALLOCA:.*]] = llvm.alloca %[[ONE]] x i32 : (i64) -> !llvm.ptr
+ %c1 = arith.constant 1 : i64
+ %alloca = llvm.alloca %c1 x i32 : (i64) -> !llvm.ptr
+ // CHECK-NEXT: llvm.return %[[ALLOCA]] : !llvm.ptr
+ llvm.return %alloca : !llvm.ptr
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_canonicalize_dynamic_array_alloca
+// CHECK-SAME: (%[[SIZE:.*]]: i64)
+llvm.func @do_not_canonicalize_dynamic_array_alloca(%size : i64) -> !llvm.ptr {
+ // CHECK-NEXT: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x i32 : (i64) -> !llvm.ptr
+ %alloca = llvm.alloca %size x i32 : (i64) -> !llvm.ptr
+ // CHECK-NEXT: llvm.return %[[ALLOCA]] : !llvm.ptr
+ llvm.return %alloca : !llvm.ptr
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_canonicalize_large_array_size
+llvm.func @do_not_canonicalize_large_array_size() -> !llvm.ptr {
+ // CHECK-NEXT: %[[SIZE:.*]] = llvm.mlir.constant(18446744073709551616 : i128) : i128
+ // CHECK-NEXT: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x i8 : (i128) -> !llvm.ptr
+ %size = llvm.mlir.constant(18446744073709551616 : i128) : i128
+ %alloca = llvm.alloca %size x i8 : (i128) -> !llvm.ptr
+ // CHECK-NEXT: llvm.return %[[ALLOCA]] : !llvm.ptr
+ llvm.return %alloca : !llvm.ptr
+}
+
+// -----
+
// CHECK-LABEL: fold_extractvalue
llvm.func @fold_extractvalue() -> i32 {
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i32
>From 3f5b0ff9292644571c0cf7a514d43ab94b16c424 Mon Sep 17 00:00:00 2001
From: GabrielWu2004 <gabrielwu2004 at gmail.com>
Date: Mon, 7 Sep 2026 12:46:45 -0400
Subject: [PATCH 2/3] address review comments
---
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp | 2 +-
mlir/test/Dialect/LLVMIR/canonicalize.mlir | 12 ++++++++++++
2 files changed, 13 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index e35785cc0cf02..bdbf49f9a3c5e 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -479,7 +479,7 @@ LogicalResult AllocaOp::canonicalize(AllocaOp op, PatternRewriter &rewriter) {
numElements.isOne() || numElements.getActiveBits() > 64)
return failure();
- Type arrayType =
+ auto arrayType =
LLVMArrayType::get(op.getElemType(), numElements.getZExtValue());
Value one = ConstantOp::create(rewriter, op.getLoc(), rewriter.getI32Type(),
/*value=*/1);
diff --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
index 42fef746fcce0..61d89928c0d87 100644
--- a/mlir/test/Dialect/LLVMIR/canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
@@ -45,6 +45,18 @@ llvm.func @canonicalize_constant_array_alloca() -> !llvm.ptr {
// -----
+// CHECK-LABEL: @canonicalize_nested_array_alloca
+llvm.func @canonicalize_nested_array_alloca() -> !llvm.ptr {
+ // CHECK-NEXT: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK-NEXT: %[[ALLOCA:.*]] = llvm.alloca %[[ONE]] x !llvm.array<4 x array<2 x i32>> : (i32) -> !llvm.ptr
+ %c4 = arith.constant 4 : i64
+ %alloca = llvm.alloca %c4 x !llvm.array<2 x i32> : (i64) -> !llvm.ptr
+ // CHECK-NEXT: llvm.return %[[ALLOCA]] : !llvm.ptr
+ llvm.return %alloca : !llvm.ptr
+}
+
+// -----
+
// CHECK-LABEL: @canonicalize_zero_array_alloca
llvm.func @canonicalize_zero_array_alloca() -> !llvm.ptr {
// CHECK-NEXT: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : i32
>From 9dc57c58a406f6f709f7c9bfd29c9663bf745a3f Mon Sep 17 00:00:00 2001
From: GabrielWu2004 <gabrielwu2004 at gmail.com>
Date: Mon, 7 Sep 2026 12:55:49 -0400
Subject: [PATCH 3/3] clang format
---
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index bdbf49f9a3c5e..0bc95d8f59203 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -483,9 +483,9 @@ LogicalResult AllocaOp::canonicalize(AllocaOp op, PatternRewriter &rewriter) {
LLVMArrayType::get(op.getElemType(), numElements.getZExtValue());
Value one = ConstantOp::create(rewriter, op.getLoc(), rewriter.getI32Type(),
/*value=*/1);
- auto newAlloca = AllocaOp::create(
- rewriter, op.getLoc(), op.getType(), one, op.getAlignmentAttr(),
- arrayType, op.getInalloca());
+ auto newAlloca =
+ AllocaOp::create(rewriter, op.getLoc(), op.getType(), one,
+ op.getAlignmentAttr(), arrayType, op.getInalloca());
newAlloca->setDiscardableAttrs(op->getDiscardableAttrDictionary());
rewriter.replaceOp(op, newAlloca);
return success();
More information about the Mlir-commits
mailing list