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

[Reinhard Stahn]

================
@@ -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()) &&
----------------
rainij wrote:

None of the three reasons apply. I was also a bit surprised, but it makes sense. Semantically the for loop has these requirements:

- The lb is either a dimension or more generally a max over multiple dimensions.
- The ub is either a dimension or more generally a min over multiple dimensions.

(And only the max/min stuff makes it really interesting for polyhedral optimization as far as my knowledge on the matter can tell. Not an expert though, just a math guy.)

And it is correct that `isValidDim` returns "no" on a `affine.max`. It really isn't a dimension in general. MLIR does not give `arith.max` any privileges beyond those given to any other pure functions: If all its args are symbols then the output is considered a symbol and hence also a dimension. But this is too weak in general. It also cannot decide "does it fit as lb or ub?" because lb and ub do not have the same requirements.

In theory one could relax e.g. the lb condition to: is either a dim or a max over multiple dims (inline in the op itself), or is an `affine.max` (its verifier already requires that the operands are dimensions). This would be semantically correct. But it would be cumbersome to check this all the time I guess.

For that reason we go the manual route and check explicitly for this operation. The reason we added this is that on discourse we agreed that we port the features of polygeist / Enzyme-JAX and they could essentially deal with this situation. The details are a bit different: the actually cannot raise the `affine.max` situation but they raise the even lower level cmp+select pattern. You recently added to the canoncializer that this will be "raised" to `arith.max`. The only gap left now is the raising to `affine.max`. But @julfarn wants to take care of this once this here is merged (or maybe earlier).

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