[Mlir-commits] [mlir] [mlir][affine] Reuse translated loop-body computations (PR #213501)
Takayuki Todokoro
llvmlistbot at llvm.org
Sat Aug 1 20:43:55 PDT 2026
https://github.com/takatodo created https://github.com/llvm/llvm-project/pull/213501
# [mlir][affine] Reuse translated loop-body computations
## Summary
- Add an opt-in `affine-loop-carried-computation-reuse` pass.
- Detect equivalent pure computation DAGs whose affine loads differ by one
loop step.
- Compute the earlier DAG before the loop, carry its result through an
`affine.for` iteration argument, and yield the later DAG for the next
iteration.
For example, the pass turns the repeated producer in
```mlir
%src, %dst = memref.distinct_objects %src0, %dst0
: memref<10xi32>, memref<8xi32>
affine.for %i = 0 to 8 {
%a = affine.load %src[%i] : memref<10xi32>
%b = affine.load %src[%i + 1] : memref<10xi32>
%pa = arith.muli %a, %a : i32
%pb = arith.muli %b, %b : i32
%sum = arith.addi %pa, %pb : i32
affine.store %sum, %dst[%i] : memref<8xi32>
}
```
into the equivalent loop-carried form:
```mlir
%a0 = affine.load %src[0] : memref<10xi32>
%initial = arith.muli %a0, %a0 : i32
affine.for %i = 0 to 8 iter_args(%previous = %initial) {
%b = affine.load %src[%i + 1] : memref<10xi32>
%next = arith.muli %b, %b : i32
%sum = arith.addi %previous, %next : i32
affine.store %sum, %dst[%i] : memref<8xi32>
affine.yield %next : i32
}
```
## Motivation
Affine fusion, short-loop unrolling, and scalar replacement can expose this
translated producer overlap, but existing passes do not carry the computed
producer result between iterations. This pass performs only that missing
semantic rewrite and leaves vectorization, unrolling, register allocation, and
instruction scheduling to downstream target passes.
The pass is deliberately conservative. It requires a side-effect-free,
speculatable, single-result DAG; exact one-step translated Affine accesses;
sources proven stable under alias analysis; at least two proven iterations;
and a safe
first-iteration preload. It rejects raw-load-only pairs, changing loop-carried
context, unsupported indexing, unknown effects, and candidates that exceed the
matcher depth bound. One whole computation is selected per loop, avoiding an
unconditional fixed-point profitability policy.
The pass is opt-in and is not added to a default pipeline.
## Validation
- `check-mlir-dialect-affine`: 74/74 tests pass on LLVM `main`
`4e0d78f97449`.
- The tests cover direct and composed accesses, non-unit steps, symbolic and
empty trip counts, existing iteration arguments, aliasing and source
mutation, unknown effects, attribute and producer mismatches, preload
speculation, and composition with fusion/unroll/scalar replacement.
- `git diff --check` passes and `git clang-format --diff` is empty.
- In a local finite-volume face-flux case, the automatic and hand-written
carried forms lower to byte-identical objects; at 1,048,576 cells the
automatic form measured 1.25x faster than the fused recomputation form.
## AI assistance
OpenAI Codex assisted with implementation, tests, and prose. I reviewed the
resulting changes and ran the checks above.
>From 2933fc70dc2c2a9e798eb44726a68e68bde86eff Mon Sep 17 00:00:00 2001
From: takatodo <takatodo1227 at gmail.com>
Date: Sun, 2 Aug 2026 09:30:02 +0900
Subject: [PATCH 1/2] [mlir][affine] Reuse translated loop-body computations
Detect structurally equivalent pure computation DAGs whose affine loads are separated by one loop step, then carry the earlier result through an affine.for iter_arg. Keep raw loads and algebraic reduction rewrites outside this pass, and fail closed on aliasing, unstable sources, unsupported indexing, and unproven trip counts.
Assisted-by: OpenAI Codex
---
.../mlir/Dialect/Affine/Transforms/Passes.h | 5 +
.../mlir/Dialect/Affine/Transforms/Passes.td | 20 +
.../Dialect/Affine/Transforms/CMakeLists.txt | 1 +
.../LoopCarriedComputationReuse.cpp | 398 ++++++++++++++++++
...oop-carried-computation-reuse-corners.mlir | 274 ++++++++++++
...op-carried-computation-reuse-pipeline.mlir | 73 ++++
.../loop-carried-computation-reuse.mlir | 229 ++++++++++
7 files changed, 1000 insertions(+)
create mode 100644 mlir/lib/Dialect/Affine/Transforms/LoopCarriedComputationReuse.cpp
create mode 100644 mlir/test/Dialect/Affine/loop-carried-computation-reuse-corners.mlir
create mode 100644 mlir/test/Dialect/Affine/loop-carried-computation-reuse-pipeline.mlir
create mode 100644 mlir/test/Dialect/Affine/loop-carried-computation-reuse.mlir
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.h b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.h
index ce764eb750583..f6481493376ef 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.h
@@ -49,6 +49,11 @@ createSimplifyAffineStructuresPass();
std::unique_ptr<OperationPass<func::FuncOp>>
createAffineLoopInvariantCodeMotionPass();
+/// Creates a pass that reuses a one-iteration translated loop-body
+/// computation through an affine.for iter_arg.
+std::unique_ptr<OperationPass<func::FuncOp>>
+createAffineLoopCarriedComputationReusePass();
+
/// Creates a pass to convert all parallel affine.for's into 1-d affine.parallel
/// ops.
std::unique_ptr<OperationPass<func::FuncOp>> createAffineParallelizePass();
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
index f2008eeae9fcd..ad6d0d2a7e373 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
@@ -183,6 +183,26 @@ def AffineLoopInvariantCodeMotion
let constructor = "mlir::affine::createAffineLoopInvariantCodeMotionPass()";
}
+def AffineLoopCarriedComputationReuse
+ : Pass<"affine-loop-carried-computation-reuse", "func::FuncOp"> {
+ let summary = "Reuse one-iteration translated computations across affine "
+ "loop iterations";
+ let description = [{
+ This pass finds two side-effect-free computation DAGs consumed in the same
+ affine loop iteration when their affine loads differ by exactly one loop
+ step. It computes the earlier DAG once before the loop, carries its result
+ through an iter_arg, and yields the later DAG for the next iteration.
+
+ The pass preserves operation structure and does not reassociate
+ floating-point operations. It does not match a pair of raw loads or
+ algebraically rewrite reductions. One candidate is materialized per loop,
+ preferring a whole computation over a translated subexpression, to bound
+ the amount of new loop-carried state.
+ }];
+ let constructor = "mlir::affine::createAffineLoopCarriedComputationReusePass()";
+ let dependentDialects = ["arith::ArithDialect"];
+}
+
def AffineLoopTiling : Pass<"affine-loop-tile", "func::FuncOp"> {
let summary = "Tile affine loop nests";
let constructor = "mlir::affine::createLoopTilingPass()";
diff --git a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
index 9d912139810b2..c32d454a81d6b 100644
--- a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
@@ -8,6 +8,7 @@ add_mlir_dialect_library(MLIRAffineTransforms
AffineScalarReplacement.cpp
DecomposeAffineOps.cpp
FoldMemRefAliasOps.cpp
+ LoopCarriedComputationReuse.cpp
LoopCoalescing.cpp
LoopFusion.cpp
LoopTiling.cpp
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopCarriedComputationReuse.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopCarriedComputationReuse.cpp
new file mode 100644
index 0000000000000..472cf071d7881
--- /dev/null
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopCarriedComputationReuse.cpp
@@ -0,0 +1,398 @@
+//===- LoopCarriedComputationReuse.cpp -----------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements reuse of a pure loop-body computation whose affine
+// accesses are translated by one iteration.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/Transforms/Passes.h"
+
+#include "mlir/Analysis/AliasAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include <cstdint>
+#include <optional>
+
+namespace mlir {
+namespace affine {
+#define GEN_PASS_DEF_AFFINELOOPCARRIEDCOMPUTATIONREUSE
+#include "mlir/Dialect/Affine/Transforms/Passes.h.inc"
+} // namespace affine
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::affine;
+
+namespace {
+
+constexpr unsigned maxComputationDepth = 256;
+
+struct ReuseCandidate {
+ Value earlierRoot;
+ Value laterRoot;
+ SmallVector<Operation *> earlierOps;
+ SmallVector<Value> sources;
+};
+
+static void canonicalizeAccess(AffineLoadOp load, AffineMap &map,
+ SmallVectorImpl<Value> &operands) {
+ map = load.getAffineMap();
+ llvm::append_range(operands, load.getMapOperands());
+ fullyComposeAffineMapAndOperands(&map, &operands);
+ map = simplifyAffineMap(map);
+ canonicalizeMapAndOperands(&map, &operands);
+}
+
+/// Return whether evaluating `earlier` one loop step later accesses exactly
+/// the location accessed by `later` in the current iteration.
+static bool areLoadsOneIterationApart(AffineLoadOp earlier, AffineLoadOp later,
+ AffineForOp loop, bool &isTranslated) {
+ if (earlier.getMemRef() != later.getMemRef() ||
+ earlier.getType() != later.getType())
+ return false;
+
+ AffineMap earlierMap, laterMap;
+ SmallVector<Value> earlierOperands, laterOperands;
+ canonicalizeAccess(earlier, earlierMap, earlierOperands);
+ canonicalizeAccess(later, laterMap, laterOperands);
+
+ auto hasUnsupportedLoopLocalOperand = [&](ArrayRef<Value> operands) {
+ return llvm::any_of(operands, [&](Value operand) {
+ return operand != loop.getInductionVar() &&
+ !loop.isDefinedOutsideOfLoop(operand);
+ });
+ };
+ if (hasUnsupportedLoopLocalOperand(earlierOperands) ||
+ hasUnsupportedLoopLocalOperand(laterOperands))
+ return false;
+
+ MLIRContext *context = loop.getContext();
+ SmallVector<AffineExpr> dimReplacements;
+ SmallVector<AffineExpr> symbolReplacements;
+ for (unsigned i = 0; i < earlierMap.getNumDims(); ++i)
+ dimReplacements.push_back(getAffineDimExpr(i, context));
+ for (unsigned i = 0; i < earlierMap.getNumSymbols(); ++i)
+ symbolReplacements.push_back(getAffineSymbolExpr(i, context));
+
+ for (auto [index, operand] : llvm::enumerate(earlierOperands)) {
+ if (operand != loop.getInductionVar())
+ continue;
+ isTranslated = true;
+ if (index < earlierMap.getNumDims())
+ dimReplacements[index] = dimReplacements[index] + loop.getStepAsInt();
+ else
+ symbolReplacements[index - earlierMap.getNumDims()] =
+ symbolReplacements[index - earlierMap.getNumDims()] +
+ loop.getStepAsInt();
+ }
+
+ AffineMap shiftedEarlierMap = earlierMap.replaceDimsAndSymbols(
+ dimReplacements, symbolReplacements, earlierMap.getNumDims(),
+ earlierMap.getNumSymbols());
+ return AffineValueMap(shiftedEarlierMap, earlierOperands) ==
+ AffineValueMap(laterMap, laterOperands);
+}
+
+/// Match two side-effect-free single-result DAGs. The only varying leaves are
+/// affine loads separated by one loop iteration. Equal values must be defined
+/// outside the loop, so an existing iter_arg cannot become a reusable context.
+class ShiftedDAGMatcher {
+public:
+ explicit ShiftedDAGMatcher(AffineForOp loop) : loop(loop) {}
+
+ LogicalResult match(Value earlier, Value later) {
+ return matchImpl(earlier, later, /*depth=*/0);
+ }
+
+ SmallVector<Operation *> takeEarlierOps() { return std::move(earlierOps); }
+
+ ArrayRef<Operation *> getLaterOps() const { return laterOps; }
+
+ bool hasTranslation() const { return hasTranslatedLoad; }
+
+ SmallVector<Value> takeSources() {
+ return SmallVector<Value>(sources.begin(), sources.end());
+ }
+
+private:
+ LogicalResult matchImpl(Value earlier, Value later, unsigned depth) {
+ if (depth > maxComputationDepth)
+ return failure();
+ if (earlier == later)
+ return success(loop.isDefinedOutsideOfLoop(earlier));
+ if (earlier.getType() != later.getType())
+ return failure();
+
+ auto knownEarlier = earlierToLater.find(earlier);
+ if (knownEarlier != earlierToLater.end())
+ return success(knownEarlier->second == later);
+ auto knownLater = laterToEarlier.find(later);
+ if (knownLater != laterToEarlier.end())
+ return success(knownLater->second == earlier);
+
+ auto earlierLoad = earlier.getDefiningOp<AffineLoadOp>();
+ auto laterLoad = later.getDefiningOp<AffineLoadOp>();
+ if (earlierLoad || laterLoad) {
+ if (!earlierLoad || !laterLoad || earlierLoad->getParentOp() != loop ||
+ laterLoad->getParentOp() != loop ||
+ !loop.isDefinedOutsideOfLoop(earlierLoad.getMemRef()) ||
+ !areLoadsOneIterationApart(earlierLoad, laterLoad, loop,
+ hasTranslatedLoad))
+ return failure();
+ for (Value operand : earlierLoad.getMapOperands())
+ if (failed(recordAffineApplyDependencies(operand, depth + 1)))
+ return failure();
+ mapValues(earlier, later);
+ record(earlierLoad, laterLoad);
+ sources.insert(earlierLoad.getMemRef());
+ return success();
+ }
+
+ Operation *earlierOp = earlier.getDefiningOp();
+ Operation *laterOp = later.getDefiningOp();
+ if (!earlierOp || !laterOp || earlierOp == laterOp ||
+ earlierOp->getParentOp() != loop || laterOp->getParentOp() != loop ||
+ earlierOp->getNumResults() != 1 || laterOp->getNumResults() != 1 ||
+ earlierOp->getNumRegions() != 0 || laterOp->getNumRegions() != 0 ||
+ !isMemoryEffectFree(earlierOp) || !isMemoryEffectFree(laterOp) ||
+ !isSpeculatable(earlierOp) || !isSpeculatable(laterOp))
+ return failure();
+
+ mapValues(earlier, later);
+ auto flags = static_cast<OperationEquivalence::Flags>(
+ OperationEquivalence::IgnoreLocations |
+ OperationEquivalence::IgnoreDiscardableAttrs |
+ OperationEquivalence::IgnoreCommutativity);
+ if (!OperationEquivalence::isEquivalentTo(
+ earlierOp, laterOp,
+ [&](Value earlierOperand, Value laterOperand) {
+ return matchImpl(earlierOperand, laterOperand, depth + 1);
+ },
+ /*markEquivalent=*/nullptr, flags))
+ return failure();
+
+ record(earlierOp, laterOp);
+ return success();
+ }
+
+ LogicalResult recordAffineApplyDependencies(Value value, unsigned depth) {
+ if (depth > maxComputationDepth)
+ return failure();
+ if (value == loop.getInductionVar() || loop.isDefinedOutsideOfLoop(value))
+ return success();
+ auto apply = value.getDefiningOp<AffineApplyOp>();
+ if (!apply || apply->getParentOp() != loop)
+ return failure();
+ if (seenEarlier.contains(apply))
+ return success();
+ for (Value operand : apply.getMapOperands())
+ if (failed(recordAffineApplyDependencies(operand, depth + 1)))
+ return failure();
+ if (seenEarlier.insert(apply).second)
+ earlierOps.push_back(apply);
+ return success();
+ }
+
+ void mapValues(Value earlier, Value later) {
+ earlierToLater.try_emplace(earlier, later);
+ laterToEarlier.try_emplace(later, earlier);
+ }
+
+ void record(Operation *earlier, Operation *later) {
+ if (seenEarlier.insert(earlier).second)
+ earlierOps.push_back(earlier);
+ if (seenLater.insert(later).second)
+ laterOps.push_back(later);
+ }
+
+ AffineForOp loop;
+ DenseMap<Value, Value> earlierToLater;
+ DenseMap<Value, Value> laterToEarlier;
+ llvm::SmallPtrSet<Operation *, 16> seenEarlier;
+ llvm::SmallPtrSet<Operation *, 16> seenLater;
+ llvm::SmallSetVector<Value, 4> sources;
+ SmallVector<Operation *> earlierOps;
+ SmallVector<Operation *> laterOps;
+ bool hasTranslatedLoad = false;
+};
+
+static bool sourceIsStable(AffineForOp loop, Value source,
+ AliasAnalysis &aliasAnalysis) {
+ for (Operation &operation : loop.getBody()->without_terminator())
+ if (aliasAnalysis.getModRef(&operation, source).isMod())
+ return false;
+ return true;
+}
+
+/// Return whether the loop is proven to execute at least twice. Handle
+/// constant bounds directly to avoid overflowing a signed bound difference.
+/// Reject a negative symbolic result returned in an APInt as well.
+static bool hasAtLeastTwoIterations(AffineForOp loop) {
+ if (loop.hasConstantBounds()) {
+ int64_t lowerBound = loop.getConstantLowerBound();
+ int64_t upperBound = loop.getConstantUpperBound();
+ if (upperBound <= lowerBound)
+ return false;
+ uint64_t span =
+ static_cast<uint64_t>(upperBound) - static_cast<uint64_t>(lowerBound);
+ return span > static_cast<uint64_t>(loop.getStepAsInt());
+ }
+
+ std::optional<APInt> tripCount = loop.getStaticTripCount();
+ return tripCount && !tripCount->isNegative() && tripCount->ugt(1);
+}
+
+static std::optional<ReuseCandidate>
+findReuseCandidate(AffineForOp loop, AliasAnalysis &aliasAnalysis) {
+ if (loop.getLowerBoundMap().getNumResults() != 1)
+ return std::nullopt;
+ if (!hasAtLeastTwoIterations(loop))
+ return std::nullopt;
+
+ SmallVector<Operation *> bodyOps;
+ for (Operation &operation : loop.getBody()->without_terminator())
+ bodyOps.push_back(&operation);
+
+ // Search consumers backwards so that a whole translated computation is
+ // chosen before a translated subexpression inside that computation.
+ for (Operation *consumer : llvm::reverse(bodyOps)) {
+ for (unsigned laterIndex = 0; laterIndex < consumer->getNumOperands();
+ ++laterIndex) {
+ for (unsigned earlierIndex = 0; earlierIndex < consumer->getNumOperands();
+ ++earlierIndex) {
+ if (earlierIndex == laterIndex)
+ continue;
+ Value earlierRoot = consumer->getOperand(earlierIndex);
+ Value laterRoot = consumer->getOperand(laterIndex);
+ if (!earlierRoot.getDefiningOp() || !laterRoot.getDefiningOp() ||
+ isa<AffineLoadOp>(earlierRoot.getDefiningOp()))
+ continue;
+
+ ShiftedDAGMatcher matcher(loop);
+ if (failed(matcher.match(earlierRoot, laterRoot)) ||
+ !matcher.hasTranslation())
+ continue;
+ if (llvm::is_contained(matcher.getLaterOps(),
+ earlierRoot.getDefiningOp()))
+ continue;
+
+ SmallVector<Operation *> earlierOps = matcher.takeEarlierOps();
+ SmallVector<Value> sources = matcher.takeSources();
+ if (earlierOps.empty() || sources.empty() ||
+ llvm::any_of(sources, [&](Value source) {
+ return !sourceIsStable(loop, source, aliasAnalysis);
+ }))
+ continue;
+
+ llvm::SmallPtrSet<Operation *, 16> earlierSet(earlierOps.begin(),
+ earlierOps.end());
+ if (llvm::any_of(earlierOps, [&](Operation *operation) {
+ return llvm::any_of(operation->getOperands(), [&](Value operand) {
+ if (operand == loop.getInductionVar() ||
+ loop.isDefinedOutsideOfLoop(operand))
+ return false;
+ Operation *definingOp = operand.getDefiningOp();
+ return !definingOp || !earlierSet.contains(definingOp);
+ });
+ }))
+ continue;
+
+ return ReuseCandidate{earlierRoot, laterRoot, std::move(earlierOps),
+ std::move(sources)};
+ }
+ }
+ }
+ return std::nullopt;
+}
+
+static LogicalResult materializeReuse(IRRewriter &rewriter, AffineForOp loop,
+ ReuseCandidate candidate) {
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPoint(loop);
+
+ Value lowerBound;
+ if (loop.hasConstantLowerBound())
+ lowerBound = arith::ConstantIndexOp::create(rewriter, loop.getLoc(),
+ loop.getConstantLowerBound());
+ else
+ lowerBound =
+ AffineApplyOp::create(rewriter, loop.getLoc(), loop.getLowerBoundMap(),
+ loop.getLowerBoundOperands());
+
+ IRMapping mapping;
+ mapping.map(loop.getInductionVar(), lowerBound);
+ SmallVector<Operation *> clonedOps;
+ clonedOps.reserve(candidate.earlierOps.size());
+ for (Operation *operation : candidate.earlierOps)
+ clonedOps.push_back(rewriter.clone(*operation, mapping));
+ Value initial = mapping.lookup(candidate.earlierRoot);
+
+ BlockArgument carried;
+ FailureOr<LoopLikeOpInterface> replacement = loop.replaceWithAdditionalYields(
+ rewriter, initial, /*replaceInitOperandUsesInLoop=*/false,
+ [&](OpBuilder &, Location, ArrayRef<BlockArgument> newArguments) {
+ carried = newArguments.front();
+ return SmallVector<Value>{candidate.laterRoot};
+ });
+ if (failed(replacement)) {
+ for (Operation *operation : llvm::reverse(clonedOps))
+ rewriter.eraseOp(operation);
+ if (lowerBound.use_empty())
+ rewriter.eraseOp(lowerBound.getDefiningOp());
+ return failure();
+ }
+
+ candidate.earlierRoot.replaceAllUsesWith(carried);
+ for (Operation *operation : llvm::reverse(candidate.earlierOps))
+ if (isOpTriviallyDead(operation))
+ rewriter.eraseOp(operation);
+ return success();
+}
+
+struct AffineLoopCarriedComputationReuse
+ : public affine::impl::AffineLoopCarriedComputationReuseBase<
+ AffineLoopCarriedComputationReuse> {
+ void runOnOperation() override {
+ AliasAnalysis &aliasAnalysis = getAnalysis<AliasAnalysis>();
+ SmallVector<AffineForOp> loops;
+ getOperation().walk<WalkOrder::PostOrder>(
+ [&](AffineForOp loop) { loops.push_back(loop); });
+
+ IRRewriter rewriter(&getContext());
+ for (AffineForOp loop : loops) {
+ std::optional<ReuseCandidate> candidate =
+ findReuseCandidate(loop, aliasAnalysis);
+ if (candidate &&
+ failed(materializeReuse(rewriter, loop, std::move(*candidate)))) {
+ signalPassFailure();
+ return;
+ }
+ }
+ }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+mlir::affine::createAffineLoopCarriedComputationReusePass() {
+ return std::make_unique<AffineLoopCarriedComputationReuse>();
+}
diff --git a/mlir/test/Dialect/Affine/loop-carried-computation-reuse-corners.mlir b/mlir/test/Dialect/Affine/loop-carried-computation-reuse-corners.mlir
new file mode 100644
index 0000000000000..4096a92ef139c
--- /dev/null
+++ b/mlir/test/Dialect/Affine/loop-carried-computation-reuse-corners.mlir
@@ -0,0 +1,274 @@
+// RUN: mlir-opt %s \
+// RUN: --pass-pipeline='builtin.module(func.func(affine-loop-carried-computation-reuse),canonicalize,cse)' \
+// RUN: | FileCheck %s
+
+#lb = affine_map<(d0) -> (d0)>
+#ub8 = affine_map<(d0) -> (d0 + 8)>
+#ub_minus1 = affine_map<(d0) -> (d0 - 1)>
+#plus0 = affine_map<(d0) -> (d0)>
+#plus1 = affine_map<(d0) -> (d0 + 1)>
+#plus2 = affine_map<(d0) -> (d0 + 2)>
+
+// The translation is one iteration, not one index unit.
+
+// CHECK-LABEL: func.func @step_two
+// CHECK: %[[INIT:.*]] = arith.addi
+// CHECK: affine.for %[[I:.*]] = 0 to 8 step 2 iter_args(%[[PREV:.*]] = %[[INIT]])
+// CHECK: %[[CUR:.*]] = arith.addi
+// CHECK: arith.subi %[[CUR]], %[[PREV]]
+// CHECK: affine.yield %[[CUR]]
+func.func @step_two(%src0: memref<12xi32>, %dst0: memref<4xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<12xi32>, memref<4xi32>
+ affine.for %i = 0 to 8 step 2 {
+ %a = affine.load %src[%i] : memref<12xi32>
+ %b = affine.load %src[%i + 2] : memref<12xi32>
+ %c = affine.load %src[%i + 4] : memref<12xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ %j = affine.apply affine_map<(d0) -> (d0 floordiv 2)>(%i)
+ affine.store %difference, %dst[%j] : memref<4xi32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @step_mismatch
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @step_mismatch(%src0: memref<10xi32>, %dst0: memref<4xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<10xi32>, memref<4xi32>
+ affine.for %i = 0 to 8 step 2 {
+ %a = affine.load %src[%i] : memref<10xi32>
+ %b = affine.load %src[%i + 1] : memref<10xi32>
+ %c = affine.load %src[%i + 2] : memref<10xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ %j = affine.apply affine_map<(d0) -> (d0 floordiv 2)>(%i)
+ affine.store %difference, %dst[%j] : memref<4xi32>
+ }
+ return
+}
+
+// affine.apply chains are composed before comparing accesses.
+
+// CHECK-LABEL: func.func @composed_accesses
+// CHECK: %[[INIT:.*]] = arith.addi
+// CHECK: affine.for %[[I:.*]] = 0 to 8 iter_args(%[[PREV:.*]] = %[[INIT]])
+// CHECK: %[[CUR:.*]] = arith.addi
+// CHECK: arith.subi %[[CUR]], %[[PREV]]
+// CHECK: affine.yield %[[CUR]]
+func.func @composed_accesses(%src0: memref<10xi32>, %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<10xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %i0 = affine.apply #plus0(%i)
+ %i1 = affine.apply #plus1(%i)
+ %i2 = affine.apply #plus2(%i)
+ %a = affine.load %src[%i0] : memref<10xi32>
+ %b = affine.load %src[%i1] : memref<10xi32>
+ %c = affine.load %src[%i2] : memref<10xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
+
+// Non-linear Affine maps are accepted when canonical substitution proves the
+// same one-iteration translation.
+
+// CHECK-LABEL: func.func @modulo_accesses
+// CHECK: %[[INIT:.*]] = arith.addi
+// CHECK: affine.for %[[I:.*]] = 0 to 8 iter_args(%[[PREV:.*]] = %[[INIT]])
+// CHECK: %[[CUR:.*]] = arith.addi
+// CHECK: arith.subi %[[CUR]], %[[PREV]]
+// CHECK: affine.yield %[[CUR]]
+func.func @modulo_accesses(%src0: memref<4xi32>, %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<4xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %a = affine.load %src[%i mod 4] : memref<4xi32>
+ %b = affine.load %src[(%i + 1) mod 4] : memref<4xi32>
+ %c = affine.load %src[(%i + 2) mod 4] : memref<4xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
+
+// A symbolic lower bound is legal when the trip count is nevertheless known.
+
+// CHECK-LABEL: func.func @symbolic_lower_fixed_trip
+// CHECK: affine.load %{{.*}}[symbol(%{{.*}})]
+// CHECK: %[[INIT:.*]] = arith.addi
+// CHECK: affine.for %[[I:.*]] = %{{.*}} to {{.*}} iter_args(%[[PREV:.*]] = %[[INIT]])
+// CHECK: %[[CUR:.*]] = arith.addi
+// CHECK: affine.yield %[[CUR]]
+func.func @symbolic_lower_fixed_trip(%src0: memref<?xi32>,
+ %dst0: memref<?xi32>, %start: index) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<?xi32>, memref<?xi32>
+ affine.for %i = #lb(%start) to #ub8(%start) {
+ %a = affine.load %src[%i] : memref<?xi32>
+ %b = affine.load %src[%i + 1] : memref<?xi32>
+ %c = affine.load %src[%i + 2] : memref<?xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<?xi32>
+ }
+ return
+}
+
+// A statically empty symbolic interval must not cause a prologue load.
+
+// CHECK-LABEL: func.func @symbolic_negative_trip
+// CHECK-NOT: affine.load %{{.*}}[symbol(%{{.*}})]
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @symbolic_negative_trip(%src0: memref<?xi32>,
+ %dst0: memref<?xi32>, %start: index) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<?xi32>, memref<?xi32>
+ affine.for %i = #lb(%start) to #ub_minus1(%start) {
+ %a = affine.load %src[%i] : memref<?xi32>
+ %b = affine.load %src[%i + 1] : memref<?xi32>
+ %c = affine.load %src[%i + 2] : memref<?xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<?xi32>
+ }
+ return
+}
+
+// A signed-overflowing bound difference is still an empty interval.
+
+// CHECK-LABEL: func.func @constant_overflow_empty
+// CHECK-NOT: affine.load
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @constant_overflow_empty(%src0: memref<?xi32>,
+ %dst0: memref<?xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<?xi32>, memref<?xi32>
+ affine.for %i = 9223372036854775807 to -9223372036854775800 {
+ %a = affine.load %src[%i] : memref<?xi32>
+ %b = affine.load %src[%i + 1] : memref<?xi32>
+ %c = affine.load %src[%i + 2] : memref<?xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<?xi32>
+ }
+ return
+}
+
+// One executing iteration cannot eliminate a repeated producer evaluation.
+
+// CHECK-LABEL: func.func @one_trip
+// CHECK-NOT: affine.load %{{.*}}[0]
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @one_trip(%src0: memref<3xi32>, %dst0: memref<1xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<3xi32>, memref<1xi32>
+ affine.for %i = 0 to 1 {
+ %a = affine.load %src[%i] : memref<3xi32>
+ %b = affine.load %src[%i + 1] : memref<3xi32>
+ %c = affine.load %src[%i + 2] : memref<3xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<1xi32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @zero_trip
+// CHECK-NOT: affine.load
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @zero_trip(%src0: memref<2xi32>, %dst0: memref<1xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<2xi32>, memref<1xi32>
+ affine.for %i = 0 to 0 {
+ %a = affine.load %src[%i] : memref<2xi32>
+ %b = affine.load %src[%i + 1] : memref<2xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %a : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<1xi32>
+ }
+ return
+}
+
+// Operation attributes are part of the producer semantics.
+
+// CHECK-LABEL: func.func @attribute_mismatch
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @attribute_mismatch(%src0: memref<10xi32>, %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<10xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %a = affine.load %src[%i] : memref<10xi32>
+ %b = affine.load %src[%i + 1] : memref<10xi32>
+ %c = affine.load %src[%i + 2] : memref<10xi32>
+ %left = arith.addi %a, %b overflow<nsw> : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
+
+// Equal-looking computations from different memory objects are not a
+// translated producer pair.
+
+// CHECK-LABEL: func.func @different_sources
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @different_sources(%a0: memref<10xi32>, %b0: memref<10xi32>,
+ %dst0: memref<8xi32>) {
+ %a, %b, %dst = memref.distinct_objects %a0, %b0, %dst0
+ : memref<10xi32>, memref<10xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %a1 = affine.load %a[%i] : memref<10xi32>
+ %a2 = affine.load %a[%i + 1] : memref<10xi32>
+ %b1 = affine.load %b[%i + 1] : memref<10xi32>
+ %b2 = affine.load %b[%i + 2] : memref<10xi32>
+ %left = arith.addi %a1, %a2 : i32
+ %right = arith.addi %b1, %b2 : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
+
+// Repeated loop-invariant computations belong to LICM/CSE, not translated
+// loop-carried reuse.
+
+// CHECK-LABEL: func.func @loop_invariant_duplicates
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @loop_invariant_duplicates(%src0: memref<1xi32>,
+ %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<1xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %a = affine.load %src[0] : memref<1xi32>
+ %b = affine.load %src[0] : memref<1xi32>
+ %left = arith.muli %a, %a : i32
+ %right = arith.muli %b, %b : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
diff --git a/mlir/test/Dialect/Affine/loop-carried-computation-reuse-pipeline.mlir b/mlir/test/Dialect/Affine/loop-carried-computation-reuse-pipeline.mlir
new file mode 100644
index 0000000000000..9f25af9f4db14
--- /dev/null
+++ b/mlir/test/Dialect/Affine/loop-carried-computation-reuse-pipeline.mlir
@@ -0,0 +1,73 @@
+// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer maximal},affine-loop-unroll{unroll-factor=-1 unroll-full-threshold=2},affine-scalrep,affine-loop-carried-computation-reuse),canonicalize,cse)' | FileCheck %s
+
+// Existing fusion, short-loop unrolling, and scalar replacement expose the
+// translated producer pair. Computation reuse is responsible only for the
+// final loop-carried SSA value.
+
+// CHECK-LABEL: func.func @fusible
+// CHECK-NOT: memref.alloc
+// CHECK: %[[A:.*]] = affine.load %[[SRC:.*]][0]
+// CHECK: %[[B:.*]] = affine.load %[[SRC]][1]
+// CHECK: %[[INIT:.*]] = arith.muli %[[A]], %[[B]]
+// CHECK: affine.for %[[I:.*]] = 0 to 16 iter_args(%[[PREV:.*]] = %[[INIT]])
+// CHECK: %[[B2:.*]] = affine.load %[[SRC]][%[[I]] + 1]
+// CHECK: %[[C:.*]] = affine.load %[[SRC]][%[[I]] + 2]
+// CHECK: %[[CURRENT:.*]] = arith.muli %[[B2]], %[[C]]
+// CHECK: arith.subi %[[CURRENT]], %[[PREV]]
+// CHECK: affine.yield %[[CURRENT]]
+// CHECK: return
+func.func @fusible(%src0: memref<18xi32>, %out0: memref<16xi32>) {
+ %src, %out = memref.distinct_objects %src0, %out0
+ : memref<18xi32>, memref<16xi32>
+ %temporary = memref.alloc() : memref<17xi32>
+ affine.for %f = 0 to 17 {
+ %a = affine.load %src[%f] : memref<18xi32>
+ %b = affine.load %src[%f + 1] : memref<18xi32>
+ %product = arith.muli %a, %b : i32
+ affine.store %product, %temporary[%f] : memref<17xi32>
+ }
+ affine.for %i = 0 to 16 {
+ %left = affine.load %temporary[%i] : memref<17xi32>
+ %right = affine.load %temporary[%i + 1] : memref<17xi32>
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %out[%i] : memref<16xi32>
+ }
+ memref.dealloc %temporary : memref<17xi32>
+ return
+}
+
+// Current affine-loop-fusion intentionally skips a destination loop with
+// results. Keep this boundary visible instead of silently growing the custom
+// reuse pass into a fusion implementation.
+
+// CHECK-LABEL: func.func @result_bearing_consumer
+// CHECK: %[[TEMP:.*]] = memref.alloc
+// CHECK: affine.for %{{.*}} = 0 to 17
+// CHECK: affine.store %{{.*}}, %[[TEMP]]
+// CHECK: %[[SUM:.*]] = affine.for %{{.*}} = 0 to 16 iter_args
+// CHECK: affine.load %[[TEMP]]
+// CHECK: affine.load %[[TEMP]]
+// CHECK: return %[[SUM]]
+func.func @result_bearing_consumer(%src0: memref<18xi32>,
+ %out0: memref<16xi32>) -> i32 {
+ %src, %out = memref.distinct_objects %src0, %out0
+ : memref<18xi32>, memref<16xi32>
+ %temporary = memref.alloc() : memref<17xi32>
+ %zero = arith.constant 0 : i32
+ affine.for %f = 0 to 17 {
+ %a = affine.load %src[%f] : memref<18xi32>
+ %b = affine.load %src[%f + 1] : memref<18xi32>
+ %product = arith.muli %a, %b : i32
+ affine.store %product, %temporary[%f] : memref<17xi32>
+ }
+ %sum = affine.for %i = 0 to 16 iter_args(%acc = %zero) -> i32 {
+ %left = affine.load %temporary[%i] : memref<17xi32>
+ %right = affine.load %temporary[%i + 1] : memref<17xi32>
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %out[%i] : memref<16xi32>
+ %next = arith.addi %acc, %difference : i32
+ affine.yield %next : i32
+ }
+ memref.dealloc %temporary : memref<17xi32>
+ return %sum : i32
+}
diff --git a/mlir/test/Dialect/Affine/loop-carried-computation-reuse.mlir b/mlir/test/Dialect/Affine/loop-carried-computation-reuse.mlir
new file mode 100644
index 0000000000000..9344036ba3967
--- /dev/null
+++ b/mlir/test/Dialect/Affine/loop-carried-computation-reuse.mlir
@@ -0,0 +1,229 @@
+// RUN: mlir-opt --allow-unregistered-dialect %s \
+// RUN: --pass-pipeline='builtin.module(func.func(affine-loop-carried-computation-reuse),canonicalize,cse)' \
+// RUN: | FileCheck %s
+
+// A non-trivial producer DAG at i+1 is identical to the DAG at i in the next
+// iteration. Carry its result without interpreting addi or muli as reductions.
+
+// CHECK-LABEL: func.func @integer_polynomial
+// CHECK: %[[A0:.*]] = affine.load %[[SRC:.*]][0]
+// CHECK: %[[A02:.*]] = arith.muli %[[A0]], %[[A0]]
+// CHECK: %[[B0:.*]] = affine.load %[[SRC]][1]
+// CHECK: %[[INITIAL:.*]] = arith.addi %[[A02]], %[[B0]]
+// CHECK: affine.for %[[I:.*]] = 0 to 8 iter_args(%[[PREVIOUS:.*]] = %[[INITIAL]])
+// CHECK-NOT: affine.load %[[SRC]][%[[I]]]
+// CHECK: %[[B:.*]] = affine.load %[[SRC]][%[[I]] + 1]
+// CHECK: %[[C:.*]] = affine.load %[[SRC]][%[[I]] + 2]
+// CHECK: %[[B2:.*]] = arith.muli %[[B]], %[[B]]
+// CHECK: %[[CURRENT:.*]] = arith.addi %[[B2]], %[[C]]
+// CHECK: arith.subi %[[CURRENT]], %[[PREVIOUS]]
+// CHECK: affine.yield %[[CURRENT]]
+func.func @integer_polynomial(%src0: memref<10xi32>,
+ %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<10xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %a = affine.load %src[%i] : memref<10xi32>
+ %b = affine.load %src[%i + 1] : memref<10xi32>
+ %c = affine.load %src[%i + 2] : memref<10xi32>
+ %a2 = arith.muli %a, %a : i32
+ %left = arith.addi %a2, %b : i32
+ %b2 = arith.muli %b, %b : i32
+ %right = arith.addi %b2, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
+
+// Existing ordered and reduction state keeps its original position. The
+// translated producer state is appended.
+
+// CHECK-LABEL: func.func @existing_iter_arg
+// CHECK: %[[RESULT:.*]]:2 = affine.for %[[I:.*]] = 0 to 8
+// CHECK-SAME: iter_args(%[[ACC:.*]] = %{{.*}}, %[[PREVIOUS:.*]] = %{{.*}})
+// CHECK: %[[CURRENT:.*]] = arith.addi
+// CHECK: %[[NEXT:.*]] = arith.addi %[[ACC]], %{{.*}}
+// CHECK: affine.yield %[[NEXT]], %[[CURRENT]]
+// CHECK: return %[[RESULT]]#0
+func.func @existing_iter_arg(%src0: memref<10xi32>,
+ %dst0: memref<8xi32>) -> i32 {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<10xi32>, memref<8xi32>
+ %zero = arith.constant 0 : i32
+ %result = affine.for %i = 0 to 8 iter_args(%acc = %zero) -> i32 {
+ %a = affine.load %src[%i] : memref<10xi32>
+ %b = affine.load %src[%i + 1] : memref<10xi32>
+ %c = affine.load %src[%i + 2] : memref<10xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ %next = arith.addi %acc, %difference : i32
+ affine.yield %next : i32
+ }
+ return %result : i32
+}
+
+// A raw-load pair is intentionally outside this computation-reuse pass.
+
+// CHECK-LABEL: func.func @raw_load_only
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @raw_load_only(%src0: memref<9xi32>, %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<9xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %left = affine.load %src[%i] : memref<9xi32>
+ %right = affine.load %src[%i + 1] : memref<9xi32>
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
+
+// Without object disjointness, the destination write may modify the source
+// before the carried value is consumed in the next iteration.
+
+// CHECK-LABEL: func.func @may_alias_output
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @may_alias_output(%src: memref<10xi32>, %dst: memref<8xi32>) {
+ affine.for %i = 0 to 8 {
+ %a = affine.load %src[%i] : memref<10xi32>
+ %b = affine.load %src[%i + 1] : memref<10xi32>
+ %c = affine.load %src[%i + 2] : memref<10xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
+
+// A write to the producer source invalidates cross-iteration reuse.
+
+// CHECK-LABEL: func.func @source_modified
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @source_modified(%src0: memref<10xi32>, %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<10xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %a = affine.load %src[%i] : memref<10xi32>
+ %b = affine.load %src[%i + 1] : memref<10xi32>
+ %c = affine.load %src[%i + 2] : memref<10xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ affine.store %right, %src[%i + 1] : memref<10xi32>
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
+
+// Structurally different roots and a non-unit translation do not match.
+
+// CHECK-LABEL: func.func @different_operator
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @different_operator(%src0: memref<10xi32>,
+ %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<10xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %a = affine.load %src[%i] : memref<10xi32>
+ %b = affine.load %src[%i + 1] : memref<10xi32>
+ %c = affine.load %src[%i + 2] : memref<10xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.muli %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @distance_two
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @distance_two(%src0: memref<11xi32>, %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<11xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %a = affine.load %src[%i] : memref<11xi32>
+ %b = affine.load %src[%i + 1] : memref<11xi32>
+ %c = affine.load %src[%i + 3] : memref<11xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
+
+// Hoisting the initial producer out of a loop that may not execute would add
+// memory reads and arithmetic.
+
+// CHECK-LABEL: func.func @dynamic_trip
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @dynamic_trip(%src0: memref<?xi32>, %dst0: memref<?xi32>,
+ %n: index) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<?xi32>, memref<?xi32>
+ affine.for %i = 0 to %n {
+ %a = affine.load %src[%i] : memref<?xi32>
+ %b = affine.load %src[%i + 1] : memref<?xi32>
+ %c = affine.load %src[%i + 2] : memref<?xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<?xi32>
+ }
+ return
+}
+
+// A producer depending on another loop-carried value has a different context
+// key in the next iteration and must not be reused.
+
+// CHECK-LABEL: func.func @context_changes
+// CHECK: affine.for
+// CHECK-NOT: iter_args({{.*}},
+func.func @context_changes(%src0: memref<10xi32>, %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<10xi32>, memref<8xi32>
+ %zero = arith.constant 0 : i32
+ affine.for %i = 0 to 8 iter_args(%state = %zero) -> i32 {
+ %a = affine.load %src[%i] : memref<10xi32>
+ %b = affine.load %src[%i + 1] : memref<10xi32>
+ %left = arith.addi %a, %state : i32
+ %right = arith.addi %b, %state : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ %next = arith.addi %state, %difference : i32
+ affine.yield %next : i32
+ }
+ return
+}
+
+// Unknown recursive effects are rejected even when the candidate loads and
+// destination are otherwise distinct.
+
+// CHECK-LABEL: func.func @unknown_effect
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @unknown_effect(%src0: memref<10xi32>, %dst0: memref<8xi32>) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<10xi32>, memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %a = affine.load %src[%i] : memref<10xi32>
+ %b = affine.load %src[%i + 1] : memref<10xi32>
+ %c = affine.load %src[%i + 2] : memref<10xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ "test.unknown_effect"() : () -> ()
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<8xi32>
+ }
+ return
+}
>From 2e3fdad8a078f6860a06b167e4d735780d35600b Mon Sep 17 00:00:00 2001
From: takatodo <takatodo1227 at gmail.com>
Date: Sun, 2 Aug 2026 11:08:07 +0900
Subject: [PATCH 2/2] [mlir][affine] Check speculation before loop-carried
preloads
Reject candidates whose first-iteration preload would cross a non-speculatable operation. Keep source stability as a whole-loop requirement while limiting the scheduling check to the prefix actually crossed by the preload.
Assisted-by: OpenAI Codex
---
.../LoopCarriedComputationReuse.cpp | 61 +++++++++++++++----
...oop-carried-computation-reuse-corners.mlir | 57 +++++++++++++++++
2 files changed, 105 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopCarriedComputationReuse.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopCarriedComputationReuse.cpp
index 472cf071d7881..d9788dadb0891 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopCarriedComputationReuse.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopCarriedComputationReuse.cpp
@@ -235,12 +235,23 @@ class ShiftedDAGMatcher {
bool hasTranslatedLoad = false;
};
-static bool sourceIsStable(AffineForOp loop, Value source,
+static bool isSourceStable(AffineForOp loop, Value source,
AliasAnalysis &aliasAnalysis) {
- for (Operation &operation : loop.getBody()->without_terminator())
- if (aliasAnalysis.getModRef(&operation, source).isMod())
- return false;
- return true;
+ bool stable = true;
+ (void)loop.getBody()->walk<WalkOrder::PostOrder>([&](Operation *operation) {
+ // Recursive-effect operations derive their effects from nested operations
+ // unless they also expose direct effects. The post-order walk has already
+ // checked those nested operations.
+ if (operation->hasTrait<OpTrait::HasRecursiveMemoryEffects>() &&
+ !isa<MemoryEffectOpInterface>(operation))
+ return WalkResult::advance();
+ if (aliasAnalysis.getModRef(operation, source).isMod()) {
+ stable = false;
+ return WalkResult::interrupt();
+ }
+ return WalkResult::advance();
+ });
+ return stable;
}
/// Return whether the loop is proven to execute at least twice. Handle
@@ -261,13 +272,39 @@ static bool hasAtLeastTwoIterations(AffineForOp loop) {
return tripCount && !tripCount->isNegative() && tripCount->ugt(1);
}
+/// Return true only when the loop executes at least twice, every source is
+/// stable, and moving `prologueOps` before the loop does not cross a blocking
+/// operation in the first iteration.
+static bool isSafeToPreload(AffineForOp loop, ValueRange sources,
+ ArrayRef<Operation *> prologueOps,
+ AliasAnalysis &aliasAnalysis) {
+ if (loop.getLowerBoundMap().getNumResults() != 1 ||
+ !hasAtLeastTwoIterations(loop) || sources.empty() || prologueOps.empty())
+ return false;
+ if (llvm::any_of(sources, [&](Value source) {
+ return !isSourceStable(loop, source, aliasAnalysis);
+ }))
+ return false;
+
+ Operation *root = prologueOps.back();
+ if (!root || root->getParentOp() != loop)
+ return false;
+ llvm::SmallPtrSet<Operation *, 16> prologueSet(prologueOps.begin(),
+ prologueOps.end());
+ for (Operation &operation : loop.getBody()->without_terminator()) {
+ bool reachedRoot = &operation == root;
+ if (!prologueSet.contains(&operation) &&
+ !isa<AffineReadOpInterface, AffineWriteOpInterface>(operation) &&
+ !isPure(&operation))
+ return false;
+ if (reachedRoot)
+ return true;
+ }
+ return false;
+}
+
static std::optional<ReuseCandidate>
findReuseCandidate(AffineForOp loop, AliasAnalysis &aliasAnalysis) {
- if (loop.getLowerBoundMap().getNumResults() != 1)
- return std::nullopt;
- if (!hasAtLeastTwoIterations(loop))
- return std::nullopt;
-
SmallVector<Operation *> bodyOps;
for (Operation &operation : loop.getBody()->without_terminator())
bodyOps.push_back(&operation);
@@ -298,9 +335,7 @@ findReuseCandidate(AffineForOp loop, AliasAnalysis &aliasAnalysis) {
SmallVector<Operation *> earlierOps = matcher.takeEarlierOps();
SmallVector<Value> sources = matcher.takeSources();
if (earlierOps.empty() || sources.empty() ||
- llvm::any_of(sources, [&](Value source) {
- return !sourceIsStable(loop, source, aliasAnalysis);
- }))
+ !isSafeToPreload(loop, sources, earlierOps, aliasAnalysis))
continue;
llvm::SmallPtrSet<Operation *, 16> earlierSet(earlierOps.begin(),
diff --git a/mlir/test/Dialect/Affine/loop-carried-computation-reuse-corners.mlir b/mlir/test/Dialect/Affine/loop-carried-computation-reuse-corners.mlir
index 4096a92ef139c..4281fe2bb9cf1 100644
--- a/mlir/test/Dialect/Affine/loop-carried-computation-reuse-corners.mlir
+++ b/mlir/test/Dialect/Affine/loop-carried-computation-reuse-corners.mlir
@@ -272,3 +272,60 @@ func.func @loop_invariant_duplicates(%src0: memref<1xi32>,
}
return
}
+
+// A first-iteration preload must not cross a potentially non-terminating
+// operation. With a dynamic memref, the original program need not execute an
+// access at all when %continue is true.
+
+// CHECK-LABEL: func.func @non_speculatable_prefix
+// CHECK-NOT: affine.load
+// CHECK: affine.for
+// CHECK-NOT: iter_args
+func.func @non_speculatable_prefix(%src0: memref<?xi32>,
+ %dst0: memref<2xi32>, %continue: i1) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<?xi32>, memref<2xi32>
+ affine.for %i = 0 to 2 {
+ scf.while : () -> () {
+ scf.condition(%continue)
+ } do {
+ scf.yield
+ }
+ %a = affine.load %src[%i] : memref<?xi32>
+ %b = affine.load %src[%i + 1] : memref<?xi32>
+ %c = affine.load %src[%i + 2] : memref<?xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<2xi32>
+ }
+ return
+}
+
+// The same operation after both producer DAGs does not block a prologue that
+// only moves work which was already executed before it.
+
+// CHECK-LABEL: func.func @non_speculatable_after_producers
+// CHECK: %[[INIT:.*]] = arith.addi
+// CHECK: affine.for {{.*}} iter_args(%{{.*}} = %[[INIT]])
+func.func @non_speculatable_after_producers(%src0: memref<4xi32>,
+ %dst0: memref<2xi32>,
+ %continue: i1) {
+ %src, %dst = memref.distinct_objects %src0, %dst0
+ : memref<4xi32>, memref<2xi32>
+ affine.for %i = 0 to 2 {
+ %a = affine.load %src[%i] : memref<4xi32>
+ %b = affine.load %src[%i + 1] : memref<4xi32>
+ %c = affine.load %src[%i + 2] : memref<4xi32>
+ %left = arith.addi %a, %b : i32
+ %right = arith.addi %b, %c : i32
+ scf.while : () -> () {
+ scf.condition(%continue)
+ } do {
+ scf.yield
+ }
+ %difference = arith.subi %right, %left : i32
+ affine.store %difference, %dst[%i] : memref<2xi32>
+ }
+ return
+}
More information about the Mlir-commits
mailing list