[Mlir-commits] [mlir] [mlir][affine] Select shared-input reductions for unroll-and-jam (PR #226270)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 24 11:50:25 PDT 2026
https://github.com/ParkHanbum created https://github.com/llvm/llvm-project/pull/226270
Add the opt-in select-shared-input-reductions option to select perfect nests with constant bounds and floating-point multiply-add reductions where exactly one input access is invariant in the outer loop.
Match the reduction through SSA definitions. Require disjoint accumulator coordinates across outer iterations and proven NoAlias between the accumulator and both inputs. Support function arguments and views while allowing the read-only inputs to alias each other.
Require the outer trip count to match the unroll-jam factor. Preserve default pass behavior and leave unsupported nests unchanged.
Add tests for selection, aliasing, unsupported patterns, multiple candidates, and option handling.
Fixes #225396
>From 019c6d3eea48405fa68876b41d98503d71877b96 Mon Sep 17 00:00:00 2001
From: hanbeom <kese111 at gmail.com>
Date: Fri, 25 Sep 2026 03:48:22 +0900
Subject: [PATCH] [mlir][affine] Select shared-input reductions for
unroll-and-jam
Add the opt-in select-shared-input-reductions option to select perfect
nests with constant bounds and floating-point multiply-add reductions
where exactly one input access is invariant in the outer loop.
Match the reduction through SSA definitions. Require disjoint accumulator
coordinates across outer iterations and proven NoAlias between the
accumulator and both inputs. Support function arguments and views while
allowing the read-only inputs to alias each other.
Require the outer trip count to match the unroll-jam factor. Preserve
default pass behavior and leave unsupported nests unchanged.
Add tests for selection, aliasing, unsupported patterns, multiple
candidates, and option handling.
Fixes #225396
---
.../mlir/Dialect/Affine/Transforms/Passes.td | 18 +++
.../Dialect/Affine/Transforms/CMakeLists.txt | 1 +
.../Affine/Transforms/LoopUnrollAndJam.cpp | 124 +++++++++++++++++-
3 files changed, 141 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
index 7d7adfb1db235..3bd985845757e 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
@@ -220,11 +220,29 @@ def AffineLoopUnroll : InterfacePass<"affine-loop-unroll", "FunctionOpInterface"
def AffineLoopUnrollAndJam : InterfacePass<"affine-loop-unroll-jam", "FunctionOpInterface"> {
let summary = "Unroll and jam affine loops";
+ let description = [{
+ By default, unroll and jam the first operation in the function entry block
+ if it is an affine loop.
+
+ With `select-shared-input-reductions`, instead select perfect constant-bound
+ nests containing independent floating-point multiply-add reductions with
+ one shared input. Alias analysis must prove that the accumulator does not
+ alias either input; it need not be a local allocation. Without such proof,
+ function arguments are not assumed to be independent. An entire accumulator
+ coordinate must be the output-loop induction variable. The output loop must
+ start at zero, have unit step, and have a trip count equal to
+ `unroll-jam-factor`. Unsupported nests are left unchanged. This is a
+ structural selection policy, not a target profitability model; it does not
+ itself eliminate redundant loads.
+ }];
let constructor = "mlir::affine::createLoopUnrollAndJamPass()";
let options = [
Option<"unrollJamFactor", "unroll-jam-factor", "unsigned",
/*default=*/"4",
"Use this unroll jam factor for all loops (default 4)">,
+ Option<"selectSharedInputReductions", "select-shared-input-reductions",
+ "bool", /*default=*/"false",
+ "Select only independent reductions sharing an input">,
];
}
diff --git a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
index 9d912139810b2..5e8076eb92689 100644
--- a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
@@ -33,6 +33,7 @@ add_mlir_dialect_library(MLIRAffineTransforms
MLIRAffineDialect
MLIRAffineAnalysis
MLIRAffineUtils
+ MLIRAnalysis
MLIRArithDialect
MLIRIR
MLIRMemRefDialect
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp
index f97174065a475..8311b56c45d6c 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp
@@ -35,10 +35,13 @@
#include "mlir/Dialect/Affine/Transforms/Passes.h"
+#include "mlir/Analysis/AliasAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/LoopUtils.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/CommandLine.h"
#include <optional>
@@ -55,8 +58,106 @@ using namespace mlir;
using namespace mlir::affine;
namespace {
-/// Loop unroll jam pass. Currently, this just unroll jams the first
-/// outer loop in a Function.
+struct MultiplyAddReduction {
+ AffineStoreOp store;
+ AffineLoadOp accLoad;
+ AffineLoadOp input0;
+ AffineLoadOp input1;
+};
+
+/// Match a body containing only three loads, a multiply, an add and a store.
+static std::optional<MultiplyAddReduction>
+matchMultiplyAddReduction(AffineForOp inner) {
+ Block *body = inner.getBody();
+ auto stores = body->getOps<AffineStoreOp>();
+ if (!llvm::hasNItems(body->without_terminator(), 6) ||
+ !llvm::hasSingleElement(stores))
+ return std::nullopt;
+
+ AffineStoreOp store = *stores.begin();
+ auto add = store.getValue().getDefiningOp<arith::AddFOp>();
+ if (!add)
+ return std::nullopt;
+
+ auto accLoad = add.getLhs().getDefiningOp<AffineLoadOp>();
+ auto mul = add.getRhs().getDefiningOp<arith::MulFOp>();
+ if (!accLoad || !mul) {
+ accLoad = add.getRhs().getDefiningOp<AffineLoadOp>();
+ mul = add.getLhs().getDefiningOp<arith::MulFOp>();
+ }
+ if (!accLoad || !mul)
+ return std::nullopt;
+
+ auto input0 = mul.getLhs().getDefiningOp<AffineLoadOp>();
+ auto input1 = mul.getRhs().getDefiningOp<AffineLoadOp>();
+ if (!input0 || !input1 || input0 == input1 || input0 == accLoad ||
+ input1 == accLoad)
+ return std::nullopt;
+
+ // These six distinct operations must all belong to this body, leaving no
+ // room for additional operations, effects or subregions.
+ Operation *matchedOps[] = {store, add, mul, accLoad, input0, input1};
+ if (llvm::any_of(matchedOps,
+ [&](Operation *op) { return op->getBlock() != body; }))
+ return std::nullopt;
+
+ return MultiplyAddReduction{store, accLoad, input0, input1};
+}
+
+static bool isSharedInputReduction(AffineForOp outer, unsigned factor,
+ AliasAnalysis &aliasAnalysis) {
+ // Full output grouping only: no cleanup/remainder or dynamic bounds.
+ if (factor < 2 || !outer.hasConstantBounds() || outer.getStepAsInt() != 1 ||
+ outer.getConstantLowerBound() != 0 ||
+ outer.getConstantUpperBound() != factor || outer.getNumResults())
+ return false;
+
+ SmallVector<AffineForOp> loops;
+ getPerfectlyNestedLoops(loops, outer);
+ if (loops.size() < 2)
+ return false;
+
+ // The outer loop was checked above.
+ for (AffineForOp loop : llvm::drop_begin(loops)) {
+ if (!loop.hasConstantBounds() || loop.getNumResults())
+ return false;
+ }
+
+ auto reduction = matchMultiplyAddReduction(loops.back());
+ if (!reduction)
+ return false;
+ auto [store, accLoad, input0, input1] = *reduction;
+ if (!(MemRefAccess(accLoad) == MemRefAccess(store)))
+ return false;
+
+ // An entire output coordinate is exactly the outer IV, so different outer
+ // iterations access distinct index tuples. Valid memref layouts do not alias
+ // distinct in-bounds tuples, regardless of allocation site, shape or strides.
+ Value outerIV = outer.getInductionVar();
+ auto operands = store.getMapOperands();
+ bool hasOuterIVCoordinate =
+ llvm::any_of(store.getAffineMap().getResults(), [&](AffineExpr expr) {
+ auto dim = dyn_cast<AffineDimExpr>(expr);
+ return dim && operands[dim.getPosition()] == outerIV;
+ });
+ if (!hasOuterIVCoordinate)
+ return false;
+
+ // Select reductions with exactly one input access invariant in the outer
+ // loop.
+ bool input0Invariant = isInvariantAccess(input0, outer);
+ bool input1Invariant = isInvariantAccess(input1, outer);
+ if (input0Invariant == input1Invariant)
+ return false;
+
+ // Require NoAlias between the accumulator and each input. The read-only
+ // inputs may alias each other. Query the accessed memrefs to retain facts
+ // from memref.distinct_objects.
+ Value accumulator = store.getMemRef();
+ return aliasAnalysis.alias(accumulator, input0.getMemRef()).isNo() &&
+ aliasAnalysis.alias(accumulator, input1.getMemRef()).isNo();
+}
+
struct LoopUnrollAndJam
: public affine::impl::AffineLoopUnrollAndJamBase<LoopUnrollAndJam> {
explicit LoopUnrollAndJam(
@@ -80,6 +181,25 @@ void LoopUnrollAndJam::runOnOperation() {
if (getOperation().isExternal())
return;
+ if (selectSharedInputReductions) {
+ // This is a structural selection policy, not a profitability model. Keep
+ // the legacy path below unchanged when the pass-local option is disabled.
+ auto &aliasAnalysis = getAnalysis<AliasAnalysis>();
+ // Finish all alias queries before any candidate is rewritten.
+ SmallVector<AffineForOp> candidates;
+ getOperation()->walk<WalkOrder::PreOrder>([&](AffineForOp loop) {
+ if (isSharedInputReduction(loop, unrollJamFactor, aliasAnalysis)) {
+ candidates.push_back(loop);
+ return WalkResult::skip(); // Never retain overlapping loop handles.
+ }
+ return WalkResult::advance();
+ });
+ for (AffineForOp loop : candidates)
+ (void)loopUnrollJamByFactor(loop, unrollJamFactor);
+ // No fallback to the legacy path if no candidate was selected.
+ return;
+ }
+
// Currently, just the outermost loop from the first loop nest is
// unroll-and-jammed by this pass. However, runOnAffineForOp can be called on
// any for operation.
More information about the Mlir-commits
mailing list