[Mlir-commits] [mlir] [SCFToAffine] Raise scf.for to affine.for (PR #200851)
Reinhard Stahn
llvmlistbot at llvm.org
Tue Jun 9 03:09:11 PDT 2026
https://github.com/rainij updated https://github.com/llvm/llvm-project/pull/200851
>From 2768998ba5628eea7ed9894353228219c975dae5 Mon Sep 17 00:00:00 2001
From: Ming Yan <nexming7 at gmail.com>
Date: Mon, 11 Aug 2025 00:14:47 +0800
Subject: [PATCH 01/37] Add a pass to raise scf to affine ops.
This patch supports the conversion from `scf.for` to `affine.for`.
---
mlir/include/mlir/Conversion/Passes.h | 1 +
mlir/include/mlir/Conversion/Passes.td | 12 ++
.../mlir/Conversion/SCFToAffine/SCFToAffine.h | 26 ++++
mlir/lib/Conversion/CMakeLists.txt | 1 +
.../lib/Conversion/SCFToAffine/CMakeLists.txt | 17 +++
.../Conversion/SCFToAffine/SCFToAffine.cpp | 136 ++++++++++++++++++
.../Conversion/SCFToAffine/scf-to-affine.mlir | 57 ++++++++
7 files changed, 250 insertions(+)
create mode 100644 mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h
create mode 100644 mlir/lib/Conversion/SCFToAffine/CMakeLists.txt
create mode 100644 mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
create mode 100644 mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
diff --git a/mlir/include/mlir/Conversion/Passes.h b/mlir/include/mlir/Conversion/Passes.h
index 82c7670296e52..577cab3a0161f 100644
--- a/mlir/include/mlir/Conversion/Passes.h
+++ b/mlir/include/mlir/Conversion/Passes.h
@@ -62,6 +62,7 @@
#include "mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h"
#include "mlir/Conversion/PDLToPDLInterp/PDLToPDLInterp.h"
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
+#include "mlir/Conversion/SCFToAffine/SCFToAffine.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Conversion/SCFToEmitC/SCFToEmitC.h"
#include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index dda756ddab152..f90765598712c 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1149,6 +1149,18 @@ def ReconcileUnrealizedCastsPass : Pass<"reconcile-unrealized-casts"> {
}];
}
+//===----------------------------------------------------------------------===//
+// SCFToAffine
+//===----------------------------------------------------------------------===//
+
+def RaiseSCFToAffinePass : Pass<"raise-scf-to-affine"> {
+ let summary = "Raise SCF to affine ops";
+ let dependentDialects = [
+ "affine::AffineDialect",
+ "scf::SCFDialect",
+ ];
+}
+
//===----------------------------------------------------------------------===//
// SCFToControlFlow
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h b/mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h
new file mode 100644
index 0000000000000..4f87ef8e6c6e4
--- /dev/null
+++ b/mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h
@@ -0,0 +1,26 @@
+//===- SCFToAffine.h - SCF to Affine Pass entrypoint ------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_SCFTOAFFINE_SCFTOAFFINE_H_
+#define MLIR_CONVERSION_SCFTOAFFINE_SCFTOAFFINE_H_
+
+#include <memory>
+
+namespace mlir {
+class Pass;
+class RewritePatternSet;
+
+#define GEN_PASS_DECL_RAISESCFTOAFFINEPASS
+#include "mlir/Conversion/Passes.h.inc"
+
+/// Collect a set of patterns to convert SCF operations to Affine operations.
+void populateSCFToAffineConversionPatterns(RewritePatternSet &patterns);
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_SCFTOAFFINE_SCFTOAFFINE_H_
diff --git a/mlir/lib/Conversion/CMakeLists.txt b/mlir/lib/Conversion/CMakeLists.txt
index f5e0bcf613e59..2e9e50e3bf67a 100644
--- a/mlir/lib/Conversion/CMakeLists.txt
+++ b/mlir/lib/Conversion/CMakeLists.txt
@@ -55,6 +55,7 @@ add_subdirectory(OpenMPToLLVM)
add_subdirectory(PDLToPDLInterp)
add_subdirectory(PtrToLLVM)
add_subdirectory(ReconcileUnrealizedCasts)
+add_subdirectory(SCFToAffine)
add_subdirectory(SCFToControlFlow)
add_subdirectory(SCFToEmitC)
add_subdirectory(SCFToGPU)
diff --git a/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt b/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt
new file mode 100644
index 0000000000000..bf1494d6f3cf0
--- /dev/null
+++ b/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt
@@ -0,0 +1,17 @@
+add_mlir_conversion_library(MLIRSCFToAffine
+ SCFToAffine.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/SCFToAffine
+
+ DEPENDS
+ MLIRConversionPassIncGen
+
+ LINK_LIBS PUBLIC
+ MLIRArithDialect
+ MLIRAffineDialect
+ MLIRLLVMDialect
+ MLIRSCFDialect
+ MLIRSCFTransforms
+ MLIRTransforms
+ )
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
new file mode 100644
index 0000000000000..35e662d88b488
--- /dev/null
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -0,0 +1,136 @@
+//===- SCFToAffine.cpp - SCF to Affine conversion -------------------------===//
+//
+// 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 a pass to raise scf.for, scf.if and loop.terminator
+// ops into affine ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/SCFToAffine/SCFToAffine.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Verifier.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/Passes.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_RAISESCFTOAFFINEPASS
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+struct SCFToAffinePass
+ : public impl::RaiseSCFToAffinePassBase<SCFToAffinePass> {
+ void runOnOperation() override;
+};
+
+bool canRaiseToAffine(scf::ForOp op) {
+ return affine::isValidDim(op.getLowerBound()) &&
+ affine::isValidDim(op.getUpperBound()) &&
+ affine::isValidSymbol(op.getStep());
+}
+
+struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
+ using OpRewritePattern<scf::ForOp>::OpRewritePattern;
+
+ std::pair<affine::AffineForOp, Value>
+ createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
+ if (auto constantStep = op.getStep().getDefiningOp<arith::ConstantOp>()) {
+ int64_t step = cast<IntegerAttr>(constantStep.getValue()).getInt();
+ if (step > 0)
+ return positiveConstantStep(op, step, rewriter);
+ }
+ return genericBounds(op, rewriter);
+ }
+
+ std::pair<affine::AffineForOp, Value>
+ positiveConstantStep(scf::ForOp op, int64_t step,
+ PatternRewriter &rewriter) const {
+ auto affineFor = affine::AffineForOp::create(
+ rewriter, op.getLoc(), ValueRange(op.getLowerBound()),
+ AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)),
+ ValueRange(op.getUpperBound()),
+ AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)), step,
+ op.getInits());
+ return std::make_pair(affineFor, affineFor.getInductionVar());
+ }
+
+ std::pair<affine::AffineForOp, Value>
+ genericBounds(scf::ForOp op, PatternRewriter &rewriter) const {
+ Value lower = op.getLowerBound();
+ Value upper = op.getUpperBound();
+ Value step = op.getStep();
+ AffineExpr lowerExpr = rewriter.getAffineDimExpr(0);
+ AffineExpr upperExpr = rewriter.getAffineDimExpr(1);
+ AffineExpr stepExpr = rewriter.getAffineSymbolExpr(0);
+ auto affineFor = affine::AffineForOp::create(
+ rewriter, op.getLoc(), ValueRange(), rewriter.getConstantAffineMap(0),
+ ValueRange({lower, upper, step}),
+ AffineMap::get(
+ 2, 1, (upperExpr - lowerExpr + stepExpr - 1).floorDiv(stepExpr)),
+ 1, op.getInits());
+
+ rewriter.setInsertionPointToStart(affineFor.getBody());
+ auto actualIndexMap = AffineMap::get(
+ 2, 1, lowerExpr + rewriter.getAffineDimExpr(1) * stepExpr);
+ auto actualIndex = affine::AffineApplyOp::create(
+ rewriter, op.getLoc(), actualIndexMap,
+ ValueRange({lower, affineFor.getInductionVar(), step}));
+ return std::make_pair(affineFor, actualIndex.getResult());
+ }
+
+ LogicalResult matchAndRewrite(scf::ForOp op,
+ PatternRewriter &rewriter) const override {
+ if (!canRaiseToAffine(op))
+ return failure();
+
+ auto [affineFor, actualIndex] = createAffineFor(op, rewriter);
+ Block *affineBody = affineFor.getBody();
+
+ if (affineBody->mightHaveTerminator())
+ rewriter.eraseOp(affineBody->getTerminator());
+
+ SmallVector<Value> argValues;
+ argValues.push_back(actualIndex);
+ llvm::append_range(argValues, affineFor.getRegionIterArgs());
+ rewriter.inlineBlockBefore(op.getBody(), affineBody, affineBody->end(),
+ argValues);
+
+ auto scfYieldOp = cast<scf::YieldOp>(affineBody->getTerminator());
+ rewriter.setInsertionPointToEnd(affineBody);
+ rewriter.replaceOpWithNewOp<affine::AffineYieldOp>(
+ scfYieldOp, scfYieldOp->getOperands());
+
+ rewriter.replaceOp(op, affineFor);
+ return success();
+ }
+};
+
+} // namespace
+
+void mlir::populateSCFToAffineConversionPatterns(RewritePatternSet &patterns) {
+ patterns.add<ForOpRewrite>(patterns.getContext());
+}
+
+void SCFToAffinePass::runOnOperation() {
+ MLIRContext &ctx = getContext();
+ RewritePatternSet patterns(&ctx);
+ populateSCFToAffineConversionPatterns(patterns);
+
+ // Configure conversion to raise SCF operations.
+ ConversionTarget target(ctx);
+ target.addDynamicallyLegalOp<scf::ForOp>(
+ [](scf::ForOp op) { return !canRaiseToAffine(op); });
+ target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+ if (failed(
+ applyPartialConversion(getOperation(), target, std::move(patterns))))
+ signalPassFailure();
+}
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
new file mode 100644
index 0000000000000..2e2649ed8ef1c
--- /dev/null
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -0,0 +1,57 @@
+// RUN: mlir-opt -raise-scf-to-affine -split-input-file %s | FileCheck %s
+
+// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1)[s0] -> ((d1 - d0 + s0 - 1) floordiv s0)>
+// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
+// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: func.func @simple_loop(
+// CHECK-SAME: %[[ARG0:.*]]: memref<?xi32>,
+// CHECK-SAME: %[[ARG1:.*]]: memref<3xindex>) {
+// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i32
+// CHECK: %[[VAL_1:.*]] = arith.constant 0 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 2 : index
+// CHECK: %[[VAL_4:.*]] = memref.load %[[ARG1]]{{\[}}%[[VAL_1]]] : memref<3xindex>
+// CHECK: %[[VAL_5:.*]] = memref.load %[[ARG1]]{{\[}}%[[VAL_2]]] : memref<3xindex>
+// CHECK: %[[VAL_6:.*]] = memref.load %[[ARG1]]{{\[}}%[[VAL_3]]] : memref<3xindex>
+// CHECK: affine.for %[[VAL_7:.*]] = 0 to #[[$ATTR_0]](%[[VAL_4]], %[[VAL_5]]){{\[}}%[[VAL_6]]] {
+// CHECK: %[[VAL_8:.*]] = affine.apply #[[$ATTR_1]](%[[VAL_4]], %[[VAL_7]]){{\[}}%[[VAL_6]]]
+// CHECK: memref.store %[[VAL_0]], %[[ARG0]]{{\[}}%[[VAL_8]]] : memref<?xi32>
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+func.func @simple_loop(%arg0: memref<?xi32>, %arg1: memref<3xindex>) {
+ %c0_i32 = arith.constant 0 : i32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %0 = memref.load %arg1[%c0] : memref<3xindex>
+ %1 = memref.load %arg1[%c1] : memref<3xindex>
+ %2 = memref.load %arg1[%c2] : memref<3xindex>
+ scf.for %arg2 = %0 to %1 step %2 {
+ memref.store %c0_i32, %arg0[%arg2] : memref<?xi32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @loop_with_constant_step(
+// CHECK-SAME: %[[ARG0:.*]]: memref<?xi32>,
+// CHECK-SAME: %[[ARG1:.*]]: index,
+// CHECK-SAME: %[[ARG2:.*]]: index) {
+// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i32
+// CHECK: %[[VAL_1:.*]] = arith.constant 3 : index
+// CHECK: affine.for %[[VAL_2:.*]] = #[[$ATTR_2]](%[[ARG1]]) to #[[$ATTR_2]](%[[ARG2]]) step 3 {
+// CHECK: memref.store %[[VAL_0]], %[[ARG0]]{{\[}}%[[VAL_2]]] : memref<?xi32>
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+func.func @loop_with_constant_step(%arg0: memref<?xi32>, %arg1: index, %arg2: index) {
+ %c0_i32 = arith.constant 0 : i32
+ %c3 = arith.constant 3 : index
+ scf.for %arg3 = %arg1 to %arg2 step %c3 {
+ memref.store %c0_i32, %arg0[%arg3] : memref<?xi32>
+ }
+ return
+}
+
>From ef394e92c3c9109784083dd36691333ef5240a90 Mon Sep 17 00:00:00 2001
From: yanming <ming.yan at terapines.com>
Date: Mon, 25 Aug 2025 13:53:06 +0800
Subject: [PATCH 02/37] Use `walkAndApplyPatterns` instead of
`applyPartialConversion`
---
mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp | 12 ++----------
1 file changed, 2 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 35e662d88b488..c8d250ab6e447 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -15,8 +15,8 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Verifier.h"
-#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/WalkPatternRewriteDriver.h"
namespace mlir {
#define GEN_PASS_DEF_RAISESCFTOAFFINEPASS
@@ -124,13 +124,5 @@ void SCFToAffinePass::runOnOperation() {
MLIRContext &ctx = getContext();
RewritePatternSet patterns(&ctx);
populateSCFToAffineConversionPatterns(patterns);
-
- // Configure conversion to raise SCF operations.
- ConversionTarget target(ctx);
- target.addDynamicallyLegalOp<scf::ForOp>(
- [](scf::ForOp op) { return !canRaiseToAffine(op); });
- target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
- if (failed(
- applyPartialConversion(getOperation(), target, std::move(patterns))))
- signalPassFailure();
+ walkAndApplyPatterns(getOperation(), std::move(patterns));
}
>From 575d4ed1551bd25a3080927ba4a6a617e41393c5 Mon Sep 17 00:00:00 2001
From: yanming <ming.yan at terapines.com>
Date: Mon, 25 Aug 2025 13:54:45 +0800
Subject: [PATCH 03/37] Add a nested loop test case.
---
.../Conversion/SCFToAffine/scf-to-affine.mlir | 33 ++++++++++++++++++-
1 file changed, 32 insertions(+), 1 deletion(-)
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 2e2649ed8ef1c..41504f987a216 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -2,7 +2,6 @@
// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1)[s0] -> ((d1 - d0 + s0 - 1) floordiv s0)>
// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
-// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func.func @simple_loop(
// CHECK-SAME: %[[ARG0:.*]]: memref<?xi32>,
// CHECK-SAME: %[[ARG1:.*]]: memref<3xindex>) {
@@ -34,6 +33,9 @@ func.func @simple_loop(%arg0: memref<?xi32>, %arg1: memref<3xindex>) {
return
}
+// -----
+
+// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func.func @loop_with_constant_step(
// CHECK-SAME: %[[ARG0:.*]]: memref<?xi32>,
// CHECK-SAME: %[[ARG1:.*]]: index,
@@ -55,3 +57,32 @@ func.func @loop_with_constant_step(%arg0: memref<?xi32>, %arg1: index, %arg2: in
return
}
+// -----
+
+// CHECK: #[[$ATTR_3:.+]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: func.func @nested_loop(
+// CHECK-SAME: %[[ARG0:.*]]: memref<?x?xi32>,
+// CHECK-SAME: %[[ARG1:.*]]: index,
+// CHECK-SAME: %[[ARG2:.*]]: index) {
+// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i32
+// CHECK: %[[VAL_1:.*]] = arith.constant 0 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 1 : index
+// CHECK: affine.for %[[VAL_3:.*]] = #[[$ATTR_3]](%[[VAL_1]]) to #[[$ATTR_3]](%[[ARG1]]) {
+// CHECK: affine.for %[[VAL_4:.*]] = #[[$ATTR_3]](%[[VAL_1]]) to #[[$ATTR_3]](%[[ARG2]]) {
+// CHECK: memref.store %[[VAL_0]], %[[ARG0]]{{\[}}%[[VAL_3]], %[[VAL_4]]] : memref<?x?xi32>
+// CHECK: }
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+func.func @nested_loop(%arg0: memref<?x?xi32>, %arg1: index, %arg2: index) {
+ %c0_i32 = arith.constant 0 : i32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.for %arg3 = %c0 to %arg1 step %c1 {
+ scf.for %arg4 = %c0 to %arg2 step %c1 {
+ memref.store %c0_i32, %arg0[%arg3, %arg4] : memref<?x?xi32>
+ }
+ }
+ return
+}
>From 043fe0301fcd6c8bf56c3be2456c72e92af296e6 Mon Sep 17 00:00:00 2001
From: yanming <ming.yan at terapines.com>
Date: Mon, 25 Aug 2025 13:58:51 +0800
Subject: [PATCH 04/37] Add debugging information.
---
mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp | 7 ++++++-
1 file changed, 6 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index c8d250ab6e447..ba47763af1486 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -23,6 +23,8 @@ namespace mlir {
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
+#define DEBUG_TYPE "raise-scf-to-affine"
+
using namespace mlir;
namespace {
@@ -89,8 +91,11 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
LogicalResult matchAndRewrite(scf::ForOp op,
PatternRewriter &rewriter) const override {
- if (!canRaiseToAffine(op))
+ if (!canRaiseToAffine(op)) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "[affine] Cannot raise scf op: " << op << "\n");
return failure();
+ }
auto [affineFor, actualIndex] = createAffineFor(op, rewriter);
Block *affineBody = affineFor.getBody();
>From 253a3125b12dde54ad2304e4ac83d5644536fbb0 Mon Sep 17 00:00:00 2001
From: yanming <ming.yan at terapines.com>
Date: Mon, 25 Aug 2025 16:46:25 +0800
Subject: [PATCH 05/37] Simplify the code
---
mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp | 9 +++++----
1 file changed, 5 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index ba47763af1486..0166552bd40f3 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -17,6 +17,7 @@
#include "mlir/IR/Verifier.h"
#include "mlir/Transforms/Passes.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"
+#include "llvm/Support/DebugLog.h"
namespace mlir {
#define GEN_PASS_DEF_RAISESCFTOAFFINEPASS
@@ -45,8 +46,9 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
std::pair<affine::AffineForOp, Value>
createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
- if (auto constantStep = op.getStep().getDefiningOp<arith::ConstantOp>()) {
- int64_t step = cast<IntegerAttr>(constantStep.getValue()).getInt();
+ IntegerAttr constAttr;
+ if (matchPattern(op.getStep(), m_Constant(&constAttr))) {
+ int64_t step = constAttr.getInt();
if (step > 0)
return positiveConstantStep(op, step, rewriter);
}
@@ -92,8 +94,7 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
LogicalResult matchAndRewrite(scf::ForOp op,
PatternRewriter &rewriter) const override {
if (!canRaiseToAffine(op)) {
- LLVM_DEBUG(llvm::dbgs()
- << "[affine] Cannot raise scf op: " << op << "\n");
+ LDBG() << "[affine] Cannot raise scf op: " << op << "\n";
return failure();
}
>From c13f1a5780bc0ac62977fced337890e9c6e25dc3 Mon Sep 17 00:00:00 2001
From: yanming <ming.yan at terapines.com>
Date: Mon, 25 Aug 2025 18:00:27 +0800
Subject: [PATCH 06/37] Fix test format.
---
.../Conversion/SCFToAffine/scf-to-affine.mlir | 91 ++++++-------------
1 file changed, 27 insertions(+), 64 deletions(-)
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 41504f987a216..22fc61ca6cca2 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -1,87 +1,50 @@
// RUN: mlir-opt -raise-scf-to-affine -split-input-file %s | FileCheck %s
-// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1)[s0] -> ((d1 - d0 + s0 - 1) floordiv s0)>
-// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
-// CHECK-LABEL: func.func @simple_loop(
-// CHECK-SAME: %[[ARG0:.*]]: memref<?xi32>,
-// CHECK-SAME: %[[ARG1:.*]]: memref<3xindex>) {
-// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i32
-// CHECK: %[[VAL_1:.*]] = arith.constant 0 : index
-// CHECK: %[[VAL_2:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_3:.*]] = arith.constant 2 : index
-// CHECK: %[[VAL_4:.*]] = memref.load %[[ARG1]]{{\[}}%[[VAL_1]]] : memref<3xindex>
-// CHECK: %[[VAL_5:.*]] = memref.load %[[ARG1]]{{\[}}%[[VAL_2]]] : memref<3xindex>
-// CHECK: %[[VAL_6:.*]] = memref.load %[[ARG1]]{{\[}}%[[VAL_3]]] : memref<3xindex>
-// CHECK: affine.for %[[VAL_7:.*]] = 0 to #[[$ATTR_0]](%[[VAL_4]], %[[VAL_5]]){{\[}}%[[VAL_6]]] {
-// CHECK: %[[VAL_8:.*]] = affine.apply #[[$ATTR_1]](%[[VAL_4]], %[[VAL_7]]){{\[}}%[[VAL_6]]]
-// CHECK: memref.store %[[VAL_0]], %[[ARG0]]{{\[}}%[[VAL_8]]] : memref<?xi32>
-// CHECK: }
-// CHECK: return
-// CHECK: }
-
-func.func @simple_loop(%arg0: memref<?xi32>, %arg1: memref<3xindex>) {
+// CHECK: #[[$UB_MAP:.+]] = affine_map<(d0, d1)[s0] -> ((d1 - d0 + s0 - 1) floordiv s0)>
+// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
+// CHECK-LABEL: @generic_loop
+// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LOWER:.*]]: index, %[[UPPER:.*]]: index, %[[STEP:.*]]: index
+func.func @generic_loop(%arr: memref<?xi32>, %lower: index, %upper: index, %step: index) {
+// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]](%[[LOWER]], %[[UPPER]])[%[[STEP]]] {
+// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LOWER]], %[[IV]])[%[[STEP]]]
+// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IDX]]] : memref<?xi32>
+// CHECK: }
%c0_i32 = arith.constant 0 : i32
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
- %c2 = arith.constant 2 : index
- %0 = memref.load %arg1[%c0] : memref<3xindex>
- %1 = memref.load %arg1[%c1] : memref<3xindex>
- %2 = memref.load %arg1[%c2] : memref<3xindex>
- scf.for %arg2 = %0 to %1 step %2 {
- memref.store %c0_i32, %arg0[%arg2] : memref<?xi32>
+ scf.for %idx = %lower to %upper step %step {
+ memref.store %c0_i32, %arr[%idx] : memref<?xi32>
}
return
}
// -----
-// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0) -> (d0)>
-// CHECK-LABEL: func.func @loop_with_constant_step(
-// CHECK-SAME: %[[ARG0:.*]]: memref<?xi32>,
-// CHECK-SAME: %[[ARG1:.*]]: index,
-// CHECK-SAME: %[[ARG2:.*]]: index) {
-// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i32
-// CHECK: %[[VAL_1:.*]] = arith.constant 3 : index
-// CHECK: affine.for %[[VAL_2:.*]] = #[[$ATTR_2]](%[[ARG1]]) to #[[$ATTR_2]](%[[ARG2]]) step 3 {
-// CHECK: memref.store %[[VAL_0]], %[[ARG0]]{{\[}}%[[VAL_2]]] : memref<?xi32>
-// CHECK: }
-// CHECK: return
-// CHECK: }
-
-func.func @loop_with_constant_step(%arg0: memref<?xi32>, %arg1: index, %arg2: index) {
+// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: @loop_with_constant_step
+// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LOWER:.*]]: index, %[[UPPER:.*]]: index
+func.func @loop_with_constant_step(%arr: memref<?xi32>, %lower: index, %upper: index) {
+// CHECK: affine.for %[[IDX:.*]] = #[[$MAP]](%[[LOWER]]) to #[[$MAP]](%[[UPPER]]) step 3 {
+// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IDX]]] : memref<?xi32>
+// CHECK: }
%c0_i32 = arith.constant 0 : i32
%c3 = arith.constant 3 : index
- scf.for %arg3 = %arg1 to %arg2 step %c3 {
- memref.store %c0_i32, %arg0[%arg3] : memref<?xi32>
+ scf.for %idx = %lower to %upper step %c3 {
+ memref.store %c0_i32, %arr[%idx] : memref<?xi32>
}
return
}
// -----
-// CHECK: #[[$ATTR_3:.+]] = affine_map<(d0) -> (d0)>
-// CHECK-LABEL: func.func @nested_loop(
-// CHECK-SAME: %[[ARG0:.*]]: memref<?x?xi32>,
-// CHECK-SAME: %[[ARG1:.*]]: index,
-// CHECK-SAME: %[[ARG2:.*]]: index) {
-// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i32
-// CHECK: %[[VAL_1:.*]] = arith.constant 0 : index
-// CHECK: %[[VAL_2:.*]] = arith.constant 1 : index
-// CHECK: affine.for %[[VAL_3:.*]] = #[[$ATTR_3]](%[[VAL_1]]) to #[[$ATTR_3]](%[[ARG1]]) {
-// CHECK: affine.for %[[VAL_4:.*]] = #[[$ATTR_3]](%[[VAL_1]]) to #[[$ATTR_3]](%[[ARG2]]) {
-// CHECK: memref.store %[[VAL_0]], %[[ARG0]]{{\[}}%[[VAL_3]], %[[VAL_4]]] : memref<?x?xi32>
-// CHECK: }
-// CHECK: }
-// CHECK: return
-// CHECK: }
-
-func.func @nested_loop(%arg0: memref<?x?xi32>, %arg1: index, %arg2: index) {
+// CHECK-LABEL: @nested_loop
+func.func @nested_loop(%arg0: memref<?x?xi32>, %upper1: index, %upper2: index) {
+// CHECK: affine.for
+// CHECK: affine.for
%c0_i32 = arith.constant 0 : i32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
- scf.for %arg3 = %c0 to %arg1 step %c1 {
- scf.for %arg4 = %c0 to %arg2 step %c1 {
- memref.store %c0_i32, %arg0[%arg3, %arg4] : memref<?x?xi32>
+ scf.for %i = %c0 to %upper1 step %c1 {
+ scf.for %j = %c0 to %upper2 step %c1 {
+ memref.store %c0_i32, %arg0[%i, %j] : memref<?x?xi32>
}
}
return
>From 0727bf3485c93f27a2d939984382b5409af760d1 Mon Sep 17 00:00:00 2001
From: Ming Yan <ming.yan at terapines.com>
Date: Mon, 25 Aug 2025 23:53:36 +0800
Subject: [PATCH 07/37] Add a description for the pass.
---
mlir/include/mlir/Conversion/Passes.td | 18 +++++++++++++++++-
1 file changed, 17 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index f90765598712c..8c14adebe3d4f 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1154,7 +1154,23 @@ def ReconcileUnrealizedCastsPass : Pass<"reconcile-unrealized-casts"> {
//===----------------------------------------------------------------------===//
def RaiseSCFToAffinePass : Pass<"raise-scf-to-affine"> {
- let summary = "Raise SCF to affine ops";
+ let summary = "Raise SCF operations to affine operations where possible";
+ let description = [{
+ This pass raises SCF operations to affine operations where possible.
+
+ Specifically:
+ - `scf.for` loops with affine-compatible bounds and steps are
+ converted to `affine.for`.
+
+ Converting SCF to affine enables affine-specific optimizations such as
+ loop tiling, unrolling, vectorization, and memory access analysis.
+
+ Note:
+ - Only loops that are statically affine can be converted;
+ non-affine loops remain in SCF form.
+ - This pass does not modify memory accesses; consider using
+ --affine-raise-from-memref for converting `memref.load`/`store`.
+ }];
let dependentDialects = [
"affine::AffineDialect",
"scf::SCFDialect",
>From a60dde862b36a487a18ec31d5e7b6f9e370b2502 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <51020828+rainij at users.noreply.github.com>
Date: Mon, 1 Jun 2026 19:37:49 +0000
Subject: [PATCH 08/37] Minor refactoring, docstrings, add some TODOs.
---
mlir/include/mlir/Conversion/Passes.td | 2 +
.../Conversion/SCFToAffine/SCFToAffine.cpp | 208 +++++++++++-------
2 files changed, 136 insertions(+), 74 deletions(-)
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 8c14adebe3d4f..4317b80f1533a 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1158,6 +1158,8 @@ def RaiseSCFToAffinePass : Pass<"raise-scf-to-affine"> {
let description = [{
This pass raises SCF operations to affine operations where possible.
+ TODO(rainij): document additional features.
+
Specifically:
- `scf.for` loops with affine-compatible bounds and steps are
converted to `affine.for`.
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 0166552bd40f3..cc9aed4ae1a21 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -5,6 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
+// TODO(rainij): update description?
//
// This file implements a pass to raise scf.for, scf.if and loop.terminator
// ops into affine ops.
@@ -30,105 +31,164 @@ using namespace mlir;
namespace {
+//===----------------------------------------------------------------------===//
+// SCFToAffinePass
+//===----------------------------------------------------------------------===//
+
struct SCFToAffinePass
: public impl::RaiseSCFToAffinePassBase<SCFToAffinePass> {
void runOnOperation() override;
};
-bool canRaiseToAffine(scf::ForOp op) {
- return affine::isValidDim(op.getLowerBound()) &&
- affine::isValidDim(op.getUpperBound()) &&
- affine::isValidSymbol(op.getStep());
-}
+//===----------------------------------------------------------------------===//
+// ForOpRewrite
+//===----------------------------------------------------------------------===//
+
+// TODO(rainij): add some patterns inspired by Enzyme-JAX to raise certain
+// scf.for ops which do not *already* satisfy canRaiseToAffine. If possible do
+// it in a way so that after some rewrite they satisfy it so that the current
+// pattern just applies.
struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
using OpRewritePattern<scf::ForOp>::OpRewritePattern;
+ LogicalResult matchAndRewrite(scf::ForOp op,
+ PatternRewriter &rewriter) const override;
+
+private:
+ /// Returns an equivalent `affine.for` skeleton - that is, whose body is
+ /// essentially empty up to a normalization of the induction variable in case
+ /// the step is not constant. Returns the normalized induction variable as
+ /// second value.
std::pair<affine::AffineForOp, Value>
- createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
- IntegerAttr constAttr;
- if (matchPattern(op.getStep(), m_Constant(&constAttr))) {
- int64_t step = constAttr.getInt();
- if (step > 0)
- return positiveConstantStep(op, step, rewriter);
- }
- return genericBounds(op, rewriter);
- }
+ createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const;
+ /// Returns `(affine.for %iv = %lb to %ub step <step> { <empty> }, %iv)`
std::pair<affine::AffineForOp, Value>
positiveConstantStep(scf::ForOp op, int64_t step,
- PatternRewriter &rewriter) const {
- auto affineFor = affine::AffineForOp::create(
- rewriter, op.getLoc(), ValueRange(op.getLowerBound()),
- AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)),
- ValueRange(op.getUpperBound()),
- AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)), step,
- op.getInits());
- return std::make_pair(affineFor, affineFor.getInductionVar());
+ PatternRewriter &rewriter) const;
+
+ /// TODO(rainij): consider a better name. Same for the other private
+ /// functions.
+ ///
+ /// Returns an equivalent `affine.for` skeleton whose step is normalized to 1.
+ /// The body contains an expression which computes the old index (old_iv = lb
+ /// + step * iv) which is also returned as second return value.
+ std::pair<affine::AffineForOp, Value>
+ genericBounds(scf::ForOp op, PatternRewriter &rewriter) const;
+
+ // TODO(rainij): should be add the body already via the two helper above? Note
+ // that the second helper already fiddles around with the body.
+
+ // TODO(rainij): I believe some docstring for the helper is needed, but it
+ // should be more concise. Maybe only document the "primary" helper and rename
+ // the two helper in a way which connects to the docstring of the primary
+ // helper.
+};
+
+bool canRaiseToAffine(scf::ForOp op) {
+ return affine::isValidDim(op.getLowerBound()) &&
+ affine::isValidDim(op.getUpperBound()) &&
+ affine::isValidSymbol(op.getStep());
+}
+
+LogicalResult ForOpRewrite::matchAndRewrite(scf::ForOp op,
+ PatternRewriter &rewriter) const {
+ if (!canRaiseToAffine(op)) {
+ // TODO(rainij): another pattern might make this raisible. We might want
+ // drop this message then, or alter it to acknowledge the possibility.
+ LDBG() << "[affine] Cannot raise scf op: " << op << "\n";
+ return failure();
}
- std::pair<affine::AffineForOp, Value>
- genericBounds(scf::ForOp op, PatternRewriter &rewriter) const {
- Value lower = op.getLowerBound();
- Value upper = op.getUpperBound();
- Value step = op.getStep();
- AffineExpr lowerExpr = rewriter.getAffineDimExpr(0);
- AffineExpr upperExpr = rewriter.getAffineDimExpr(1);
- AffineExpr stepExpr = rewriter.getAffineSymbolExpr(0);
- auto affineFor = affine::AffineForOp::create(
- rewriter, op.getLoc(), ValueRange(), rewriter.getConstantAffineMap(0),
- ValueRange({lower, upper, step}),
- AffineMap::get(
- 2, 1, (upperExpr - lowerExpr + stepExpr - 1).floorDiv(stepExpr)),
- 1, op.getInits());
-
- rewriter.setInsertionPointToStart(affineFor.getBody());
- auto actualIndexMap = AffineMap::get(
- 2, 1, lowerExpr + rewriter.getAffineDimExpr(1) * stepExpr);
- auto actualIndex = affine::AffineApplyOp::create(
- rewriter, op.getLoc(), actualIndexMap,
- ValueRange({lower, affineFor.getInductionVar(), step}));
- return std::make_pair(affineFor, actualIndex.getResult());
+ auto [affineFor, actualIndex] = createAffineFor(op, rewriter);
+ Block *affineBody = affineFor.getBody();
+
+ if (affineBody->mightHaveTerminator()) {
+ Operation *terminator = affineBody->getTerminator();
+ assert(isa<affine::AffineYieldOp>(terminator) &&
+ "expected affine.yield if there *might* be terminator");
+ rewriter.eraseOp(terminator);
}
- LogicalResult matchAndRewrite(scf::ForOp op,
- PatternRewriter &rewriter) const override {
- if (!canRaiseToAffine(op)) {
- LDBG() << "[affine] Cannot raise scf op: " << op << "\n";
- return failure();
- }
-
- auto [affineFor, actualIndex] = createAffineFor(op, rewriter);
- Block *affineBody = affineFor.getBody();
-
- if (affineBody->mightHaveTerminator())
- rewriter.eraseOp(affineBody->getTerminator());
-
- SmallVector<Value> argValues;
- argValues.push_back(actualIndex);
- llvm::append_range(argValues, affineFor.getRegionIterArgs());
- rewriter.inlineBlockBefore(op.getBody(), affineBody, affineBody->end(),
- argValues);
-
- auto scfYieldOp = cast<scf::YieldOp>(affineBody->getTerminator());
- rewriter.setInsertionPointToEnd(affineBody);
- rewriter.replaceOpWithNewOp<affine::AffineYieldOp>(
- scfYieldOp, scfYieldOp->getOperands());
-
- rewriter.replaceOp(op, affineFor);
- return success();
+ SmallVector<Value> argValues;
+ argValues.push_back(actualIndex);
+ llvm::append_range(argValues, affineFor.getRegionIterArgs());
+ rewriter.inlineBlockBefore(op.getBody(), affineBody, affineBody->end(),
+ argValues);
+
+ auto scfYieldOp = cast<scf::YieldOp>(affineBody->getTerminator());
+ rewriter.setInsertionPointToEnd(affineBody);
+ rewriter.replaceOpWithNewOp<affine::AffineYieldOp>(scfYieldOp,
+ scfYieldOp->getOperands());
+
+ rewriter.replaceOp(op, affineFor);
+ return success();
+}
+
+std::pair<affine::AffineForOp, Value>
+ForOpRewrite::createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
+ IntegerAttr constAttr;
+ if (matchPattern(op.getStep(), m_Constant(&constAttr))) {
+ int64_t step = constAttr.getInt();
+ if (step > 0)
+ return positiveConstantStep(op, step, rewriter);
}
-};
+ return genericBounds(op, rewriter);
+}
-} // namespace
+std::pair<affine::AffineForOp, Value>
+ForOpRewrite::positiveConstantStep(scf::ForOp op, int64_t step,
+ PatternRewriter &rewriter) const {
+ auto affineFor = affine::AffineForOp::create(
+ rewriter, op.getLoc(), ValueRange(op.getLowerBound()),
+ AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)),
+ ValueRange(op.getUpperBound()),
+ AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)), step, op.getInits());
+ return std::make_pair(affineFor, affineFor.getInductionVar());
+}
-void mlir::populateSCFToAffineConversionPatterns(RewritePatternSet &patterns) {
- patterns.add<ForOpRewrite>(patterns.getContext());
+std::pair<affine::AffineForOp, Value>
+ForOpRewrite::genericBounds(scf::ForOp op, PatternRewriter &rewriter) const {
+ Value lower = op.getLowerBound();
+ Value upper = op.getUpperBound();
+ Value step = op.getStep();
+
+ AffineExpr lowerExpr = rewriter.getAffineDimExpr(0);
+ AffineExpr upperExpr = rewriter.getAffineDimExpr(1);
+ AffineExpr stepExpr = rewriter.getAffineSymbolExpr(0);
+
+ auto affineFor = affine::AffineForOp::create(
+ rewriter, op.getLoc(), ValueRange(), rewriter.getConstantAffineMap(0),
+ ValueRange({lower, upper, step}),
+ AffineMap::get(2, 1,
+ (upperExpr - lowerExpr + stepExpr - 1).floorDiv(stepExpr)),
+ 1, op.getInits());
+
+ rewriter.setInsertionPointToStart(affineFor.getBody());
+ auto actualIndexMap =
+ AffineMap::get(2, 1, lowerExpr + rewriter.getAffineDimExpr(1) * stepExpr);
+ auto actualIndex = affine::AffineApplyOp::create(
+ rewriter, op.getLoc(), actualIndexMap,
+ ValueRange({lower, affineFor.getInductionVar(), step}));
+ return std::make_pair(affineFor, actualIndex.getResult());
}
void SCFToAffinePass::runOnOperation() {
MLIRContext &ctx = getContext();
RewritePatternSet patterns(&ctx);
populateSCFToAffineConversionPatterns(patterns);
+ // TODO(rainij): we might need a different rewriter (which tries to converge)
+ // if we add more features.
walkAndApplyPatterns(getOperation(), std::move(patterns));
}
+
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// API
+//===----------------------------------------------------------------------===//
+
+void mlir::populateSCFToAffineConversionPatterns(RewritePatternSet &patterns) {
+ patterns.add<ForOpRewrite>(patterns.getContext());
+}
>From 65b04bf45973dac2cf3d7182ff266631259abf57 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Wed, 3 Jun 2026 11:21:42 +0000
Subject: [PATCH 09/37] Better docstring
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 45 +++++++++----------
1 file changed, 20 insertions(+), 25 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index cc9aed4ae1a21..7c833873215a4 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -56,34 +56,28 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
PatternRewriter &rewriter) const override;
private:
- /// Returns an equivalent `affine.for` skeleton - that is, whose body is
- /// essentially empty up to a normalization of the induction variable in case
- /// the step is not constant. Returns the normalized induction variable as
- /// second value.
+ /// Returns an equivalent `affine.for` skeleton. There are two cases.
+ ///
+ /// (1) If the step is a constant we trivially raise the `scf.for` by
+ /// essentially keeping lb, ub, iv as is. The body is left empty. The second
+ /// return value is the induction variable in this case.
+ ///
+ /// (2) Otherwise (generic step) we normalize the loop by setting step = 1, lb
+ /// = 0, ub = ceil((old_ub - old_lb) / old_step). Moreover we insert ops to
+ /// compute old_iv = lb + step * iv in the body and return old_iv as second
+ /// result. Apart from that the body is empty.
+ ///
+ /// The resulting `affine.for` is valid (satisfies affine constraints) if lb
+ /// and ub of the `scf.for` are dimensions and its step is a symbol.
std::pair<affine::AffineForOp, Value>
createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const;
- /// Returns `(affine.for %iv = %lb to %ub step <step> { <empty> }, %iv)`
std::pair<affine::AffineForOp, Value>
- positiveConstantStep(scf::ForOp op, int64_t step,
+ caseConstantStep(scf::ForOp op, int64_t step,
PatternRewriter &rewriter) const;
- /// TODO(rainij): consider a better name. Same for the other private
- /// functions.
- ///
- /// Returns an equivalent `affine.for` skeleton whose step is normalized to 1.
- /// The body contains an expression which computes the old index (old_iv = lb
- /// + step * iv) which is also returned as second return value.
std::pair<affine::AffineForOp, Value>
- genericBounds(scf::ForOp op, PatternRewriter &rewriter) const;
-
- // TODO(rainij): should be add the body already via the two helper above? Note
- // that the second helper already fiddles around with the body.
-
- // TODO(rainij): I believe some docstring for the helper is needed, but it
- // should be more concise. Maybe only document the "primary" helper and rename
- // the two helper in a way which connects to the docstring of the primary
- // helper.
+ caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const;
};
bool canRaiseToAffine(scf::ForOp op) {
@@ -132,13 +126,14 @@ ForOpRewrite::createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
if (matchPattern(op.getStep(), m_Constant(&constAttr))) {
int64_t step = constAttr.getInt();
if (step > 0)
- return positiveConstantStep(op, step, rewriter);
+ return caseConstantStep(op, step, rewriter);
+ // TODO(rainij): what about step <= 0? Is this possible?
}
- return genericBounds(op, rewriter);
+ return caseGenericStep(op, rewriter);
}
std::pair<affine::AffineForOp, Value>
-ForOpRewrite::positiveConstantStep(scf::ForOp op, int64_t step,
+ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
PatternRewriter &rewriter) const {
auto affineFor = affine::AffineForOp::create(
rewriter, op.getLoc(), ValueRange(op.getLowerBound()),
@@ -149,7 +144,7 @@ ForOpRewrite::positiveConstantStep(scf::ForOp op, int64_t step,
}
std::pair<affine::AffineForOp, Value>
-ForOpRewrite::genericBounds(scf::ForOp op, PatternRewriter &rewriter) const {
+ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
Value lower = op.getLowerBound();
Value upper = op.getUpperBound();
Value step = op.getStep();
>From 16b5975961167e6709d2d544b1c31b5865f2028d Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Wed, 3 Jun 2026 11:27:37 +0000
Subject: [PATCH 10/37] Turning if-check for positive steps into assert
---
mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp | 5 ++---
1 file changed, 2 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 7c833873215a4..e314344a104f1 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -125,9 +125,8 @@ ForOpRewrite::createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
IntegerAttr constAttr;
if (matchPattern(op.getStep(), m_Constant(&constAttr))) {
int64_t step = constAttr.getInt();
- if (step > 0)
- return caseConstantStep(op, step, rewriter);
- // TODO(rainij): what about step <= 0? Is this possible?
+ assert(step > 0 && "scf.for has positive step");
+ return caseConstantStep(op, step, rewriter);
}
return caseGenericStep(op, rewriter);
}
>From 29cf0c62c309fa4cea94c5ad58ec011dade6e2ae Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Wed, 3 Jun 2026 12:30:07 +0000
Subject: [PATCH 11/37] Refactor for readability
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 40 +++++++++----------
1 file changed, 20 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index e314344a104f1..4fa24f77d9bfe 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -49,6 +49,8 @@ struct SCFToAffinePass
// it in a way so that after some rewrite they satisfy it so that the current
// pattern just applies.
+/// Raise an `scf.for` to an equivalent `affine.for` if lb, ub are dimensions
+/// and step is a symbol.
struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
using OpRewritePattern<scf::ForOp>::OpRewritePattern;
@@ -56,7 +58,7 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
PatternRewriter &rewriter) const override;
private:
- /// Returns an equivalent `affine.for` skeleton. There are two cases.
+ /// Returns an equivalent `affine.for` skeleton. There are two cases.
///
/// (1) If the step is a constant we trivially raise the `scf.for` by
/// essentially keeping lb, ub, iv as is. The body is left empty. The second
@@ -64,8 +66,8 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
///
/// (2) Otherwise (generic step) we normalize the loop by setting step = 1, lb
/// = 0, ub = ceil((old_ub - old_lb) / old_step). Moreover we insert ops to
- /// compute old_iv = lb + step * iv in the body and return old_iv as second
- /// result. Apart from that the body is empty.
+ /// compute old_iv = old_lb + old_step * new_iv in the body and return old_iv
+ /// as second result. Apart from that the body is empty.
///
/// The resulting `affine.for` is valid (satisfies affine constraints) if lb
/// and ub of the `scf.for` are dimensions and its step is a symbol.
@@ -74,7 +76,7 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
std::pair<affine::AffineForOp, Value>
caseConstantStep(scf::ForOp op, int64_t step,
- PatternRewriter &rewriter) const;
+ PatternRewriter &rewriter) const;
std::pair<affine::AffineForOp, Value>
caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const;
@@ -133,39 +135,37 @@ ForOpRewrite::createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
std::pair<affine::AffineForOp, Value>
ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
- PatternRewriter &rewriter) const {
+ PatternRewriter &rewriter) const {
auto affineFor = affine::AffineForOp::create(
rewriter, op.getLoc(), ValueRange(op.getLowerBound()),
AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)),
ValueRange(op.getUpperBound()),
AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)), step, op.getInits());
+
return std::make_pair(affineFor, affineFor.getInductionVar());
}
std::pair<affine::AffineForOp, Value>
ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
- Value lower = op.getLowerBound();
- Value upper = op.getUpperBound();
+ Value lb = op.getLowerBound();
+ Value ub = op.getUpperBound();
Value step = op.getStep();
- AffineExpr lowerExpr = rewriter.getAffineDimExpr(0);
- AffineExpr upperExpr = rewriter.getAffineDimExpr(1);
- AffineExpr stepExpr = rewriter.getAffineSymbolExpr(0);
+ AffineExpr d0 = rewriter.getAffineDimExpr(0);
+ AffineExpr d1 = rewriter.getAffineDimExpr(1);
+ AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
auto affineFor = affine::AffineForOp::create(
rewriter, op.getLoc(), ValueRange(), rewriter.getConstantAffineMap(0),
- ValueRange({lower, upper, step}),
- AffineMap::get(2, 1,
- (upperExpr - lowerExpr + stepExpr - 1).floorDiv(stepExpr)),
- 1, op.getInits());
+ ValueRange({lb, ub, step}),
+ AffineMap::get(2, 1, (d1 - d0 + s0 - 1).floorDiv(s0)), 1, op.getInits());
rewriter.setInsertionPointToStart(affineFor.getBody());
- auto actualIndexMap =
- AffineMap::get(2, 1, lowerExpr + rewriter.getAffineDimExpr(1) * stepExpr);
- auto actualIndex = affine::AffineApplyOp::create(
- rewriter, op.getLoc(), actualIndexMap,
- ValueRange({lower, affineFor.getInductionVar(), step}));
- return std::make_pair(affineFor, actualIndex.getResult());
+ auto oldIV = affine::AffineApplyOp::create(
+ rewriter, op.getLoc(), AffineMap::get(2, 1, d0 + d1 * s0),
+ ValueRange({lb, affineFor.getInductionVar(), step}));
+
+ return std::make_pair(affineFor, oldIV);
}
void SCFToAffinePass::runOnOperation() {
>From 333ba0476b0c687dd56cf77dfa3bea7930b96b7b Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Wed, 3 Jun 2026 15:10:17 +0000
Subject: [PATCH 12/37] Remove dependent dialect which we actually not depend
on
---
mlir/include/mlir/Conversion/Passes.td | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 4317b80f1533a..478ddaa3af698 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1175,7 +1175,6 @@ def RaiseSCFToAffinePass : Pass<"raise-scf-to-affine"> {
}];
let dependentDialects = [
"affine::AffineDialect",
- "scf::SCFDialect",
];
}
>From 004b211e82d810b8d78ce816c692fa2803dd1b60 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Wed, 3 Jun 2026 15:18:17 +0000
Subject: [PATCH 13/37] Switch to greedy pass driver (anticipating more
patterns)
---
mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp | 12 +++++++++---
mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir | 7 +++----
2 files changed, 12 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 4fa24f77d9bfe..5cbf892e6b9a7 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -16,6 +16,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Verifier.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/Passes.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"
#include "llvm/Support/DebugLog.h"
@@ -168,13 +169,16 @@ ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
return std::make_pair(affineFor, oldIV);
}
+//===----------------------------------------------------------------------===//
+// Pass implementation
+//===----------------------------------------------------------------------===//
+
void SCFToAffinePass::runOnOperation() {
MLIRContext &ctx = getContext();
RewritePatternSet patterns(&ctx);
populateSCFToAffineConversionPatterns(patterns);
- // TODO(rainij): we might need a different rewriter (which tries to converge)
- // if we add more features.
- walkAndApplyPatterns(getOperation(), std::move(patterns));
+
+ (void)applyPatternsGreedily(getOperation(), std::move(patterns));
}
} // namespace
@@ -183,6 +187,8 @@ void SCFToAffinePass::runOnOperation() {
// API
//===----------------------------------------------------------------------===//
+// TODO(rainij): assign right *benefits* to the patterns.
+
void mlir::populateSCFToAffineConversionPatterns(RewritePatternSet &patterns) {
patterns.add<ForOpRewrite>(patterns.getContext());
}
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 22fc61ca6cca2..a3cc00f45d0a6 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -1,11 +1,11 @@
// RUN: mlir-opt -raise-scf-to-affine -split-input-file %s | FileCheck %s
-// CHECK: #[[$UB_MAP:.+]] = affine_map<(d0, d1)[s0] -> ((d1 - d0 + s0 - 1) floordiv s0)>
+// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1, s2] -> ((s0 - s1 + s2 - 1) floordiv s0)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
// CHECK-LABEL: @generic_loop
// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LOWER:.*]]: index, %[[UPPER:.*]]: index, %[[STEP:.*]]: index
func.func @generic_loop(%arr: memref<?xi32>, %lower: index, %upper: index, %step: index) {
-// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]](%[[LOWER]], %[[UPPER]])[%[[STEP]]] {
+// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]]()[%[[STEP]], %[[LOWER]], %[[UPPER]]] {
// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LOWER]], %[[IV]])[%[[STEP]]]
// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IDX]]] : memref<?xi32>
// CHECK: }
@@ -18,11 +18,10 @@ func.func @generic_loop(%arr: memref<?xi32>, %lower: index, %upper: index, %step
// -----
-// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @loop_with_constant_step
// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LOWER:.*]]: index, %[[UPPER:.*]]: index
func.func @loop_with_constant_step(%arr: memref<?xi32>, %lower: index, %upper: index) {
-// CHECK: affine.for %[[IDX:.*]] = #[[$MAP]](%[[LOWER]]) to #[[$MAP]](%[[UPPER]]) step 3 {
+// CHECK: affine.for %[[IDX:.*]] = %[[LOWER]] to %[[UPPER]] step 3 {
// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IDX]]] : memref<?xi32>
// CHECK: }
%c0_i32 = arith.constant 0 : i32
>From 302859aff425829274c441d8e5ba1467fb7394e9 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Wed, 3 Jun 2026 17:48:36 +0000
Subject: [PATCH 14/37] Rewrite pattern for index casts
Co-authored-by: Julian Farnsteiner <jcf96 at proton.me>
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 53 ++++++++++++++++++-
.../Conversion/SCFToAffine/scf-to-affine.mlir | 46 ++++++++++++++++
2 files changed, 98 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 5cbf892e6b9a7..6504c883693fe 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -169,6 +169,57 @@ ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
return std::make_pair(affineFor, oldIV);
}
+//===----------------------------------------------------------------------===//
+// Index casts
+//===----------------------------------------------------------------------===//
+
+/// Cast lb, ub, and iv of `scf.for` ops to `index` type.
+struct ForBoundsIndexCast : public OpRewritePattern<scf::ForOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(scf::ForOp loop,
+ PatternRewriter &rewriter) const override {
+ Value lb = loop.getLowerBound();
+ Value ub = loop.getUpperBound();
+ Value step = loop.getStep();
+ Type originalType = step.getType();
+
+ assert(lb.getType() == originalType && ub.getType() == originalType &&
+ "expected lb, ub, and step to have the same type");
+
+ if (isa<IndexType>(originalType)) {
+ return rewriter.notifyMatchFailure(
+ loop, "bounds and step are already index-typed");
+ }
+
+ auto createIndexCast = [&](Value value, Type targetType) -> Value {
+ Location loc = loop.getLoc();
+ if (loop.getUnsignedCmp()) {
+ return arith::IndexCastUIOp::create(rewriter, loc, targetType, value);
+ }
+ return arith::IndexCastOp::create(rewriter, loc, targetType, value);
+ };
+
+ Value newLb = createIndexCast(lb, rewriter.getIndexType());
+ Value newUb = createIndexCast(ub, rewriter.getIndexType());
+ Value newStep = createIndexCast(step, rewriter.getIndexType());
+
+ loop.setLowerBound(newLb);
+ loop.setUpperBound(newUb);
+ loop.setStep(newStep);
+
+ Value iv = loop.getInductionVar();
+ iv.setType(rewriter.getIndexType()); // TODO(rainij): setType advocates for
+ // not using itself.
+ rewriter.setInsertionPointToStart(loop.getBody());
+ Value castIv = createIndexCast(iv, originalType);
+
+ iv.replaceAllUsesExcept(castIv, castIv.getDefiningOp());
+
+ return success();
+ }
+};
+
//===----------------------------------------------------------------------===//
// Pass implementation
//===----------------------------------------------------------------------===//
@@ -190,5 +241,5 @@ void SCFToAffinePass::runOnOperation() {
// TODO(rainij): assign right *benefits* to the patterns.
void mlir::populateSCFToAffineConversionPatterns(RewritePatternSet &patterns) {
- patterns.add<ForOpRewrite>(patterns.getContext());
+ patterns.add<ForOpRewrite, ForBoundsIndexCast>(patterns.getContext());
}
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index a3cc00f45d0a6..9c4e67bb73967 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -48,3 +48,49 @@ func.func @nested_loop(%arg0: memref<?x?xi32>, %upper1: index, %upper2: index) {
}
return
}
+
+// -----
+
+func.func private @some_func(%arg: i32)
+
+func.func @no_index_type(%lb: i32, %ub: i32) {
+ %step = arith.constant 1 : i32
+ scf.for %i = %lb to %ub step %step : i32 {
+ func.call @some_func(%i) : (i32) -> ()
+ }
+ return
+}
+
+// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1] -> (-s0 + s1)>
+// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
+// CHECK-LABEL: func.func private @some_func(i32)
+
+// CHECK-LABEL: func.func @no_index_type(
+// CHECK-SAME: %[[LB:.*]]: i32,
+// CHECK-SAME: %[[UB:.*]]: i32) {
+// CHECK: %[[STEP:.*]] = arith.constant 1 : index
+// CHECK: %[[LB_1:.*]] = arith.index_cast %[[LB]] : i32 to index
+// CHECK: %[[UB_1:.*]] = arith.index_cast %[[UB]] : i32 to index
+// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]](){{\[}}%[[LB_1]], %[[UB_1]]] {
+// CHECK: %[[IV_OLD:.*]] = affine.apply #[[$IV_MAP]](%[[LB_1]], %[[IV]]){{\[}}%[[STEP]]]
+// CHECK: %[[IV_OLD_1:.*]] = arith.index_cast %[[IV_OLD]] : index to i32
+// CHECK: func.call @some_func(%[[IV_OLD_1]]) : (i32) -> ()
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+// -----
+
+func.func private @some_func(%arg: i32)
+
+// CHECK-LABEL: func.func @no_index_type_unsigned(
+func.func @no_index_type_unsigned(%lb: i32, %ub: i32) {
+// CHECK: %{{.*}} = arith.index_castui
+// CHECK: %{{.*}} = arith.index_castui
+ %step = arith.constant 1 : i32
+ scf.for unsigned %i = %lb to %ub step %step : i32 {
+// CHECK: %{{.*}} = arith.index_castui
+ func.call @some_func(%i) : (i32) -> ()
+ }
+ return
+}
\ No newline at end of file
>From 8d2a0597e36bd49b4a019f72ccae7b070aac914c Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Thu, 4 Jun 2026 13:20:53 +0000
Subject: [PATCH 15/37] Set the right benefit
---
mlir/include/mlir/Conversion/Passes.td | 1 +
mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp | 5 ++---
2 files changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 478ddaa3af698..261d36b82e461 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1153,6 +1153,7 @@ def ReconcileUnrealizedCastsPass : Pass<"reconcile-unrealized-casts"> {
// SCFToAffine
//===----------------------------------------------------------------------===//
+// TODO(rainij): reconsider the pass name. I like it, but it feels inconsistent with -affine-raise-from-memref
def RaiseSCFToAffinePass : Pass<"raise-scf-to-affine"> {
let summary = "Raise SCF operations to affine operations where possible";
let description = [{
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 6504c883693fe..0d20d030f4466 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -238,8 +238,7 @@ void SCFToAffinePass::runOnOperation() {
// API
//===----------------------------------------------------------------------===//
-// TODO(rainij): assign right *benefits* to the patterns.
-
void mlir::populateSCFToAffineConversionPatterns(RewritePatternSet &patterns) {
- patterns.add<ForOpRewrite, ForBoundsIndexCast>(patterns.getContext());
+ patterns.add<ForBoundsIndexCast>(patterns.getContext(), /*benefit=*/2);
+ patterns.add<ForOpRewrite>(patterns.getContext(), /*benefit=*/1);
}
>From a957c2607e89ebc3e0d1444a7f71a5ba004d09ec Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Thu, 4 Jun 2026 13:21:11 +0000
Subject: [PATCH 16/37] test on unsigned loop index
---
mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 9c4e67bb73967..2965fa179fa8d 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -88,7 +88,7 @@ func.func @no_index_type_unsigned(%lb: i32, %ub: i32) {
// CHECK: %{{.*}} = arith.index_castui
// CHECK: %{{.*}} = arith.index_castui
%step = arith.constant 1 : i32
- scf.for unsigned %i = %lb to %ub step %step : i32 {
+ scf.for unsigned %i = %lb to %ub step %step : i32 { // NOTE: "unsigned" matters
// CHECK: %{{.*}} = arith.index_castui
func.call @some_func(%i) : (i32) -> ()
}
>From a6f7fe1f02c57063255c242059c90e1384379997 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Thu, 4 Jun 2026 13:35:06 +0000
Subject: [PATCH 17/37] autogenerated tests to unify the style
---
.../Conversion/SCFToAffine/scf-to-affine.mlir | 101 ++++++++++++------
1 file changed, 68 insertions(+), 33 deletions(-)
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 2965fa179fa8d..452bea9ba3079 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -1,14 +1,22 @@
// RUN: mlir-opt -raise-scf-to-affine -split-input-file %s | FileCheck %s
+// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
+
// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1, s2] -> ((s0 - s1 + s2 - 1) floordiv s0)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
-// CHECK-LABEL: @generic_loop
-// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LOWER:.*]]: index, %[[UPPER:.*]]: index, %[[STEP:.*]]: index
+// CHECK-LABEL: func.func @generic_loop(
+// CHECK-SAME: %[[ARR:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<?xi32>,
+// CHECK-SAME: %[[LOWER:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index,
+// CHECK-SAME: %[[UPPER:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index,
+// CHECK-SAME: %[[STEP:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) {
+// CHECK: %[[C0:.*]] = arith.constant 0 : i32
+// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]](){{\[}}%[[STEP]], %[[LOWER]], %[[UPPER]]] {
+// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LOWER]], %[[IV]]){{\[}}%[[STEP]]]
+// CHECK: memref.store %[[C0]], %[[ARR]]{{\[}}%[[IDX]]] : memref<?xi32>
+// CHECK: }
+// CHECK: return
+// CHECK: }
func.func @generic_loop(%arr: memref<?xi32>, %lower: index, %upper: index, %step: index) {
-// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]]()[%[[STEP]], %[[LOWER]], %[[UPPER]]] {
-// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LOWER]], %[[IV]])[%[[STEP]]]
-// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IDX]]] : memref<?xi32>
-// CHECK: }
%c0_i32 = arith.constant 0 : i32
scf.for %idx = %lower to %upper step %step {
memref.store %c0_i32, %arr[%idx] : memref<?xi32>
@@ -18,12 +26,17 @@ func.func @generic_loop(%arr: memref<?xi32>, %lower: index, %upper: index, %step
// -----
-// CHECK-LABEL: @loop_with_constant_step
-// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LOWER:.*]]: index, %[[UPPER:.*]]: index
+// CHECK-LABEL: func.func @loop_with_constant_step(
+// CHECK-SAME: %[[ARR:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<?xi32>,
+// CHECK-SAME: %[[LOWER:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index,
+// CHECK-SAME: %[[UPPER:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) {
+// CHECK: %[[C0:.*]] = arith.constant 0 : i32
+// CHECK: affine.for %[[IV:.*]] = %[[LOWER]] to %[[UPPER]] step 3 {
+// CHECK: memref.store %[[C0]], %[[ARR]]{{\[}}%[[IV]]] : memref<?xi32>
+// CHECK: }
+// CHECK: return
+// CHECK: }
func.func @loop_with_constant_step(%arr: memref<?xi32>, %lower: index, %upper: index) {
-// CHECK: affine.for %[[IDX:.*]] = %[[LOWER]] to %[[UPPER]] step 3 {
-// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IDX]]] : memref<?xi32>
-// CHECK: }
%c0_i32 = arith.constant 0 : i32
%c3 = arith.constant 3 : index
scf.for %idx = %lower to %upper step %c3 {
@@ -34,10 +47,19 @@ func.func @loop_with_constant_step(%arr: memref<?xi32>, %lower: index, %upper: i
// -----
-// CHECK-LABEL: @nested_loop
+// CHECK-LABEL: func.func @nested_loop(
+// CHECK-SAME: %[[ARR:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<?x?xi32>,
+// CHECK-SAME: %[[UPPER1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index,
+// CHECK-SAME: %[[UPPER2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) {
+// CHECK: %[[C0:.*]] = arith.constant 0 : i32
+// CHECK: affine.for %[[I:.*]] = 0 to %[[UPPER1]] {
+// CHECK: affine.for %[[J:.*]] = 0 to %[[UPPER2]] {
+// CHECK: memref.store %[[C0]], %[[ARR]]{{\[}}%[[I]], %[[J]]] : memref<?x?xi32>
+// CHECK: }
+// CHECK: }
+// CHECK: return
+// CHECK: }
func.func @nested_loop(%arg0: memref<?x?xi32>, %upper1: index, %upper2: index) {
-// CHECK: affine.for
-// CHECK: affine.for
%c0_i32 = arith.constant 0 : i32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
@@ -51,8 +73,25 @@ func.func @nested_loop(%arg0: memref<?x?xi32>, %upper1: index, %upper2: index) {
// -----
+// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1] -> (-s0 + s1)>
+// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
+// CHECK-LABEL: func.func private @some_func(i32)
+
func.func private @some_func(%arg: i32)
+// CHECK-LABEL: func.func @no_index_type(
+// CHECK-SAME: %[[LB:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32,
+// CHECK-SAME: %[[UB:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32) {
+// CHECK: %[[STEP:.*]] = arith.constant 1 : index
+// CHECK: %[[LB_IDX:.*]] = arith.index_cast %[[LB]] : i32 to index
+// CHECK: %[[UB_IDX:.*]] = arith.index_cast %[[UB]] : i32 to index
+// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]](){{\[}}%[[LB_IDX]], %[[UB_IDX]]] {
+// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LB_IDX]], %[[IV]]){{\[}}%[[STEP]]]
+// CHECK: %[[IV_I32:.*]] = arith.index_cast %[[IDX]] : index to i32
+// CHECK: func.call @some_func(%[[IV_I32]]) : (i32) -> ()
+// CHECK: }
+// CHECK: return
+// CHECK: }
func.func @no_index_type(%lb: i32, %ub: i32) {
%step = arith.constant 1 : i32
scf.for %i = %lb to %ub step %step : i32 {
@@ -61,36 +100,32 @@ func.func @no_index_type(%lb: i32, %ub: i32) {
return
}
+// -----
+
// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1] -> (-s0 + s1)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
// CHECK-LABEL: func.func private @some_func(i32)
-// CHECK-LABEL: func.func @no_index_type(
-// CHECK-SAME: %[[LB:.*]]: i32,
-// CHECK-SAME: %[[UB:.*]]: i32) {
+func.func private @some_func(%arg: i32)
+
+// CHECK-LABEL: func.func @no_index_type_unsigned(
+// CHECK-SAME: %[[LB:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32,
+// CHECK-SAME: %[[UB:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32) {
// CHECK: %[[STEP:.*]] = arith.constant 1 : index
-// CHECK: %[[LB_1:.*]] = arith.index_cast %[[LB]] : i32 to index
-// CHECK: %[[UB_1:.*]] = arith.index_cast %[[UB]] : i32 to index
-// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]](){{\[}}%[[LB_1]], %[[UB_1]]] {
-// CHECK: %[[IV_OLD:.*]] = affine.apply #[[$IV_MAP]](%[[LB_1]], %[[IV]]){{\[}}%[[STEP]]]
-// CHECK: %[[IV_OLD_1:.*]] = arith.index_cast %[[IV_OLD]] : index to i32
-// CHECK: func.call @some_func(%[[IV_OLD_1]]) : (i32) -> ()
+// CHECK: %[[LB_IDX:.*]] = arith.index_castui %[[LB]] : i32 to index
+// CHECK: %[[UB_IDX:.*]] = arith.index_castui %[[UB]] : i32 to index
+// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]](){{\[}}%[[LB_IDX]], %[[UB_IDX]]] {
+// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LB_IDX]], %[[IV]]){{\[}}%[[STEP]]]
+// CHECK: %[[IV_I32:.*]] = arith.index_castui %[[IDX]] : index to i32
+// CHECK: func.call @some_func(%[[IV_I32]]) : (i32) -> ()
// CHECK: }
// CHECK: return
// CHECK: }
-
-// -----
-
-func.func private @some_func(%arg: i32)
-
-// CHECK-LABEL: func.func @no_index_type_unsigned(
func.func @no_index_type_unsigned(%lb: i32, %ub: i32) {
-// CHECK: %{{.*}} = arith.index_castui
-// CHECK: %{{.*}} = arith.index_castui
%step = arith.constant 1 : i32
scf.for unsigned %i = %lb to %ub step %step : i32 { // NOTE: "unsigned" matters
-// CHECK: %{{.*}} = arith.index_castui
func.call @some_func(%i) : (i32) -> ()
}
return
-}
\ No newline at end of file
+}
+
>From c0fed3527187a07af7066cf6164fa7b35312f9f8 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Thu, 4 Jun 2026 14:00:32 +0000
Subject: [PATCH 18/37] Shorten the check lines
---
.../Conversion/SCFToAffine/scf-to-affine.mlir | 110 +++++++-----------
1 file changed, 39 insertions(+), 71 deletions(-)
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 452bea9ba3079..861a07f4631c4 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -1,24 +1,15 @@
// RUN: mlir-opt -raise-scf-to-affine -split-input-file %s | FileCheck %s
-// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-
// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1, s2] -> ((s0 - s1 + s2 - 1) floordiv s0)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
-// CHECK-LABEL: func.func @generic_loop(
-// CHECK-SAME: %[[ARR:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<?xi32>,
-// CHECK-SAME: %[[LOWER:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index,
-// CHECK-SAME: %[[UPPER:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index,
-// CHECK-SAME: %[[STEP:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) {
-// CHECK: %[[C0:.*]] = arith.constant 0 : i32
-// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]](){{\[}}%[[STEP]], %[[LOWER]], %[[UPPER]]] {
-// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LOWER]], %[[IV]]){{\[}}%[[STEP]]]
-// CHECK: memref.store %[[C0]], %[[ARR]]{{\[}}%[[IDX]]] : memref<?xi32>
-// CHECK: }
-// CHECK: return
-// CHECK: }
-func.func @generic_loop(%arr: memref<?xi32>, %lower: index, %upper: index, %step: index) {
+// CHECK-LABEL: @generic_loop
+// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
+// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]]()[%[[STEP]], %[[LB]], %[[UB]]] {
+// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LB]], %[[IV]])[%[[STEP]]]
+// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IDX]]]
+func.func @generic_loop(%arr: memref<?xi32>, %lb: index, %ub: index, %step: index) {
%c0_i32 = arith.constant 0 : i32
- scf.for %idx = %lower to %upper step %step {
+ scf.for %idx = %lb to %ub step %step {
memref.store %c0_i32, %arr[%idx] : memref<?xi32>
}
return
@@ -26,20 +17,14 @@ func.func @generic_loop(%arr: memref<?xi32>, %lower: index, %upper: index, %step
// -----
-// CHECK-LABEL: func.func @loop_with_constant_step(
-// CHECK-SAME: %[[ARR:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<?xi32>,
-// CHECK-SAME: %[[LOWER:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index,
-// CHECK-SAME: %[[UPPER:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) {
-// CHECK: %[[C0:.*]] = arith.constant 0 : i32
-// CHECK: affine.for %[[IV:.*]] = %[[LOWER]] to %[[UPPER]] step 3 {
-// CHECK: memref.store %[[C0]], %[[ARR]]{{\[}}%[[IV]]] : memref<?xi32>
-// CHECK: }
-// CHECK: return
-// CHECK: }
-func.func @loop_with_constant_step(%arr: memref<?xi32>, %lower: index, %upper: index) {
+// CHECK-LABEL: @loop_with_constant_step
+// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LB:.*]]: index, %[[UB:.*]]: index
+// CHECK: affine.for %[[IV:.*]] = %[[LB]] to %[[UB]] step 3 {
+// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IV]]]
+func.func @loop_with_constant_step(%arr: memref<?xi32>, %lb: index, %ub: index) {
%c0_i32 = arith.constant 0 : i32
%c3 = arith.constant 3 : index
- scf.for %idx = %lower to %upper step %c3 {
+ scf.for %idx = %lb to %ub step %c3 {
memref.store %c0_i32, %arr[%idx] : memref<?xi32>
}
return
@@ -47,24 +32,17 @@ func.func @loop_with_constant_step(%arr: memref<?xi32>, %lower: index, %upper: i
// -----
-// CHECK-LABEL: func.func @nested_loop(
-// CHECK-SAME: %[[ARR:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<?x?xi32>,
-// CHECK-SAME: %[[UPPER1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index,
-// CHECK-SAME: %[[UPPER2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) {
-// CHECK: %[[C0:.*]] = arith.constant 0 : i32
-// CHECK: affine.for %[[I:.*]] = 0 to %[[UPPER1]] {
-// CHECK: affine.for %[[J:.*]] = 0 to %[[UPPER2]] {
-// CHECK: memref.store %[[C0]], %[[ARR]]{{\[}}%[[I]], %[[J]]] : memref<?x?xi32>
-// CHECK: }
-// CHECK: }
-// CHECK: return
-// CHECK: }
-func.func @nested_loop(%arg0: memref<?x?xi32>, %upper1: index, %upper2: index) {
+// CHECK-LABEL: @nested_loop
+// CHECK-SAME: %[[ARR:.*]]: memref<?x?xi32>, %[[UB1:.*]]: index, %[[UB2:.*]]: index
+// CHECK: affine.for %[[I:.*]] = 0 to %[[UB1]] {
+// CHECK: affine.for %[[J:.*]] = 0 to %[[UB2]] {
+// CHECK: memref.store %{{.*}}, %[[ARR]][%[[I]], %[[J]]]
+func.func @nested_loop(%arg0: memref<?x?xi32>, %ub1: index, %ub2: index) {
%c0_i32 = arith.constant 0 : i32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
- scf.for %i = %c0 to %upper1 step %c1 {
- scf.for %j = %c0 to %upper2 step %c1 {
+ scf.for %i = %c0 to %ub1 step %c1 {
+ scf.for %j = %c0 to %ub2 step %c1 {
memref.store %c0_i32, %arg0[%i, %j] : memref<?x?xi32>
}
}
@@ -75,23 +53,18 @@ func.func @nested_loop(%arg0: memref<?x?xi32>, %upper1: index, %upper2: index) {
// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1] -> (-s0 + s1)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
-// CHECK-LABEL: func.func private @some_func(i32)
+// CHECK-LABEL: @no_index_type
+// CHECK-SAME: %[[LB:.*]]: i32, %[[UB:.*]]: i32
+// CHECK: %[[STEP:.*]] = arith.constant 1 : index
+// CHECK: %[[LB_IDX:.*]] = arith.index_cast %[[LB]] : i32 to index
+// CHECK: %[[UB_IDX:.*]] = arith.index_cast %[[UB]] : i32 to index
+// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]]()[%[[LB_IDX]], %[[UB_IDX]]] {
+// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LB_IDX]], %[[IV]])[%[[STEP]]]
+// CHECK: %[[IV_I32:.*]] = arith.index_cast %[[IDX]] : index to i32
+// CHECK: func.call @some_func(%[[IV_I32]])
func.func private @some_func(%arg: i32)
-// CHECK-LABEL: func.func @no_index_type(
-// CHECK-SAME: %[[LB:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32,
-// CHECK-SAME: %[[UB:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32) {
-// CHECK: %[[STEP:.*]] = arith.constant 1 : index
-// CHECK: %[[LB_IDX:.*]] = arith.index_cast %[[LB]] : i32 to index
-// CHECK: %[[UB_IDX:.*]] = arith.index_cast %[[UB]] : i32 to index
-// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]](){{\[}}%[[LB_IDX]], %[[UB_IDX]]] {
-// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LB_IDX]], %[[IV]]){{\[}}%[[STEP]]]
-// CHECK: %[[IV_I32:.*]] = arith.index_cast %[[IDX]] : index to i32
-// CHECK: func.call @some_func(%[[IV_I32]]) : (i32) -> ()
-// CHECK: }
-// CHECK: return
-// CHECK: }
func.func @no_index_type(%lb: i32, %ub: i32) {
%step = arith.constant 1 : i32
scf.for %i = %lb to %ub step %step : i32 {
@@ -104,23 +77,18 @@ func.func @no_index_type(%lb: i32, %ub: i32) {
// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1] -> (-s0 + s1)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
-// CHECK-LABEL: func.func private @some_func(i32)
+// CHECK-LABEL: @no_index_type_unsigned
+// CHECK-SAME: %[[LB:.*]]: i32, %[[UB:.*]]: i32
+// CHECK: %[[STEP:.*]] = arith.constant 1 : index
+// CHECK: %[[LB_IDX:.*]] = arith.index_castui %[[LB]] : i32 to index
+// CHECK: %[[UB_IDX:.*]] = arith.index_castui %[[UB]] : i32 to index
+// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]]()[%[[LB_IDX]], %[[UB_IDX]]] {
+// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LB_IDX]], %[[IV]])[%[[STEP]]]
+// CHECK: %[[IV_I32:.*]] = arith.index_castui %[[IDX]] : index to i32
+// CHECK: func.call @some_func(%[[IV_I32]])
func.func private @some_func(%arg: i32)
-// CHECK-LABEL: func.func @no_index_type_unsigned(
-// CHECK-SAME: %[[LB:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32,
-// CHECK-SAME: %[[UB:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32) {
-// CHECK: %[[STEP:.*]] = arith.constant 1 : index
-// CHECK: %[[LB_IDX:.*]] = arith.index_castui %[[LB]] : i32 to index
-// CHECK: %[[UB_IDX:.*]] = arith.index_castui %[[UB]] : i32 to index
-// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]](){{\[}}%[[LB_IDX]], %[[UB_IDX]]] {
-// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LB_IDX]], %[[IV]]){{\[}}%[[STEP]]]
-// CHECK: %[[IV_I32:.*]] = arith.index_castui %[[IDX]] : index to i32
-// CHECK: func.call @some_func(%[[IV_I32]]) : (i32) -> ()
-// CHECK: }
-// CHECK: return
-// CHECK: }
func.func @no_index_type_unsigned(%lb: i32, %ub: i32) {
%step = arith.constant 1 : i32
scf.for unsigned %i = %lb to %ub step %step : i32 { // NOTE: "unsigned" matters
>From c8de746d2f962d70cf27029729aed1fa44cc5a81 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Thu, 4 Jun 2026 17:31:15 +0000
Subject: [PATCH 19/37] Non rectangular nests for constants steps
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 44 +++++++++++++++----
.../Conversion/SCFToAffine/scf-to-affine.mlir | 42 ++++++++++++++++--
2 files changed, 73 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 0d20d030f4466..088bd45a608ea 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -16,9 +16,10 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Verifier.h"
+#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/Passes.h"
-#include "mlir/Transforms/WalkPatternRewriteDriver.h"
+#include "llvm/Support/Casting.h"
#include "llvm/Support/DebugLog.h"
namespace mlir {
@@ -59,6 +60,7 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
PatternRewriter &rewriter) const override;
private:
+ /// TODO(rainij): adjust docstring
/// Returns an equivalent `affine.for` skeleton. There are two cases.
///
/// (1) If the step is a constant we trivially raise the `scf.for` by
@@ -84,9 +86,16 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
};
bool canRaiseToAffine(scf::ForOp op) {
- return affine::isValidDim(op.getLowerBound()) &&
- affine::isValidDim(op.getUpperBound()) &&
- affine::isValidSymbol(op.getStep());
+ auto lb = op.getLowerBound();
+ auto ub = op.getUpperBound();
+
+ bool lbOK =
+ affine::isValidDim(lb) || isa<affine::AffineMaxOp>(lb.getDefiningOp());
+ bool ubOK =
+ affine::isValidDim(ub) || isa<affine::AffineMinOp>(ub.getDefiningOp());
+ bool stepOK = affine::isValidSymbol(op.getStep());
+
+ return lbOK && ubOK && stepOK;
}
LogicalResult ForOpRewrite::matchAndRewrite(scf::ForOp op,
@@ -137,15 +146,32 @@ ForOpRewrite::createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
std::pair<affine::AffineForOp, Value>
ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
PatternRewriter &rewriter) const {
- auto affineFor = affine::AffineForOp::create(
- rewriter, op.getLoc(), ValueRange(op.getLowerBound()),
- AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)),
- ValueRange(op.getUpperBound()),
- AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)), step, op.getInits());
+ auto lb = op.getLowerBound();
+ auto ub = op.getUpperBound();
+
+ auto lbOperands = ValueRange(lb);
+ auto ubOperands = ValueRange(ub);
+ auto lbMap = AffineMap::get(1, 0, rewriter.getAffineDimExpr(0));
+ auto ubMap = AffineMap::get(1, 0, rewriter.getAffineDimExpr(0));
+
+ if (auto ubMin = ub.getDefiningOp<affine::AffineMinOp>()) {
+ ubOperands = ubMin->getOperands();
+ ubMap = ubMin.getAffineMap();
+ }
+
+ if (auto lbMax = lb.getDefiningOp<affine::AffineMaxOp>()) {
+ lbOperands = lbMax->getOperands();
+ lbMap = lbMax.getAffineMap();
+ }
+
+ auto affineFor =
+ affine::AffineForOp::create(rewriter, op.getLoc(), lbOperands, lbMap,
+ ubOperands, ubMap, step, op.getInits());
return std::make_pair(affineFor, affineFor.getInductionVar());
}
+// TODO(rainij): add support for max/min ops.
std::pair<affine::AffineForOp, Value>
ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
Value lb = op.getLowerBound();
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 861a07f4631c4..c88e38e0434df 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -2,12 +2,12 @@
// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1, s2] -> ((s0 - s1 + s2 - 1) floordiv s0)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
-// CHECK-LABEL: @generic_loop
+// CHECK-LABEL: @generic_step
// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]]()[%[[STEP]], %[[LB]], %[[UB]]] {
// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LB]], %[[IV]])[%[[STEP]]]
// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IDX]]]
-func.func @generic_loop(%arr: memref<?xi32>, %lb: index, %ub: index, %step: index) {
+func.func @generic_step(%arr: memref<?xi32>, %lb: index, %ub: index, %step: index) {
%c0_i32 = arith.constant 0 : i32
scf.for %idx = %lb to %ub step %step {
memref.store %c0_i32, %arr[%idx] : memref<?xi32>
@@ -17,11 +17,11 @@ func.func @generic_loop(%arr: memref<?xi32>, %lb: index, %ub: index, %step: inde
// -----
-// CHECK-LABEL: @loop_with_constant_step
+// CHECK-LABEL: @constant_step
// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LB:.*]]: index, %[[UB:.*]]: index
// CHECK: affine.for %[[IV:.*]] = %[[LB]] to %[[UB]] step 3 {
// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IV]]]
-func.func @loop_with_constant_step(%arr: memref<?xi32>, %lb: index, %ub: index) {
+func.func @constant_step(%arr: memref<?xi32>, %lb: index, %ub: index) {
%c0_i32 = arith.constant 0 : i32
%c3 = arith.constant 3 : index
scf.for %idx = %lb to %ub step %c3 {
@@ -97,3 +97,37 @@ func.func @no_index_type_unsigned(%lb: i32, %ub: i32) {
return
}
+// -----
+
+// CHECK: #[[$LB_MAP:.+]] = affine_map<(d0) -> (0, -d0 + 3)>
+// CHECK: #[[$UB_MAP:.+]] = affine_map<(d0) -> (3, -d0 + 10)>
+// CHECK-LABEL: func.func @constant_step_non_rectangular_nest
+// CHECK: affine.for %[[I:.*]] = 0 to 10 {
+// CHECK: affine.for %[[J:.*]] = max #[[$LB_MAP]](%[[I]]) to min #[[$UB_MAP]](%[[I]]) {
+// CHECK: func.call @some_func(%[[I]], %[[J]]) : (index, index) -> ()
+
+#lbs = affine_map<(i)[K, N] -> (0, K - i)>
+#ubs = affine_map<(i)[K, N] -> (K, N - i)>
+
+func.func private @some_func(%i: index, %j: index)
+
+func.func @constant_step_non_rectangular_nest() {
+ %zero = arith.constant 0 : index
+ %one = arith.constant 1 : index
+
+ %N = arith.constant 10 : index
+ %K = arith.constant 3 : index
+ %step = arith.constant 1 : index
+
+ scf.for %i = %zero to %N step %one {
+ %lb = affine.max #lbs(%i)[%K, %N] // NOTE: %lb is *not* a dimension.
+ %ub = affine.min #ubs(%i)[%K, %N] // NOTE: %ub is *not* a dimension.
+ scf.for %j = %lb to %ub step %one {
+ func.call @some_func(%i, %j) : (index, index) -> ()
+ }
+ }
+
+ return
+}
+
+// TODO(rainij): similar test with generic step
\ No newline at end of file
>From 881190956e9bd0e9b7f8ff9cf56c51d52e456cb6 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 12:53:41 +0000
Subject: [PATCH 20/37] Allow ub = affine.min with generic step.
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 85 ++++++++++++++++++-
.../Conversion/SCFToAffine/scf-to-affine.mlir | 35 +++++++-
2 files changed, 114 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 088bd45a608ea..62b0dfb504a7a 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -15,11 +15,14 @@
#include "mlir/Conversion/SCFToAffine/SCFToAffine.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Value.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/Passes.h"
-#include "llvm/Support/Casting.h"
+#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/DebugLog.h"
namespace mlir {
@@ -82,13 +85,17 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
PatternRewriter &rewriter) const;
std::pair<affine::AffineForOp, Value>
- caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const;
+ caseGenericStepOld(scf::ForOp op, PatternRewriter &rewriter) const;
+
+ std::pair<affine::AffineForOp, Value>
+ caseGenericStepNew(scf::ForOp op, PatternRewriter &rewriter) const;
};
bool canRaiseToAffine(scf::ForOp op) {
auto lb = op.getLowerBound();
auto ub = op.getUpperBound();
+ // TODO(rainij): lb can be affine.min only if step is constant.
bool lbOK =
affine::isValidDim(lb) || isa<affine::AffineMaxOp>(lb.getDefiningOp());
bool ubOK =
@@ -140,7 +147,7 @@ ForOpRewrite::createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
assert(step > 0 && "scf.for has positive step");
return caseConstantStep(op, step, rewriter);
}
- return caseGenericStep(op, rewriter);
+ return caseGenericStepNew(op, rewriter);
}
std::pair<affine::AffineForOp, Value>
@@ -173,7 +180,8 @@ ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
// TODO(rainij): add support for max/min ops.
std::pair<affine::AffineForOp, Value>
-ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
+ForOpRewrite::caseGenericStepOld(scf::ForOp op,
+ PatternRewriter &rewriter) const {
Value lb = op.getLowerBound();
Value ub = op.getUpperBound();
Value step = op.getStep();
@@ -195,6 +203,75 @@ ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
return std::make_pair(affineFor, oldIV);
}
+std::pair<affine::AffineForOp, Value>
+ForOpRewrite::caseGenericStepNew(scf::ForOp op,
+ PatternRewriter &rewriter) const {
+ Value lb = op.getLowerBound();
+ Value ub = op.getUpperBound();
+ Value step = op.getStep();
+
+ AffineExpr d0 = rewriter.getAffineDimExpr(0);
+ AffineExpr d1 = rewriter.getAffineDimExpr(1);
+ AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
+
+ // NOTE: ubs transformed with floor((x - lb + step - 1) / step)
+ // where x goes over all ub_i. lb is essentially also transformed this way but
+ // this will end up being 0 anyway.
+
+ llvm::SmallVector<Value, 0> lbOperands = {};
+ llvm::SmallVector<Value, 3> ubOperands = {lb, ub, step};
+
+ AffineMap zeroMap = rewriter.getConstantAffineMap(0);
+ AffineMap ubMap = AffineMap::get(2, 1, (d1 - d0 + s0 - 1).floorDiv(s0));
+
+ assert(!mlir::isa_and_present<affine::AffineMaxOp>(lb.getDefiningOp()) &&
+ "not raisible in general if step is generic");
+
+ if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
+ AffineMap origUbMap = ubMinOp.getAffineMap();
+ unsigned ubDims = origUbMap.getNumDims();
+ unsigned ubSyms = origUbMap.getNumSymbols();
+
+ // Combined space: dims = [ub dims, lb]
+ // syms = [ub syms, step]
+ AffineExpr lbDim = rewriter.getAffineDimExpr(ubDims);
+ AffineExpr stepSym = rewriter.getAffineSymbolExpr(ubSyms);
+
+ SmallVector<AffineExpr> ubExprs;
+ ubExprs.reserve(origUbMap.getNumResults());
+ for (AffineExpr ubI : origUbMap.getResults()) {
+ ubExprs.push_back((ubI - lbDim + stepSym - 1).floorDiv(stepSym));
+ }
+
+ ubMap =
+ AffineMap::get(ubDims + 1, ubSyms + 1, ubExprs, rewriter.getContext());
+
+ // Operand order consistent with "combined space" above:
+ ValueRange ubOps = ubMinOp->getOperands();
+ SmallVector<Value> combined;
+ combined.append(ubOps.begin(), ubOps.begin() + ubDims); // ub dims
+ combined.push_back(lb); // lb (single dim)
+ combined.append(ubOps.begin() + ubDims, ubOps.end()); // ub syms
+ combined.push_back(op.getStep()); // step (single sym)
+ ubOperands = std::move(combined);
+ }
+
+ auto affineFor = affine::AffineForOp::create(
+ rewriter, op.getLoc(), {}, zeroMap, ubOperands, ubMap, 1, op.getInits());
+
+ // NOTE: old_iv = old_lb + new_iv * step
+ AffineMap ivMap = AffineMap::get(2, 1, d0 + d1 * s0);
+
+ llvm::SmallVector<Value, 3> ivOperands = {lb, affineFor.getInductionVar(),
+ step};
+
+ rewriter.setInsertionPointToStart(affineFor.getBody());
+ auto oldIV =
+ affine::AffineApplyOp::create(rewriter, op.getLoc(), ivMap, ivOperands);
+
+ return std::make_pair(affineFor, oldIV);
+}
+
//===----------------------------------------------------------------------===//
// Index casts
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index c88e38e0434df..1448be2554bc2 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -117,7 +117,6 @@ func.func @constant_step_non_rectangular_nest() {
%N = arith.constant 10 : index
%K = arith.constant 3 : index
- %step = arith.constant 1 : index
scf.for %i = %zero to %N step %one {
%lb = affine.max #lbs(%i)[%K, %N] // NOTE: %lb is *not* a dimension.
@@ -130,4 +129,36 @@ func.func @constant_step_non_rectangular_nest() {
return
}
-// TODO(rainij): similar test with generic step
\ No newline at end of file
+// -----
+
+// CHECK: #[[$LB_MAP:.+]] = affine_map<(d0)[s0] -> ((s0 + 5) floordiv s0, (-d0 + s0 + 98) floordiv s0)>
+// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
+// CHECK-LABEL: func.func @generic_step_non_rectangular_nest(
+// CHECK-SAME: %[[INNER_STEP:.*]]: index) {
+// CHECK: %[[ONE:.*]] = arith.constant 1 : index
+// CHECK: affine.for %[[I:.*]] = 0 to 100 {
+// CHECK: affine.for %[[J:.*]] = 0 to min #[[$LB_MAP]](%[[I]])[%[[INNER_STEP]]] {
+// CHECK: %[[OLD_IV:.*]] = affine.apply #[[$IV_MAP]](%[[ONE]], %[[J]])[%[[INNER_STEP]]]
+// CHECK: func.call @some_func(%[[I]], %[[OLD_IV]]) : (index, index) -> ()
+
+#ub_map = affine_map<(i)[K, N] -> (K, N - i)>
+
+func.func private @some_func(%i: index, %j: index)
+
+func.func @generic_step_non_rectangular_nest(%inner_step: index) {
+ %zero = arith.constant 0 : index
+ %one = arith.constant 1 : index
+
+ %N = arith.constant 100 : index
+ %K = arith.constant 7 : index
+
+ scf.for %i = %zero to %N step %one {
+ // NOTE: lower bounds cannot be a max in general if step is not constant.
+ %ub = affine.min #ub_map(%i)[%K, %N] // NOTE: %ub is *not* a dimension.
+ scf.for %j = %one to %ub step %inner_step {
+ func.call @some_func(%i, %j) : (index, index) -> ()
+ }
+ }
+
+ return
+}
\ No newline at end of file
>From 4df02ef51eb0b3725ff365bd76496a4c8445b9c3 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 12:55:22 +0000
Subject: [PATCH 21/37] Removed old code
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 35 ++-----------------
1 file changed, 3 insertions(+), 32 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 62b0dfb504a7a..e0cb80e390bc2 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -85,10 +85,7 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
PatternRewriter &rewriter) const;
std::pair<affine::AffineForOp, Value>
- caseGenericStepOld(scf::ForOp op, PatternRewriter &rewriter) const;
-
- std::pair<affine::AffineForOp, Value>
- caseGenericStepNew(scf::ForOp op, PatternRewriter &rewriter) const;
+ caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const;
};
bool canRaiseToAffine(scf::ForOp op) {
@@ -147,7 +144,7 @@ ForOpRewrite::createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
assert(step > 0 && "scf.for has positive step");
return caseConstantStep(op, step, rewriter);
}
- return caseGenericStepNew(op, rewriter);
+ return caseGenericStep(op, rewriter);
}
std::pair<affine::AffineForOp, Value>
@@ -178,34 +175,8 @@ ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
return std::make_pair(affineFor, affineFor.getInductionVar());
}
-// TODO(rainij): add support for max/min ops.
-std::pair<affine::AffineForOp, Value>
-ForOpRewrite::caseGenericStepOld(scf::ForOp op,
- PatternRewriter &rewriter) const {
- Value lb = op.getLowerBound();
- Value ub = op.getUpperBound();
- Value step = op.getStep();
-
- AffineExpr d0 = rewriter.getAffineDimExpr(0);
- AffineExpr d1 = rewriter.getAffineDimExpr(1);
- AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
-
- auto affineFor = affine::AffineForOp::create(
- rewriter, op.getLoc(), ValueRange(), rewriter.getConstantAffineMap(0),
- ValueRange({lb, ub, step}),
- AffineMap::get(2, 1, (d1 - d0 + s0 - 1).floorDiv(s0)), 1, op.getInits());
-
- rewriter.setInsertionPointToStart(affineFor.getBody());
- auto oldIV = affine::AffineApplyOp::create(
- rewriter, op.getLoc(), AffineMap::get(2, 1, d0 + d1 * s0),
- ValueRange({lb, affineFor.getInductionVar(), step}));
-
- return std::make_pair(affineFor, oldIV);
-}
-
std::pair<affine::AffineForOp, Value>
-ForOpRewrite::caseGenericStepNew(scf::ForOp op,
- PatternRewriter &rewriter) const {
+ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
Value lb = op.getLowerBound();
Value ub = op.getUpperBound();
Value step = op.getStep();
>From aec39baa1937b275324fe2953edd4d4e583b33ce Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 15:48:18 +0000
Subject: [PATCH 22/37] Update docstrings
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 59 ++++++++++---------
1 file changed, 32 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index e0cb80e390bc2..8de5ce0357076 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -5,10 +5,8 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
-// TODO(rainij): update description?
//
-// This file implements a pass to raise scf.for, scf.if and loop.terminator
-// ops into affine ops.
+// This file implements a pass to raise scf ops to affine ops.
//
//===----------------------------------------------------------------------===//
@@ -49,13 +47,8 @@ struct SCFToAffinePass
// ForOpRewrite
//===----------------------------------------------------------------------===//
-// TODO(rainij): add some patterns inspired by Enzyme-JAX to raise certain
-// scf.for ops which do not *already* satisfy canRaiseToAffine. If possible do
-// it in a way so that after some rewrite they satisfy it so that the current
-// pattern just applies.
-
-/// Raise an `scf.for` to an equivalent `affine.for` if lb, ub are dimensions
-/// and step is a symbol.
+/// Raise an `scf.for` to an equivalent `affine.for` if lb, ub and step satisfy
+/// certain constraints making this possible.
struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
using OpRewritePattern<scf::ForOp>::OpRewritePattern;
@@ -63,20 +56,25 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
PatternRewriter &rewriter) const override;
private:
- /// TODO(rainij): adjust docstring
- /// Returns an equivalent `affine.for` skeleton. There are two cases.
+ /// Returns an equivalent `affine.for` skeleton and the *old* induction
+ /// variable to be used by the body to be inserted later on. The affine loop
+ /// body is left empty except for an operation computing the old induction
+ /// variable from the new one *iff* it differs from the new one.
+ ///
+ /// Assumes: `canRaiseToAffine(op) == true`.
///
- /// (1) If the step is a constant we trivially raise the `scf.for` by
- /// essentially keeping lb, ub, iv as is. The body is left empty. The second
- /// return value is the induction variable in this case.
+ /// There are two cases:
///
- /// (2) Otherwise (generic step) we normalize the loop by setting step = 1, lb
- /// = 0, ub = ceil((old_ub - old_lb) / old_step). Moreover we insert ops to
- /// compute old_iv = old_lb + old_step * new_iv in the body and return old_iv
- /// as second result. Apart from that the body is empty.
+ /// 1. step is constant
+ /// 2. step is generic (not constant)
///
- /// The resulting `affine.for` is valid (satisfies affine constraints) if lb
- /// and ub of the `scf.for` are dimensions and its step is a symbol.
+ /// In case (1) and if lb, ub are dimensions `scf.for` is trivially raised
+ /// (leaving lb, ub, iv as is). If lb is an `affine.max` we "inline" it into
+ /// the loop's lb attribute position. Similarly if ub is an `affine.min`.
+ ///
+ /// In case (2) we apply *step normalization* and rescale lb and ub so that
+ /// the new lb starts at 0 and the new step is 1. In this case we require that
+ /// lb is a dimension (ub can still be a min).
std::pair<affine::AffineForOp, Value>
createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const;
@@ -84,17 +82,25 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
caseConstantStep(scf::ForOp op, int64_t step,
PatternRewriter &rewriter) const;
+ /// TODO(rainij): rename generic to dynamic?
std::pair<affine::AffineForOp, Value>
caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const;
};
+/// An `scf.for` can trivially be raised if lb, ub are dimensions and step is a
+/// constant. With some more work one can raise under relaxed constraints as
+/// expressed by this function.
bool canRaiseToAffine(scf::ForOp op) {
auto lb = op.getLowerBound();
auto ub = op.getUpperBound();
+ IntegerAttr constAttr;
- // TODO(rainij): lb can be affine.min only if step is constant.
- bool lbOK =
- affine::isValidDim(lb) || isa<affine::AffineMaxOp>(lb.getDefiningOp());
+ // The asymmetry between lb and ub comes from the fact that the step
+ // normalization (for non-constant (generic) steps) does not work with
+ // multiple *lower* bounds (max).
+ bool lbOK = affine::isValidDim(lb) ||
+ (isa<affine::AffineMaxOp>(lb.getDefiningOp()) &&
+ matchPattern(op.getStep(), m_Constant(&constAttr)));
bool ubOK =
affine::isValidDim(ub) || isa<affine::AffineMinOp>(ub.getDefiningOp());
bool stepOK = affine::isValidSymbol(op.getStep());
@@ -105,8 +111,6 @@ bool canRaiseToAffine(scf::ForOp op) {
LogicalResult ForOpRewrite::matchAndRewrite(scf::ForOp op,
PatternRewriter &rewriter) const {
if (!canRaiseToAffine(op)) {
- // TODO(rainij): another pattern might make this raisible. We might want
- // drop this message then, or alter it to acknowledge the possibility.
LDBG() << "[affine] Cannot raise scf op: " << op << "\n";
return failure();
}
@@ -114,6 +118,7 @@ LogicalResult ForOpRewrite::matchAndRewrite(scf::ForOp op,
auto [affineFor, actualIndex] = createAffineFor(op, rewriter);
Block *affineBody = affineFor.getBody();
+ // TODO(rainij): why *exactly* is this necessary?
if (affineBody->mightHaveTerminator()) {
Operation *terminator = affineBody->getTerminator();
assert(isa<affine::AffineYieldOp>(terminator) &&
@@ -230,7 +235,7 @@ ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
auto affineFor = affine::AffineForOp::create(
rewriter, op.getLoc(), {}, zeroMap, ubOperands, ubMap, 1, op.getInits());
- // NOTE: old_iv = old_lb + new_iv * step
+ // old_iv = old_lb + new_iv * step
AffineMap ivMap = AffineMap::get(2, 1, d0 + d1 * s0);
llvm::SmallVector<Value, 3> ivOperands = {lb, affineFor.getInductionVar(),
>From 935acf62cc0b724c0821f86404dfbc7e53f850f2 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 15:48:29 +0000
Subject: [PATCH 23/37] Update pass description
---
mlir/include/mlir/Conversion/Passes.td | 39 +++++++++++++++-----------
1 file changed, 22 insertions(+), 17 deletions(-)
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 261d36b82e461..03cad88781e7c 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1154,25 +1154,30 @@ def ReconcileUnrealizedCastsPass : Pass<"reconcile-unrealized-casts"> {
//===----------------------------------------------------------------------===//
// TODO(rainij): reconsider the pass name. I like it, but it feels inconsistent with -affine-raise-from-memref
+// TODO: extend to scf.if, scf.parallel, possibly more.
def RaiseSCFToAffinePass : Pass<"raise-scf-to-affine"> {
- let summary = "Raise SCF operations to affine operations where possible";
+ let summary = "Raise SCF operations to affine operations (best effort)";
let description = [{
- This pass raises SCF operations to affine operations where possible.
-
- TODO(rainij): document additional features.
-
- Specifically:
- - `scf.for` loops with affine-compatible bounds and steps are
- converted to `affine.for`.
-
- Converting SCF to affine enables affine-specific optimizations such as
- loop tiling, unrolling, vectorization, and memory access analysis.
-
- Note:
- - Only loops that are statically affine can be converted;
- non-affine loops remain in SCF form.
- - This pass does not modify memory accesses; consider using
- --affine-raise-from-memref for converting `memref.load`/`store`.
+ Raising to the affine dialect enables affine analyses and transforms such
+ as loop tiling, unrolling, vectorization, and dependence analysis.
+ Currently only `scf.for` is handled; loops that cannot be raised are left
+ unchanged.
+
+ A `scf.for` is raised when its lower and upper bounds are valid affine
+ dimensions (or an `affine.max` / `affine.min` respectively) and its step
+ is a valid affine symbol. Two cases are handled:
+
+ - Constant step: raised directly, preserving the step and any
+ `affine.min` / `affine.max` bounds as the loop's bounds.
+ - Dynamic (symbolic) step: the loop is normalized to unit step and the
+ original induction variable is reconstructed in the body. An
+ `affine.min` upper bound is supported here; an `affine.max` lower bound
+ is not, as it cannot be raised soundly without a constant step.
+
+ Bounds and step of non-`index` type are cast to `index` first.
+
+ This pass does not modify memory accesses; use `--affine-raise-from-memref`
+ to convert `memref.load` / `memref.store`.
}];
let dependentDialects = [
"affine::AffineDialect",
>From ca70573752c418952bdc20a0d499389a58b56733 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 16:09:43 +0000
Subject: [PATCH 24/37] Refined docstring of createAffineFor
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 20 ++++++++++---------
1 file changed, 11 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 8de5ce0357076..a4df07477982e 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -57,9 +57,9 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
private:
/// Returns an equivalent `affine.for` skeleton and the *old* induction
- /// variable to be used by the body to be inserted later on. The affine loop
- /// body is left empty except for an operation computing the old induction
- /// variable from the new one *iff* it differs from the new one.
+ /// variable for use by the body that is inlined later. The affine loop body
+ /// is left empty except for an operation computing the old induction variable
+ /// from the new one *iff* it differs from the new one.
///
/// Assumes: `canRaiseToAffine(op) == true`.
///
@@ -68,13 +68,15 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
/// 1. step is constant
/// 2. step is generic (not constant)
///
- /// In case (1) and if lb, ub are dimensions `scf.for` is trivially raised
- /// (leaving lb, ub, iv as is). If lb is an `affine.max` we "inline" it into
- /// the loop's lb attribute position. Similarly if ub is an `affine.min`.
+ /// In case (1) and if lb, ub are (valid) dimensions `scf.for` is trivially
+ /// raised (leaving lb, ub, iv as is). If lb is an `affine.max` we "inline" it
+ /// into the loop's lower bound map. Similarly if ub is an `affine.min`.
///
- /// In case (2) we apply *step normalization* and rescale lb and ub so that
- /// the new lb starts at 0 and the new step is 1. In this case we require that
- /// lb is a dimension (ub can still be a min).
+ /// In case (2) we *normalize* the loop to run from 0 with step 1: the new
+ /// upper bound is `ceil((ub - lb) / step)` and the original induction
+ /// variable is recovered in the body as `lb + step * new_iv`. Here we require
+ /// lb to be a dimension; ub may still be an `affine.min`, which is rescaled
+ /// accordingly.
std::pair<affine::AffineForOp, Value>
createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const;
>From 0185b1df720b6201c5edd3d170337ed4389fd443 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 16:26:13 +0000
Subject: [PATCH 25/37] cleanup
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 21 ++++++++++---------
.../Conversion/SCFToAffine/scf-to-affine.mlir | 8 +++----
2 files changed, 15 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index a4df07477982e..c3eb2f30d476b 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -66,7 +66,7 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
/// There are two cases:
///
/// 1. step is constant
- /// 2. step is generic (not constant)
+ /// 2. step is dynamic (not constant)
///
/// In case (1) and if lb, ub are (valid) dimensions `scf.for` is trivially
/// raised (leaving lb, ub, iv as is). If lb is an `affine.max` we "inline" it
@@ -84,9 +84,8 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
caseConstantStep(scf::ForOp op, int64_t step,
PatternRewriter &rewriter) const;
- /// TODO(rainij): rename generic to dynamic?
std::pair<affine::AffineForOp, Value>
- caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const;
+ caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const;
};
/// An `scf.for` can trivially be raised if lb, ub are dimensions and step is a
@@ -98,7 +97,7 @@ bool canRaiseToAffine(scf::ForOp op) {
IntegerAttr constAttr;
// The asymmetry between lb and ub comes from the fact that the step
- // normalization (for non-constant (generic) steps) does not work with
+ // normalization (for non-constant (dynamic) steps) does not work with
// multiple *lower* bounds (max).
bool lbOK = affine::isValidDim(lb) ||
(isa<affine::AffineMaxOp>(lb.getDefiningOp()) &&
@@ -151,7 +150,7 @@ ForOpRewrite::createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
assert(step > 0 && "scf.for has positive step");
return caseConstantStep(op, step, rewriter);
}
- return caseGenericStep(op, rewriter);
+ return caseDynamicStep(op, rewriter);
}
std::pair<affine::AffineForOp, Value>
@@ -183,7 +182,7 @@ ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
}
std::pair<affine::AffineForOp, Value>
-ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
+ForOpRewrite::caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const {
Value lb = op.getLowerBound();
Value ub = op.getUpperBound();
Value step = op.getStep();
@@ -192,9 +191,9 @@ ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
AffineExpr d1 = rewriter.getAffineDimExpr(1);
AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
- // NOTE: ubs transformed with floor((x - lb + step - 1) / step)
- // where x goes over all ub_i. lb is essentially also transformed this way but
- // this will end up being 0 anyway.
+ // ubs transformed with (x - lb + step - 1) floorDiv step where x goes over
+ // all ub_i. lb is essentially also transformed this way but this will end up
+ // being 0 anyway.
llvm::SmallVector<Value, 0> lbOperands = {};
llvm::SmallVector<Value, 3> ubOperands = {lb, ub, step};
@@ -203,7 +202,7 @@ ForOpRewrite::caseGenericStep(scf::ForOp op, PatternRewriter &rewriter) const {
AffineMap ubMap = AffineMap::get(2, 1, (d1 - d0 + s0 - 1).floorDiv(s0));
assert(!mlir::isa_and_present<affine::AffineMaxOp>(lb.getDefiningOp()) &&
- "not raisible in general if step is generic");
+ "not raisible in general if step is not constant");
if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
AffineMap origUbMap = ubMinOp.getAffineMap();
@@ -320,6 +319,8 @@ void SCFToAffinePass::runOnOperation() {
//===----------------------------------------------------------------------===//
void mlir::populateSCFToAffineConversionPatterns(RewritePatternSet &patterns) {
+ // ForBoundsIndexCast runs first (higher benefit): it rewrites non-index
+ // bounds to index so ForOpRewrite can then match them.
patterns.add<ForBoundsIndexCast>(patterns.getContext(), /*benefit=*/2);
patterns.add<ForOpRewrite>(patterns.getContext(), /*benefit=*/1);
}
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 1448be2554bc2..edfb3d632caea 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -2,12 +2,12 @@
// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1, s2] -> ((s0 - s1 + s2 - 1) floordiv s0)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
-// CHECK-LABEL: @generic_step
+// CHECK-LABEL: @dynamic_step
// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
// CHECK: affine.for %[[IV:.*]] = 0 to #[[$UB_MAP]]()[%[[STEP]], %[[LB]], %[[UB]]] {
// CHECK: %[[IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LB]], %[[IV]])[%[[STEP]]]
// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IDX]]]
-func.func @generic_step(%arr: memref<?xi32>, %lb: index, %ub: index, %step: index) {
+func.func @dynamic_step(%arr: memref<?xi32>, %lb: index, %ub: index, %step: index) {
%c0_i32 = arith.constant 0 : i32
scf.for %idx = %lb to %ub step %step {
memref.store %c0_i32, %arr[%idx] : memref<?xi32>
@@ -133,7 +133,7 @@ func.func @constant_step_non_rectangular_nest() {
// CHECK: #[[$LB_MAP:.+]] = affine_map<(d0)[s0] -> ((s0 + 5) floordiv s0, (-d0 + s0 + 98) floordiv s0)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
-// CHECK-LABEL: func.func @generic_step_non_rectangular_nest(
+// CHECK-LABEL: func.func @dynamic_step_non_rectangular_nest(
// CHECK-SAME: %[[INNER_STEP:.*]]: index) {
// CHECK: %[[ONE:.*]] = arith.constant 1 : index
// CHECK: affine.for %[[I:.*]] = 0 to 100 {
@@ -145,7 +145,7 @@ func.func @constant_step_non_rectangular_nest() {
func.func private @some_func(%i: index, %j: index)
-func.func @generic_step_non_rectangular_nest(%inner_step: index) {
+func.func @dynamic_step_non_rectangular_nest(%inner_step: index) {
%zero = arith.constant 0 : index
%one = arith.constant 1 : index
>From 0ae0036dab5454d2c3edeee745a2158fdf1c0fdb Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 16:33:50 +0000
Subject: [PATCH 26/37] Reorder tests a bit
---
.../Conversion/SCFToAffine/scf-to-affine.mlir | 30 +++++++++----------
1 file changed, 15 insertions(+), 15 deletions(-)
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index edfb3d632caea..08acec82c205f 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -1,5 +1,20 @@
// RUN: mlir-opt -raise-scf-to-affine -split-input-file %s | FileCheck %s
+// CHECK-LABEL: @constant_step
+// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LB:.*]]: index, %[[UB:.*]]: index
+// CHECK: affine.for %[[IV:.*]] = %[[LB]] to %[[UB]] step 3 {
+// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IV]]]
+func.func @constant_step(%arr: memref<?xi32>, %lb: index, %ub: index) {
+ %c0_i32 = arith.constant 0 : i32
+ %c3 = arith.constant 3 : index
+ scf.for %idx = %lb to %ub step %c3 {
+ memref.store %c0_i32, %arr[%idx] : memref<?xi32>
+ }
+ return
+}
+
+// -----
+
// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1, s2] -> ((s0 - s1 + s2 - 1) floordiv s0)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
// CHECK-LABEL: @dynamic_step
@@ -17,21 +32,6 @@ func.func @dynamic_step(%arr: memref<?xi32>, %lb: index, %ub: index, %step: inde
// -----
-// CHECK-LABEL: @constant_step
-// CHECK-SAME: %[[ARR:.*]]: memref<?xi32>, %[[LB:.*]]: index, %[[UB:.*]]: index
-// CHECK: affine.for %[[IV:.*]] = %[[LB]] to %[[UB]] step 3 {
-// CHECK: memref.store %{{.*}}, %[[ARR]][%[[IV]]]
-func.func @constant_step(%arr: memref<?xi32>, %lb: index, %ub: index) {
- %c0_i32 = arith.constant 0 : i32
- %c3 = arith.constant 3 : index
- scf.for %idx = %lb to %ub step %c3 {
- memref.store %c0_i32, %arr[%idx] : memref<?xi32>
- }
- return
-}
-
-// -----
-
// CHECK-LABEL: @nested_loop
// CHECK-SAME: %[[ARR:.*]]: memref<?x?xi32>, %[[UB1:.*]]: index, %[[UB2:.*]]: index
// CHECK: affine.for %[[I:.*]] = 0 to %[[UB1]] {
>From 70bb8ba9d325812c1b5874ab93d8128399b17f72 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 16:53:53 +0000
Subject: [PATCH 27/37] cleanup
---
mlir/include/mlir/Conversion/Passes.td | 1 +
mlir/lib/Conversion/SCFToAffine/CMakeLists.txt | 4 +---
mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp | 4 ++--
3 files changed, 4 insertions(+), 5 deletions(-)
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 03cad88781e7c..8848f9b47ffd6 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1181,6 +1181,7 @@ def RaiseSCFToAffinePass : Pass<"raise-scf-to-affine"> {
}];
let dependentDialects = [
"affine::AffineDialect",
+ "arith::ArithDialect",
];
}
diff --git a/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt b/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt
index bf1494d6f3cf0..f978d8309fc39 100644
--- a/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt
+++ b/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt
@@ -8,10 +8,8 @@ add_mlir_conversion_library(MLIRSCFToAffine
MLIRConversionPassIncGen
LINK_LIBS PUBLIC
- MLIRArithDialect
MLIRAffineDialect
- MLIRLLVMDialect
+ MLIRArithDialect
MLIRSCFDialect
- MLIRSCFTransforms
MLIRTransforms
)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index c3eb2f30d476b..272bfdf380e05 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -12,14 +12,14 @@
#include "mlir/Conversion/SCFToAffine/SCFToAffine.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Value.h"
-#include "mlir/IR/Verifier.h"
+#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "mlir/Transforms/Passes.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/DebugLog.h"
>From a99bf704caae6f788f7eaed77929337c9bfffff8 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 17:04:04 +0000
Subject: [PATCH 28/37] Added test with iter args
---
.../Conversion/SCFToAffine/scf-to-affine.mlir | 24 ++++++++++++++++---
1 file changed, 21 insertions(+), 3 deletions(-)
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 08acec82c205f..d0dc3d2a73d01 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -131,13 +131,13 @@ func.func @constant_step_non_rectangular_nest() {
// -----
-// CHECK: #[[$LB_MAP:.+]] = affine_map<(d0)[s0] -> ((s0 + 5) floordiv s0, (-d0 + s0 + 98) floordiv s0)>
+// CHECK: #[[$UB_MAP:.+]] = affine_map<(d0)[s0] -> ((s0 + 5) floordiv s0, (-d0 + s0 + 98) floordiv s0)>
// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
// CHECK-LABEL: func.func @dynamic_step_non_rectangular_nest(
// CHECK-SAME: %[[INNER_STEP:.*]]: index) {
// CHECK: %[[ONE:.*]] = arith.constant 1 : index
// CHECK: affine.for %[[I:.*]] = 0 to 100 {
-// CHECK: affine.for %[[J:.*]] = 0 to min #[[$LB_MAP]](%[[I]])[%[[INNER_STEP]]] {
+// CHECK: affine.for %[[J:.*]] = 0 to min #[[$UB_MAP]](%[[I]])[%[[INNER_STEP]]] {
// CHECK: %[[OLD_IV:.*]] = affine.apply #[[$IV_MAP]](%[[ONE]], %[[J]])[%[[INNER_STEP]]]
// CHECK: func.call @some_func(%[[I]], %[[OLD_IV]]) : (index, index) -> ()
@@ -161,4 +161,22 @@ func.func @dynamic_step_non_rectangular_nest(%inner_step: index) {
}
return
-}
\ No newline at end of file
+}
+
+// -----
+
+// CHECK-LABEL: func.func @with_iter_args(
+// CHECK-SAME: %[[LB:.*]]: index, %[[UB:.*]]: index, %[[INIT:.*]]: f32
+// CHECK: %[[RESULT:.*]] = affine.for %[[I:.*]] = %[[LB]] to %[[UB]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (f32) {
+// CHECK: %[[NEXT_ACC:.*]] = arith.addf %[[ACC]], %[[ACC]] : f32
+// CHECK: affine.yield %[[NEXT_ACC]] : f32
+// CHECK: }
+// CHECK: return %[[RESULT]] : f32
+func.func @with_iter_args(%lb: index, %ub: index, %init: f32) -> f32 {
+ %c1 = arith.constant 1 : index
+ %r = scf.for %i = %lb to %ub step %c1 iter_args(%acc = %init) -> (f32) {
+ %v = arith.addf %acc, %acc : f32
+ scf.yield %v : f32
+ }
+ return %r : f32
+}
>From cedd9fd768da8a440eabf45f216ca1cca819f76e Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 17:46:30 +0000
Subject: [PATCH 29/37] Add max_lb_symbol_dynamic_step
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 4 +--
.../Conversion/SCFToAffine/scf-to-affine.mlir | 25 +++++++++++++++++++
2 files changed, 27 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 272bfdf380e05..22dd894eb7f63 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -201,8 +201,8 @@ ForOpRewrite::caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const {
AffineMap zeroMap = rewriter.getConstantAffineMap(0);
AffineMap ubMap = AffineMap::get(2, 1, (d1 - d0 + s0 - 1).floorDiv(s0));
- assert(!mlir::isa_and_present<affine::AffineMaxOp>(lb.getDefiningOp()) &&
- "not raisible in general if step is not constant");
+ assert(affine::isValidDim(lb) &&
+ "dynamic-step lower bound must be a valid affine dim");
if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
AffineMap origUbMap = ubMinOp.getAffineMap();
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index d0dc3d2a73d01..af70586216c8c 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -180,3 +180,28 @@ func.func @with_iter_args(%lb: index, %ub: index, %init: f32) -> f32 {
}
return %r : f32
}
+
+// -----
+
+// CHECK: #[[$LB_MAP:.+]] = affine_map<()[s0] -> (0, s0)>
+// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1] -> ((s0 - s1 + 99) floordiv s0)>
+// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
+// CHECK-LABEL: func.func @max_lb_symbol_dynamic_step(
+// CHECK-SAME: %[[STEP:.*]]: index, %[[K:.*]]: index) {
+// CHECK: %[[LB:.*]] = affine.max #[[$LB_MAP]]()[%[[K]]]
+// CHECK: affine.for %[[NEW_I:.*]] = 0 to #[[$UB_MAP]]()[%[[STEP]], %[[LB]]] {
+// CHECK: %[[OLD_I:.*]] = affine.apply #[[$IV_MAP]](%[[LB]], %[[NEW_I]])[%[[STEP]]]
+// CHECK: func.call @some_func(%[[OLD_I]]) : (index) -> ()
+
+#lb_map = affine_map<()[K] -> (0, K)>
+
+func.func private @some_func(%i: index)
+
+func.func @max_lb_symbol_dynamic_step(%step: index, %K: index) {
+ %ub = arith.constant 100 : index
+ %lb = affine.max #lb_map()[%K]
+ scf.for %i = %lb to %ub step %step {
+ func.call @some_func(%i) : (index) -> ()
+ }
+ return
+}
\ No newline at end of file
>From 6bb30cf2a07a648815fb0324b36bc3d95685ae16 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Fri, 5 Jun 2026 17:53:52 +0000
Subject: [PATCH 30/37] Added test max_lb_iv_dynamic_step_not_raised
---
.../Conversion/SCFToAffine/scf-to-affine.mlir | 24 +++++++++++++++++++
1 file changed, 24 insertions(+)
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index af70586216c8c..25a14174e7bf8 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -204,4 +204,28 @@ func.func @max_lb_symbol_dynamic_step(%step: index, %K: index) {
func.call @some_func(%i) : (index) -> ()
}
return
+}
+
+// -----
+
+// CHECK-LABEL: @max_lb_iv_dynamic_step_not_raised
+// CHECK: affine.for %{{.*}} = 0 to %{{.*}} {
+// CHECK: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK-NOT: affine.for
+
+#lb_map = affine_map<(i)[K] -> (0, K - i)>
+
+func.func private @some_func(%i: index, %j: index)
+
+func.func @max_lb_iv_dynamic_step_not_raised(%n: index, %ub: index, %k: index,
+ %step: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.for %i = %c0 to %n step %c1 {
+ %lb = affine.max #lb_map(%i)[%k] // this lb prevents the inner scf.for from raising
+ scf.for %j = %lb to %ub step %step {
+ func.call @some_func(%i, %j) : (index, index) -> ()
+ }
+ }
+ return
}
\ No newline at end of file
>From af54fdcb95b4a169c61cf9c0fb93da6f980d4f35 Mon Sep 17 00:00:00 2001
From: Julian Farnsteiner <jcf96 at proton.me>
Date: Sun, 7 Jun 2026 19:02:00 +0000
Subject: [PATCH 31/37] simplify terminator check
---
mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp | 12 +++++-------
1 file changed, 5 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 22dd894eb7f63..997a24623be41 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -119,13 +119,11 @@ LogicalResult ForOpRewrite::matchAndRewrite(scf::ForOp op,
auto [affineFor, actualIndex] = createAffineFor(op, rewriter);
Block *affineBody = affineFor.getBody();
- // TODO(rainij): why *exactly* is this necessary?
- if (affineBody->mightHaveTerminator()) {
- Operation *terminator = affineBody->getTerminator();
- assert(isa<affine::AffineYieldOp>(terminator) &&
- "expected affine.yield if there *might* be terminator");
- rewriter.eraseOp(terminator);
- }
+ // TODO(julfarn): We control the creation of the affine loop, can we guarantee
+ // the absence of a terminator and skip this step?
+ if (!affineBody->empty())
+ if (auto terminator = dyn_cast<affine::AffineYieldOp>(affineBody->back()))
+ rewriter.eraseOp(terminator);
SmallVector<Value> argValues;
argValues.push_back(actualIndex);
>From bd0be787b16ffb255b341a679c89a4c783668f84 Mon Sep 17 00:00:00 2001
From: Julian Farnsteiner <jcf96 at proton.me>
Date: Mon, 8 Jun 2026 13:26:09 +0000
Subject: [PATCH 32/37] legibility, avoid setType
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 84 ++++++++++---------
1 file changed, 44 insertions(+), 40 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 997a24623be41..86ca746a6dad0 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -159,8 +159,9 @@ ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
auto lbOperands = ValueRange(lb);
auto ubOperands = ValueRange(ub);
- auto lbMap = AffineMap::get(1, 0, rewriter.getAffineDimExpr(0));
- auto ubMap = AffineMap::get(1, 0, rewriter.getAffineDimExpr(0));
+
+ auto lbMap = AffineMap::getMultiDimIdentityMap(1, rewriter.getContext());
+ auto ubMap = AffineMap::getMultiDimIdentityMap(1, rewriter.getContext());
if (auto ubMin = ub.getDefiningOp<affine::AffineMinOp>()) {
ubOperands = ubMin->getOperands();
@@ -185,50 +186,52 @@ ForOpRewrite::caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const {
Value ub = op.getUpperBound();
Value step = op.getStep();
+ assert(affine::isValidDim(lb) &&
+ "dynamic-step lower bound must be a valid affine dim");
+
AffineExpr d0 = rewriter.getAffineDimExpr(0);
AffineExpr d1 = rewriter.getAffineDimExpr(1);
AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
+ AffineMap zeroMap = rewriter.getConstantAffineMap(0);
- // ubs transformed with (x - lb + step - 1) floorDiv step where x goes over
- // all ub_i. lb is essentially also transformed this way but this will end up
- // being 0 anyway.
-
- llvm::SmallVector<Value, 0> lbOperands = {};
llvm::SmallVector<Value, 3> ubOperands = {lb, ub, step};
- AffineMap zeroMap = rewriter.getConstantAffineMap(0);
- AffineMap ubMap = AffineMap::get(2, 1, (d1 - d0 + s0 - 1).floorDiv(s0));
+ // ub is transformed with (x - lb + step - 1) floorDiv step where x ranges
+ // over all ub_i.
+ // lb is transformed to zero.
- assert(affine::isValidDim(lb) &&
- "dynamic-step lower bound must be a valid affine dim");
+ AffineMap ubMap;
+ {
+ if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
+ AffineMap origUbMap = ubMinOp.getAffineMap();
+ unsigned ubDims = origUbMap.getNumDims();
+ unsigned ubSyms = origUbMap.getNumSymbols();
- if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
- AffineMap origUbMap = ubMinOp.getAffineMap();
- unsigned ubDims = origUbMap.getNumDims();
- unsigned ubSyms = origUbMap.getNumSymbols();
+ AffineExpr lbDim = rewriter.getAffineDimExpr(ubDims);
+ AffineExpr stepSym = rewriter.getAffineSymbolExpr(ubSyms);
- // Combined space: dims = [ub dims, lb]
- // syms = [ub syms, step]
- AffineExpr lbDim = rewriter.getAffineDimExpr(ubDims);
- AffineExpr stepSym = rewriter.getAffineSymbolExpr(ubSyms);
+ SmallVector<AffineExpr> ubExprs;
+ ubExprs.reserve(origUbMap.getNumResults());
+ for (AffineExpr ubI : origUbMap.getResults()) {
+ ubExprs.push_back((ubI - lbDim + stepSym - 1).floorDiv(stepSym));
+ }
- SmallVector<AffineExpr> ubExprs;
- ubExprs.reserve(origUbMap.getNumResults());
- for (AffineExpr ubI : origUbMap.getResults()) {
- ubExprs.push_back((ubI - lbDim + stepSym - 1).floorDiv(stepSym));
+ // Combined space: dims = [ub dims, lb]
+ // syms = [ub syms, step]
+ ubMap = AffineMap::get(ubDims + 1, ubSyms + 1, ubExprs,
+ rewriter.getContext());
+
+ // Operand order consistent with "combined space" above:
+ ValueRange ubOps = ubMinOp->getOperands();
+ SmallVector<Value> combined;
+ combined.append(ubOps.begin(), ubOps.begin() + ubDims); // ub dims
+ combined.push_back(lb); // lb (single dim)
+ combined.append(ubOps.begin() + ubDims, ubOps.end()); // ub syms
+ combined.push_back(op.getStep()); // step (single sym)
+ ubOperands = std::move(combined);
+ } else {
+ ubMap = AffineMap::get(2, 1, (d1 - d0 + s0 - 1).floorDiv(s0));
}
-
- ubMap =
- AffineMap::get(ubDims + 1, ubSyms + 1, ubExprs, rewriter.getContext());
-
- // Operand order consistent with "combined space" above:
- ValueRange ubOps = ubMinOp->getOperands();
- SmallVector<Value> combined;
- combined.append(ubOps.begin(), ubOps.begin() + ubDims); // ub dims
- combined.push_back(lb); // lb (single dim)
- combined.append(ubOps.begin() + ubDims, ubOps.end()); // ub syms
- combined.push_back(op.getStep()); // step (single sym)
- ubOperands = std::move(combined);
}
auto affineFor = affine::AffineForOp::create(
@@ -286,13 +289,14 @@ struct ForBoundsIndexCast : public OpRewritePattern<scf::ForOp> {
loop.setUpperBound(newUb);
loop.setStep(newStep);
- Value iv = loop.getInductionVar();
- iv.setType(rewriter.getIndexType()); // TODO(rainij): setType advocates for
- // not using itself.
+ Value originalIV = loop.getInductionVar();
+ Value iv = loop.getBody()->insertArgument(
+ (unsigned)0, rewriter.getIndexType(), loop.getLoc());
rewriter.setInsertionPointToStart(loop.getBody());
Value castIv = createIndexCast(iv, originalType);
-
- iv.replaceAllUsesExcept(castIv, castIv.getDefiningOp());
+ originalIV.replaceAllUsesWith(castIv);
+ loop.getBody()->eraseArgument(
+ 1); // Original induction var is now at index 1.
return success();
}
>From 1dd3f3b622c7e50edd9e8504516b54dc23c138e6 Mon Sep 17 00:00:00 2001
From: Julian Farnsteiner <jcf96 at proton.me>
Date: Mon, 8 Jun 2026 13:27:05 +0000
Subject: [PATCH 33/37] Revert "simplify terminator check"
This reverts commit af54fdcb95b4a169c61cf9c0fb93da6f980d4f35.
---
mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp | 12 +++++++-----
1 file changed, 7 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 86ca746a6dad0..234d812b3c851 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -119,11 +119,13 @@ LogicalResult ForOpRewrite::matchAndRewrite(scf::ForOp op,
auto [affineFor, actualIndex] = createAffineFor(op, rewriter);
Block *affineBody = affineFor.getBody();
- // TODO(julfarn): We control the creation of the affine loop, can we guarantee
- // the absence of a terminator and skip this step?
- if (!affineBody->empty())
- if (auto terminator = dyn_cast<affine::AffineYieldOp>(affineBody->back()))
- rewriter.eraseOp(terminator);
+ // TODO(rainij): why *exactly* is this necessary?
+ if (affineBody->mightHaveTerminator()) {
+ Operation *terminator = affineBody->getTerminator();
+ assert(isa<affine::AffineYieldOp>(terminator) &&
+ "expected affine.yield if there *might* be terminator");
+ rewriter.eraseOp(terminator);
+ }
SmallVector<Value> argValues;
argValues.push_back(actualIndex);
>From 56c59981a38f84e5acbd16db620a154772c80d7b Mon Sep 17 00:00:00 2001
From: Julian Farnsteiner <jcf96 at proton.me>
Date: Mon, 8 Jun 2026 13:33:39 +0000
Subject: [PATCH 34/37] clarifying comment for mightHaveTerminator
---
mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 234d812b3c851..b988691130269 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -119,8 +119,8 @@ LogicalResult ForOpRewrite::matchAndRewrite(scf::ForOp op,
auto [affineFor, actualIndex] = createAffineFor(op, rewriter);
Block *affineBody = affineFor.getBody();
- // TODO(rainij): why *exactly* is this necessary?
if (affineBody->mightHaveTerminator()) {
+ // No unregistered ops in the body, so this is definitive.
Operation *terminator = affineBody->getTerminator();
assert(isa<affine::AffineYieldOp>(terminator) &&
"expected affine.yield if there *might* be terminator");
>From 5c93874c6f086374d15b1d6ab11acfd3ef44f540 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Tue, 9 Jun 2026 08:08:56 +0000
Subject: [PATCH 35/37] cleanup
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 96 +++++++++----------
1 file changed, 47 insertions(+), 49 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index b988691130269..a276e20eef39d 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -165,14 +165,14 @@ ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
auto lbMap = AffineMap::getMultiDimIdentityMap(1, rewriter.getContext());
auto ubMap = AffineMap::getMultiDimIdentityMap(1, rewriter.getContext());
- if (auto ubMin = ub.getDefiningOp<affine::AffineMinOp>()) {
- ubOperands = ubMin->getOperands();
- ubMap = ubMin.getAffineMap();
+ if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
+ ubOperands = ubMinOp->getOperands();
+ ubMap = ubMinOp.getAffineMap();
}
- if (auto lbMax = lb.getDefiningOp<affine::AffineMaxOp>()) {
- lbOperands = lbMax->getOperands();
- lbMap = lbMax.getAffineMap();
+ if (auto lbMaxOp = lb.getDefiningOp<affine::AffineMaxOp>()) {
+ lbOperands = lbMaxOp->getOperands();
+ lbMap = lbMaxOp.getAffineMap();
}
auto affineFor =
@@ -199,41 +199,37 @@ ForOpRewrite::caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const {
llvm::SmallVector<Value, 3> ubOperands = {lb, ub, step};
// ub is transformed with (x - lb + step - 1) floorDiv step where x ranges
- // over all ub_i.
- // lb is transformed to zero.
-
- AffineMap ubMap;
- {
- if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
- AffineMap origUbMap = ubMinOp.getAffineMap();
- unsigned ubDims = origUbMap.getNumDims();
- unsigned ubSyms = origUbMap.getNumSymbols();
-
- AffineExpr lbDim = rewriter.getAffineDimExpr(ubDims);
- AffineExpr stepSym = rewriter.getAffineSymbolExpr(ubSyms);
-
- SmallVector<AffineExpr> ubExprs;
- ubExprs.reserve(origUbMap.getNumResults());
- for (AffineExpr ubI : origUbMap.getResults()) {
- ubExprs.push_back((ubI - lbDim + stepSym - 1).floorDiv(stepSym));
- }
+ // over all ub_i. lb is transformed to zero.
+
+ AffineMap ubMap = AffineMap::get(2, 1, (d1 - d0 + s0 - 1).floorDiv(s0));
+
+ if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
+ AffineMap origUbMap = ubMinOp.getAffineMap();
+ unsigned ubDims = origUbMap.getNumDims();
+ unsigned ubSyms = origUbMap.getNumSymbols();
+
+ AffineExpr lbDim = rewriter.getAffineDimExpr(ubDims);
+ AffineExpr stepSym = rewriter.getAffineSymbolExpr(ubSyms);
- // Combined space: dims = [ub dims, lb]
- // syms = [ub syms, step]
- ubMap = AffineMap::get(ubDims + 1, ubSyms + 1, ubExprs,
- rewriter.getContext());
-
- // Operand order consistent with "combined space" above:
- ValueRange ubOps = ubMinOp->getOperands();
- SmallVector<Value> combined;
- combined.append(ubOps.begin(), ubOps.begin() + ubDims); // ub dims
- combined.push_back(lb); // lb (single dim)
- combined.append(ubOps.begin() + ubDims, ubOps.end()); // ub syms
- combined.push_back(op.getStep()); // step (single sym)
- ubOperands = std::move(combined);
- } else {
- ubMap = AffineMap::get(2, 1, (d1 - d0 + s0 - 1).floorDiv(s0));
+ SmallVector<AffineExpr> ubExprs;
+ ubExprs.reserve(origUbMap.getNumResults());
+ for (AffineExpr ubI : origUbMap.getResults()) {
+ ubExprs.push_back((ubI - lbDim + stepSym - 1).floorDiv(stepSym));
}
+
+ // Combined space: dims = [ub dims, lb]
+ // syms = [ub syms, step]
+ ubMap =
+ AffineMap::get(ubDims + 1, ubSyms + 1, ubExprs, rewriter.getContext());
+
+ // Operand order consistent with "combined space" above:
+ ValueRange ubOps = ubMinOp->getOperands();
+ SmallVector<Value> combined;
+ combined.append(ubOps.begin(), ubOps.begin() + ubDims); // ub dims
+ combined.push_back(lb); // lb (single dim)
+ combined.append(ubOps.begin() + ubDims, ubOps.end()); // ub syms
+ combined.push_back(op.getStep()); // step (single sym)
+ ubOperands = std::move(combined);
}
auto affineFor = affine::AffineForOp::create(
@@ -275,17 +271,17 @@ struct ForBoundsIndexCast : public OpRewritePattern<scf::ForOp> {
loop, "bounds and step are already index-typed");
}
- auto createIndexCast = [&](Value value, Type targetType) -> Value {
+ auto createIndexCast = [&](Type out, Value in) -> Value {
Location loc = loop.getLoc();
if (loop.getUnsignedCmp()) {
- return arith::IndexCastUIOp::create(rewriter, loc, targetType, value);
+ return arith::IndexCastUIOp::create(rewriter, loc, out, in);
}
- return arith::IndexCastOp::create(rewriter, loc, targetType, value);
+ return arith::IndexCastOp::create(rewriter, loc, out, in);
};
- Value newLb = createIndexCast(lb, rewriter.getIndexType());
- Value newUb = createIndexCast(ub, rewriter.getIndexType());
- Value newStep = createIndexCast(step, rewriter.getIndexType());
+ Value newLb = createIndexCast(rewriter.getIndexType(), lb);
+ Value newUb = createIndexCast(rewriter.getIndexType(), ub);
+ Value newStep = createIndexCast(rewriter.getIndexType(), step);
loop.setLowerBound(newLb);
loop.setUpperBound(newUb);
@@ -294,11 +290,13 @@ struct ForBoundsIndexCast : public OpRewritePattern<scf::ForOp> {
Value originalIV = loop.getInductionVar();
Value iv = loop.getBody()->insertArgument(
(unsigned)0, rewriter.getIndexType(), loop.getLoc());
+
rewriter.setInsertionPointToStart(loop.getBody());
- Value castIv = createIndexCast(iv, originalType);
- originalIV.replaceAllUsesWith(castIv);
- loop.getBody()->eraseArgument(
- 1); // Original induction var is now at index 1.
+ Value castIV = createIndexCast(originalType, iv);
+ originalIV.replaceAllUsesWith(castIV);
+
+ // Original induction var is now at index 1.
+ loop.getBody()->eraseArgument(1);
return success();
}
>From c0eaa17e079362fa1dd7de89ba0426721a0f1879 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Tue, 9 Jun 2026 08:58:21 +0000
Subject: [PATCH 36/37] cleanup
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 30 +++++++++---------
.../Conversion/SCFToAffine/scf-to-affine.mlir | 31 +++++++++++++++++--
2 files changed, 45 insertions(+), 16 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index a276e20eef39d..f8be4be3011cd 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -283,20 +283,22 @@ struct ForBoundsIndexCast : public OpRewritePattern<scf::ForOp> {
Value newUb = createIndexCast(rewriter.getIndexType(), ub);
Value newStep = createIndexCast(rewriter.getIndexType(), step);
- loop.setLowerBound(newLb);
- loop.setUpperBound(newUb);
- loop.setStep(newStep);
-
- Value originalIV = loop.getInductionVar();
- Value iv = loop.getBody()->insertArgument(
- (unsigned)0, rewriter.getIndexType(), loop.getLoc());
-
- rewriter.setInsertionPointToStart(loop.getBody());
- Value castIV = createIndexCast(originalType, iv);
- originalIV.replaceAllUsesWith(castIV);
-
- // Original induction var is now at index 1.
- loop.getBody()->eraseArgument(1);
+ rewriter.modifyOpInPlace(loop, [&] {
+ loop.setLowerBound(newLb);
+ loop.setUpperBound(newUb);
+ loop.setStep(newStep);
+
+ Value originalIV = loop.getInductionVar();
+ Value iv = loop.getBody()->insertArgument(
+ (unsigned)0, rewriter.getIndexType(), loop.getLoc());
+
+ rewriter.setInsertionPointToStart(loop.getBody());
+ Value castIV = createIndexCast(originalType, iv);
+ rewriter.replaceAllUsesWith(originalIV, castIV);
+
+ // Original induction var is now at index 1.
+ loop.getBody()->eraseArgument(1);
+ });
return success();
}
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 25a14174e7bf8..5740f0628a80d 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -165,14 +165,14 @@ func.func @dynamic_step_non_rectangular_nest(%inner_step: index) {
// -----
-// CHECK-LABEL: func.func @with_iter_args(
+// CHECK-LABEL: func.func @with_iter_args_simple(
// CHECK-SAME: %[[LB:.*]]: index, %[[UB:.*]]: index, %[[INIT:.*]]: f32
// CHECK: %[[RESULT:.*]] = affine.for %[[I:.*]] = %[[LB]] to %[[UB]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (f32) {
// CHECK: %[[NEXT_ACC:.*]] = arith.addf %[[ACC]], %[[ACC]] : f32
// CHECK: affine.yield %[[NEXT_ACC]] : f32
// CHECK: }
// CHECK: return %[[RESULT]] : f32
-func.func @with_iter_args(%lb: index, %ub: index, %init: f32) -> f32 {
+func.func @with_iter_args_simple(%lb: index, %ub: index, %init: f32) -> f32 {
%c1 = arith.constant 1 : index
%r = scf.for %i = %lb to %ub step %c1 iter_args(%acc = %init) -> (f32) {
%v = arith.addf %acc, %acc : f32
@@ -228,4 +228,31 @@ func.func @max_lb_iv_dynamic_step_not_raised(%n: index, %ub: index, %k: index,
}
}
return
+}
+
+// -----
+
+// CHECK: #[[$UB_MAP:.+]] = affine_map<()[s0, s1, s2] -> ((s0 - s1 + s2 - 1) floordiv s0)>
+// CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
+// CHECK-LABEL: func.func @index_cast_with_iter_args(
+// CHECK-SAME: %[[LB:.*]]: i32, %[[UB:.*]]: i32, %[[STEP:.*]]: i32, %[[INIT:.*]]: f32) -> f32 {
+// CHECK: %[[LB_IDX:.*]] = arith.index_cast %[[LB]] : i32 to index
+// CHECK: %[[UB_IDX:.*]] = arith.index_cast %[[UB]] : i32 to index
+// CHECK: %[[STEP_IDX:.*]] = arith.index_cast %[[STEP]] : i32 to index
+// CHECK: %[[RESULT:.*]] = affine.for %[[IV_NEW:.*]] = 0 to #[[$UB_MAP]]()[%[[STEP_IDX]], %[[LB_IDX]], %[[UB_IDX]]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (f32) {
+// CHECK: %[[IV_OLD_IDX:.*]] = affine.apply #[[$IV_MAP]](%[[LB_IDX]], %[[IV_NEW]])[%[[STEP_IDX]]]
+// CHECK: %[[IV_OLD:.*]] = arith.index_cast %[[IV_OLD_IDX]] : index to i32
+// CHECK: %[[VAL:.*]] = arith.sitofp %[[IV_OLD]] : i32 to f32
+// CHECK: %[[ACC_NEXT:.*]] = arith.addf %[[ACC]], %[[VAL]] : f32
+// CHECK: affine.yield %[[ACC_NEXT]] : f32
+// CHECK: }
+// CHECK: return %[[RESULT]] : f32
+
+func.func @index_cast_with_iter_args(%lb: i32, %ub: i32, %step: i32, %init: f32) -> f32 {
+ %r = scf.for %i = %lb to %ub step %step iter_args(%acc = %init) -> (f32) : i32 {
+ %f = arith.sitofp %i : i32 to f32
+ %v = arith.addf %acc, %f : f32
+ scf.yield %v : f32
+ }
+ return %r : f32
}
\ No newline at end of file
>From 1f1d833a2b9eda2cf6f46f069d8aee36b9622e15 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Tue, 9 Jun 2026 10:03:02 +0000
Subject: [PATCH 37/37] Fix bug related to truncating index cast
---
mlir/include/mlir/Conversion/Passes.td | 3 ++-
.../Conversion/SCFToAffine/SCFToAffine.cpp | 20 +++++++++++++++----
.../Conversion/SCFToAffine/scf-to-affine.mlir | 18 +++++++++++++++++
3 files changed, 36 insertions(+), 5 deletions(-)
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 8848f9b47ffd6..685744bd84150 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1174,7 +1174,8 @@ def RaiseSCFToAffinePass : Pass<"raise-scf-to-affine"> {
`affine.min` upper bound is supported here; an `affine.max` lower bound
is not, as it cannot be raised soundly without a constant step.
- Bounds and step of non-`index` type are cast to `index` first.
+ Bounds and step of non-`index` type are cast to `index` first if the source
+ type is not wider than `index`.
This pass does not modify memory accesses; use `--affine-raise-from-memref`
to convert `memref.load` / `memref.store`.
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index f8be4be3011cd..e07e43234d616 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -17,6 +17,7 @@
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -100,10 +101,10 @@ bool canRaiseToAffine(scf::ForOp op) {
// normalization (for non-constant (dynamic) steps) does not work with
// multiple *lower* bounds (max).
bool lbOK = affine::isValidDim(lb) ||
- (isa<affine::AffineMaxOp>(lb.getDefiningOp()) &&
+ (isa_and_present<affine::AffineMaxOp>(lb.getDefiningOp()) &&
matchPattern(op.getStep(), m_Constant(&constAttr)));
- bool ubOK =
- affine::isValidDim(ub) || isa<affine::AffineMinOp>(ub.getDefiningOp());
+ bool ubOK = affine::isValidDim(ub) ||
+ isa_and_present<affine::AffineMinOp>(ub.getDefiningOp());
bool stepOK = affine::isValidSymbol(op.getStep());
return lbOK && ubOK && stepOK;
@@ -252,7 +253,8 @@ ForOpRewrite::caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const {
// Index casts
//===----------------------------------------------------------------------===//
-/// Cast lb, ub, and iv of `scf.for` ops to `index` type.
+/// Cast lb, ub, and iv of `scf.for` ops to `index` type. Allow widening
+/// (semantic preserving) but dissallow truncation (not semantic preserving).
struct ForBoundsIndexCast : public OpRewritePattern<scf::ForOp> {
using OpRewritePattern::OpRewritePattern;
@@ -271,6 +273,16 @@ struct ForBoundsIndexCast : public OpRewritePattern<scf::ForOp> {
loop, "bounds and step are already index-typed");
}
+ if (auto intType = dyn_cast<IntegerType>(originalType)) {
+ llvm::TypeSize indexWidth =
+ DataLayout::closest(loop).getTypeSizeInBits(rewriter.getIndexType());
+ if (intType.getWidth() > indexWidth.getFixedValue()) {
+ return rewriter.notifyMatchFailure(
+ loop, "source integer type is wider than index; cast would "
+ "truncate the bounds");
+ }
+ }
+
auto createIndexCast = [&](Type out, Value in) -> Value {
Location loc = loop.getLoc();
if (loop.getUnsignedCmp()) {
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 5740f0628a80d..70fa16e702345 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -255,4 +255,22 @@ func.func @index_cast_with_iter_args(%lb: i32, %ub: i32, %step: i32, %init: f32)
scf.yield %v : f32
}
return %r : f32
+}
+
+// -----
+
+// CHECK-LABEL: @wider_than_index_not_raised
+// CHECK: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} : i64
+// CHECK-NOT: affine.for
+
+// Use dlti dialect to pin width(index) == 32.
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<index, 32>>} {
+ func.func private @some_func(%arg: i64)
+
+ func.func @wider_than_index_not_raised(%lb: i64, %ub: i64, %step: i64) {
+ scf.for %i = %lb to %ub step %step : i64 {
+ func.call @some_func(%i) : (i64) -> ()
+ }
+ return
+ }
}
\ No newline at end of file
More information about the Mlir-commits
mailing list