[SCFToAffine] Add a pass to raise scf to affine ops.

[NexMing]

This patch adds an initial pass to convert eligible scf.for loops into affine.for. The Affine dialect enables stronger optimization and analysis capabilities compared to SCF, including dependence analysis, loop tiling/fusion, and parallelization. Promoting analyzable SCF loops to Affine provides opportunities for polyhedral-style optimizations at the core dialect level.

This SCF-to-Affine pass is inspired by the Polygeist project, which implemented similar transformations. However, this implementation is independently developed in MLIR project. This allows frontends generating SCF loops, such as ClangIR or Flang, to benefit from Affine-based transformations.

This patch is the first step in an incremental effort:

• Extend support for converting other scf operations.
• Extend support for more complex loop forms and bounds.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Ming Yan (NexMing)

Changes

This patch supports the conversion from scf.for to affine.for.

---

Full diff: https://github.com/llvm/llvm-project/pull/152925.diff

7 Files Affected:

• (modified) mlir/include/mlir/Conversion/Passes.h (+1)
• (modified) mlir/include/mlir/Conversion/Passes.td (+12)
• (added) mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h (+26)
• (modified) mlir/lib/Conversion/CMakeLists.txt (+1)
• (added) mlir/lib/Conversion/SCFToAffine/CMakeLists.txt (+17)
• (added) mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp (+100)
• (added) mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir (+34)

diff --git a/mlir/include/mlir/Conversion/Passes.h b/mlir/include/mlir/Conversion/Passes.h
index 3dc48b2201cf2..2507ef2834dc5 100644
--- a/mlir/include/mlir/Conversion/Passes.h
+++ b/mlir/include/mlir/Conversion/Passes.h
@@ -58,6 +58,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 6e1baaf23fcf7..38f35b2dadd94 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1025,6 +1025,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 785cb8293810c..b8059fcbfb028 100644
--- a/mlir/lib/Conversion/CMakeLists.txt
+++ b/mlir/lib/Conversion/CMakeLists.txt
@@ -51,6 +51,7 @@ add_subdirectory(OpenACCToSCF)
 add_subdirectory(OpenMPToLLVM)
 add_subdirectory(PDLToPDLInterp)
 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..e68bb2123cadc
--- /dev/null
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -0,0 +1,100 @@
+//===- 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;
+};
+
+struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
+  using OpRewritePattern<scf::ForOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(scf::ForOp op,
+                                PatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto lower = op.getLowerBound();
+    auto upper = op.getUpperBound();
+    auto step = op.getStep();
+
+    if (!affine::isValidDim(lower) || !affine::isValidDim(upper) ||
+        !affine::isValidSymbol(step))
+      return llvm::failure();
+
+    auto lowerDim = rewriter.getAffineDimExpr(0);
+    auto upperDim = rewriter.getAffineDimExpr(1);
+    auto stepSym = rewriter.getAffineSymbolExpr(0);
+    auto affineFor = affine::AffineForOp::create(
+        rewriter, loc, ValueRange(), rewriter.getConstantAffineMap(0),
+        ValueRange({lower, upper, step}),
+        AffineMap::get(2, 1,
+                       (upperDim - lowerDim + stepSym - 1).floorDiv(stepSym)),
+        1, op.getInits());
+    auto affineBody = affineFor.getBody();
+
+    if (affineBody->mightHaveTerminator())
+      rewriter.eraseOp(affineBody->getTerminator());
+
+    rewriter.setInsertionPointToStart(affineBody);
+    auto actualIndexMap =
+        AffineMap::get(2, 1, lowerDim + rewriter.getAffineDimExpr(1) * stepSym);
+    Value newIndVar =
+        affine::AffineApplyOp::create(
+            rewriter, op.getLoc(), actualIndexMap,
+            ValueRange({lower, affineFor.getInductionVar(), step}))
+            .getResult();
+
+    SmallVector<Value> argValues;
+    argValues.push_back(newIndVar);
+    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() {
+  RewritePatternSet patterns(&getContext());
+  populateSCFToAffineConversionPatterns(patterns);
+
+  if (failed(applyPatternsGreedily(getOperation(), 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..6f419bc8ee9ce
--- /dev/null
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -0,0 +1,34 @@
+// RUN: mlir-opt -raise-scf-to-affine -split-input-file %s | FileCheck %s
+
+// CHECK: #[[$ATTR_0:.+]] = affine_map<()[s0, s1, s2] -> ((s0 - s1 + s2 - 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_6]], %[[VAL_4]], %[[VAL_5]]] {
+// 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
+}

[ftynse]

Please update the commit message to explain why this is added, https://mlir.llvm.org/getting_started/Contributing/#commit-messages. It is also a good idea to describe whether this a one-off patch or whether further work is planned.

[NexMing]

@ftynse Could you help me review this PR?

[linuxlonelyeagle]

I left some simple comments.

[rscottmanley]

I would suggest adding a test case with 2d loops where inner can be converted but outer cannot be. And another where the outer loop is affine on its own but the inner is not and prevents the outer from being converted.

[NexMing]

Port the Polygeist code to the main project instead of reimplementing it.
Refer to this PR: #158267
Related discussion: https://discourse.llvm.org/t/rfc-add-scf-to-affine-conversion-pass-in-mlir/88036.