LoopRangeFolding.cpp

Go to the documentation of this file.

//===- LoopRangeFolding.cpp - Code to perform loop range folding-----------===//
//
// 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 loop range folding.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/SCF/Transforms/Passes.h"
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/IRMapping.h"
 
namespace mlir {
#define GEN_PASS_DEF_SCFFORLOOPRANGEFOLDING
#include "mlir/Dialect/SCF/Transforms/Passes.h.inc"
} // namespace mlir
 
using namespace mlir;
using namespace mlir::scf;
 
namespace {
struct ForLoopRangeFolding
    : public impl::SCFForLoopRangeFoldingBase<ForLoopRangeFolding> {
  void runOnOperation() override;
};
} // namespace
 
void ForLoopRangeFolding::runOnOperation() {
  getOperation()->walk([&](ForOp op) {
    Value indVar = op.getInductionVar();
 
    auto canBeFolded = [&](Value value) {
      return op.isDefinedOutsideOfLoop(value) || value == indVar;
    };
 
    // Fold until a fixed point is reached
    while (true) {
 
      // If the induction variable is used by more than one operation, we can't
      // fold its arith ops into the loop range. Note that a single operation
      // may use the induction variable several times (e.g. `arith.addi %i,
      // %i`), so check for a single user rather than a single use.
      if (indVar.use_empty())
        break;
 
      Operation *user = *indVar.getUsers().begin();
      if (!llvm::all_of(indVar.getUsers(),
                        [&](Operation *u) { return u == user; }))
        break;
 
      if (!isa<arith::AddIOp, arith::MulIOp>(user))
        break;
 
      if (!llvm::all_of(user->getOperands(), canBeFolded))
        break;
 
      OpBuilder b(op);
      IRMapping lbMap;
      lbMap.map(indVar, op.getLowerBound());
      IRMapping ubMap;
      ubMap.map(indVar, op.getUpperBound());
      IRMapping stepMap;
      stepMap.map(indVar, op.getStep());
 
      if (auto addOp = dyn_cast<arith::AddIOp>(user)) {
        Operation *lbFold = b.clone(*user, lbMap);
        Operation *ubFold = b.clone(*user, ubMap);
 
        op.setLowerBound(lbFold->getResult(0));
        op.setUpperBound(ubFold->getResult(0));
 
        // `arith.addi %i, %i` is `2 * %i`, so the step has to be doubled too.
        if (addOp.getLhs() == indVar && addOp.getRhs() == indVar) {
          Operation *stepFold = b.clone(*user, stepMap);
          op.setStep(stepFold->getResult(0));
        }
 
      } else if (auto mulOp = dyn_cast<arith::MulIOp>(user)) {
        // Only fold if the multiplier is a known strictly positive constant.
        // Multiplying by zero or a negative value would produce an invalid
        // step (scf.for requires a strictly positive step).
        Value multiplier =
            (mulOp.getLhs() == indVar) ? mulOp.getRhs() : mulOp.getLhs();
        std::optional<int64_t> multiplierVal = getConstantIntValue(multiplier);
        if (!multiplierVal || *multiplierVal <= 0)
          break;
 
        Operation *lbFold = b.clone(*user, lbMap);
        Operation *ubFold = b.clone(*user, ubMap);
        Operation *stepFold = b.clone(*user, stepMap);
 
        op.setLowerBound(lbFold->getResult(0));
        op.setUpperBound(ubFold->getResult(0));
        op.setStep(stepFold->getResult(0));
      }
 
      ValueRange wrapIndvar(indVar);
      user->replaceAllUsesWith(wrapIndvar);
      user->erase();
    }
  });
}
 
std::unique_ptr<Pass> mlir::createForLoopRangeFoldingPass() {
  return std::make_unique<ForLoopRangeFolding>();
}