[Mlir-commits] [mlir] [mlir][scf] Fuse parallel loops with private allocations (PR #227766)

[Ivan Butygin]

Hardcode84 wrote:

The current change is LGTM, but AI agent review exposed a couple more gaps in the existing alias checks, adding there for documentation purposes:
```mlir
// DEFINE: %{lower} = -scf-parallel-for-to-nested-fors -convert-scf-to-cf \
// DEFINE:   -convert-arith-to-llvm -finalize-memref-to-llvm \
// DEFINE:   -convert-func-to-llvm -convert-cf-to-llvm -reconcile-unrealized-casts
// RUN: mlir-opt %s %{lower} | \
// RUN:   mlir-runner -e main -entry-point-result=i32 | FileCheck %s
// RUN: mlir-opt %s -scf-parallel-loop-fusion %{lower} | \
// RUN:   mlir-runner -e main -entry-point-result=i32 | FileCheck %s

// The helper writes to a different element in each iteration. All calls must
// finish before the second loop reads the buffer in reverse order.
// If the loops fuse, the first read can occur before the second helper call.
// Sequential execution of the fused loop then returns the initial value, 0.
// CHECK: {{^1$}}

func.func private @set(%buffer: memref<2xi32>, %i: index, %value: i32) {
  memref.store %value, %buffer[%i] : memref<2xi32>
  return
}

func.func @main() -> i32 {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c2 = arith.constant 2 : index
  %zero = arith.constant 0 : i32
  %one = arith.constant 1 : i32
  %buffer = memref.alloc() : memref<2xi32>
  %out = memref.alloc() : memref<2xi32>
  memref.store %zero, %buffer[%c0] : memref<2xi32>
  memref.store %zero, %buffer[%c1] : memref<2xi32>
  scf.parallel (%i) = (%c0) to (%c2) step (%c1) {
    %private = memref.alloca() : memref<1xi32>
    memref.store %one, %private[%c0] : memref<1xi32>
    %value = memref.load %private[%c0] : memref<1xi32>
    func.call @set(%buffer, %i, %value) : (memref<2xi32>, index, i32) -> ()
  }
  scf.parallel (%i) = (%c0) to (%c2) step (%c1) {
    %j = arith.subi %c1, %i : index
    %value = memref.load %buffer[%j] : memref<2xi32>
    memref.store %value, %out[%i] : memref<2xi32>
  }
  %answer = memref.load %out[%c0] : memref<2xi32>
  memref.dealloc %out : memref<2xi32>
  memref.dealloc %buffer : memref<2xi32>
  return %answer : i32
}
```

```mlir
// DEFINE: %{lower} = -scf-parallel-for-to-nested-fors -convert-scf-to-cf \
// DEFINE:   -convert-arith-to-llvm -finalize-memref-to-llvm \
// DEFINE:   -convert-func-to-llvm -convert-cf-to-llvm -reconcile-unrealized-casts
// RUN: mlir-opt %s %{lower} | \
// RUN:   mlir-runner -e main -entry-point-result=i32 | FileCheck %s
// RUN: mlir-opt %s -scf-parallel-loop-fusion %{lower} | \
// RUN:   mlir-runner -e main -entry-point-result=i32 | FileCheck %s

// Each loop writes to a different element in each iteration. The second loop
// must run after all writes from the first loop. Both elements must contain 2.
// If the loops fuse, writes from different iterations conflict. Sequential
// execution of the fused loop leaves 1 in the second element.
// CHECK: {{^2$}}

func.func @main() -> i32 {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c2 = arith.constant 2 : index
  %one = arith.constant 1 : i32
  %two = arith.constant 2 : i32
  %buffer = memref.alloc() : memref<2xi32>
  scf.parallel (%i) = (%c0) to (%c2) step (%c1) {
    %private = memref.alloca() : memref<1xi32>
    memref.store %one, %private[%c0] : memref<1xi32>
    %value = memref.load %private[%c0] : memref<1xi32>
    memref.store %value, %buffer[%i] : memref<2xi32>
  }
  scf.parallel (%i) = (%c0) to (%c2) step (%c1) {
    %j = arith.subi %c1, %i : index
    memref.store %two, %buffer[%j] : memref<2xi32>
  }
  %answer = memref.load %buffer[%c1] : memref<2xi32>
  memref.dealloc %buffer : memref<2xi32>
  return %answer : i32
}
```
We can fix in follow up PRs.

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