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

[llvmlistbot at llvm.org]

GabrielWu2004 wrote:

> Just out of curiosity. Did you experiment with MLIR's SROA pass? I would assume that it benefits from this canonicalization as well (in any case LLVM IR's SROA is probably more robust).

Just ran an experiment, and indeed MLIR's SROA benefits from this canonicalization because previously it doesn't handle array allocas.

Input IR:
```
llvm.func @array() -> i32 {
  %0 = llvm.mlir.constant(10 : i32) : i32
  %1 = llvm.alloca %0 x i32 {alignment = 8 : i64} : (i32) -> !llvm.ptr
  %2 = llvm.getelementptr inbounds %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<10 x i32>
  %3 = llvm.load %2 : !llvm.ptr -> i32
  llvm.return %3 : i32
}
```

Result from running SROA only (no change):
```
module {
  llvm.func @array() -> i32 {
    %0 = llvm.mlir.constant(10 : i32) : i32
    %1 = llvm.alloca %0 x i32 {alignment = 8 : i64} : (i32) -> !llvm.ptr
    %2 = llvm.getelementptr inbounds %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<10 x i32>
    %3 = llvm.load %2 : !llvm.ptr -> i32
    llvm.return %3 : i32
  }
}
```

Result from running canonicalize + SROA:
```
module {
  llvm.func @array() -> i32 {
    %0 = llvm.mlir.constant(1 : i32) : i32
    %1 = llvm.alloca %0 x i32 : (i32) -> !llvm.ptr
    %2 = llvm.load %1 : !llvm.ptr -> i32
    llvm.return %2 : i32
  }
}
```

This probably doesn't need a separate lit test, but it's good to know that this canonicalization enabled more SROA opportunities in MLIR :). But as you said, we typically rely on LLVM's SROA pass anyway. 

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