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

[Ming Yan]

================
@@ -0,0 +1,368 @@
+//===- 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"
+
+namespace mlir {
+#define GEN_PASS_DEF_RAISESCFTOAFFINEPASS
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+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>
+  createAffineForWithConstantStep(scf::ForOp op, int64_t step,
+                                  PatternRewriter &rewriter) const;
+
+  std::pair<affine::AffineForOp, Value>
+  createAffineForWithDynamicStep(scf::ForOp op,
+                                 PatternRewriter &rewriter) const;
+};
+
+bool indexBoundsRaisable(scf::ForOp op) {
+  Value lb = op.getLowerBound();
+  Value 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()) &&
----------------
NexMing wrote:

Why is special-case handling needed here for `affine.min/max`? `isValidDim` is a prerequisite for constructing an `affine.for`. If the result doesn't match your expectation, the reason may be:

1. The transformation itself is invalid.
2. There is an issue with the `isValidDim` implementation, causing it to fail to achieve the expected behavior.
3. `isValidDim` is not being used correctly.

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