SCFToAffine.cpp

Go to the documentation of this file.

//===- 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;
};
 
static bool areValidAffineMapOperands(AffineMap map, ValueRange operands,
                                      Region *scope) {
  assert(map.getNumInputs() == operands.size() &&
         "expected one operand per affine map input");
  return llvm::all_of(
             operands.take_front(map.getNumDims()),
             [&](Value value) { return affine::isValidDim(value, scope); }) &&
         llvm::all_of(operands.drop_front(map.getNumDims()), [&](Value value) {
           return affine::isValidSymbol(value, scope);
         });
}
 
bool indexBoundsRaisable(scf::ForOp op) {
  Value lb = op.getLowerBound();
  Value ub = op.getUpperBound();
  IntegerAttr constAttr;
  Region *scope = affine::getAffineScope(op);
  if (!scope)
    return false;
 
  // 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).
  auto lbMaxOp = lb.getDefiningOp<affine::AffineMaxOp>();
  bool lbOK = affine::isValidDim(lb, scope) ||
              (lbMaxOp && matchPattern(op.getStep(), m_Constant(&constAttr)) &&
               areValidAffineMapOperands(lbMaxOp.getAffineMap(),
                                         lbMaxOp->getOperands(), scope));
  auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>();
  bool ubOK =
      affine::isValidDim(ub, scope) ||
      (ubMinOp && areValidAffineMapOperands(ubMinOp.getAffineMap(),
                                            ubMinOp->getOperands(), scope));
  bool stepOK = affine::isValidSymbol(op.getStep(), scope);
 
  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 requiredWidth = intType.getWidth() + (op.getUnsignedCmp() ? 1 : 0);
  if (requiredWidth > 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)) {
    return rewriter.notifyMatchFailure(op, "cannot raise scf op to affine");
  }
 
  if (!isa<IndexType>(op.getInductionVar().getType()))
    castBoundsToIndex(op, rewriter);
 
  auto [affineFor, oldIV] = createAffineFor(op, rewriter);
  Block *affineBody = affineFor.getBody();
 
  if (affineBody->mightHaveTerminator()) {
    // No unregistered ops in the body, so this is definitive.
    Operation *terminator = affineBody->getTerminator();
    assert(isa<affine::AffineYieldOp>(terminator) &&
           "expected affine.yield if there *might* be terminator");
    rewriter.eraseOp(terminator);
  }
 
  SmallVector<Value> argValues;
  argValues.push_back(oldIV);
  llvm::append_range(argValues, affineFor.getRegionIterArgs());
  rewriter.inlineBlockBefore(op.getBody(), affineBody, affineBody->end(),
                             argValues);
 
  auto scfYieldOp = cast<scf::YieldOp>(affineBody->getTerminator());
  rewriter.setInsertionPointToEnd(affineBody);
  rewriter.replaceOpWithNewOp<affine::AffineYieldOp>(scfYieldOp,
                                                     scfYieldOp->getOperands());
 
  rewriter.replaceOp(op, affineFor);
  return success();
}
 
std::pair<affine::AffineForOp, Value>
ForOpRewrite::createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const {
  IntegerAttr constAttr;
  if (matchPattern(op.getStep(), m_Constant(&constAttr))) {
    int64_t step = constAttr.getInt();
    assert(step > 0 && "scf.for has positive step");
    return createAffineForWithConstantStep(op, step, rewriter);
  }
  return createAffineForWithDynamicStep(op, rewriter);
}
 
std::pair<affine::AffineForOp, Value>
ForOpRewrite::createAffineForWithConstantStep(scf::ForOp op, int64_t step,
                                              PatternRewriter &rewriter) const {
  Value lb = op.getLowerBound();
  Value ub = op.getUpperBound();
 
  auto lbOperands = ValueRange(lb);
  auto ubOperands = ValueRange(ub);
 
  auto lbMap = AffineMap::getMultiDimIdentityMap(1, rewriter.getContext());
  auto ubMap = AffineMap::getMultiDimIdentityMap(1, rewriter.getContext());
 
  if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
    ubOperands = ubMinOp->getOperands();
    ubMap = ubMinOp.getAffineMap();
  }
 
  if (auto lbMaxOp = lb.getDefiningOp<affine::AffineMaxOp>()) {
    lbOperands = lbMaxOp->getOperands();
    lbMap = lbMaxOp.getAffineMap();
  }
 
  auto affineFor =
      affine::AffineForOp::create(rewriter, op.getLoc(), lbOperands, lbMap,
                                  ubOperands, ubMap, step, op.getInits());
 
  return std::make_pair(affineFor, affineFor.getInductionVar());
}
 
