[Mlir-commits] [mlir] [mlir][Affine] Account for eliminated stores in fusion cost (PR #220711)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/220711

>From 4af0cb11bf27aa2cdc8421dfaa4a4270b5eb4397 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 2 Sep 2026 07:20:13 -0700
Subject: [PATCH] [mlir][Affine] Account for eliminated stores in fusion cost

Charge eliminated stores to their enclosing source loops when computing a
single-iteration fusion slice. The extracted cost utility instead charged
the destination insertion loop, where the adjustment was overwritten before
use.

Pre-fix regression:

  error: unexpected remark: fusion compute cost: 2
  error: expected remark "fusion compute cost: 1" was not produced

Found by Coverity.

Assisted-by: Codex
---
 .../Dialect/Affine/Utils/LoopFusionUtils.cpp  | 15 ++++----
 mlir/test/Dialect/Affine/loop-fusion-3.mlir   | 10 +++---
 .../Affine/loop-fusion-compute-cost.mlir      | 16 +++++++++
 .../lib/Dialect/Affine/TestLoopFusion.cpp     | 35 +++++++++++++++++++
 4 files changed, 64 insertions(+), 12 deletions(-)
 create mode 100644 mlir/test/Dialect/Affine/loop-fusion-compute-cost.mlir

diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 68296ea3368a1f..fc0bb23138f7bc 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -596,15 +596,17 @@ bool mlir::affine::getFusionComputeCost(AffineForOp srcForOp,
   if (storeLoadFwdGuaranteed) {
     // Subtract from operation count the loads/store we expect load/store
     // forwarding to remove.
-    unsigned storeCount = 0;
+    auto subtractFromEnclosingLoop = [&](Operation &op) {
+      SmallVector<AffineForOp, 4> loops;
+      getAffineForIVs(op, &loops);
+      if (!loops.empty())
+        --computeCostMap[loops.back()];
+    };
     llvm::SmallDenseSet<Value, 4> storeMemrefs;
     srcForOp.walk([&](AffineWriteOpInterface storeOp) {
       storeMemrefs.insert(storeOp.getMemRef());
-      ++storeCount;
+      subtractFromEnclosingLoop(*storeOp);
     });
-    // Subtract out any store ops in single-iteration src slice loop nest.
-    if (storeCount > 0)
-      computeCostMap[insertPointParent] = -storeCount;
     // Subtract out any load users of 'storeMemrefs' nested below
     // 'insertPointParent'.
     for (Value memref : storeMemrefs) {
@@ -616,8 +618,7 @@ bool mlir::affine::getFusionComputeCost(AffineForOp srcForOp,
         // 'insertPointParent'.
         getAffineForIVs(*user, &loops);
         if (llvm::is_contained(loops, cast<AffineForOp>(insertPointParent))) {
-          if (auto forOp = dyn_cast_or_null<AffineForOp>(user->getParentOp()))
-            --computeCostMap[forOp];
+          subtractFromEnclosingLoop(*user);
         }
       }
     }
diff --git a/mlir/test/Dialect/Affine/loop-fusion-3.mlir b/mlir/test/Dialect/Affine/loop-fusion-3.mlir
index 70d6c82105543c..6f9bde51ee022a 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-3.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-3.mlir
@@ -1056,11 +1056,11 @@ func.func @reduce_add_non_innermost(%arg0: memref<64x64xf32, 1>, %arg1: memref<1
   }
   return
 }
-// Test checks the loop structure is preserved after sibling fusion.
-// CHECK:         affine.for
-// CHECK-NEXT:      affine.for
-// CHECK-NEXT:        affine.for
-// CHECK:            affine.for
+// The cost model accounts for eliminated stores and chooses maximal sibling
+// fusion.
+// CHECK-COUNT-3: affine.for
+// CHECK-NOT:     affine.for
+// CHECK:         return
 
 
 
diff --git a/mlir/test/Dialect/Affine/loop-fusion-compute-cost.mlir b/mlir/test/Dialect/Affine/loop-fusion-compute-cost.mlir
new file mode 100644
index 00000000000000..b99fd9a1ec697e
--- /dev/null
+++ b/mlir/test/Dialect/Affine/loop-fusion-compute-cost.mlir
@@ -0,0 +1,16 @@
+// RUN: mlir-opt %s -test-loop-fusion=test-fusion-compute-cost \
+// RUN:   -verify-diagnostics
+
+func.func @single_iteration_store_load(%input: memref<1xf32>, %value: f32) {
+  affine.for %i = 0 to 1 {
+    affine.store %value, %input[%i] : memref<1xf32>
+  }
+  // expected-remark at below {{fusion compute cost: 1}}
+  affine.for %j = 0 to 1 {
+    affine.for %k = 0 to 1 {
+      %loaded = affine.load %input[%k] : memref<1xf32>
+      %unused = arith.addf %loaded, %loaded : f32
+    }
+  }
+  return
+}
diff --git a/mlir/test/lib/Dialect/Affine/TestLoopFusion.cpp b/mlir/test/lib/Dialect/Affine/TestLoopFusion.cpp
index bf11d94596fa7f..59e8a64801fa2a 100644
--- a/mlir/test/lib/Dialect/Affine/TestLoopFusion.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestLoopFusion.cpp
@@ -52,6 +52,11 @@ struct TestLoopFusion
       *this, "test-loop-fusion-utilities",
       llvm::cl::desc("Enable testing of loop fusion transformation utilities"),
       llvm::cl::init(false)};
+
+  Option<bool> clTestFusionComputeCost{
+      *this, "test-fusion-compute-cost",
+      llvm::cl::desc("Enable testing of the fusion compute cost"),
+      llvm::cl::init(false)};
 };
 
 } // namespace
@@ -154,6 +159,28 @@ static bool testLoopFusionUtilities(AffineForOp forOpA, AffineForOp forOpB,
   return false;
 }
 
+/// Computes and emits the cost of fusing `srcForOp` into `dstForOp`.
+static LogicalResult testFusionComputeCost(AffineForOp srcForOp,
+                                           AffineForOp dstForOp) {
+  LoopNestStats srcStats;
+  LoopNestStats dstStats;
+  if (!getLoopNestStats(srcForOp, &srcStats) ||
+      !getLoopNestStats(dstForOp, &dstStats))
+    return failure();
+
+  ComputationSliceState slice;
+  if (canFuseLoops(srcForOp, dstForOp, /*dstLoopDepth=*/1, &slice).value !=
+      FusionResult::Success)
+    return failure();
+
+  int64_t computeCost;
+  if (!getFusionComputeCost(srcForOp, srcStats, dstForOp, dstStats, slice,
+                            &computeCost))
+    return failure();
+  dstForOp->emitRemark() << "fusion compute cost: " << computeCost;
+  return success();
+}
+
 using LoopFunc = function_ref<bool(AffineForOp, AffineForOp, unsigned, unsigned,
                                    unsigned, unsigned)>;
 
@@ -198,6 +225,14 @@ void TestLoopFusion::runOnOperation() {
   // Gather all AffineForOps by loop depth.
   gatherLoops(getOperation(), depthToLoops);
 
+  if (clTestFusionComputeCost) {
+    if (depthToLoops.empty() || depthToLoops.front().size() != 2 ||
+        failed(testFusionComputeCost(depthToLoops.front()[0],
+                                     depthToLoops.front()[1])))
+      return signalPassFailure();
+    return;
+  }
+
   // Run tests on all combinations of src/dst loop nests in 'depthToLoops'.
   if (clTestDependenceCheck)
     iterateLoops(depthToLoops, testDependenceCheck);