llvm / llvm/llvm-project

[MLIR][Affine] affine-loop-fusion truncates the producer iteration space to the consumer domain

Open
#212,028 2 comments 0 reactions 0 assignees View on GitHub
mlir
Dominant language
LLVM
Stars
40.5k
Forks
18.7k
PR merge metrics
PR metrics pending

Description

## Summary

Applying `affine-loop-fusion` to the verifier-valid program below changes the result from `-4` to `0`.

The producer covers `j = [0, 4)`, whereas the consumer covers `j = [0, 3)`. The generated fused nest uses the consumer's smaller `j` domain for both bodies, so the producer iteration at `j = 3` is omitted.

Public `mlir-opt` reproducer: https://godbolt.org/z/8aW11aYT6

## Version

Tested at `llvm/llvm-project` commit `450cf0aafeab9fed6666b112b018d3913a36527c` (`LLVM 24.0.0git`).

The same incorrect target was produced by Release builds with assertions both enabled and disabled.

## Reproducer

```mlir
module {
func.func private @kernel(%ey: memref<4x4xf64>, %hz: memref<4x4xf64>) {
affine.for %t = 0 to 2 {
affine.for %i = 1 to 4 {
affine.for %j = 0 to 4 {
%old = affine.load %ey[%i, %j] : memref<4x4xf64>
%a = affine.load %hz[%i, %j] : memref<4x4xf64>
%b = affine.load %hz[%i - 1, %j] : memref<4x4xf64>
%d = arith.subf %a, %b : f64
%new = arith.subf %old, %d : f64
affine.store %new, %ey[%i, %j] : memref<4x4xf64>
}
}
affine.for %i = 0 to 3 {
affine.for %j = 0 to 3 {
%old = affine.load %hz[%i, %j] : memref<4x4xf64>
%y1 = affine.load %ey[%i + 1, %j] : memref<4x4xf64>
%y0 = affine.load %ey[%i, %j] : memref<4x4xf64>
%d = arith.subf %y1, %y0 : f64
%new = arith.subf %old, %d : f64
affine.store %new, %hz[%i, %j] : memref<4x4xf64>
}
}
}
return
}

func.func @main() -> i64 {
%ey = memref.alloca() : memref<4x4xf64>
%hz = memref.alloca() : memref<4x4xf64>
%zero = arith.constant 0.0 : f64
%two = arith.constant 2.0 : f64

affine.for %i = 0 to 4 {
affine.for %j = 0 to 4 {
affine.store %zero, %ey[%i, %j] : memref<4x4xf64>
affine.store %zero, %hz[%i, %j] : memref<4x4xf64>
}
}

affine.store %two, %hz[1, 3] : memref<4x4xf64>
func.call @kernel(%ey, %hz)
: (memref<4x4xf64>, memref<4x4xf64>) -> ()

%result = affine.load %ey[1, 3] : memref<4x4xf64>
%readable = arith.fptosi %result : f64 to i64
return %readable : i64
}
}
```

## Commands

```sh
LLVM_BUILD=/path/to/llvm-build

LOWER='builtin.module(lower-affine,convert-scf-to-cf,expand-strided-metadata,finalize-memref-to-llvm,convert-func-to-llvm,convert-arith-to-llvm,convert-cf-to-llvm,reconcile-unrealized-casts)'

SHARED_LIBS="$LLVM_BUILD/lib/libmlir_runner_utils.so,$LLVM_BUILD/lib/libmlir_c_runner_utils.so"

# Baseline.
"$LLVM_BUILD/bin/mlir-opt" input.mlir \
-pass-pipeline="$LOWER" \
-o baseline.mlir

"$LLVM_BUILD/bin/mlir-runner" baseline.mlir \
-e main \
--entry-point-result=i64 \
--shared-libs="$SHARED_LIBS"

# Apply affine-loop-fusion to the same input, then lower and execute it.
"$LLVM_BUILD/bin/mlir-opt" input.mlir \
--affine-loop-fusion \
-o fused.mlir

"$LLVM_BUILD/bin/mlir-opt" fused.mlir \
-pass-pipeline="$LOWER" \
-o fused-lowered.mlir

"$LLVM_BUILD/bin/mlir-runner" fused-lowered.mlir \
-e main \
--entry-point-result=i64 \
--shared-libs="$SHARED_LIBS"
```

## Actual behavior

The baseline and optimized executions produce:

```text
without affine-loop-fusion: -4
with affine-loop-fusion: 0
```

The generated function contains the following loop structure:

```mlir
affine.for %t = 0 to 2 {
affine.for %j = 0 to 3 {
affine.for %i = 1 to 4 {
// producer body
}
affine.for %i = 0 to 3 {
// consumer body
}
}
}
```

The producer body is therefore executed only for `j = [0, 3)`, even though its source iteration space is `j = [0, 4)`. The update of `ey[1, 3]` is lost, and the initialized value `0.0` is returned instead of `-4.0`.

## Expected behavior

`affine-loop-fusion` should preserve the complete producer iteration space, including producer iterations outside the consumer domain, or decline the fusion.

Applying the pass must not change the result of this program.

## Validity

Both the input and generated target pass MLIR verification and lower successfully. The two memrefs are distinct and fully initialized, all memory accesses are in bounds, and the arithmetic and runner entry point are defined. No undefined behavior explains the result difference.

Contributor guide

Open the contributing guide

Research direction

Start with the provided MLIR reproducer and run it through mlir-opt with --affine-loop-fusion, comparing the generated loop structure and results against the baseline commands. Trace the affine-loop-fusion pass behavior around mismatched producer and consumer domains; done means preserving the producer iteration at j = 3 or declining fusion so execution remains -4.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
50/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.