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

[Reinhard Stahn]

================
@@ -0,0 +1,371 @@
+//===- 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"
+#include "llvm/Support/DebugLog.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_RAISESCFTOAFFINEPASS
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+#define DEBUG_TYPE "raise-scf-to-affine"
+
+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>
+  caseConstantStep(scf::ForOp op, int64_t step,
+                   PatternRewriter &rewriter) const;
+
+  std::pair<affine::AffineForOp, Value>
+  caseDynamicStep(scf::ForOp op, PatternRewriter &rewriter) const;
+};
+
+bool indexBoundsRaisable(scf::ForOp op) {
+  auto lb = op.getLowerBound();
+  auto 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()) &&
+               matchPattern(op.getStep(), m_Constant(&constAttr)));
+  bool ubOK = affine::isValidDim(ub) ||
+              isa_and_present<affine::AffineMinOp>(ub.getDefiningOp());
+  bool stepOK = affine::isValidSymbol(op.getStep());
+
+  return lbOK && ubOK && stepOK;
+}
+
+/// Decide whether an integer-typed loop can be raised by first casting its
+/// bounds (lb, ub, step) to `index`. Requires the cast to be lossless under
+/// affine's *signed* `index` interpretation, and every bound to be available at
+/// the top level of the affine scope (so the inserted casts are valid symbols).
----------------
rainij wrote:

Affine imposes certain restriction on what it allows as e.g. lb and ub. In general lb and ub have to be _dimensions_ (or max/min of something respectively). A special case of a dimension is a _symbol_. They are values which are (required to be) invariant throughout the loop nest. (A dimension allows also e.g. the induction variable of a further out loop inside the same nest.)

Now in this particular function we make sure that certain conditions are met so that after the cast the resulting lb and ub are indeed (recognized as) symbols. Being a top-level value in a scope which counts as _affine scope_ (a trait) ensures that.

Concerning the word "valid". I believe it is redundant. _Valid symbol_ means the same as _(affine) symbol_. I just added it because there is an API function `isValidSymbol(...)`.

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