std::pair<affine::AffineForOp, Value>
ForOpRewrite::createAffineForWithDynamicStep(scf::ForOp op,
                                             PatternRewriter &rewriter) const {
  Value lb = op.getLowerBound();
  Value ub = op.getUpperBound();
  Value step = op.getStep();
 
  assert(affine::isValidDim(lb) &&
         "dynamic-step lower bound must be a valid affine dim");
 
  AffineExpr d0 = rewriter.getAffineDimExpr(0);
  AffineExpr d1 = rewriter.getAffineDimExpr(1);
  AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
  AffineMap zeroMap = rewriter.getConstantAffineMap(0);
 
  llvm::SmallVector<Value, 3> ubOperands = {lb, ub, step};
 
  // ub is transformed with (x - lb + step - 1) floorDiv step where x ranges
  // over all ub_i. lb is transformed to zero.
 
  AffineMap ubMap = AffineMap::get(2, 1, (d1 - d0 + s0 - 1).floorDiv(s0));
 
  if (auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>()) {
    AffineMap origUbMap = ubMinOp.getAffineMap();
    unsigned ubDims = origUbMap.getNumDims();
    unsigned ubSyms = origUbMap.getNumSymbols();
 
    AffineExpr lbDim = rewriter.getAffineDimExpr(ubDims);
    AffineExpr stepSym = rewriter.getAffineSymbolExpr(ubSyms);
 
    SmallVector<AffineExpr> ubExprs;
    ubExprs.reserve(origUbMap.getNumResults());
    for (AffineExpr ubI : origUbMap.getResults()) {
      ubExprs.push_back((ubI - lbDim + stepSym - 1).floorDiv(stepSym));
    }
 
    // Combined space: dims = [ub dims, lb]
    //                 syms = [ub syms, step]
    ubMap =
        AffineMap::get(ubDims + 1, ubSyms + 1, ubExprs, rewriter.getContext());
 
    // Operand order consistent with "combined space" above:
    ValueRange ubOps = ubMinOp->getOperands();
    SmallVector<Value> combined;
    combined.append(ubOps.begin(), ubOps.begin() + ubDims); // ub dims
    combined.push_back(lb);                                 // lb (single dim)
    combined.append(ubOps.begin() + ubDims, ubOps.end());   // ub syms
    combined.push_back(op.getStep());                       // step (single sym)
    ubOperands = std::move(combined);
  }
 
  auto affineFor = affine::AffineForOp::create(
      rewriter, op.getLoc(), {}, zeroMap, ubOperands, ubMap, 1, op.getInits());
 
  // old_iv = old_lb + new_iv * step
  AffineMap ivMap = AffineMap::get(2, 1, d0 + d1 * s0);
 
  llvm::SmallVector<Value, 3> ivOperands = {lb, affineFor.getInductionVar(),
                                            step};
 
  rewriter.setInsertionPointToStart(affineFor.getBody());
  auto oldIV =
      affine::AffineApplyOp::create(rewriter, op.getLoc(), ivMap, ivOperands);
 
  return std::make_pair(affineFor, oldIV);
}
 
void ForOpRewrite::castBoundsToIndex(scf::ForOp loop,
                                     PatternRewriter &rewriter) const {
  OpBuilder::InsertionGuard guard(rewriter);
 
  Value lb = loop.getLowerBound();
  Value ub = loop.getUpperBound();
  Value step = loop.getStep();
  Type originalType = step.getType();
 
  assert(lb.getType() == originalType && ub.getType() == originalType &&
         "expected lb, ub, and step to have the same type");
 
  auto createIndexCast = [&](Type out, Value in) -> Value {
    Location loc = loop.getLoc();
    if (loop.getUnsignedCmp())
      return arith::IndexCastUIOp::create(rewriter, loc, out, in);
    return arith::IndexCastOp::create(rewriter, loc, out, in);
  };
 
  // We place the bound casts at the top level of the affine scope so that they
  // are identified as valid affine symbols.
 
  Region *scope = affine::getAffineScope(loop);
  Operation *anchor = loop;
  while (anchor->getParentRegion() != scope)
    anchor = anchor->getParentOp();
  rewriter.setInsertionPoint(anchor);
 
  Value newLb = createIndexCast(rewriter.getIndexType(), lb);
  Value newUb = createIndexCast(rewriter.getIndexType(), ub);
  Value newStep = createIndexCast(rewriter.getIndexType(), step);
 
  rewriter.modifyOpInPlace(loop, [&] {
    loop.setLowerBound(newLb);
    loop.setUpperBound(newUb);
    loop.setStep(newStep);
 
    Value originalIV = loop.getInductionVar();
    Value iv = loop.getBody()->insertArgument(
        (unsigned)0, rewriter.getIndexType(), loop.getLoc());
 
    rewriter.setInsertionPointToStart(loop.getBody());
    Value castIV = createIndexCast(originalType, iv);
    rewriter.replaceAllUsesWith(originalIV, castIV);
 
    // Original induction var is now at index 1.
    loop.getBody()->eraseArgument(1);
  });
}
 
//===----------------------------------------------------------------------===//
// Pass implementation
//===----------------------------------------------------------------------===//
 
void SCFToAffinePass::runOnOperation() {
  MLIRContext &ctx = getContext();
  RewritePatternSet patterns(&ctx);
  populateSCFToAffineConversionPatterns(patterns);
 
  (void)applyPatternsGreedily(getOperation(), std::move(patterns));
}
 
} // namespace
 
//===----------------------------------------------------------------------===//
// API
//===----------------------------------------------------------------------===//
 
void mlir::populateSCFToAffineConversionPatterns(RewritePatternSet &patterns) {
  patterns.add<ForOpRewrite>(patterns.getContext());
}