[Mlir-commits] [mlir] [MLIR][Affine] Fix result-bearing unroll-and-jam legality (PR #236076)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 27 04:34:36 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-affine
Author: Omnishambles (HowieYin)
<details>
<summary>Changes</summary>
Fixes #<!-- -->235800.
## Summary
`loopUnrollJamByFactor` currently duplicates every target-loop `iter_args`
initializer into each unrolled lane and then combines supported reduction
lanes. For a factor of `F`, this applies a non-identity initializer `F` times
instead of once.
There is a related legality gap for result-bearing loops: if a target
`iter_arg` is not recognized as a supported reduction, its additional lane
state is created but never combined. For example, an integer `subi` recurrence
with `init = 10` and `trip_count = 4` returns `8` after factor-2 unroll-and-jam
instead of `6`.
This change:
- keeps the original initializer on lane 0 and initializes additional
supported reduction lanes with the corresponding reduction identity;
- requires every target-loop `iter_arg` to be a supported reduction before
mutating the IR;
- initially limits result-bearing targets to scalar integer reductions without
`nsw`/`nuw`, conservatively rejecting floating-point, `index`, shaped, and
flagged integer reductions;
- preserves the existing mapped initializer behavior for nested loops being
jammed;
- validates identity attributes and cleanup-bound construction before mutating
the IR, then materializes identity constants only after cleanup generation
succeeds.
## Reproducer
For:
```mlir
func.func @<!-- -->sum(%init: i32) -> i32 {
%result = affine.for %i = 0 to 2
iter_args(%acc = %init) -> (i32) {
%one = arith.constant 1 : i32
%next = arith.addi %acc, %one : i32
affine.yield %next : i32
}
return %result : i32
}
```
factor-2 unroll-and-jam currently produces the equivalent of:
```text
(init + 1) + (init + 1)
```
With `init = 5`, the baseline result is `7`, while the transformed result is
`12`.
## Why
The target loop is extended with one `iter_args` group per unrolled lane.
Today, each group receives the original initializer. The lane results are then
combined with `arith::getReductionOp`, so the initializer is counted once per
lane:
```text
current: (I op C0) op (I op C1) ... op (I op C(F-1))
expected: I op C0 op C1 ... op C(F-1)
```
The correct split is:
```text
lane 0: original initializer
lane 1..F-1: reduction identity
```
Only the selected target loop uses identities. Nested loops represent
independent computations for different outer lanes and keep their mapped
initializers.
## Failure behavior
All new rejection conditions are checked before mutating the IR. Reduction
identities are first computed as `TypedAttr`s, and cleanup bounds are checked
for affine expressibility. Identity constants are created only after cleanup
generation succeeds.
For unsupported and mixed target `iter_args`, the transformation now fails
before mutation and leaves the IR unchanged instead of dropping lane state.
## Test
- Added focused structural coverage in
`mlir/test/Dialect/Affine/unroll-jam-reductions.mlir` for:
- factor 2 and factor 4;
- `addi` identity `0` and `muli` identity `1`;
- multiple reduction `iter_args`;
- selected target reduction with a nested jammed loop;
- final loop operands retaining identities after operand fix-up;
- rejection of mixed/unsupported `iter_args`, floating-point, `index`,
shaped types, and `nsw`/`nuw`;
- failure atomicity for an inexpressible cleanup bound.
- Added `mlir/test/Integration/Dialect/Affine/CPU/unroll-jam-reduction.mlir`.
Before the fix, the dynamic trip-count matrix returned total error `20`.
After the fix, trip counts `0`, `1`, `F-1`, `F`, and `F+1` all satisfy
`result = init + trip_count`, with total error `0`.
- `ninja -C build-mlir check-mlir-dialect-affine
check-mlir-integration-dialect-affine-cpu`
- Affine dialect: 73 passed.
- Affine CPU integration: 1 passed.
- `ninja -C build-mlir check-mlir`
- 4,438 passed.
- 404 unsupported.
- 1 skipped.
- 1 expected failure.
- 0 unexpected failures.
## Related work
#<!-- -->196554 changes unroll-and-jam dependence legality and touches the same utility
and test file, so this patch may require a rebase. #<!-- -->226270 changes pass-level
loop selection but does not change the utility's result-handling semantics.
---
Patch is 24.49 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/236076.diff
5 Files Affected:
- (modified) mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp (+69-18)
- (added) mlir/test/Dialect/Affine/unroll-jam-reductions.mlir (+204)
- (modified) mlir/test/Dialect/Affine/unroll-jam.mlir (+10-22)
- (modified) mlir/test/Dialect/SCF/transform-ops.mlir (+19-18)
- (added) mlir/test/Integration/Dialect/Affine/CPU/unroll-jam-reduction.mlir (+46)
``````````diff
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index b1efda47945304..11b2053f3f7a7b 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -12,10 +12,12 @@
#include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/Utils.h"
#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
@@ -1155,6 +1157,54 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
if (!areInnerBoundsInvariant(forOp))
return failure();
+ // Result-bearing target loops are only supported when all iter_args are
+ // unflagged scalar integer reductions. Compute their identities without
+ // mutating the IR.
+ SmallVector<LoopReduction> reductions;
+ SmallVector<TypedAttr> reductionIdentityAttrs(forOp.getNumIterOperands());
+ if (forOp.getNumIterOperands() > 0) {
+ getSupportedReductions(forOp, reductions);
+ if (reductions.size() != forOp.getNumIterOperands())
+ return failure();
+
+ OpBuilder attrBuilder(forOp.getContext());
+ for (const LoopReduction &reduction : reductions) {
+ unsigned pos = reduction.iterArgPosition;
+ Type resultType = forOp.getResult(pos).getType();
+ if (!isa<IntegerType>(resultType))
+ return failure();
+
+ Operation *combinerOp =
+ *forOp.getRegionIterArgs()[pos].getUsers().begin();
+ if (auto overflowOp =
+ dyn_cast<arith::ArithIntegerOverflowFlagsInterface>(combinerOp);
+ overflowOp && overflowOp.getOverflowAttr().getValue() !=
+ arith::IntegerOverflowFlags::none)
+ return failure();
+
+ TypedAttr identity = arith::getIdentityValueAttr(
+ reduction.kind, resultType, attrBuilder, forOp.getLoc());
+ if (!identity)
+ return failure();
+ reductionIdentityAttrs[pos] = identity;
+ }
+ }
+
+ bool needsCleanup =
+ getLargestDivisorOfTripCount(forOp) % unrollJamFactor != 0;
+ if (needsCleanup) {
+ // Multi-result bounds cannot express the cleanup loop lower bound.
+ if (forOp.getLowerBoundMap().getNumResults() != 1 ||
+ forOp.getUpperBoundMap().getNumResults() != 1)
+ return failure();
+
+ AffineMap tripCountMap;
+ SmallVector<Value> tripCountOperands;
+ getTripCountMapAndOperands(forOp, &tripCountMap, &tripCountOperands);
+ if (!tripCountMap)
+ return failure();
+ }
+
// Gather all sub-blocks to jam upon the loop being unrolled.
JamBlockGatherer<AffineForOp> jbg;
jbg.walk(forOp);
@@ -1167,21 +1217,9 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
loopsWithIterArgs.push_back(aForOp);
});
- // Get supported reductions to be used for creating reduction ops at the end.
- SmallVector<LoopReduction> reductions;
- if (forOp.getNumIterOperands() > 0)
- getSupportedReductions(forOp, reductions);
-
// Generate the cleanup loop if trip count isn't a multiple of
// unrollJamFactor.
- if (getLargestDivisorOfTripCount(forOp) % unrollJamFactor != 0) {
- // Loops where the lower bound is a max expression or the upper bound is
- // a min expression and the trip count doesn't divide the unroll factor
- // can't be unrolled since the lower bound of the cleanup loop in such cases
- // cannot be expressed as an affine function or a max over affine functions.
- if (forOp.getLowerBoundMap().getNumResults() != 1 ||
- forOp.getUpperBoundMap().getNumResults() != 1)
- return failure();
+ if (needsCleanup) {
if (failed(generateCleanupLoopForUnroll(forOp, unrollJamFactor)))
assert(false && "cleanup loop lower bound map for single result lower "
"and upper bound maps can always be determined");
@@ -1196,6 +1234,15 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
// operands.
SmallVector<AffineForOp, 4> newLoopsWithIterArgs;
IRRewriter rewriter(forOp.getContext());
+ SmallVector<Value> reductionIdentities(reductionIdentityAttrs.size());
+ if (!reductionIdentityAttrs.empty()) {
+ rewriter.setInsertionPoint(forOp);
+ for (const LoopReduction &reduction : reductions) {
+ unsigned pos = reduction.iterArgPosition;
+ reductionIdentities[pos] = arith::ConstantOp::create(
+ rewriter, forOp.getLoc(), reductionIdentityAttrs[pos]);
+ }
+ }
for (AffineForOp oldForOp : loopsWithIterArgs) {
SmallVector<Value> dupIterOperands, dupYieldOperands;
ValueRange oldIterOperands = oldForOp.getInits();
@@ -1204,13 +1251,16 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
cast<AffineYieldOp>(oldForOp.getBody()->getTerminator()).getOperands();
// Get additional iterOperands, iterArgs, and yield operands. We will
// fix iterOperands and yield operands after cloning of sub-blocks.
+ bool forOpReplaced = oldForOp == forOp;
for (unsigned i = unrollJamFactor - 1; i >= 1; --i) {
- dupIterOperands.append(oldIterOperands.begin(), oldIterOperands.end());
+ if (forOpReplaced)
+ dupIterOperands.append(reductionIdentities);
+ else
+ dupIterOperands.append(oldIterOperands.begin(), oldIterOperands.end());
dupYieldOperands.append(oldYieldOperands.begin(), oldYieldOperands.end());
}
// Create a new loop with additional iterOperands, iter_args and yield
// operands. This new loop will take the loop body of the original loop.
- bool forOpReplaced = oldForOp == forOp;
AffineForOp newForOp =
cast<AffineForOp>(*oldForOp.replaceWithAdditionalYields(
rewriter, dupIterOperands, /*replaceInitOperandUsesInLoop=*/false,
@@ -1280,9 +1330,10 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
// The `i`th duplication of an old iterOperand or yield op operand
// needs to be replaced with a mapped value from `operandMaps[i - 1]`
// if such mapped value exists.
- newForOp.setOperand(numControlOperands + i * oldNumIterOperands + j,
- operandMaps[i - 1].lookupOrDefault(
- newForOp.getOperand(numControlOperands + j)));
+ if (newForOp != forOp)
+ newForOp.setOperand(numControlOperands + i * oldNumIterOperands + j,
+ operandMaps[i - 1].lookupOrDefault(
+ newForOp.getOperand(numControlOperands + j)));
yieldOp.setOperand(
i * oldNumYieldOperands + j,
operandMaps[i - 1].lookupOrDefault(yieldOp.getOperand(j)));
diff --git a/mlir/test/Dialect/Affine/unroll-jam-reductions.mlir b/mlir/test/Dialect/Affine/unroll-jam-reductions.mlir
new file mode 100644
index 00000000000000..23259acefc549b
--- /dev/null
+++ b/mlir/test/Dialect/Affine/unroll-jam-reductions.mlir
@@ -0,0 +1,204 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=2}))" | FileCheck %s --check-prefixes=CHECK,FACTOR2
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=4}))" | FileCheck %s --check-prefixes=CHECK,FACTOR4
+
+// FACTOR2-LABEL: func.func @addi_reduction
+// FACTOR2-SAME: %[[INIT:.*]]: i32
+// FACTOR2: %[[ZERO:.*]] = arith.constant 0 : i32
+// FACTOR2: %[[RESULT:.*]]:2 = affine.for {{.*}} step 2 iter_args(%[[ACC0:.*]] = %[[INIT]], %[[ACC1:.*]] = %[[ZERO]]) -> (i32, i32) {
+// FACTOR2: %[[ADD0:.*]] = arith.addi %[[ACC0]],
+// FACTOR2: %[[ADD1:.*]] = arith.addi %[[ACC1]],
+// FACTOR2: affine.yield %[[ADD0]], %[[ADD1]]
+// FACTOR2: }
+// FACTOR2: %[[COMBINED:.*]] = arith.addi %[[RESULT]]#0, %[[RESULT]]#1
+// FACTOR2: return %[[COMBINED]]
+
+// FACTOR4-LABEL: func.func @addi_reduction
+// FACTOR4-SAME: %[[INIT:.*]]: i32
+// FACTOR4: %[[ZERO:.*]] = arith.constant 0 : i32
+// FACTOR4: %[[RESULT:.*]]:4 = affine.for {{.*}} step 4 iter_args(%{{.*}} = %[[INIT]], %{{.*}} = %[[ZERO]], %{{.*}} = %[[ZERO]], %{{.*}} = %[[ZERO]]) -> (i32, i32, i32, i32) {
+// FACTOR4: }
+// FACTOR4: %[[COMBINED0:.*]] = arith.addi %[[RESULT]]#0, %[[RESULT]]#3
+// FACTOR4: %[[COMBINED1:.*]] = arith.addi %[[COMBINED0]], %[[RESULT]]#2
+// FACTOR4: %[[COMBINED2:.*]] = arith.addi %[[COMBINED1]], %[[RESULT]]#1
+// FACTOR4: return %[[COMBINED2]]
+
+func.func @addi_reduction(%init: i32) -> i32 {
+ %result = affine.for %i = 0 to 8 iter_args(%acc = %init) -> (i32) {
+ %one = arith.constant 1 : i32
+ %next = arith.addi %acc, %one : i32
+ affine.yield %next : i32
+ }
+ return %result : i32
+}
+
+// FACTOR2-LABEL: func.func @muli_reduction
+// FACTOR2-SAME: %[[INIT:.*]]: i32
+// FACTOR2: %[[ONE:.*]] = arith.constant 1 : i32
+// FACTOR2: affine.for {{.*}} step 2 iter_args(%{{.*}} = %[[INIT]], %{{.*}} = %[[ONE]])
+
+func.func @muli_reduction(%init: i32) -> i32 {
+ %result = affine.for %i = 0 to 8 iter_args(%acc = %init) -> (i32) {
+ %two = arith.constant 2 : i32
+ %next = arith.muli %acc, %two : i32
+ affine.yield %next : i32
+ }
+ return %result : i32
+}
+
+// FACTOR2-LABEL: func.func @multiple_reductions
+// FACTOR2-SAME: %[[ADD_INIT:.*]]: i32, %[[MUL_INIT:.*]]: i32
+// FACTOR2: %[[ZERO:.*]] = arith.constant 0 : i32
+// FACTOR2: %[[ONE:.*]] = arith.constant 1 : i32
+// FACTOR2: affine.for {{.*}} step 2 iter_args(%{{.*}} = %[[ADD_INIT]], %{{.*}} = %[[MUL_INIT]], %{{.*}} = %[[ZERO]], %{{.*}} = %[[ONE]])
+
+func.func @multiple_reductions(%add_init: i32, %mul_init: i32) -> (i32, i32) {
+ %result:2 = affine.for %i = 0 to 8
+ iter_args(%add_acc = %add_init, %mul_acc = %mul_init) -> (i32, i32) {
+ %one = arith.constant 1 : i32
+ %two = arith.constant 2 : i32
+ %next_add = arith.addi %add_acc, %one : i32
+ %next_mul = arith.muli %mul_acc, %two : i32
+ affine.yield %next_add, %next_mul : i32, i32
+ }
+ return %result#0, %result#1 : i32, i32
+}
+
+// The target loop uses an identity for its extra reduction lane. The nested
+// loop still uses the mapped initializer for each independent outer lane.
+// FACTOR2-LABEL: func.func @nested_iter_args
+// FACTOR2-SAME: %[[INIT:.*]]: i32, %[[INNER_INIT:.*]]: i32
+// FACTOR2: %[[ZERO:.*]] = arith.constant 0 : i32
+// FACTOR2: affine.for {{.*}} step 2 iter_args(%[[OUTER0:.*]] = %[[INIT]], %[[OUTER1:.*]] = %[[ZERO]]) -> (i32, i32) {
+// FACTOR2: affine.for {{.*}} iter_args(%{{.*}} = %[[INNER_INIT]], %{{.*}} = %[[INNER_INIT]])
+
+func.func @nested_iter_args(%init: i32, %inner_init: i32) -> i32 {
+ %result = affine.for %i = 0 to 8 iter_args(%acc = %init) -> (i32) {
+ %inner = affine.for %j = 0 to 4
+ iter_args(%inner_acc = %inner_init) -> (i32) {
+ %one = arith.constant 1 : i32
+ %next_inner = arith.addi %inner_acc, %one : i32
+ affine.yield %next_inner : i32
+ }
+ %next = arith.addi %acc, %inner : i32
+ affine.yield %next : i32
+ }
+ return %result : i32
+}
+
+// CHECK-LABEL: func.func @unsupported_subi
+// CHECK: affine.for {{.*}} = 0 to 8 iter_args({{.*}}) -> (i32) {
+// CHECK-NOT: step 2
+// CHECK-NOT: step 4
+// CHECK: arith.subi
+// CHECK: }
+
+func.func @unsupported_subi(%init: i32) -> i32 {
+ %result = affine.for %i = 0 to 8 iter_args(%acc = %init) -> (i32) {
+ %one = arith.constant 1 : i32
+ %next = arith.subi %acc, %one : i32
+ affine.yield %next : i32
+ }
+ return %result : i32
+}
+
+// CHECK-LABEL: func.func @mixed_reductions
+// CHECK: affine.for {{.*}} = 0 to 8 iter_args({{.*}}) -> (i32, i32) {
+// CHECK-NOT: step 2
+// CHECK-NOT: step 4
+// CHECK: arith.addi
+// CHECK: arith.subi
+// CHECK: }
+
+func.func @mixed_reductions(%add_init: i32, %sub_init: i32) -> (i32, i32) {
+ %result:2 = affine.for %i = 0 to 8
+ iter_args(%add_acc = %add_init, %sub_acc = %sub_init) -> (i32, i32) {
+ %one = arith.constant 1 : i32
+ %next_add = arith.addi %add_acc, %one : i32
+ %next_sub = arith.subi %sub_acc, %one : i32
+ affine.yield %next_add, %next_sub : i32, i32
+ }
+ return %result#0, %result#1 : i32, i32
+}
+
+// CHECK-LABEL: func.func @float_reduction
+// CHECK: affine.for {{.*}} = 0 to 8 iter_args({{.*}}) -> (f32) {
+// CHECK-NOT: step 2
+// CHECK-NOT: step 4
+// CHECK: arith.addf
+// CHECK: }
+
+func.func @float_reduction(%init: f32) -> f32 {
+ %result = affine.for %i = 0 to 8 iter_args(%acc = %init) -> (f32) {
+ %one = arith.constant 1.0 : f32
+ %next = arith.addf %acc, %one : f32
+ affine.yield %next : f32
+ }
+ return %result : f32
+}
+
+// CHECK-LABEL: func.func @index_reduction
+// CHECK: affine.for {{.*}} = 0 to 8 iter_args({{.*}}) -> (index) {
+// CHECK-NOT: step 2
+// CHECK-NOT: step 4
+// CHECK: arith.addi
+// CHECK: }
+
+func.func @index_reduction(%init: index) -> index {
+ %result = affine.for %i = 0 to 8 iter_args(%acc = %init) -> (index) {
+ %one = arith.constant 1 : index
+ %next = arith.addi %acc, %one : index
+ affine.yield %next : index
+ }
+ return %result : index
+}
+
+// CHECK-LABEL: func.func @vector_reduction
+// CHECK: affine.for {{.*}} = 0 to 8 iter_args({{.*}}) -> (vector<2xi32>) {
+// CHECK-NOT: step 2
+// CHECK-NOT: step 4
+// CHECK: arith.addi
+// CHECK: }
+
+func.func @vector_reduction(%init: vector<2xi32>, %value: vector<2xi32>)
+ -> vector<2xi32> {
+ %result = affine.for %i = 0 to 8
+ iter_args(%acc = %init) -> (vector<2xi32>) {
+ %next = arith.addi %acc, %value : vector<2xi32>
+ affine.yield %next : vector<2xi32>
+ }
+ return %result : vector<2xi32>
+}
+
+// CHECK-LABEL: func.func @flagged_reduction
+// CHECK: affine.for {{.*}} = 0 to 8 iter_args({{.*}}) -> (i32) {
+// CHECK-NOT: step 2
+// CHECK-NOT: step 4
+// CHECK: arith.addi {{.*}} overflow<nsw>
+// CHECK: }
+
+func.func @flagged_reduction(%init: i32) -> i32 {
+ %result = affine.for %i = 0 to 8 iter_args(%acc = %init) -> (i32) {
+ %one = arith.constant 1 : i32
+ %next = arith.addi %acc, %one overflow<nsw> : i32
+ affine.yield %next : i32
+ }
+ return %result : i32
+}
+
+// CHECK-LABEL: func.func @inexpressible_cleanup
+// CHECK-NOT: arith.constant 0 : i32
+// CHECK: affine.for {{.*}} = max
+// CHECK-NOT: step 2
+// CHECK-NOT: step 4
+// CHECK: return
+
+func.func @inexpressible_cleanup(%init: i32, %lb0: index, %lb1: index,
+ %ub: index) -> i32 {
+ %result = affine.for %i = max affine_map<()[s0, s1] -> (s0, s1)>()
+ [%lb0, %lb1] to %ub iter_args(%acc = %init) -> (i32) {
+ %one = arith.constant 1 : i32
+ %next = arith.addi %acc, %one : i32
+ affine.yield %next : i32
+ }
+ return %result : i32
+}
diff --git a/mlir/test/Dialect/Affine/unroll-jam.mlir b/mlir/test/Dialect/Affine/unroll-jam.mlir
index 36054ba7852a77..4c7945ef6cc996 100644
--- a/mlir/test/Dialect/Affine/unroll-jam.mlir
+++ b/mlir/test/Dialect/Affine/unroll-jam.mlir
@@ -499,30 +499,17 @@ func.func @unroll_jam_nested_iter_args_mulf(%arg0: memref<21x30xf32, 1>, %init :
return
}
-// CHECK: %[[CONST0:[a-zA-Z0-9_]*]] = arith.constant 20 : index
-// CHECK-NEXT: [[RES:%[0-9]+]]:2 = affine.for %[[IV0:arg[0-9]+]] = 0 to 20 step 2 iter_args([[ACC0:%arg[0-9]+]] = [[INIT0]], [[ACC1:%arg[0-9]+]] = [[INIT0]]) -> (f32, f32) {
-// CHECK-NEXT: [[RES1:%[0-9]+]]:2 = affine.for %[[IV1:arg[0-9]+]] = 0 to 30 iter_args([[ACC2:%arg[0-9]+]] = [[INIT1]], [[ACC3:%arg[0-9]+]] = [[INIT1]]) -> (f32, f32) {
+// Floating-point target reductions are left unchanged because unroll-and-jam
+// reassociates the reduction.
+// CHECK: affine.for %[[IV0:arg[0-9]+]] = 0 to 21 iter_args([[ACC0:%arg[0-9]+]] = [[INIT0]]) -> (f32) {
+// CHECK-NEXT: [[RES1:%[0-9]+]] = affine.for %[[IV1:arg[0-9]+]] = 0 to 30 iter_args([[ACC1:%arg[0-9]+]] = [[INIT1]]) -> (f32) {
// CHECK-NEXT: [[LOAD1:%[0-9]+]] = affine.load {{.*}}[%[[IV0]], %[[IV1]]]
-// CHECK-NEXT: [[ADD1:%[0-9]+]] = arith.addf [[ACC2]], [[LOAD1]] : f32
-// CHECK-NEXT: %[[INC1:[0-9]+]] = affine.apply [[$MAP_PLUS_1]](%[[IV0]])
-// CHECK-NEXT: [[LOAD2:%[0-9]+]] = affine.load {{.*}}[%[[INC1]], %[[IV1]]]
-// CHECK-NEXT: [[ADD2:%[0-9]+]] = arith.addf [[ACC3]], [[LOAD2]] : f32
-// CHECK-NEXT: affine.yield [[ADD1]], [[ADD2]]
+// CHECK-NEXT: [[ADD1:%[0-9]+]] = arith.addf [[ACC1]], [[LOAD1]] : f32
+// CHECK-NEXT: affine.yield [[ADD1]] : f32
// CHECK-NEXT: }
-// CHECK-NEXT: [[MUL1:%[0-9]+]] = arith.mulf [[ACC0]], [[RES1]]#0 : f32
-// CHECK-NEXT: affine.apply
-// CHECK-NEXT: [[MUL2:%[0-9]+]] = arith.mulf [[ACC1]], [[RES1]]#1 : f32
-// CHECK-NEXT: affine.yield [[MUL1]], [[MUL2]]
-// CHECK-NEXT: }
-// Reduction op.
-// CHECK-NEXT: [[MUL3:%[0-9]+]] = arith.mulf [[RES]]#0, [[RES]]#1 : f32
-// Cleanup loop (single iteration).
-// CHECK-NEXT: [[RES2:%[0-9]+]] = affine.for %[[IV2:arg[0-9]+]] = 0 to 30 iter_args([[ACC4:%arg[0-9]+]] = [[INIT1]]) -> (f32) {
-// CHECK-NEXT: [[LOAD3:%[0-9]+]] = affine.load {{.*}}[%[[CONST0]], %[[IV2]]]
-// CHECK-NEXT: [[ADD3:%[0-9]+]] = arith.addf [[ACC4]], [[LOAD3]] : f32
-// CHECK-NEXT: affine.yield [[ADD3]] : f32
+// CHECK-NEXT: [[MUL1:%[0-9]+]] = arith.mulf [[ACC0]], [[RES1]] : f32
+// CHECK-NEXT: affine.yield [[MUL1]] : f32
// CHECK-NEXT: }
-// CHECK-NEXT: [[MUL4:%[0-9]+]] = arith.mulf [[MUL3]], [[RES2]] : f32
// CHECK-NEXT: return
// CHECK-LABEL: func @unroll_jam_iter_args_addi
@@ -537,7 +524,8 @@ func.func @unroll_jam_iter_args_addi(%arg0: memref<21xi32, 1>, %init : i32) {
}
// CHECK: %[[CONST0:[a-zA-Z0-9_]*]] = arith.constant 20 : index
-// CHECK-NEXT: [[RES:%[0-9]+]]:2 = affine.for %[[IV0:arg[0-9]+]] = 0 to 20 step 2 iter_args([[ACC0:%arg[0-9]+]] = [[INIT0]], [[ACC1:%arg[0-9]+]] = [[INIT0]]) -> (i32, i32) {
+// CHECK-NEXT: %[[ZERO:[a-zA-Z0-9_]*]] = arith.constant 0 : i32
+// CHECK-NEXT: [[RES:%[0-9]+]]:2 = affine.for %[[IV0:arg[0-9]+]] = 0 to 20 step 2 iter_args([[ACC0:%arg[0-9]+]] = [[INIT0]], [[ACC1:%arg[0-9]+]] = %[[ZERO]]) -> (i32, i32) {
// CHECK-NEXT: [[LOAD1:%[0-9]+]] = affine.load {{.*}}[%[[IV0]]]
// CHECK-NEXT: [[ADD1:%[0-9]+]] = arith.addi [[ACC0]], [[LOAD1]] : i32
// CHECK-NEXT: %[[INC1:[0-9]+]] = affine.apply [[$MAP_PLUS_1]](%[[IV0]])
diff --git a/mlir/test/Dialect/SCF/transform-ops.mlir b/mlir/test/Dialect/SCF/transform-ops.mlir
index 96078c09c1eebe..eed66cc19c3c89 100644
--- a/mlir/test/Dialect/SCF/transform-ops.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops.mlir
@@ -384,43 +384,44 @@ module attributes {transform.with_named_sequence} {
// -----
// CHECK-LABEL: @loop_unroll_and_jam_op
-// CHECK: %[[VAL_0:.*]]: memref<21x30xf32, 1>, %[[INIT0:.*]]: f32, %[[INIT1:.*]]: f32) {
-func.func @loop_unroll_and_jam_op(%arg0: memref<21x30xf32, 1>, %init : f32, %init1 : f32) {
+// CHECK: %[[VAL_0:.*]]: memref<21x30xi32, 1>, %[[INIT0:.*]]: i32, %[[INIT1:.*]]: i32) {
+func.func @loop_unroll_and_jam_op(%arg0: memref<21x30xi32, 1>, %init : i32, %init1 : i32) {
// CHECK: %[[LAST_OUT_ITER:.*]] = arith.constant 20
- // CHECK: %[[VAL_4:.*]]:2 = affine.for %[[OUTER_I:.*]] = 0 to 20 step 2 iter_args(%[[VAL_6:.*]] = %[[INIT0]], %[[VAL_7:.*]] = %[[INIT0]])
- %0 = affine.for %arg3 = 0 to 21 iter_args(%arg4 = %init) -> (f32) {
+ // CHECK: %[[ONE:.*]] = arith.constant 1 : i32
+ // CHECK: %[[VAL_4:.*]]:2 = affine.for %[[OUTER_I:.*]] = 0 to 20 step 2 iter_args(%[[VAL_6:.*]] = %[[INIT0]], %[[VAL_7:.*]] = %[[ONE]])
+ %0 = affine.for %arg3 = 0 to 21 iter_args(%arg4 = %init) -> (i32) {
// CHECK: %[[VAL_8:.*]]:2 = affine.for %[[INNER_I:.*]] = 0 to 30 iter_args(%[[SUM0:.*]] = %[[INIT1]], %[[SUM1:.*]] = %[[INIT1]])
- %1 = affine.for %arg5 = 0 to 30 iter_args(%arg6 = %init1) -> (f32) {
+ %1 = affine.for %arg5 = 0 to 30 iter_args(%arg6 = %init1) -> (i32) {
// CHECK: %[[LOAD0:.*]] = affine.load %[[VAL_0]]{{\[}}%[[OUTER_I]], %[[INNER_I]]]
- // CHECK: %[[ITER_SUM0:.*]] = arith.addf %[[SUM0]], %[[LOAD0]]
+ // CHECK: %[[ITER_SUM0:.*]] = arith.addi %[[SUM0]], %[[LOAD0]]
// CHECK: %[[APPLY_OUTER_I:.*]] = affine.apply #map(%[[OUTER_I]])
// CHECK: %[[LOAD1:.*]] = affine.load %[[VAL_0]]{{\[}}%[[APPLY_OUTER_I]], %[[INNER_I]]]
- // CHECK: ...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/236076
More information about the Mlir-commits
mailing list