[Mlir-commits] [mlir] [SCFToAffine] Raise scf.for to affine.for (PR #200851)

Reinhard Stahn llvmlistbot at llvm.org
Wed Jun 3 12:53:16 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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



More information about the Mlir-commits mailing list