[Mlir-commits] [mlir] [SCFToAffine] Add a pass to raise scf to affine ops. (PR #152925)

[Oleksandr Alex Zinenko]

================
@@ -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());
----------------
ftynse wrote:

Why all this complexity? Affine `for`s support having values as lower and upper bounds (by using a 1D identity map), and a non-unit step. We could then "normalize" loops in a separate pass when desired, and we may already have such a pass.

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