ControlFlowToSCF.cpp

Go to the documentation of this file.

//===- ControlFlowToSCF.h - ControlFlow to SCF -------------*- C++ ------*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Define conversions from the ControlFlow dialect to the SCF dialect.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Conversion/ControlFlowToSCF/ControlFlowToSCF.h"
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/CFGToSCF.h"
 
namespace mlir {
#define GEN_PASS_DEF_LIFTCONTROLFLOWTOSCFPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
 
using namespace mlir;
 
bool ControlFlowToSCFTransformation::canConvertMultiSuccessorBranchOp(
    Operation *op) {
  return isa<cf::CondBranchOp, cf::SwitchOp>(op);
}
 
FailureOr<Operation *>
ControlFlowToSCFTransformation::createStructuredBranchRegionOp(
    OpBuilder &builder, Operation *controlFlowCondOp, TypeRange resultTypes,
    MutableArrayRef<Region> regions) {
  if (auto condBrOp = dyn_cast<cf::CondBranchOp>(controlFlowCondOp)) {
    assert(regions.size() == 2);
    auto ifOp = scf::IfOp::create(builder, controlFlowCondOp->getLoc(),
                                  resultTypes, condBrOp.getCondition());
    ifOp.getThenRegion().takeBody(regions[0]);
    ifOp.getElseRegion().takeBody(regions[1]);
    return ifOp.getOperation();
  }
 
  if (auto switchOp = dyn_cast<cf::SwitchOp>(controlFlowCondOp)) {
    // `getCFGSwitchValue` returns an i32 that we need to convert to index
    // fist.
    auto cast = arith::IndexCastUIOp::create(
        builder, controlFlowCondOp->getLoc(), builder.getIndexType(),
        switchOp.getFlag());
    SmallVector<int64_t> cases;
    if (auto caseValues = switchOp.getCaseValues())
      llvm::append_range(
          cases, llvm::map_range(*caseValues, [](const llvm::APInt &apInt) {
            return apInt.getZExtValue();
          }));
 
    assert(regions.size() == cases.size() + 1);
 
    auto indexSwitchOp =
        scf::IndexSwitchOp::create(builder, controlFlowCondOp->getLoc(),
                                   resultTypes, cast, cases, cases.size());
 
    indexSwitchOp.getDefaultRegion().takeBody(regions[0]);
    for (auto &&[targetRegion, sourceRegion] :
         llvm::zip(indexSwitchOp.getCaseRegions(), llvm::drop_begin(regions)))
      targetRegion.takeBody(sourceRegion);
 
    return indexSwitchOp.getOperation();
  }
 
  controlFlowCondOp->emitOpError(
      "Cannot convert unknown control flow op to structured control flow");
  return failure();
}
 
LogicalResult
ControlFlowToSCFTransformation::createStructuredBranchRegionTerminatorOp(
    Location loc, OpBuilder &builder, Operation *branchRegionOp,
    Operation *replacedControlFlowOp, ValueRange results) {
  scf::YieldOp::create(builder, loc, results);
  return success();
}
 
FailureOr<Operation *>
ControlFlowToSCFTransformation::createStructuredDoWhileLoopOp(
    OpBuilder &builder, Operation *replacedOp, ValueRange loopVariablesInit,
    Value condition, ValueRange loopVariablesNextIter, Region &&loopBody) {
  Location loc = replacedOp->getLoc();
  auto whileOp = scf::WhileOp::create(
      builder, loc, loopVariablesInit.getTypes(), loopVariablesInit);
 
  whileOp.getBefore().takeBody(loopBody);
 
  builder.setInsertionPointToEnd(&whileOp.getBefore().back());
  // `getCFGSwitchValue` returns a i32. We therefore need to truncate the
  // condition to i1 first. It is guaranteed to be either 0 or 1 already.
  scf::ConditionOp::create(
      builder, loc,
      arith::TruncIOp::create(builder, loc, builder.getI1Type(), condition),
      loopVariablesNextIter);
 
  Block *afterBlock = builder.createBlock(&whileOp.getAfter());
  afterBlock->addArguments(
      loopVariablesInit.getTypes(),
      SmallVector<Location>(loopVariablesInit.size(), loc));
  scf::YieldOp::create(builder, loc, afterBlock->getArguments());
 
  return whileOp.getOperation();
}
 
Value ControlFlowToSCFTransformation::getCFGSwitchValue(Location loc,
                                                        OpBuilder &builder,
                                                        unsigned int value) {
  return arith::ConstantOp::create(builder, loc,
                                   builder.getI32IntegerAttr(value));
}
 
void ControlFlowToSCFTransformation::createCFGSwitchOp(
    Location loc, OpBuilder &builder, Value flag,
    ArrayRef<unsigned int> caseValues, BlockRange caseDestinations,
    ArrayRef<ValueRange> caseArguments, Block *defaultDest,
    ValueRange defaultArgs) {
  cf::SwitchOp::create(builder, loc, flag, defaultDest, defaultArgs,
                       llvm::to_vector_of<int32_t>(caseValues),
                       caseDestinations, caseArguments);
}
 
Value ControlFlowToSCFTransformation::getUndefValue(Location loc,
                                                    OpBuilder &builder,
                                                    Type type) {
  return ub::PoisonOp::create(builder, loc, type, nullptr);
}
 
FailureOr<Operation *>
ControlFlowToSCFTransformation::createUnreachableTerminator(Location loc,
                                                            OpBuilder &builder,
                                                            Region &region) {
 
  // TODO: This should create a `ub.unreachable` op. Once such an operation
  //       exists to make the pass independent of the func dialect. For now just
  //       return poison values.
  Operation *parentOp = region.getParentOp();
  auto funcOp = dyn_cast<func::FuncOp>(parentOp);
  if (!funcOp)
    return emitError(loc, "Cannot create unreachable terminator for '")
           << parentOp->getName() << "'";
 
  return func::ReturnOp::create(
             builder, loc,
             llvm::map_to_vector(
                 funcOp.getResultTypes(),
                 [&](Type type) { return getUndefValue(loc, builder, type); }))
      .getOperation();
}
 
namespace {
 
struct LiftControlFlowToSCF
    : public impl::LiftControlFlowToSCFPassBase<LiftControlFlowToSCF> {
 
  using Base::Base;
 
  void runOnOperation() override {
    ControlFlowToSCFTransformation transformation;
 
    bool changed = false;
    Operation *op = getOperation();
    WalkResult result = op->walk([&](func::FuncOp funcOp) {
      if (funcOp.getBody().empty())
        return WalkResult::advance();
 
      auto &domInfo = funcOp != op ? getChildAnalysis<DominanceInfo>(funcOp)
                                   : getAnalysis<DominanceInfo>();
 
      auto visitor = [&](Operation *innerOp) -> WalkResult {
        for (Region &reg : innerOp->getRegions()) {
          FailureOr<bool> changedFunc =
              transformCFGToSCF(reg, transformation, domInfo);
          if (failed(changedFunc))
            return WalkResult::interrupt();
 
          changed |= *changedFunc;
        }
        return WalkResult::advance();
      };
 
      if (funcOp->walk<WalkOrder::PostOrder>(visitor).wasInterrupted())
        return WalkResult::interrupt();
 
      return WalkResult::advance();
    });
    if (result.wasInterrupted())
      return signalPassFailure();
 
    if (!changed)
      markAllAnalysesPreserved();
  }
};
} // namespace