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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: GabrielWu2004

<details>
<summary>Changes</summary>

# Description
Add the following canonicalization pattern to AllocaOp in MLIR's LLVM dialect:

```
%n = constant(N)
%alloca = llvm.alloca %n x Ty
```
->
```
%one = constant(1)
%alloca = llvm.alloca %one x !llvm.array<N x Ty>
```

Here N is a constant representable with 64-bits. With this canonicalization, the AllocaOp will be translated to `alloca [N x Ty]` in LLVM IR. Without this canonicalization, the AllocaOp was translated to `alloca Ty, N`. We prefer `alloca [N x Ty]` over `alloca Ty, N` for fixed-size array allocation.

See [this RFC](https://discourse.llvm.org/t/rfc-allocaop-canonicalization-in-mlirs-llvm-dialect/91734) for more detailed problem motivation and description.

Assisted by: Codex

---
Full diff: https://github.com/llvm/llvm-project/pull/221508.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td (+1) 
- (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp (+20-1) 
- (modified) mlir/test/Dialect/LLVMIR/canonicalize.mlir (+59) 


``````````diff
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

``````````

</details>


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