[Mlir-commits] [mlir] [mlir][bufferization] Skip parallel analysis without parallel regions (PR #227202)

[Ziliang Zhang]

ziliangzl wrote:

# Stress Test Construction

Each chain starts with `bufferization.alloc_tensor`, alternates between statically
and dynamically shaped `tensor.cast` operations, and ends with one `linalg.fill`
and one `tensor.extract`. The entire input contains no parallel region. While
analyzing in-place candidates, the old implementation repeatedly follows these
long chains to find definitions. The fast-exit removes the entire parallel
privatization phase.

The test generator is shown below.

```python
#!/usr/bin/env python3
import sys

length = int(sys.argv[1]) if len(sys.argv) > 1 else 256
chains = int(sys.argv[2]) if len(sys.argv) > 2 else 16

print("module {\n  func.func @stress() {")
print("    %cst = arith.constant 0.0 : f32\n    %c0 = arith.constant 0 : index")
for chain in range(chains):
    value, ty = f"%base_{chain}", "tensor<1xf32>"
    print(f"    {value} = bufferization.alloc_tensor() : {ty}")
    for index in range(length):
        result = f"%cast_{chain}_{index}"
        result_ty = "tensor<?xf32>" if ty == "tensor<1xf32>" else "tensor<1xf32>"
        print(f"    {result} = tensor.cast {value} : {ty} to {result_ty}")
        value, ty = result, result_ty
    print(f"    %filled_{chain} = linalg.fill ins(%cst : f32) outs({value} : {ty}) -> {ty}")
    print(f"    %read_{chain} = tensor.extract %filled_{chain}[%c0] : {ty}")
    print(f'    "test.consume"(%read_{chain}) : (f32) -> ()')
print("    return\n  }\n}")
```

For example, with `length=3` and `chains=1`, the script generates:

```mlir
module {
  func.func @stress() {
    %cst = arith.constant 0.0 : f32
    %c0 = arith.constant 0 : index
    %base_0 = bufferization.alloc_tensor() : tensor<1xf32>
    %cast_0_0 = tensor.cast %base_0 : tensor<1xf32> to tensor<?xf32>
    %cast_0_1 = tensor.cast %cast_0_0 : tensor<?xf32> to tensor<1xf32>
    %cast_0_2 = tensor.cast %cast_0_1 : tensor<1xf32> to tensor<?xf32>
    %filled_0 = linalg.fill ins(%cst : f32) outs(%cast_0_2 : tensor<?xf32>) -> tensor<?xf32>
    %read_0 = tensor.extract %filled_0[%c0] : tensor<?xf32>
    "test.consume"(%read_0) : (f32) -> ()
    return
  }
}
```

For the measured inputs:

```bash
python3 generate.py 256 16 > stress.mlir
mlir-opt stress.mlir -allow-unregistered-dialect \
  -one-shot-bufferize=test-analysis-only -o /dev/null
```

## Test Environment and Baseline

The test ran on arm64 macOS. Both `mlir-opt` binaries were built in Release mode
with assertions disabled from the same checkout:

* Baseline: `555d8a58b85f`.
* Optimized: `f74c2cdc4dd5`, adding the no-parallel-region fast-exit and hoisting
  the write-side parallel-region lookup out of the definition loop.

The stress input contains no parallel region, so hoisting the write-side lookup
does not affect these measurements.

Each configuration was warmed up once, then run five times in interleaved order.
The table reports median wall time. In addition to the default configurations
before and after the optimization, the baseline binary was measured with
`check-parallel-regions=false` to isolate the cost of parallel analysis itself.

## Results

| Casts per chain | Total casts | Baseline, default | Baseline, parallel check disabled | Optimized, default | Reduction |
| ---: | ---: | ---: | ---: | ---: | ---: |
| 64 | 1024 | 0.20 s | 0.08 s | 0.08 s | 60.0% |
| 128 | 2048 | 1.44 s | 0.53 s | 0.53 s | 63.2% |
| 256 | 4096 | 11.62 s | 4.06 s | 4.13 s | 64.5% |


## Conclusion

For the largest input, the baseline default configuration was 7.56 seconds
slower than the baseline with parallel checking disabled, accounting for about
65.1% of total runtime. 
This stress test shows that parallel analysis can become the dominant cost for
long alias chains without parallel regions, and that a global negative fast-exit
has clear value. 

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