[Mlir-commits] [mlir] [mlir][affine] Select shared-input reductions for unroll-and-jam (PR #226270)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 11 09:28:26 PDT 2026
https://github.com/ParkHanbum updated https://github.com/llvm/llvm-project/pull/226270
>From 8d5c600392c26b2306381d85a6e2e189600d18bc Mon Sep 17 00:00:00 2001
From: hanbeom <kese111 at gmail.com>
Date: Fri, 2 Oct 2026 00:00:08 +0900
Subject: [PATCH 1/2] [mlir][affine] Add opt-in unroll and jam for entry-block
loops
Add unroll-jam-all-loops to process all affine loops in the function entry
block, including loops after setup operations. The option is disabled by
default, preserving the existing first-operation selection.
---
.../mlir/Dialect/Affine/Transforms/Passes.td | 10 +
.../Affine/Transforms/LoopUnrollAndJam.cpp | 19 +-
.../Affine/unroll-jam-entry-block.mlir | 193 ++++++++++++++++++
3 files changed, 215 insertions(+), 7 deletions(-)
create mode 100644 mlir/test/Dialect/Affine/unroll-jam-entry-block.mlir
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
index 7d7adfb1db235..b5151ac8cb49c 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
@@ -220,11 +220,21 @@ 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 unroll-jam-all-loops enabled, select
+ every affine loop in the entry block, including loops preceded or separated
+ by non-loop operations. Use the requested unroll-jam-factor for each selected
+ loop. Nested loops are handled as part of their containing loop nest.
+ }];
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<"unrollJamAllLoops", "unroll-jam-all-loops", "bool",
+ /*default=*/"false",
+ "Unroll and jam all affine loops in the function entry block">,
];
}
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp
index f97174065a475..55ab41a4c8d87 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp
@@ -39,6 +39,7 @@
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/LoopUtils.h"
+#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/CommandLine.h"
#include <optional>
@@ -55,8 +56,7 @@ using namespace mlir;
using namespace mlir::affine;
namespace {
-/// Loop unroll jam pass. Currently, this just unroll jams the first
-/// outer loop in a Function.
+/// Unroll and jam the affine loops in the function entry block.
struct LoopUnrollAndJam
: public affine::impl::AffineLoopUnrollAndJamBase<LoopUnrollAndJam> {
explicit LoopUnrollAndJam(
@@ -80,10 +80,15 @@ void LoopUnrollAndJam::runOnOperation() {
if (getOperation().isExternal())
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.
- auto &entryBlock = getOperation().front();
- if (auto forOp = dyn_cast<AffineForOp>(entryBlock.front()))
+ // Collect entry-block loops before rewriting: the utility can erase a
+ // selected loop or create cleanup loops, which must not be revisited.
+ SmallVector<AffineForOp> loops;
+ for (Operation &op : getOperation().front()) {
+ if (auto forOp = dyn_cast<AffineForOp>(op))
+ loops.push_back(forOp);
+ if (!unrollJamAllLoops)
+ break;
+ }
+ for (auto forOp : loops)
(void)loopUnrollJamByFactor(forOp, unrollJamFactor);
}
diff --git a/mlir/test/Dialect/Affine/unroll-jam-entry-block.mlir b/mlir/test/Dialect/Affine/unroll-jam-entry-block.mlir
new file mode 100644
index 0000000000000..296b741080144
--- /dev/null
+++ b/mlir/test/Dialect/Affine/unroll-jam-entry-block.mlir
@@ -0,0 +1,193 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=2 unroll-jam-all-loops}))" | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=4 unroll-jam-all-loops=true}))" | FileCheck %s --check-prefix=FOUR
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=2}))" | FileCheck %s --check-prefix=DEFAULT
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=2 unroll-jam-all-loops=false}))" | FileCheck %s --check-prefix=DEFAULT
+
+// With the option enabled, setup operations must not prevent selection. No
+// shared-input condition is required: the load depends on both loop IVs.
+// DEFAULT-LABEL: func.func @nonloop_prefix
+// DEFAULT: memref.distinct_objects
+// DEFAULT-NEXT: affine.for {{.*}} = 0 to 4 {
+// DEFAULT-NEXT: affine.for {{.*}} = 0 to 8 {
+// DEFAULT: affine.store
+// DEFAULT-NOT: affine.store
+// DEFAULT: return
+// CHECK-LABEL: func.func @nonloop_prefix
+// CHECK: arith.constant 1 : i32
+// CHECK: memref.distinct_objects
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 step 2 {
+// CHECK-NEXT: affine.for {{.*}} = 0 to 8 {
+// CHECK-COUNT-2: affine.store
+// CHECK-NOT: affine.store
+// CHECK: return
+// FOUR-LABEL: func.func @nonloop_prefix
+// FOUR: memref.distinct_objects
+// FOUR-NEXT: affine.for {{.*}} = 0 to 8 {
+// FOUR-COUNT-4: affine.store
+// FOUR-NOT: affine.store
+// FOUR: return
+func.func @nonloop_prefix(%x: memref<4x8xi32>, %y: memref<4x8xi32>) {
+ %one = arith.constant 1 : i32
+ %input, %output = memref.distinct_objects %x, %y : memref<4x8xi32>, memref<4x8xi32>
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 8 {
+ %v = affine.load %input[%i, %j] : memref<4x8xi32>
+ %sum = arith.addi %v, %one : i32
+ affine.store %sum, %output[%i, %j] : memref<4x8xi32>
+ }
+ }
+ return
+}
+
+// The initialization loop must not hide the output loop containing two
+// reduction dimensions. Factor 2 partially jams the output loop; factor 4
+// removes it, retaining both reduction loops and each output's strict FP chain.
+// DEFAULT-LABEL: func.func @factor_two_reduction_dims
+// DEFAULT: memref.alloc() : memref<4xf32>
+// DEFAULT-NEXT: affine.for {{.*}} = 0 to 4 {
+// DEFAULT: affine.store
+// DEFAULT-NEXT: }
+// DEFAULT-NEXT: affine.for {{.*}} = 0 to 4 {
+// DEFAULT-NEXT: affine.for {{.*}} = 0 to 3 {
+// DEFAULT-NEXT: affine.for {{.*}} = 0 to 5 {
+// DEFAULT: arith.mulf
+// DEFAULT-NOT: arith.mulf
+// DEFAULT: return
+// CHECK-LABEL: func.func @factor_two_reduction_dims(
+// CHECK-SAME: [[X:%.*]]: memref<3x5xf32>, [[W:%.*]]: memref<4x3x5xf32>
+// CHECK: [[ACC:%.*]] = memref.alloc() : memref<4xf32>
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 step 2 {
+// CHECK-COUNT-2: affine.store
+// CHECK-NEXT: }
+// CHECK-NEXT: affine.for [[F:%.*]] = 0 to 4 step 2 {
+// CHECK-NEXT: affine.for [[C:%.*]] = 0 to 3 {
+// CHECK-NEXT: affine.for [[K:%.*]] = 0 to 5 {
+// CHECK-NEXT: [[X0:%.*]] = affine.load [[X]][[[C]], [[K]]] : memref<3x5xf32>
+// CHECK-NEXT: [[W0:%.*]] = affine.load [[W]][[[F]], [[C]], [[K]]] : memref<4x3x5xf32>
+// CHECK-NEXT: [[A0:%.*]] = affine.load [[ACC]][[[F]]] : memref<4xf32>
+// CHECK-NEXT: [[P0:%.*]] = arith.mulf [[X0]], [[W0]] : f32
+// CHECK-NEXT: [[S0:%.*]] = arith.addf [[A0]], [[P0]] : f32
+// CHECK-NEXT: affine.store [[S0]], [[ACC]][[[F]]] : memref<4xf32>
+// CHECK-NEXT: [[F1:%.*]] = affine.apply {{.*}}([[F]])
+// CHECK-NEXT: [[X1:%.*]] = affine.load [[X]][[[C]], [[K]]] : memref<3x5xf32>
+// CHECK-NEXT: [[W1:%.*]] = affine.load [[W]][[[F1]], [[C]], [[K]]] : memref<4x3x5xf32>
+// CHECK-NEXT: [[A1:%.*]] = affine.load [[ACC]][[[F1]]] : memref<4xf32>
+// CHECK-NEXT: [[P1:%.*]] = arith.mulf [[X1]], [[W1]] : f32
+// CHECK-NEXT: [[S1:%.*]] = arith.addf [[A1]], [[P1]] : f32
+// CHECK-NEXT: affine.store [[S1]], [[ACC]][[[F1]]] : memref<4xf32>
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: }
+// CHECK-NEXT: return [[ACC]] : memref<4xf32>
+// FOUR-LABEL: func.func @factor_two_reduction_dims
+// FOUR: memref.alloc() : memref<4xf32>
+// FOUR-COUNT-4: affine.store
+// FOUR-NEXT: affine.for {{.*}} = 0 to 3 {
+// FOUR-NEXT: affine.for {{.*}} = 0 to 5 {
+// FOUR-COUNT-4: arith.mulf
+// FOUR-NOT: affine.for
+// FOUR: return
+func.func @factor_two_reduction_dims(%x: memref<3x5xf32>, %w: memref<4x3x5xf32>) -> memref<4xf32> {
+ %zero = arith.constant 0.0 : f32
+ %acc = memref.alloc() : memref<4xf32>
+ affine.for %i = 0 to 4 {
+ affine.store %zero, %acc[%i] : memref<4xf32>
+ }
+ affine.for %f = 0 to 4 {
+ affine.for %c = 0 to 3 {
+ affine.for %k = 0 to 5 {
+ %xv = affine.load %x[%c, %k] : memref<3x5xf32>
+ %wv = affine.load %w[%f, %c, %k] : memref<4x3x5xf32>
+ %old = affine.load %acc[%f] : memref<4xf32>
+ %product = arith.mulf %xv, %wv : f32
+ %sum = arith.addf %old, %product : f32
+ affine.store %sum, %acc[%f] : memref<4xf32>
+ }
+ }
+ }
+ return %acc : memref<4xf32>
+}
+
+// Select both entry-block nests, including a loop after an intervening op.
+// Without the option, only the first nest is selected.
+// Nested loops must not be selected again after their outer loop is jammed.
+// DEFAULT-LABEL: func.func @separate_nests
+// DEFAULT: affine.for {{.*}} = 0 to 4 step 2 {
+// DEFAULT-NEXT: affine.for {{.*}} = 0 to 8 {
+// DEFAULT-COUNT-2: affine.store
+// DEFAULT: arith.addi
+// DEFAULT-NEXT: affine.for {{.*}} = 0 to 4 {
+// DEFAULT-NEXT: affine.for {{.*}} = 0 to 8 {
+// DEFAULT: affine.store
+// DEFAULT-NOT: affine.store
+// DEFAULT: return
+// CHECK-LABEL: func.func @separate_nests
+// CHECK: affine.for {{.*}} = 0 to 4 step 2 {
+// CHECK-NEXT: affine.for {{.*}} = 0 to 8 {
+// CHECK-COUNT-2: affine.store
+// CHECK: arith.addi
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 step 2 {
+// CHECK-NEXT: affine.for {{.*}} = 0 to 8 {
+// CHECK-COUNT-2: affine.store
+// CHECK-NOT: affine.store
+// CHECK: return
+// FOUR-LABEL: func.func @separate_nests
+// FOUR: affine.for {{.*}} = 0 to 8 {
+// FOUR-COUNT-4: affine.store
+// FOUR: arith.addi
+// FOUR-NEXT: affine.for {{.*}} = 0 to 8 {
+// FOUR-COUNT-4: affine.store
+// FOUR-NOT: affine.store
+// FOUR: return
+func.func @separate_nests(%a: memref<4x8xi32>, %b: memref<4x8xi32>, %value: i32) {
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 8 {
+ affine.store %value, %a[%i, %j] : memref<4x8xi32>
+ }
+ }
+ %next = arith.addi %value, %value : i32
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 8 {
+ affine.store %next, %b[%i, %j] : memref<4x8xi32>
+ }
+ }
+ return
+}
+
+// Promoting a single-iteration cleanup exposes another entry-block loop.
+// Neither that loop nor the main nest's inner loop should be revisited.
+// CHECK-LABEL: func.func @cleanup_not_revisited
+// CHECK: affine.for {{.*}} = 0 to 4 step 2 {
+// CHECK-NEXT: affine.for {{.*}} = 0 to 8 {
+// CHECK-COUNT-2: affine.store
+// CHECK: affine.for {{.*}} = 0 to 8 {
+// CHECK: affine.store
+// CHECK-NOT: affine.store
+// CHECK: return
+// FOUR-LABEL: func.func @cleanup_not_revisited
+// FOUR: affine.for {{.*}} = 0 to 8 {
+// FOUR-COUNT-4: affine.store
+// FOUR: affine.for {{.*}} = 0 to 8 {
+// FOUR: affine.store
+// FOUR-NOT: affine.store
+// FOUR: return
+func.func @cleanup_not_revisited(%a: memref<5x8xi32>, %value: i32) {
+ affine.for %i = 0 to 5 {
+ affine.for %j = 0 to 8 {
+ affine.store %value, %a[%i, %j] : memref<5x8xi32>
+ }
+ }
+ return
+}
+
+// DEFAULT-LABEL: func.func @no_loop
+// DEFAULT-NEXT: return
+// CHECK-LABEL: func.func @no_loop
+// CHECK-NEXT: return
+// FOUR-LABEL: func.func @no_loop
+// FOUR-NEXT: return
+func.func @no_loop() {
+ return
+}
+
+func.func private @external()
>From 12ce958bf5604ab6f3b75b054c2c022cb87ab032 Mon Sep 17 00:00:00 2001
From: hanbeom <kese111 at gmail.com>
Date: Sun, 11 Oct 2026 15:59:00 +0900
Subject: [PATCH 2/2] [mlir][affine] Check legality for opt-in loop unroll and
jam
Require proven iteration independence for additional loops selected by
unroll-jam-all-loops. Reuse isLoopParallel and require NoAlias for distinct
memrefs involving a write, conservatively skipping unsupported operations
and accesses. Preserve the original target's behavior.
---
.../mlir/Dialect/Affine/Transforms/Passes.td | 19 +-
.../Dialect/Affine/Transforms/CMakeLists.txt | 1 +
.../Affine/Transforms/LoopUnrollAndJam.cpp | 67 ++++-
.../Dialect/Affine/unroll-jam-legality.mlir | 275 ++++++++++++++++++
4 files changed, 349 insertions(+), 13 deletions(-)
create mode 100644 mlir/test/Dialect/Affine/unroll-jam-legality.mlir
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
index b5151ac8cb49c..0e4031853dc41 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
@@ -221,11 +221,18 @@ 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 unroll-jam-all-loops enabled, select
- every affine loop in the entry block, including loops preceded or separated
- by non-loop operations. Use the requested unroll-jam-factor for each selected
- loop. Nested loops are handled as part of their containing loop nest.
+ By default, select the first operation in the function entry block if it
+ is an affine loop, and unroll and jam it by `unroll-jam-factor`.
+
+ With `unroll-jam-all-loops` enabled, also consider additional top-level
+ affine loops in the entry block, including those following setup operations.
+ These additional candidates must have proven iteration independence. The
+ proof conservatively accepts affine loads/stores with invariant memref
+ bases and pure, region-free operations, and requires non-aliasing between
+ distinct memrefs involved in a write.
+
+ The original target uses the existing transformation path with either
+ option setting. Nested loops are handled within their containing loop nest.
}];
let constructor = "mlir::affine::createLoopUnrollAndJamPass()";
let options = [
@@ -234,7 +241,7 @@ def AffineLoopUnrollAndJam : InterfacePass<"affine-loop-unroll-jam", "FunctionOp
"Use this unroll jam factor for all loops (default 4)">,
Option<"unrollJamAllLoops", "unroll-jam-all-loops", "bool",
/*default=*/"false",
- "Unroll and jam all affine loops in the function entry block">,
+ "Consider additional top-level affine loops in the function entry block">,
];
}
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 55ab41a4c8d87..c02035af0eb53 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/Interfaces/SideEffectInterfaces.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/CommandLine.h"
#include <optional>
@@ -55,6 +58,50 @@ namespace affine {
using namespace mlir;
using namespace mlir::affine;
+/// Iteration independence is a sufficient condition for unroll-and-jam.
+/// Conservatively reject operations and aliases that the dependence analysis
+/// cannot model. The utility still checks the factor and loop structure.
+static bool isUnrollJamProvenSafe(AffineForOp loop,
+ AliasAnalysis &aliasAnalysis) {
+ SmallVector<Operation *> accesses;
+ auto walkResult = loop.walk([&](Operation *op) {
+ if (isa<AffineForOp, AffineYieldOp>(op))
+ return WalkResult::advance();
+ if (isa<AffineLoadOp, AffineStoreOp>(op)) {
+ // Access relations require a memref base fixed across iterations.
+ if (!loop.isDefinedOutsideOfLoop(MemRefAccess(op).memref))
+ return WalkResult::interrupt();
+ accesses.push_back(op);
+ return WalkResult::advance();
+ }
+ // Memory-effect freedom alone does not rule out nontermination or other
+ // implicit behavior. Initially accept only pure, region-free operations.
+ if (op->getNumRegions() == 0 && isPure(op))
+ return WalkResult::advance();
+ return WalkResult::interrupt();
+ });
+ if (walkResult.wasInterrupted())
+ return false;
+
+ // Affine dependence analysis assumes different memref SSA values do not
+ // alias. Establish that precondition for each pair involving a write.
+ for (Operation *src : accesses) {
+ for (Operation *dst : accesses) {
+ if (!isa<AffineStoreOp>(src) && !isa<AffineStoreOp>(dst))
+ continue;
+ Value srcMemref = MemRefAccess(src).memref;
+ Value dstMemref = MemRefAccess(dst).memref;
+ if (srcMemref != dstMemref &&
+ !aliasAnalysis.alias(srcMemref, dstMemref).isNo())
+ return false;
+ }
+ }
+
+ // Reject outer SSA recurrences and memory dependences. Inner-loop iter_args
+ // retain their per-iteration order when the containing loop is jammed.
+ return isLoopParallel(loop);
+}
+
namespace {
/// Unroll and jam the affine loops in the function entry block.
struct LoopUnrollAndJam
@@ -80,14 +127,20 @@ void LoopUnrollAndJam::runOnOperation() {
if (getOperation().isExternal())
return;
- // Collect entry-block loops before rewriting: the utility can erase a
- // selected loop or create cleanup loops, which must not be revisited.
+ // Keep the legacy target and prove legality for additional loops before
+ // rewriting. The utility can erase selected loops or create cleanup loops,
+ // which must not be revisited.
+ auto &entryBlock = getOperation().front();
SmallVector<AffineForOp> loops;
- for (Operation &op : getOperation().front()) {
- if (auto forOp = dyn_cast<AffineForOp>(op))
- loops.push_back(forOp);
- if (!unrollJamAllLoops)
- break;
+ if (auto forOp = dyn_cast<AffineForOp>(entryBlock.front()))
+ loops.push_back(forOp);
+
+ if (unrollJamAllLoops) {
+ for (Operation &op : llvm::drop_begin(entryBlock)) {
+ if (auto forOp = dyn_cast<AffineForOp>(op))
+ if (isUnrollJamProvenSafe(forOp, getAnalysis<AliasAnalysis>()))
+ loops.push_back(forOp);
+ }
}
for (auto forOp : loops)
(void)loopUnrollJamByFactor(forOp, unrollJamFactor);
diff --git a/mlir/test/Dialect/Affine/unroll-jam-legality.mlir b/mlir/test/Dialect/Affine/unroll-jam-legality.mlir
new file mode 100644
index 0000000000000..8beb92b277141
--- /dev/null
+++ b/mlir/test/Dialect/Affine/unroll-jam-legality.mlir
@@ -0,0 +1,275 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=2 unroll-jam-all-loops}))" | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=2}))" | FileCheck %s --check-prefix=LEGACY
+
+// CHECK-DAG: [[PLUS_ONE:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>
+
+// The first operation's loop keeps the legacy behavior, even when the option
+// is enabled. The new independence proof applies only to additional targets.
+// CHECK-LABEL: func.func @carried_memory_dependence
+// CHECK-NEXT: affine.for {{.*}} = 1 to 5 step 2 {
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK-COUNT-2: affine.store
+// CHECK: return
+// LEGACY-LABEL: func.func @carried_memory_dependence
+// LEGACY-NEXT: affine.for {{.*}} = 1 to 5 step 2 {
+// LEGACY-NEXT: affine.for {{.*}} = 0 to 4 {
+// LEGACY-COUNT-2: affine.store
+// LEGACY: return
+func.func @carried_memory_dependence(%a: memref<5x5xi32>) {
+ affine.for %i = 1 to 5 {
+ affine.for %j = 0 to 4 {
+ %v = affine.load %a[%i - 1, %j + 1] : memref<5x5xi32>
+ affine.store %v, %a[%i, %j] : memref<5x5xi32>
+ }
+ }
+ return
+}
+
+// The existing target is retained even when its distinct memref arguments
+// cannot be proven NoAlias. Additional targets must establish that proof.
+// CHECK-LABEL: func.func @may_alias
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 step 2 {
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK-COUNT-2: affine.store
+// CHECK: return
+// LEGACY-LABEL: func.func @may_alias
+// LEGACY-NEXT: affine.for {{.*}} = 0 to 4 step 2 {
+// LEGACY-COUNT-2: affine.store
+// LEGACY: return
+func.func @may_alias(%input: memref<5xi32>, %output: memref<4xi32>) {
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 4 {
+ %v = affine.load %input[%j + 1] : memref<5xi32>
+ affine.store %v, %output[%i] : memref<4xi32>
+ }
+ }
+ return
+}
+
+// The first nest is the legacy target. The second nest is an additional
+// target: jamming would move the next row's read before the previous row's
+// write at j + 1, so it must remain unchanged with the option enabled.
+// CHECK-LABEL: func.func @additional_carried_memory_dependence
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 step 2 {
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK-COUNT-2: affine.store
+// CHECK: affine.for {{.*}} = 1 to 5 {
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK: affine.store
+// CHECK-NOT: affine.store
+// CHECK: return
+// LEGACY-LABEL: func.func @additional_carried_memory_dependence
+// LEGACY-NEXT: affine.for {{.*}} = 0 to 4 step 2 {
+// LEGACY-NEXT: affine.for {{.*}} = 0 to 4 {
+// LEGACY-COUNT-2: affine.store
+// LEGACY: affine.for {{.*}} = 1 to 5 {
+// LEGACY-NEXT: affine.for {{.*}} = 0 to 4 {
+// LEGACY: affine.store
+// LEGACY-NOT: affine.store
+// LEGACY: return
+func.func @additional_carried_memory_dependence(%a: memref<5x5xi32>, %setup: memref<4x4xi32>, %value: i32) {
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 4 {
+ affine.store %value, %setup[%i, %j] : memref<4x4xi32>
+ }
+ }
+ affine.for %i = 1 to 5 {
+ affine.for %j = 0 to 4 {
+ %v = affine.load %a[%i - 1, %j + 1] : memref<5x5xi32>
+ affine.store %v, %a[%i, %j] : memref<5x5xi32>
+ }
+ }
+ return
+}
+
+// A loop after setup is also an additional target. Distinct SSA values do
+// not prove NoAlias, so this loop must remain unchanged when the option is on.
+// CHECK-LABEL: func.func @additional_may_alias
+// CHECK-NEXT: arith.constant 1 : i32
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK: affine.store
+// CHECK-NOT: affine.store
+// CHECK: return
+func.func @additional_may_alias(%input: memref<5xi32>, %output: memref<4xi32>) {
+ %one = arith.constant 1 : i32
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 4 {
+ %v = affine.load %input[%j + 1] : memref<5xi32>
+ %next = arith.addi %v, %one : i32
+ affine.store %next, %output[%i] : memref<4xi32>
+ }
+ }
+ return
+}
+
+// Views of one allocation can alias even when their SSA values differ.
+// CHECK-LABEL: func.func @aliasing_cast
+// CHECK: memref.cast
+// CHECK-NEXT: affine.for {{.*}} = 1 to 5 {
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK: affine.store
+// CHECK-NOT: affine.store
+// CHECK: return
+func.func @aliasing_cast(%a: memref<5x5xi32>) {
+ %view = memref.cast %a : memref<5x5xi32> to memref<?x?xi32>
+ affine.for %i = 1 to 5 {
+ affine.for %j = 0 to 4 {
+ %v = affine.load %a[%i - 1, %j + 1] : memref<5x5xi32>
+ affine.store %v, %view[%i, %j] : memref<?x?xi32>
+ }
+ }
+ return
+}
+
+// Outer iter_args carry an SSA dependence. In particular, do not reassociate
+// a strict floating-point recurrence across the iterations being jammed.
+// CHECK-LABEL: func.func @outer_iter_args
+// CHECK-NEXT: arith.constant {{.*}} : f32
+// CHECK-NEXT: {{.*}} = affine.for {{.*}} = 0 to 4 iter_args
+// CHECK: arith.addf
+// CHECK-NOT: arith.addf
+// CHECK: return
+func.func @outer_iter_args(%a: memref<4x4xf32>, %delta: f32) -> f32 {
+ %initial = arith.constant 1.0 : f32
+ %r = affine.for %i = 0 to 4 iter_args(%acc = %initial) -> f32 {
+ affine.for %j = 0 to 4 {
+ affine.store %acc, %a[%i, %j] : memref<4x4xf32>
+ }
+ %next = arith.addf %acc, %delta : f32
+ affine.yield %next : f32
+ }
+ return %r : f32
+}
+
+// Calls with unknown effects prevent the independence proof.
+// CHECK-LABEL: func.func @unknown_effects
+// CHECK-NEXT: arith.constant 1 : i32
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK: call @effect
+// CHECK-NOT: call @effect
+// CHECK: return
+func.func @unknown_effects(%a: memref<4x4xi32>, %value: i32) {
+ %one = arith.constant 1 : i32
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 4 {
+ func.call @effect(%i, %j) : (index, index) -> ()
+ %next = arith.addi %value, %one : i32
+ affine.store %next, %a[%i, %j] : memref<4x4xi32>
+ }
+ }
+ return
+}
+func.func private @effect(index, index)
+
+// Memory-effect freedom alone is insufficient for arbitrary operations.
+// Keep this first check conservative for non-speculatable computations.
+// CHECK-LABEL: func.func @non_speculatable
+// CHECK-NEXT: arith.constant 1 : i32
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK: arith.divsi
+// CHECK-NOT: arith.divsi
+// CHECK: return
+func.func @non_speculatable(%a: memref<4x4xi32>, %value: i32, %divisor: i32) {
+ %one = arith.constant 1 : i32
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 4 {
+ %v = arith.divsi %value, %divisor : i32
+ %next = arith.addi %v, %one : i32
+ affine.store %next, %a[%i, %j] : memref<4x4xi32>
+ }
+ }
+ return
+}
+
+// Do not infer execution properties for regions other than affine loops,
+// even if all their nested operations are memory-effect-free.
+// CHECK-LABEL: func.func @unsupported_region
+// CHECK-NEXT: arith.constant 1 : i32
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK: scf.execute_region
+// CHECK-NOT: scf.execute_region
+// CHECK: return
+func.func @unsupported_region(%a: memref<4x4xi32>, %value: i32) {
+ %one = arith.constant 1 : i32
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 4 {
+ %v = scf.execute_region -> i32 {
+ %sum = arith.addi %value, %one : i32
+ scf.yield %sum : i32
+ }
+ affine.store %v, %a[%i, %j] : memref<4x4xi32>
+ }
+ }
+ return
+}
+
+// The varying view offset cancels i: every iteration accesses row 4. Modeling
+// the view as a fixed base would incorrectly infer independence from its i index.
+// CHECK-LABEL: func.func @varying_memref_base
+// CHECK-NEXT: arith.constant 4 : index
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK: memref.subview
+// CHECK: affine.store
+// CHECK-NOT: affine.store
+// CHECK: return
+func.func @varying_memref_base(%a: memref<8x4xi32>) {
+ %four = arith.constant 4 : index
+ affine.for %i = 0 to 4 {
+ %offset = affine.apply affine_map<(d0)[s0] -> (s0 - d0)>(%i)[%four]
+ %view = memref.subview %a[%offset, 0] [4, 4] [1, 1] : memref<8x4xi32> to memref<4x4xi32, strided<[4, 1], offset: ?>>
+ affine.for %j = 0 to 3 {
+ %v = affine.load %view[%i, %j + 1] : memref<4x4xi32, strided<[4, 1], offset: ?>>
+ affine.store %v, %view[%i, %j] : memref<4x4xi32, strided<[4, 1], offset: ?>>
+ }
+ }
+ return
+}
+
+// A semi-affine access outside the dependence analysis's supported relations
+// must be skipped when the analysis fails to prove independence.
+// CHECK-LABEL: func.func @unanalyzable_access
+// CHECK-NEXT: arith.constant 1 : i32
+// CHECK-NEXT: affine.for {{.*}} = 0 to 4 {
+// CHECK: affine.store
+// CHECK-NOT: affine.store
+// CHECK: return
+func.func @unanalyzable_access(%a: memref<?x4xi32>, %value: i32, %stride: index) {
+ %one = arith.constant 1 : i32
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 4 {
+ %next = arith.addi %value, %one : i32
+ affine.store %next, %a[%i * symbol(%stride), %j] : memref<?x4xi32>
+ }
+ }
+ return
+}
+
+// Inner iter_args do not carry a dependence across the selected outer loop.
+// Keep each output's nonzero initial value and strict FP chain separate.
+// CHECK-LABEL: func.func @inner_iter_args(
+// CHECK-SAME: [[X:%.*]]: memref<5xf32>, [[A:%.*]]: memref<4xf32>, [[INIT:%.*]]: f32
+// CHECK: memref.distinct_objects
+// CHECK-NEXT: affine.for [[I:%.*]] = 0 to 4 step 2 {
+// CHECK-NEXT: [[R:%.*]]:2 = affine.for {{.*}} = 0 to 5 iter_args([[ACC0:%.*]] = [[INIT]], [[ACC1:%.*]] = [[INIT]]) -> (f32, f32) {
+// CHECK: [[SUM0:%.*]] = arith.addf [[ACC0]], {{.*}} : f32
+// CHECK: [[SUM1:%.*]] = arith.addf [[ACC1]], {{.*}} : f32
+// CHECK-NEXT: affine.yield [[SUM0]], [[SUM1]] : f32, f32
+// CHECK-NEXT: }
+// CHECK-NEXT: affine.store [[R]]#0, [[OUT:%.*]][[[I]]] : memref<4xf32>
+// CHECK-NEXT: [[I1:%.*]] = affine.apply [[PLUS_ONE]]([[I]])
+// CHECK-NEXT: affine.store [[R]]#1, [[OUT]][[[I1]]] : memref<4xf32>
+// CHECK-NOT: arith.addf
+// CHECK: return
+func.func @inner_iter_args(%x: memref<5xf32>, %a: memref<4xf32>, %initial: f32) {
+ %input, %output = memref.distinct_objects %x, %a : memref<5xf32>, memref<4xf32>
+ affine.for %i = 0 to 4 {
+ %r = affine.for %k = 0 to 5 iter_args(%acc = %initial) -> f32 {
+ %v = affine.load %input[%k] : memref<5xf32>
+ %next = arith.addf %acc, %v : f32
+ affine.yield %next : f32
+ }
+ affine.store %r, %output[%i] : memref<4xf32>
+ }
+ return
+}
More information about the Mlir-commits
mailing list