[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).
+bool intBoundsRaisable(scf::ForOp op, IntegerType intType) {
+  uint64_t indexWidth = DataLayout::closest(op)
+                            .getTypeSizeInBits(IndexType::get(op.getContext()))
+                            .getFixedValue();
+  // Lossless under signed index: sign-extend needs width <= indexWidth;
+  // zero-extend (unsigned) needs a spare sign bit, i.e. width < indexWidth.
+  uint64_t need = intType.getWidth() + (op.getUnsignedCmp() ? 1 : 0);
+  if (need > indexWidth)
+    return false;
+
+  Region *scope = affine::getAffineScope(op);
+  if (!scope)
+    return false;
+
+  // Being top-level implies the value is a symbol once it is casted to index.
+  return affine::isTopLevelValue(op.getLowerBound(), scope) &&
+         affine::isTopLevelValue(op.getUpperBound(), scope) &&
+         affine::isTopLevelValue(op.getStep(), scope);
+}
+
+bool ForOpRewrite::canRaiseToAffine(scf::ForOp op) const {
+  Type type = op.getInductionVar().getType();
+  if (isa<IndexType>(type))
+    return indexBoundsRaisable(op);
+  if (auto intType = dyn_cast<IntegerType>(type))
+    return intBoundsRaisable(op, intType);
+  return false;
+}
+
+LogicalResult ForOpRewrite::matchAndRewrite(scf::ForOp op,
+                                            PatternRewriter &rewriter) const {
+  if (!canRaiseToAffine(op)) {
+    LDBG() << "[affine] Cannot raise scf op: " << op << "\n";
+    return failure();
----------------
rainij wrote:

Agree, changed as suggested.

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