[Mlir-commits] [mlir] [MLIR][Affine] Preserve WAR dependences in loop fusion (PR #212045)
Mingfei Guo
llvmlistbot at llvm.org
Sat Jul 25 10:34:35 PDT 2026
https://github.com/guoriyue created https://github.com/llvm/llvm-project/pull/212045
## Summary
Preserve affine write-after-read dependences when fusing a producer slice into a consumer loop.
The legality check composes the fused schedule (T -> S) with the original dependence (S -> D) and rejects fusion when the producer would execute lexicographically after its destination (T > D). Unsupported relation shapes are conservatively rejected.
Fixes #212032.
## Testing
- ninja -C build-212032 -j8 mlir-opt FileCheck count not
- llvm-lit: loop-fusion-3.mlir (1 passed)
- llvm-lit: full Affine loop-fusion filter (10 passed)
- Same-build A/B with the new guard disabled: 5 negative checks fail without the fix
>From 2cb8501fb4ef979931043d866e86baaa168b172d Mon Sep 17 00:00:00 2001
From: Mingfei Guo <1800012773 at pku.edu.cn>
Date: Sat, 25 Jul 2026 10:29:20 -0700
Subject: [PATCH] [MLIR][Affine] Preserve WAR dependences in loop fusion
---
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 148 +++++++++++++
mlir/test/Dialect/Affine/loop-fusion-3.mlir | 198 ++++++++++++++++++
2 files changed, 346 insertions(+)
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 68296ea3368a1..2a44ef3faaac2 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -11,6 +11,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/LoopFusionUtils.h"
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Analysis/TopologicalSortUtils.h"
#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
@@ -242,6 +243,146 @@ static unsigned getMaxLoopDepth(ArrayRef<Operation *> srcOps,
return loopDepth;
}
+/// Return whether `dependenceConstraints` contains a dependence that becomes
+/// lexicographically backward after executing the source slice in the
+/// destination loop nest.
+static std::optional<bool> hasBackwardDependenceAfterFusion(
+ const FlatAffineValueConstraints &dependenceConstraints,
+ const FlatAffineValueConstraints &sliceConstraints,
+ ArrayRef<Value> srcIVs, ArrayRef<Value> dstIVs, unsigned dstLoopDepth,
+ unsigned numCommonLoops) {
+ if (dstLoopDepth < numCommonLoops || dstIVs.size() < dstLoopDepth)
+ return std::nullopt;
+
+ unsigned numSrcIVs = srcIVs.size();
+ if (sliceConstraints.getNumDimVars() != numSrcIVs ||
+ dependenceConstraints.getNumDimVars() !=
+ numSrcIVs + dstIVs.size())
+ return std::nullopt;
+ for (auto [index, iv] : llvm::enumerate(srcIVs)) {
+ if (!sliceConstraints.hasValue(index) ||
+ sliceConstraints.getValue(index) != iv ||
+ !dependenceConstraints.hasValue(index) ||
+ dependenceConstraints.getValue(index) != iv)
+ return std::nullopt;
+ }
+ for (auto [index, iv] : llvm::enumerate(dstIVs)) {
+ unsigned pos = numSrcIVs + index;
+ if (!dependenceConstraints.hasValue(pos) ||
+ dependenceConstraints.getValue(pos) != iv)
+ return std::nullopt;
+ }
+
+ // View the slice as a relation from its destination insertion point to the
+ // source iteration, and the dependence as source-to-destination. Their
+ // composition directly relates the transformed source schedule to the
+ // original dependence destination.
+ unsigned firstScheduleSymbol = 0;
+ unsigned numScheduleDims = dstLoopDepth - numCommonLoops;
+ for (unsigned depth = numCommonLoops; depth < dstLoopDepth; ++depth) {
+ unsigned pos;
+ if (!sliceConstraints.findVar(dstIVs[depth], &pos,
+ sliceConstraints.getNumDimVars()))
+ return std::nullopt;
+ unsigned symbolPos = pos - sliceConstraints.getNumDimVars();
+ if (depth == numCommonLoops)
+ firstScheduleSymbol = symbolPos;
+ else if (symbolPos != firstScheduleSymbol + depth - numCommonLoops)
+ return std::nullopt;
+ }
+
+ presburger::IntegerRelation sliceRelation(sliceConstraints);
+ sliceRelation.convertVarKind(presburger::VarKind::Symbol,
+ firstScheduleSymbol,
+ firstScheduleSymbol + numScheduleDims,
+ presburger::VarKind::Domain);
+ presburger::IntegerRelation fusedDependence(dependenceConstraints);
+ fusedDependence.convertVarKind(presburger::VarKind::Range, 0, numSrcIVs,
+ presburger::VarKind::Domain);
+ fusedDependence.mergeAndCompose(sliceRelation);
+ if (fusedDependence.getNumDomainVars() != numScheduleDims ||
+ fusedDependence.getNumRangeVars() != dstIVs.size())
+ return std::nullopt;
+
+ unsigned dstOffset =
+ fusedDependence.getVarKindOffset(presburger::VarKind::Range);
+ for (unsigned depth = 0; depth < numScheduleDims; ++depth) {
+ presburger::IntegerRelation backwardConstraints(fusedDependence);
+ for (unsigned outerDepth = 0; outerDepth < depth; ++outerDepth) {
+ SmallVector<int64_t> equality(backwardConstraints.getNumCols(), 0);
+ equality[outerDepth] = 1;
+ equality[dstOffset + numCommonLoops + outerDepth] = -1;
+ backwardConstraints.addEquality(equality);
+ }
+
+ SmallVector<int64_t> inequality(backwardConstraints.getNumCols(), 0);
+ inequality[depth] = 1;
+ inequality[dstOffset + numCommonLoops + depth] = -1;
+ inequality.back() = -1;
+ backwardConstraints.addInequality(inequality);
+ if (!backwardConstraints.isIntegerEmpty())
+ return true;
+ }
+ return false;
+}
+
+/// Check affine WAR dependences between the source and destination that did
+/// not define the producer-consumer slice.
+static bool preservesAdditionalAffineWARDependences(
+ ArrayRef<Operation *> srcOps, ArrayRef<Operation *> dstOps,
+ unsigned dstLoopDepth, unsigned numCommonLoops,
+ const ComputationSliceState &srcSlice) {
+ FlatAffineValueConstraints sliceConstraints;
+ if (failed(srcSlice.getAsConstraints(&sliceConstraints)))
+ return false;
+
+ for (Operation *srcOp : srcOps) {
+ MemRefAccess srcAccess(srcOp);
+ for (Operation *dstOp : dstOps) {
+ MemRefAccess dstAccess(dstOp);
+ if (srcAccess.memref != dstAccess.memref)
+ continue;
+ // Source stores already participate in the producer-consumer slice.
+ // Only source-read/destination-write dependences are additional to it.
+ if (srcAccess.isStore() || !dstAccess.isStore())
+ continue;
+
+ FlatAffineValueConstraints dependenceConstraints;
+ DependenceResult result = checkMemrefAccessDependence(
+ srcAccess, dstAccess, numCommonLoops + 1, &dependenceConstraints);
+ if (result.value == DependenceResult::NoDependence)
+ continue;
+ if (result.value == DependenceResult::Failure)
+ return false;
+
+ SmallVector<AffineForOp> dstLoops;
+ getAffineForIVs(*dstOp, &dstLoops);
+ SmallVector<AffineForOp> srcLoops;
+ getAffineForIVs(*srcOp, &srcLoops);
+ SmallVector<Value> srcIVs;
+ for (AffineForOp loop : srcLoops)
+ srcIVs.push_back(loop.getInductionVar());
+ SmallVector<Value> dstIVs;
+ for (AffineForOp loop : dstLoops)
+ dstIVs.push_back(loop.getInductionVar());
+ std::optional<bool> hasBackward = hasBackwardDependenceAfterFusion(
+ dependenceConstraints, sliceConstraints, srcIVs, dstIVs,
+ dstLoopDepth, numCommonLoops);
+ if (!hasBackward) {
+ LDBG() << "Could not compare the affine dependence with the fused "
+ "schedule";
+ return false;
+ }
+ if (*hasBackward) {
+ LDBG() << "Affine WAR dependence becomes backward between " << *srcOp
+ << " and " << *dstOp;
+ return false;
+ }
+ }
+ }
+ return true;
+}
+
// TODO: This pass performs some computation that is the same for all the depths
// (e.g., getMaxLoopDepth). Implement a version of this utility that processes
// all the depths at once or only the legal maximal depth for maximal fusion.
@@ -349,6 +490,13 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
return FusionResult::FailIncorrectSlice;
}
+ if (fusionStrategy.getStrategy() == FusionStrategy::ProducerConsumer &&
+ !preservesAdditionalAffineWARDependences(
+ opsA, opsB, dstLoopDepth, numCommonLoops, *srcSlice)) {
+ LDBG() << "Fusion would reverse an additional affine WAR dependence";
+ return FusionResult::FailFusionDependence;
+ }
+
return FusionResult::Success;
}
diff --git a/mlir/test/Dialect/Affine/loop-fusion-3.mlir b/mlir/test/Dialect/Affine/loop-fusion-3.mlir
index 70d6c82105543..edd802e0a6b44 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-3.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-3.mlir
@@ -1,5 +1,6 @@
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion))' -split-input-file | FileCheck %s
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{maximal}))' -split-input-file | FileCheck %s --check-prefix=MAXIMAL
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer maximal}))' -split-input-file | FileCheck %s --check-prefix=PC-MAXIMAL
// Part I of fusion tests in mlir/test/Transforms/loop-fusion.mlir.
// Part II of fusion tests in mlir/test/Transforms/loop-fusion-2.mlir
@@ -1294,5 +1295,202 @@ func.func @unknown_memref_def_op() {
}
func.func private @bar() -> memref<10xf32>
+// -----
+
+// A producer-read/consumer-write dependence becomes backward if the producer
+// is sliced pointwise: source iteration 2 would read %a[1] after destination
+// iteration 1 writes it.
+
+// CHECK-LABEL: func.func @producer_read_consumer_write_backward
+// MAXIMAL-LABEL: func.func @producer_read_consumer_write_backward
+// PC-MAXIMAL-LABEL: func.func @producer_read_consumer_write_backward
+func.func @producer_read_consumer_write_backward(
+ %a: memref<3xf32>, %scratch: memref<3xf32>) {
+ affine.for %s = 1 to 3 {
+ %value = affine.load %a[1] : memref<3xf32>
+ affine.store %value, %scratch[%s] : memref<3xf32>
+ }
+ // PC-MAXIMAL: affine.for %[[S:.*]] = 1 to 3
+ // PC-MAXIMAL: affine.load %{{.*}}[1]
+ // PC-MAXIMAL: affine.for %[[D:.*]] = 1 to 3
+ // PC-MAXIMAL: affine.load %{{.*}}[%[[D]]]
+ affine.for %d = 1 to 3 {
+ %value = affine.load %scratch[%d] : memref<3xf32>
+ affine.store %value, %a[%d] : memref<3xf32>
+ }
+ return
+}
+
+// -----
+
+// The producer iteration s is scheduled at destination iteration s + 1. The
+// additional dependence also targets s + 1, so it remains forward.
+
+// CHECK-LABEL: func.func @offset_slice_forward_war
+// MAXIMAL-LABEL: func.func @offset_slice_forward_war
+// PC-MAXIMAL-LABEL: func.func @offset_slice_forward_war
+// PC-MAXIMAL-SAME: %[[A:[a-zA-Z0-9_]+]]: memref<10xf32>
+func.func @offset_slice_forward_war(
+ %a: memref<10xf32>, %scratch: memref<10xf32>) {
+ affine.for %s = 0 to 8 {
+ %value = affine.load %a[%s + 1] : memref<10xf32>
+ affine.store %value, %scratch[%s + 1] : memref<10xf32>
+ }
+ // The source slice is cloned into the destination loop.
+ // PC-MAXIMAL: affine.for %{{.*}} = 0 to 8
+ // PC-MAXIMAL: affine.for %[[D:.*]] = 1 to 9
+ // PC-MAXIMAL: affine.load %[[A]][%{{.*}} + 1]
+ // PC-MAXIMAL: affine.store %{{.*}}, %[[A]][%[[D]]]
+ affine.for %d = 1 to 9 {
+ %value = affine.load %scratch[%d] : memref<10xf32>
+ affine.store %value, %a[%d] : memref<10xf32>
+ }
+ return
+}
+
+// -----
+
+// Here the same source iteration is scheduled at s + 1, but its read must
+// precede destination iteration s. Fusion would reverse that dependence.
+
+// CHECK-LABEL: func.func @offset_slice_backward_war
+// MAXIMAL-LABEL: func.func @offset_slice_backward_war
+// PC-MAXIMAL-LABEL: func.func @offset_slice_backward_war
+// PC-MAXIMAL-SAME: %[[A:[a-zA-Z0-9_]+]]: memref<10xf32>, %[[SCRATCH:[a-zA-Z0-9_]+]]: memref<10xf32>
+func.func @offset_slice_backward_war(
+ %a: memref<10xf32>, %scratch: memref<10xf32>) {
+ affine.for %s = 0 to 8 {
+ %value = affine.load %a[%s] : memref<10xf32>
+ affine.store %value, %scratch[%s + 1] : memref<10xf32>
+ }
+ // PC-MAXIMAL: affine.for %{{.*}} = 0 to 8
+ // PC-MAXIMAL: affine.for %[[D:.*]] = 1 to 9
+ // PC-MAXIMAL-NOT: affine.load %[[A]]
+ // PC-MAXIMAL: affine.load %[[SCRATCH]][%[[D]]]
+ // PC-MAXIMAL: affine.store %{{.*}}, %[[A]][%[[D]]]
+ affine.for %d = 1 to 9 {
+ %value = affine.load %scratch[%d] : memref<10xf32>
+ affine.store %value, %a[%d] : memref<10xf32>
+ }
+ return
+}
+
+// -----
+
+// Verify that common surrounding loops remain parameters of the composed
+// schedule relation and do not prevent a safe pointwise fusion.
+
+// CHECK-LABEL: func.func @common_outer_pointwise_war
+// MAXIMAL-LABEL: func.func @common_outer_pointwise_war
+// PC-MAXIMAL-LABEL: func.func @common_outer_pointwise_war
+// PC-MAXIMAL-SAME: %[[A:[a-zA-Z0-9_]+]]: memref<4x10xf32>
+func.func @common_outer_pointwise_war(
+ %a: memref<4x10xf32>, %scratch: memref<4x10xf32>) {
+ affine.for %outer = 0 to 4 {
+ affine.for %s = 1 to 9 {
+ %value = affine.load %a[%outer, %s] : memref<4x10xf32>
+ affine.store %value, %scratch[%outer, %s] : memref<4x10xf32>
+ }
+ // The source slice is cloned into the destination loop.
+ // PC-MAXIMAL: affine.for %[[OUTER:.*]] = 0 to 4
+ // PC-MAXIMAL: affine.for %{{.*}} = 1 to 9
+ // PC-MAXIMAL: affine.for %[[D:.*]] = 1 to 9
+ // PC-MAXIMAL: affine.load %[[A]][%[[OUTER]], %[[D]]]
+ affine.for %d = 1 to 9 {
+ %value = affine.load %scratch[%outer, %d] : memref<4x10xf32>
+ affine.store %value, %a[%outer, %d] : memref<4x10xf32>
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @common_outer_backward_war
+// MAXIMAL-LABEL: func.func @common_outer_backward_war
+// PC-MAXIMAL-LABEL: func.func @common_outer_backward_war
+// PC-MAXIMAL-SAME: %[[A:[a-zA-Z0-9_]+]]: memref<4x10xf32>, %[[SCRATCH:[a-zA-Z0-9_]+]]: memref<4x10xf32>
+func.func @common_outer_backward_war(
+ %a: memref<4x10xf32>, %scratch: memref<4x10xf32>) {
+ affine.for %outer = 0 to 4 {
+ affine.for %s = 1 to 9 {
+ %value = affine.load %a[%outer, 1] : memref<4x10xf32>
+ affine.store %value, %scratch[%outer, %s] : memref<4x10xf32>
+ }
+ // PC-MAXIMAL: affine.for %[[OUTER:.*]] = 0 to 4
+ // PC-MAXIMAL: affine.for %{{.*}} = 1 to 9
+ // PC-MAXIMAL: affine.for %[[D:.*]] = 1 to 9
+ // PC-MAXIMAL-NOT: affine.load %[[A]]
+ // PC-MAXIMAL: affine.load %[[SCRATCH]][%[[OUTER]], %[[D]]]
+ // PC-MAXIMAL: affine.store %{{.*}}, %[[A]][%[[OUTER]], %[[D]]]
+ affine.for %d = 1 to 9 {
+ %value = affine.load %scratch[%outer, %d] : memref<4x10xf32>
+ affine.store %value, %a[%outer, %d] : memref<4x10xf32>
+ }
+ }
+ return
+}
+
+// -----
+
+// An inner-dimension backward dependence prevents depth-2 fusion, while
+// depth-1 fusion remains legal because each full source inner loop still
+// executes before the corresponding destination inner loop.
+
+// CHECK-LABEL: func.func @two_dimensional_inner_backward_war
+// MAXIMAL-LABEL: func.func @two_dimensional_inner_backward_war
+// PC-MAXIMAL-LABEL: func.func @two_dimensional_inner_backward_war
+func.func @two_dimensional_inner_backward_war(
+ %a: memref<6x6xf32>, %scratch: memref<6x6xf32>) {
+ affine.for %s0 = 1 to 5 {
+ affine.for %s1 = 1 to 5 {
+ %value = affine.load %a[%s0, %s1 - 1] : memref<6x6xf32>
+ affine.store %value, %scratch[%s0, %s1] : memref<6x6xf32>
+ }
+ }
+ // PC-MAXIMAL: affine.for %[[OUTER:.*]] = 1 to 5
+ // PC-MAXIMAL: affine.for %{{.*}} = 1 to 5
+ // PC-MAXIMAL: affine.for %{{.*}} = 1 to 5
+ // PC-MAXIMAL-NOT: affine.for
+ // PC-MAXIMAL: return
+ affine.for %d0 = 1 to 5 {
+ affine.for %d1 = 1 to 5 {
+ %value = affine.load %scratch[%d0, %d1] : memref<6x6xf32>
+ affine.store %value, %a[%d0, %d1] : memref<6x6xf32>
+ }
+ }
+ return
+}
+
+// -----
+
+// A backward dependence in the outer dimension prevents fusion at every
+// depth, irrespective of the inner-dimension direction.
+
+// CHECK-LABEL: func.func @two_dimensional_outer_backward_war
+// MAXIMAL-LABEL: func.func @two_dimensional_outer_backward_war
+// PC-MAXIMAL-LABEL: func.func @two_dimensional_outer_backward_war
+func.func @two_dimensional_outer_backward_war(
+ %a: memref<6x6xf32>, %scratch: memref<6x6xf32>) {
+ affine.for %s0 = 1 to 5 {
+ affine.for %s1 = 1 to 5 {
+ %value = affine.load %a[%s0 - 1, %s1 + 1] : memref<6x6xf32>
+ affine.store %value, %scratch[%s0, %s1] : memref<6x6xf32>
+ }
+ }
+ // PC-MAXIMAL: affine.for %{{.*}} = 1 to 5
+ // PC-MAXIMAL: affine.for %{{.*}} = 1 to 5
+ // PC-MAXIMAL: affine.for %{{.*}} = 1 to 5
+ // PC-MAXIMAL: affine.for %{{.*}} = 1 to 5
+ // PC-MAXIMAL-NOT: affine.for
+ // PC-MAXIMAL: return
+ affine.for %d0 = 1 to 5 {
+ affine.for %d1 = 1 to 5 {
+ %value = affine.load %scratch[%d0, %d1] : memref<6x6xf32>
+ affine.store %value, %a[%d0, %d1] : memref<6x6xf32>
+ }
+ }
+ return
+}
// Add further tests in mlir/test/Transforms/loop-fusion-4.mlir
More information about the Mlir-commits
mailing list