[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 09:46:58 PDT 2026
https://github.com/GabrielWu2004 updated https://github.com/llvm/llvm-project/pull/221508
>From 0a96cd60ac2a3f4774d4228acdc82a460f2751c6 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/2] 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 ab641a5da8c22..48d6d0f8a8f3d 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 cd6425f4a63fa..cdd4b5a9a4138 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -394,7 +394,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) {
@@ -469,6 +469,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 9bc476537bc01fbae858808dacc5b3c1aec893fd 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/2] 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 cdd4b5a9a4138..5f7b1b975f393 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -476,7 +476,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
More information about the Mlir-commits
mailing list