[Mlir-commits] [mlir] add canonicalization for array alloca in MLIR LLVM dialect (PR #221508)

[llvmlistbot at llvm.org]

https://github.com/GabrielWu2004 created https://github.com/llvm/llvm-project/pull/221508

Will add description later

>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] 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