[MLIR][Bufferization] Quadratic compile time when simplifying wide deallocs with distinct allocation roots
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
### Problem
`-buffer-deallocation-simplification` shows near-quadratic compile-time scaling
for a wide `bufferization.dealloc` whose operands have distinct allocation
roots.
This occurs in `SplitDeallocWhenNotAliasingAnyOther`. For each memref, the pass
copies the full operand list, removes that memref, and checks it against every
remaining operand:
```c++
for (int64_t i = 0, e = deallocOp.getMemrefs().size(); i < e; ++i) {
Value memref = deallocOp.getMemrefs()[i];
SmallVector otherMemrefs(deallocOp.getMemrefs());
otherMemrefs.erase(otherMemrefs.begin() + i);
if (potentiallyAliasesMemref(analysis, otherMemrefs, memref))
continue;
// Split memref into a separate dealloc.
}
```
If all operands come from different `memref.alloc` operations, none aliases and
there is no early exit. The pass therefore performs `N * (N - 1)` directed pair
checks and copies an `N`-element list `N` times. For 2,547 operands, that is
6,484,662 directed pair checks.
### Real-world impact
This was found while lowering the public
[ONNX Model Zoo DenseNet-121 opset 12 model](https://huggingface.co/onnxmodelzoo/densenet-12/tree/60b30cf31d0c0ed06fbd9f695ce3da6f71450ded).
Its actual pass input contains one `bufferization.dealloc` with 2,547 memrefs
and 2,547 conditions.
On that captured input, seven alternating baseline/prototype pairs changed the
complete simplification command from a median of **3.30365 s to 0.69447 s**
(paired median **-78.88%**). The parse-subtracted pass cost changed from
2.75902 s to 0.15902 s.
As an end-to-end check, five alternating full DenseNet lowering pairs changed
from **15.24 s to 12.74 s** (paired median **-17.06%**). The pass-input result
was measured against upstream main; the end-to-end result used the same source
change relinked into the ABI-compatible LLVM revision used by ONNX-MLIR.
### Upstream-only reproducer
This is a synthetic reduction of the shape observed in DenseNet, not the
literal DenseNet IR. It uses only upstream common dialects.
```python
#!/usr/bin/env python3
import sys
n = int(sys.argv[1])
names = [f"%a{i}" for i in range(n)]
print("module { func.func @f() {")
print(" %true = arith.constant true")
for name in names:
print(f" {name} = memref.alloc() : memref<1xi8>")
print(" bufferization.dealloc (" + ", ".join(names) + " : " +
", ".join(["memref<1xi8>"] * n) + ") if (" +
", ".join(["%true"] * n) + ")")
print(" return } }")
```
```bash
python3 make-wide-dealloc.py 4096 > wide-4096.mlir
/usr/bin/time -f 'elapsed=%e maxrss_kib=%M' \
mlir-opt wide-4096.mlir --mlir-disable-threading \
--buffer-deallocation-simplification -o /dev/null
```
Release `-O3`, assertions-enabled builds were pinned to one CPU and measured in
seven alternating pairs after warmup:
| operands | baseline | allocation-root prototype | paired median change |
|---:|---:|---:|---:|
| 512 | 0.13454 s | 0.02518 s | -80.29% |
| 1,024 | 0.45973 s | 0.03081 s | -93.05% |
| 2,048 | 1.66288 s | 0.03827 s | -97.71% |
| 4,096 | 6.47635 s | 0.06017 s | -99.07% |
After subtracting parse time, scaling from 512 to 4,096 operands increased the
baseline pass cost by 58.72x (observed exponent 1.96), versus 6.28x (exponent
0.88) for the prototype.
Measurements used LLVM revision
`902c28258ab67d9095e5b8dbf51dce7d0c90d9a7`; the relevant source and focused
test were unchanged through main
`399a8ce19e2c8b6c5b5916204ab46eab7158cc04`.
Contributor guide
Assessment
This issue has not been assessed yet.