[Mlir-commits] [mlir] [mlir][SCFToAffine] Raise scf.for to affine.for (PR #200851)
Reinhard Stahn
llvmlistbot at llvm.org
Wed Jun 17 08:38:09 PDT 2026
https://github.com/rainij updated https://github.com/llvm/llvm-project/pull/200851
>From 94f8daaed3eaeff3b31a997bc5c1fa919cba8b54 Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Wed, 10 Jun 2026 11:29:54 +0000
Subject: [PATCH 1/2] [mlir][SCFToAffine] Raise scf.for to affine.for
Add a pass `-raise-scf-to-affine` that rewrites `scf.for` into
`affine.for` when the bounds and step are valid affine quantities. It
handles constant and dynamic steps, and integer-typed loops (by a
lossless cast of the bounds to `index`).
This is a first step; raising further scf ops to affine will follow.
Co-authored-by: Ming Yan <nexming7 at gmail.com>
Co-authored-by: Julian Farnsteiner <jcf96 at proton.me>
Assisted-by: Claude Code (Anthropic)
---
mlir/include/mlir/Conversion/Passes.h | 1 +
mlir/include/mlir/Conversion/Passes.td | 37 ++
.../mlir/Conversion/SCFToAffine/SCFToAffine.h | 26 ++
mlir/lib/Conversion/CMakeLists.txt | 1 +
.../lib/Conversion/SCFToAffine/CMakeLists.txt | 15 +
.../Conversion/SCFToAffine/SCFToAffine.cpp | 371 ++++++++++++++++++
.../Conversion/SCFToAffine/scf-to-affine.mlir | 330 ++++++++++++++++
7 files changed, 781 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 c30dd3b07d028..7a4cef880b930 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1149,6 +1149,43 @@ def ReconcileUnrealizedCastsPass : Pass<"reconcile-unrealized-casts"> {
}];
}
+//===----------------------------------------------------------------------===//
+// SCFToAffine
+//===----------------------------------------------------------------------===//
+
+// TODO: extend to scf.if, scf.parallel, possibly more.
+def RaiseSCFToAffinePass : Pass<"raise-scf-to-affine"> {
+ let summary = "Raise SCF operations to affine operations (best effort)";
+ let description = [{
+ 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 if the source
+ type allows lossless conversion (otherwise we do not touch the loop). In any
+ case we preserve semantics.
+
+ This pass does not modify memory accesses; use `--affine-raise-from-memref`
+ to convert `memref.load` / `memref.store`.
+ }];
+ let dependentDialects = [
+ "affine::AffineDialect",
+ "arith::ArithDialect",
+ ];
+}
+
//===----------------------------------------------------------------------===//
// 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..f978d8309fc39
--- /dev/null
+++ b/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt
@@ -0,0 +1,15 @@
+add_mlir_conversion_library(MLIRSCFToAffine
+ SCFToAffine.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/SCFToAffine
+
+ DEPENDS
+ MLIRConversionPassIncGen
+
+ LINK_LIBS PUBLIC
+ MLIRAffineDialect
+ MLIRArithDialect
+ MLIRSCFDialect
+ MLIRTransforms
+ )
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
new file mode 100644
index 0000000000000..ca97b0b72be2f
--- /dev/null
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -0,0 +1,371 @@
+//===- 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 ops to affine ops.
+//
+//===----------------------------------------------------------------------===//
+
+#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/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/DebugLog.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_RAISESCFTOAFFINEPASS
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+#define DEBUG_TYPE "raise-scf-to-affine"
+
+using namespace mlir;
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// SCFToAffinePass
+//===----------------------------------------------------------------------===//
+
+struct SCFToAffinePass
+ : public impl::RaiseSCFToAffinePassBase<SCFToAffinePass> {
+ void runOnOperation() override;
+};
+
+//===----------------------------------------------------------------------===//
+// ForOpRewrite
+//===----------------------------------------------------------------------===//
+
+/// 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;
+
+ LogicalResult matchAndRewrite(scf::ForOp op,
+ PatternRewriter &rewriter) const override;
+
+private:
+ /// Definitively decide whether we are going to raise or not.
+ ///
+ /// 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) const;
+
+ /// Cast lb, ub, step and the induction variable of an integer-typed `op` to
+ /// `index`, in place. The bound and step casts are placed at the top level of
+ /// the affine scope so they are valid affine symbols; the induction variable
+ /// is cast back to its original type at the start of the body so the body is
+ /// left unchanged. Assumes `canRaiseToAffine(op) == true`.
+ void castBoundsToIndex(scf::ForOp op, PatternRewriter &rewriter) const;
+
+ /// Returns an equivalent `affine.for` skeleton and the *old* induction
+ /// 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` and index casts were performed (if
+ /// necessary).
+ ///
+ /// There are two cases:
+ ///
+ /// 1. step is 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
+ /// into the loop's lower bound map. Similarly if ub is an `affine.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;
+
+ std::pair<affine::AffineForOp, Value>
+ caseConstantStep(scf::ForOp op, int64_t step,
+ PatternRewriter &rewriter) const;
+
+ std::pair<affine::AffineForOp, Value>
+ caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const;
+};
+
+bool indexBoundsRaisable(scf::ForOp op) {
+ auto lb = op.getLowerBound();
+ auto ub = op.getUpperBound();
+ IntegerAttr constAttr;
+
+ // The asymmetry between lb and ub comes from the fact that the step
+ // normalization (for non-constant (dynamic) steps) does not work with
+ // multiple *lower* bounds (max).
+ bool lbOK = affine::isValidDim(lb) ||
+ (isa_and_present<affine::AffineMaxOp>(lb.getDefiningOp()) &&
+ matchPattern(op.getStep(), m_Constant(&constAttr)));
+ bool ubOK = affine::isValidDim(ub) ||
+ isa_and_present<affine::AffineMinOp>(ub.getDefiningOp());
+ bool stepOK = affine::isValidSymbol(op.getStep());
+
+ return lbOK && ubOK && stepOK;
+}
+
+/// Decide whether an integer-typed loop can be raised by first casting its
+/// bounds (lb, ub, step) to `index`. Requires the cast to be lossless under
+/// affine's *signed* `index` interpretation, and every bound to be available at
+/// the top level of the affine scope (so the inserted casts are valid symbols).
+bool intBoundsRaisable(scf::ForOp op, IntegerType intType) {
+ uint64_t indexWidth = DataLayout::closest(op)
+ .getTypeSizeInBits(IndexType::get(op.getContext()))
+ .getFixedValue();
+ // Lossless under signed index: sign-extend needs width <= indexWidth;
+ // zero-extend (unsigned) needs a spare sign bit, i.e. width < indexWidth.
+ uint64_t need = intType.getWidth() + (op.getUnsignedCmp() ? 1 : 0);
+ if (need > indexWidth)
+ return false;
+
+ Region *scope = affine::getAffineScope(op);
+ if (!scope)
+ return false;
+
+ // Being top-level implies the value is a symbol once it is casted to index.
+ return affine::isTopLevelValue(op.getLowerBound(), scope) &&
+ affine::isTopLevelValue(op.getUpperBound(), scope) &&
+ affine::isTopLevelValue(op.getStep(), scope);
+}
+
+bool ForOpRewrite::canRaiseToAffine(scf::ForOp op) const {
+ Type type = op.getInductionVar().getType();
+ if (isa<IndexType>(type))
+ return indexBoundsRaisable(op);
+ if (auto intType = dyn_cast<IntegerType>(type))
+ return intBoundsRaisable(op, intType);
+ return false;
+}
+
+LogicalResult ForOpRewrite::matchAndRewrite(scf::ForOp op,
+ PatternRewriter &rewriter) const {
+ if (!canRaiseToAffine(op)) {
+ LDBG() << "[affine] Cannot raise scf op: " << op << "\n";
+ return failure();
+ }
+
+ if (!isa<IndexType>(op.getInductionVar().getType()))
+ castBoundsToIndex(op, rewriter);
+
+ auto [affineFor, oldIV] = createAffineFor(op, rewriter);
+ Block *affineBody = affineFor.getBody();
+
+ 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");
+ rewriter.eraseOp(terminator);
+ }
+
+ SmallVector<Value> argValues;
+ argValues.push_back(oldIV);
+ 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();
+ assert(step > 0 && "scf.for has positive step");
+ return caseConstantStep(op, step, rewriter);
+ }
+ return caseDynamicStep(op, rewriter);
+}
+
+std::pair<affine::AffineForOp, Value>
+ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
+ PatternRewriter &rewriter) const {
+ auto lb = op.getLowerBound();
+ auto ub = op.getUpperBound();
+
+ auto lbOperands = ValueRange(lb);
+ auto ubOperands = ValueRange(ub);
+
+ auto lbMap = AffineMap::getMultiDimIdentityMap(1, rewriter.getContext());
+ auto ubMap = AffineMap::getMultiDimIdentityMap(1, rewriter.getContext());
+
+ if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
+ ubOperands = ubMinOp->getOperands();
+ ubMap = ubMinOp.getAffineMap();
+ }
+
+ if (auto lbMaxOp = lb.getDefiningOp<affine::AffineMaxOp>()) {
+ lbOperands = lbMaxOp->getOperands();
+ lbMap = lbMaxOp.getAffineMap();
+ }
+
+ auto affineFor =
+ affine::AffineForOp::create(rewriter, op.getLoc(), lbOperands, lbMap,
+ ubOperands, ubMap, step, op.getInits());
+
+ return std::make_pair(affineFor, affineFor.getInductionVar());
+}
+
+std::pair<affine::AffineForOp, Value>
+ForOpRewrite::caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const {
+ Value lb = op.getLowerBound();
+ 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);
+
+ 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 = 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);
+
+ 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(
+ rewriter, op.getLoc(), {}, zeroMap, ubOperands, ubMap, 1, op.getInits());
+
+ // 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);
+}
+
+void ForOpRewrite::castBoundsToIndex(scf::ForOp loop,
+ PatternRewriter &rewriter) const {
+ OpBuilder::InsertionGuard guard(rewriter);
+
+ 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");
+
+ auto createIndexCast = [&](Type out, Value in) -> Value {
+ Location loc = loop.getLoc();
+ if (loop.getUnsignedCmp()) {
+ return arith::IndexCastUIOp::create(rewriter, loc, out, in);
+ }
+ return arith::IndexCastOp::create(rewriter, loc, out, in);
+ };
+
+ // We place the bound casts at the top level of the affine scope so that they
+ // are identified as valid affine symbols.
+
+ Region *scope = affine::getAffineScope(loop);
+ Operation *anchor = loop;
+ while (anchor->getParentRegion() != scope)
+ anchor = anchor->getParentOp();
+ rewriter.setInsertionPoint(anchor);
+
+ Value newLb = createIndexCast(rewriter.getIndexType(), lb);
+ Value newUb = createIndexCast(rewriter.getIndexType(), ub);
+ Value newStep = createIndexCast(rewriter.getIndexType(), step);
+
+ 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);
+ });
+}
+
+//===----------------------------------------------------------------------===//
+// Pass implementation
+//===----------------------------------------------------------------------===//
+
+void SCFToAffinePass::runOnOperation() {
+ MLIRContext &ctx = getContext();
+ RewritePatternSet patterns(&ctx);
+ populateSCFToAffineConversionPatterns(patterns);
+
+ (void)applyPatternsGreedily(getOperation(), std::move(patterns));
+}
+
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// API
+//===----------------------------------------------------------------------===//
+
+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
new file mode 100644
index 0000000000000..543527ef32093
--- /dev/null
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -0,0 +1,330 @@
+// 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
+// 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 @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>
+ }
+ return
+}
+
+// -----
+
+// 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 %ub1 step %c1 {
+ scf.for %j = %c0 to %ub2 step %c1 {
+ memref.store %c0_i32, %arg0[%i, %j] : memref<?x?xi32>
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @index_cast_simple
+// CHECK-SAME: %[[LB:.*]]: i32, %[[UB:.*]]: i32
+// 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 #{{.*}}%[[LB_IDX]]{{.*}}%[[UB_IDX]]
+// CHECK: %[[IDX:.*]] = affine.apply #{{.*}}%[[LB_IDX]]{{.*}}%[[IV]]
+// CHECK: %[[IV_I32:.*]] = arith.index_cast %[[IDX]] : index to i32
+// CHECK: func.call @some_func(%[[IV_I32]])
+
+func.func private @some_func(%arg: i32)
+
+func.func @index_cast_simple(%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-LABEL: @index_cast_unsigned
+// CHECK-SAME: %[[LB:.*]]: i32, %[[UB:.*]]: i32
+// CHECK: %[[LB_IDX:.*]] = arith.index_castui %[[LB]] : i32 to index
+// CHECK: %[[UB_IDX:.*]] = arith.index_castui %[[UB]] : i32 to index
+// CHECK: affine.for
+// CHECK: %[[IV_I32:.*]] = arith.index_castui %{{.*}} : index to i32
+// CHECK: func.call @some_func(%[[IV_I32]])
+
+func.func private @some_func(%arg: i32)
+
+func.func @index_cast_unsigned(%lb: i32, %ub: i32) {
+ %step = arith.constant 1 : i32
+ scf.for unsigned %i = %lb to %ub step %step : i32 {
+ func.call @some_func(%i) : (i32) -> ()
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @nested_loop_index_cast(
+// CHECK-SAME: %[[UB1:.*]]: i16, %[[UB2:.*]]: i16) {
+// NOTE: index casts for *all* upper bounds deliberately hoisted to top-level
+// CHECK: %[[UB1_IDX:.*]] = arith.index_cast %[[UB1]] : i16 to index
+// CHECK: %[[UB2_IDX:.*]] = arith.index_cast %[[UB2]] : i16 to index
+// CHECK: affine.for %[[I_NEW:.*]] = 0 to %[[UB1_IDX]] {
+// CHECK: %[[I_OLD_IDX:.*]] = affine.apply #{{.*}}%[[I_NEW]]
+// CHECK: %[[I_OLD:.*]] = arith.index_cast %[[I_OLD_IDX]] : index to i16
+// CHECK: affine.for %[[J_NEW:.*]] = 0 to %[[UB2_IDX]] {
+// CHECK: %[[J_OLD_IDX:.*]] = affine.apply #{{.*}}%[[J_NEW]]
+// CHECK: %[[J_OLD:.*]] = arith.index_cast %[[J_OLD_IDX]] : index to i16
+// CHECK: func.call @some_func(%[[I_OLD]], %[[J_OLD]]) : (i16, i16) -> ()
+
+func.func private @some_func(%i: i16, %j: i16)
+
+func.func @nested_loop_index_cast(%ub1: i16, %ub2: i16) {
+ %c0_i32 = arith.constant 0 : i32
+ %c0 = arith.constant 0 : i16
+ %c1 = arith.constant 1 : i16
+ scf.for %i = %c0 to %ub1 step %c1 : i16{
+ scf.for %j = %c0 to %ub2 step %c1 : i16 {
+ func.call @some_func(%i, %j) : (i16, i16) -> ()
+ }
+ }
+ 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
+
+ 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
+}
+
+// -----
+
+// 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 #[[$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) -> ()
+
+#ub_map = affine_map<(i)[K, N] -> (K, N - i)>
+
+func.func private @some_func(%i: index, %j: index)
+
+func.func @dynamic_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
+}
+
+// -----
+
+// 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_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
+ scf.yield %v : 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
+}
+
+// -----
+
+// 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
+}
+
+// -----
+
+// 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 {{.*}} iter_args(%[[ACC:.*]] = %[[INIT]]) -> (f32) {
+// CHECK: %[[IV_OLD_IDX:.*]] = affine.apply #{{.*}}%[[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
+}
+
+// -----
+
+// CHECK-LABEL: @wider_than_index_not_raised
+// CHECK: scf.for %{{.*}} : 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
+ }
+}
+
+// -----
+
+// CHECK-LABEL: @unsigned_same_width_not_raised
+// CHECK: scf.for unsigned %{{.*}} : i32
+// CHECK-NOT: affine.for
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<index, 32>>} {
+ func.func private @some_func(%arg: i32)
+
+ func.func @unsigned_same_width_not_raised(%lb: i32, %ub: i32, %step: i32) {
+ scf.for unsigned %i = %lb to %ub step %step : i32 {
+ func.call @some_func(%i) : (i32) -> ()
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK-LABEL: @signed_same_width_raised
+// CHECK: affine.for
+// CHECK-NOT: scf.for
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<index, 32>>} {
+ func.func private @some_func(%arg: i32)
+
+ func.func @signed_same_width_raised(%lb: i32, %ub: i32, %step: i32) {
+ scf.for %i = %lb to %ub step %step : i32 {
+ func.call @some_func(%i) : (i32) -> ()
+ }
+ return
+ }
+}
\ No newline at end of file
>From bef46100576e9a654e35e885225611ee60a5314a Mon Sep 17 00:00:00 2001
From: Reinhard Stahn <rainij36 at proton.me>
Date: Wed, 17 Jun 2026 15:34:34 +0000
Subject: [PATCH 2/2] Addressed review comments by IgWod
---
.../Conversion/SCFToAffine/SCFToAffine.cpp | 39 +++++++++----------
.../Conversion/SCFToAffine/scf-to-affine.mlir | 2 +-
2 files changed, 19 insertions(+), 22 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index ca97b0b72be2f..488aae4c714c5 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -22,15 +22,12 @@
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/DebugLog.h"
namespace mlir {
#define GEN_PASS_DEF_RAISESCFTOAFFINEPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
-#define DEBUG_TYPE "raise-scf-to-affine"
-
using namespace mlir;
namespace {
@@ -97,16 +94,17 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const;
std::pair<affine::AffineForOp, Value>
- caseConstantStep(scf::ForOp op, int64_t step,
- PatternRewriter &rewriter) const;
+ createAffineForWithConstantStep(scf::ForOp op, int64_t step,
+ PatternRewriter &rewriter) const;
std::pair<affine::AffineForOp, Value>
- caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const;
+ createAffineForWithDynamicStep(scf::ForOp op,
+ PatternRewriter &rewriter) const;
};
bool indexBoundsRaisable(scf::ForOp op) {
- auto lb = op.getLowerBound();
- auto ub = op.getUpperBound();
+ Value lb = op.getLowerBound();
+ Value ub = op.getUpperBound();
IntegerAttr constAttr;
// The asymmetry between lb and ub comes from the fact that the step
@@ -132,8 +130,8 @@ bool intBoundsRaisable(scf::ForOp op, IntegerType intType) {
.getFixedValue();
// Lossless under signed index: sign-extend needs width <= indexWidth;
// zero-extend (unsigned) needs a spare sign bit, i.e. width < indexWidth.
- uint64_t need = intType.getWidth() + (op.getUnsignedCmp() ? 1 : 0);
- if (need > indexWidth)
+ uint64_t requiredWidth = intType.getWidth() + (op.getUnsignedCmp() ? 1 : 0);
+ if (requiredWidth > indexWidth)
return false;
Region *scope = affine::getAffineScope(op);
@@ -158,8 +156,7 @@ bool ForOpRewrite::canRaiseToAffine(scf::ForOp op) const {
LogicalResult ForOpRewrite::matchAndRewrite(scf::ForOp op,
PatternRewriter &rewriter) const {
if (!canRaiseToAffine(op)) {
- LDBG() << "[affine] Cannot raise scf op: " << op << "\n";
- return failure();
+ return rewriter.notifyMatchFailure(op, "cannot raise scf op to affine");
}
if (!isa<IndexType>(op.getInductionVar().getType()))
@@ -197,16 +194,16 @@ ForOpRewrite::createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
if (matchPattern(op.getStep(), m_Constant(&constAttr))) {
int64_t step = constAttr.getInt();
assert(step > 0 && "scf.for has positive step");
- return caseConstantStep(op, step, rewriter);
+ return createAffineForWithConstantStep(op, step, rewriter);
}
- return caseDynamicStep(op, rewriter);
+ return createAffineForWithDynamicStep(op, rewriter);
}
std::pair<affine::AffineForOp, Value>
-ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
- PatternRewriter &rewriter) const {
- auto lb = op.getLowerBound();
- auto ub = op.getUpperBound();
+ForOpRewrite::createAffineForWithConstantStep(scf::ForOp op, int64_t step,
+ PatternRewriter &rewriter) const {
+ Value lb = op.getLowerBound();
+ Value ub = op.getUpperBound();
auto lbOperands = ValueRange(lb);
auto ubOperands = ValueRange(ub);
@@ -232,7 +229,8 @@ ForOpRewrite::caseConstantStep(scf::ForOp op, int64_t step,
}
std::pair<affine::AffineForOp, Value>
-ForOpRewrite::caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const {
+ForOpRewrite::createAffineForWithDynamicStep(scf::ForOp op,
+ PatternRewriter &rewriter) const {
Value lb = op.getLowerBound();
Value ub = op.getUpperBound();
Value step = op.getStep();
@@ -311,9 +309,8 @@ void ForOpRewrite::castBoundsToIndex(scf::ForOp loop,
auto createIndexCast = [&](Type out, Value in) -> Value {
Location loc = loop.getLoc();
- if (loop.getUnsignedCmp()) {
+ if (loop.getUnsignedCmp())
return arith::IndexCastUIOp::create(rewriter, loc, out, in);
- }
return arith::IndexCastOp::create(rewriter, loc, out, in);
};
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index 543527ef32093..67e56c45a4b77 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -raise-scf-to-affine -split-input-file %s | FileCheck %s
+// 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
More information about the Mlir-commits
mailing list