[Mlir-commits] [mlir] [MLIR][Affine] Fix fusion across ops with unknown memory effects (PR #203231)

[Federico Bruzzone]

https://github.com/FedericoBruzzone commented:

Of course, thank you, this is a much appreciated contribution 🫶


The problem is reproducible locally via this AI-generated script:
<details>
<summary>test_zero_tolerance.sh</summary>

```bash
#!/bin/bash
# Script: test_zero_tolerance.sh
# Verifies that with compute-tolerance=0.0, loops are NOT fused.
# Usage: ./test_zero_tolerance.sh [path to mlir-opt]

MLIROPT="${1:-mlir-opt}"
if ! command -v "$MLIROPT" &> /dev/null; then
  echo "ERROR: '$MLIROPT' could not be found. Please ensure it is in your PATH or provide the correct path as an argument."
  exit 1
fi

INPUT=$(cat << 'EOF'
#map = affine_map<(d0) -> (d0 mod 65536)>
module {
func.func @zero_tolerance(%arg0: memref<65536xcomplex<f64>>, %arg1: memref<30x131072xi64>,
%arg2: memref<30xi64>, %arg3: memref<30xi64>, %arg4: memref<30xi64>, %arg5: memref<30xi64>) {
  %c65536 = arith.constant 65536 : index
  %cst = arith.constant 0.000000e+00 : f64
  %cst_0 = arith.constant 0x4320000000380004 : f64
  %cst_1 = arith.constant 5.000000e-01 : f64
  %alloc = memref.alloc() {alignment = 128 : i64} : memref<30x131072xi64>
  %alloc_1 = memref.alloc() {alignment = 128 : i64} : memref<131072xi1>
  %alloc_2 = memref.alloc() {alignment = 128 : i64} : memref<131072xi128>

  // First nest: produces values in alloc_2 and alloc_1
  affine.for %i = 0 to 131072 {
    %idx = affine.apply #map(%i)
    %v0 = affine.load %arg0[%idx] : memref<65536xcomplex<f64>>
    %cmp = arith.cmpi ult, %i, %c65536 : index
    %im = complex.im %v0 : complex<f64>
    %re = complex.re %v0 : complex<f64>
    %sel = arith.select %cmp, %re, %im : f64
    %olt = arith.cmpf olt, %sel, %cst : f64
    %neg = arith.negf %sel : f64
    %sel2 = arith.select %olt, %neg, %sel : f64
    %mul = arith.mulf %sel2, %cst_0 : f64
    %add = arith.addf %mul, %cst_1 : f64
    %i128 = arith.fptosi %add : f64 to i128
    affine.store %i128, %alloc_2[%i] : memref<131072xi128>
    affine.store %olt, %alloc_1[%i] : memref<131072xi1>
  }

  // Second nest: consumes alloc_2/alloc_1, calls external function
  affine.for %j = 0 to 30 {
    affine.for %k = 0 to 131072 {
      %a0 = affine.load %arg5[%j] : memref<30xi64>
      %a1 = affine.load %arg2[%j] : memref<30xi64>
      %a2 = affine.load %arg4[%j] : memref<30xi64>
      %a3 = affine.load %arg3[%j] : memref<30xi64>
      %buf = affine.load %alloc_2[%k] : memref<131072xi128>
      %flag = affine.load %alloc_1[%k] : memref<131072xi1>
      %ret = func.call @__external_reduce_barrett(%a0, %a1, %a2, %a3, %buf) {outputModFac = 1 : i64} : (i64, i64, i64, i64, i128) -> i64
      %sub = arith.subi %a0, %ret : i64
      %sel = arith.select %flag, %sub, %ret : i64
      affine.store %sel, %alloc[%j, %k] : memref<30x131072xi64>
    }
  }

  // Opaque external call (writes to alloc in-place)
  func.call @__external_levelwise_forward_ntt(%alloc) : (memref<30x131072xi64>) -> ()

  // Third nest: reads from alloc, writes to arg1
  affine.for %l = 0 to 30 {
    affine.for %m = 0 to 131072 {
      %v = affine.load %alloc[%l, %m] : memref<30x131072xi64>
      affine.store %v, %arg1[%l, %m] : memref<30x131072xi64>
    }
  }

  memref.dealloc %alloc_2 : memref<131072xi128>
  memref.dealloc %alloc_1 : memref<131072xi1>
  memref.dealloc %alloc : memref<30x131072xi64>
  return
}
func.func private @__external_levelwise_forward_ntt(memref<30x131072xi64>)
func.func private @__external_reduce_barrett(i64, i64, i64, i64, i128) -> i64
}
EOF
)

echo "=== INPUT (before fusion) ==="
echo "$INPUT"
echo

echo "=== OUTPUT (after affine-loop-fusion with compute-tolerance=0.0) ==="
RESULT=$(echo "$INPUT" | "$MLIROPT" --allow-unregistered-dialect -mlir-disable-threading \
  -pass-pipeline='builtin.module(func.func(affine-loop-fusion{compute-tolerance=0.0}))' 2>&1)

# Display the output module/function
echo "$RESULT"

echo
echo "=== VERIFICATION: Did loops remain separate? ==="
# Count how many affine.for operations are present in the output
NUM_LOOPS=$(echo "$RESULT" | grep -c 'affine.for')
echo "Number of affine.for loops in the output: $NUM_LOOPS"
if [ "$NUM_LOOPS" -eq 5 ]; then
  echo "✅ CORRECT: 5 loops (3 nests, NONE fused) — tolerance 0.0 prevents fusion"
else
  echo "⚠️  Warning: $NUM_LOOPS loops found (expected 5)"
fi
```

</details>

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