[Mlir-commits] [mlir] [mlir][affine] Fix consumer IV ordering in ComputationSliceState::isMaximal (PR #213513)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Aug 2 00:42:17 PDT 2026
https://github.com/xxxxbc created https://github.com/llvm/llvm-project/pull/213513
isMaximal() collected the consumer IVs in the order they appear in lbOperands[0], but addDomainFromSliceMaps expects dimension i of the constraint system to correspond to the loop at position i of the nest. When the two orders disagree, some dimensions receive no constraints; an unconstrained dimension looks unbounded, so a non-maximal slice is reported as maximal. Loop fusion then removes a producer whose extra iterations are still live, changing the result of the program.
Build the Values in dimension order instead: for slice bound i that is an equality on a consumer IV, use that IV, and leave the entry null for slice loops not yet materialized in the IR.
Fixes #212028.
>From 825debf5fc39c48365c839a1fa70c2ce86f86164 Mon Sep 17 00:00:00 2001
From: hehuan <2128534713 at qq.com>
Date: Sun, 2 Aug 2026 15:39:19 +0800
Subject: [PATCH] [mlir][affine] Fix consumer IV ordering in
ComputationSliceState::isMaximal
isMaximal() collected the consumer IVs in the order they appear in
lbOperands[0], but addDomainFromSliceMaps expects dimension i of the
constraint system to correspond to the loop at position i of the nest.
When the two orders disagree, some dimensions receive no constraints;
an unconstrained dimension looks unbounded, so a non-maximal slice is
reported as maximal. Loop fusion then removes a producer whose extra
iterations are still live, changing the result of the program.
Build the Values in dimension order instead: for slice bound i that is
an equality on a consumer IV, use that IV, and leave the entry null for
slice loops not yet materialized in the IR.
Fixes #212028.
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 35 +++++++++----
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 56 +++++++++++++++++++--
2 files changed, 78 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index cac305df8ba75..7ef16815e96d9 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -1120,15 +1120,32 @@ std::optional<bool> ComputationSliceState::isMaximal() const {
// Create constraints for the slice using the dst loop nest information. We
// retrieve existing dst loops from the lbOperands.
- SmallVector<Value> consumerIVs;
- for (Value lbOp : lbOperands[0])
- if (getForInductionVarOwner(lbOp))
- consumerIVs.push_back(lbOp);
-
- // Add empty IV Values for those new loops that are not equalities and,
- // therefore, are not yet materialized in the IR.
- for (int i = consumerIVs.size(), end = ivs.size(); i < end; ++i)
- consumerIVs.push_back(Value());
+ //
+ // `addDomainFromSliceMaps` assumes that dimension `i` of the system is the
+ // loop at position `i` of the nest (see its doc comment). Build the Values
+ // in that order: dimension `i` describes source IV `ivs[i]`, so use the
+ // consumer IV that slice bound `i` is an equality on, and leave the rest
+ // null for slice loops that aren't materialized in the IR yet. Collecting
+ // the operands in their own order instead would put the dimensions out of
+ // step with `lbs`/`ubs`, leaving some dimensions unconstrained -- and an
+ // unconstrained dimension makes the slice look unbounded, which in turn
+ // makes a non-maximal slice test as maximal.
+ SmallVector<Value> consumerIVs(ivs.size(), Value());
+ for (unsigned i = 0, e = std::min<size_t>(lbs.size(), ivs.size()); i < e;
+ ++i) {
+ AffineMap lbMap = lbs[i], ubMap = ubs[i];
+ if (!lbMap || !ubMap || lbMap.getNumResults() != 1 ||
+ ubMap.getNumResults() != 1 ||
+ lbMap.getResult(0) + 1 != ubMap.getResult(0) ||
+ isa<AffineConstantExpr>(lbMap.getResult(0)))
+ continue;
+ auto dim = dyn_cast<AffineDimExpr>(lbMap.getResult(0));
+ if (!dim || dim.getPosition() >= lbOperands[i].size())
+ continue;
+ Value operand = lbOperands[i][dim.getPosition()];
+ if (getForInductionVarOwner(operand))
+ consumerIVs[i] = operand;
+ }
FlatAffineValueConstraints sliceConstraints(/*numDims=*/consumerIVs.size(),
/*numSymbols=*/0,
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index cf530016c201a..a532920f0daa9 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -831,9 +831,21 @@ func.func @fusion_non_constant_bounds_1(%N: index, %M: memref<?xf32>, %cst: f32)
return
}
-// No fusion here as the cost models computing slice costs run out of 64-bit precision.
+// Both producer nests fuse into the consumer here. The two reduction loops
+// keep accumulating into the same `%alloc` element, so the slice covers the
+// full source iteration space.
// PRODUCER-CONSUMER-LABEL: func @high_trip_count
+// PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 16 {
+// PRODUCER-CONSUMER-NEXT: affine.for %{{.*}} = 0 to 512 {
+// PRODUCER-CONSUMER-NEXT: affine.for %{{.*}} = 0 to 64 {
+// PRODUCER-CONSUMER-NEXT: affine.for %{{.*}} = 0 to 16 {
+// PRODUCER-CONSUMER-NEXT: affine.for %{{.*}} = 0 to 2048 {
+// PRODUCER-CONSUMER: affine.store
+// PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 2048 {
+// PRODUCER-CONSUMER: affine.store
+// PRODUCER-CONSUMER: affine.load
+// PRODUCER-CONSUMER-NEXT: affine.store
func.func @high_trip_count(%arg0: memref<1024x4096xf32>, %arg1: memref<8192x4096xf32>) -> memref<1024x8192xf32> {
%cst_0 = arith.constant 1.000000e+00 : f32
%alloc = memref.alloc() : memref<1024x8192xf32>
@@ -847,7 +859,6 @@ func.func @high_trip_count(%arg0: memref<1024x4096xf32>, %arg1: memref<8192x4096
%2 = affine.load %alloc[%arg5 + %arg2 * 64, %arg6 + %arg4 * 16] : memref<1024x8192xf32>
%3 = arith.mulf %0, %1 : f32
%4 = arith.addf %2, %3 : f32
- // PRODUCER-CONSUMER: affine.store
affine.store %4, %alloc[%arg5 + %arg2 * 64, %arg6 + %arg4 * 16] : memref<1024x8192xf32>
}
}
@@ -863,7 +874,6 @@ func.func @high_trip_count(%arg0: memref<1024x4096xf32>, %arg1: memref<8192x4096
%2 = affine.load %alloc[%arg5 + %arg2 * 64, %arg6 + %arg4 * 16] : memref<1024x8192xf32>
%3 = arith.mulf %0, %1 : f32
%4 = arith.addf %2, %3 : f32
- // PRODUCER-CONSUMER: affine.store
affine.store %4, %alloc[%arg5 + %arg2 * 64, %arg6 + %arg4 * 16] : memref<1024x8192xf32>
}
}
@@ -871,7 +881,6 @@ func.func @high_trip_count(%arg0: memref<1024x4096xf32>, %arg1: memref<8192x4096
}
}
- // PRODUCER-CONSUMER: affine.for {{.*}} = 0 to 16
affine.for %arg2 = 0 to 16 {
affine.for %arg3 = 0 to 512 {
affine.for %arg4 = 0 to 64 {
@@ -884,3 +893,42 @@ func.func @high_trip_count(%arg0: memref<1024x4096xf32>, %arg1: memref<8192x4096
}
return %alloc : memref<1024x8192xf32>
}
+
+// -----
+
+// The producer covers j = [0, 4) while the consumer only covers j = [0, 3).
+// Whatever fusion happens, the producer's j = 3 iteration must survive:
+// dropping it loses the update to %ey[1, 3] and changes the result of the
+// program. Check that a loop covering the full j = [0, 4) range remains, and
+// that the producer body is still guarded by it.
+// See https://github.com/llvm/llvm-project/issues/212028.
+
+// PRODUCER-CONSUMER-LABEL: func @non_maximal_slice_keeps_producer
+// PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 2 {
+// PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 4 {
+// PRODUCER-CONSUMER: affine.for %{{.*}} = 1 to 4 {
+func.func @non_maximal_slice_keeps_producer(%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
+}
More information about the Mlir-commits
mailing list