CIRSimplify.cpp

Go to the documentation of this file.

//===----------------------------------------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#include "PassDetail.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Block.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Region.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "clang/CIR/Dialect/IR/CIRDialect.h"
#include "clang/CIR/Dialect/Passes.h"
#include "llvm/ADT/SmallVector.h"
 
using namespace mlir;
using namespace cir;
 
namespace mlir {
#define GEN_PASS_DEF_CIRSIMPLIFY
#include "clang/CIR/Dialect/Passes.h.inc"
} // namespace mlir
 
//===----------------------------------------------------------------------===//
// Rewrite patterns
//===----------------------------------------------------------------------===//
 
namespace {
 
/// Dominance information that stays coherent while patterns rewrite the IR.
/// DominanceInfo caches per region and never notices IR edits, so this doubles
/// as the rewrite listener and drops a region's cached tree whenever its block
/// list changes. Successor edits would need notifyOperationModified as well; no
/// pattern here makes any.
class RewriteDominanceInfo final : public RewriterBase::Listener {
public:
  DominanceInfo &getDominanceInfo() { return domInfo; }
 
  void notifyBlockInserted(Block *block, Region *previous,
                           Region::iterator previousIt) override {
    domInfo.invalidate(block->getParent(), /*recursive=*/false);
    if (previous)
      domInfo.invalidate(previous, /*recursive=*/false);
  }
 
  void notifyBlockErased(Block *block) override {
    domInfo.invalidate(block->getParent(), /*recursive=*/false);
  }
 
private:
  DominanceInfo domInfo;
};
 
/// Find the `cir.store` operation that stores to the given alloca and dominates
/// the given load operation. Dominance calculation is done through the given
/// DominanceInfo object.
///
/// Return nullptr if no such store operation exists or if multiple store
/// operations satisfy the criteria.
cir::StoreOp findDominatingInitOp(cir::AllocaOp alloca, cir::LoadOp load,
                                  const DominanceInfo &domInfo) {
  cir::StoreOp result;
 
  // Walk through all uses of the alloca and visit the store operations that
  // store to the alloca
  for (const mlir::OpOperand &use : alloca->getUses()) {
    auto store = mlir::dyn_cast<cir::StoreOp>(use.getOwner());
    if (!store)
      continue;
 
    // `cir.store` has two operands, we're only interested if the store is
    // storing into the alloca, not if the store is storing the address of the
    // alloca slot into somewhere else
    if (use.getOperandNumber() != cir::StoreOp::odsIndex_addr)
      continue;
 
    if (domInfo.dominates(store, load)) {
      if (result) {
        // If we have already found a dominating store, then there are multiple
        // dominating stores, we intentionally don't simplify the load.
        return nullptr;
      }
      result = store;
    }
  }
 
  return result;
}
 
/// Simplify `cir.load` that loads from an alloca marked as "constant".
///
/// For example:
///
///   %0 = cir.alloca "x" align(4) const : !cir.ptr<!s32i>
///   cir.store %init, %0 : !s32i, !cir.ptr<!s32i>
///   %1 = cir.load %0 : !cir.ptr<!s32i>
///
/// All uses of the load above could be replaced with the SSA value `%init`.
struct SimplifyConstantLoad final : public OpRewritePattern<LoadOp> {
  SimplifyConstantLoad(MLIRContext *context, RewriteDominanceInfo &domInfo)
      : OpRewritePattern(context), domInfo(domInfo) {}
 
  LogicalResult matchAndRewrite(LoadOp op,
                                PatternRewriter &rewriter) const override {
    // Volatile or atomic loads should not be simplified.
    if (op.getIsVolatile() || op.getMemOrder())
      return mlir::failure();
 
    auto allocaOp = op.getAddr().getDefiningOp<cir::AllocaOp>();
    if (!allocaOp || !allocaOp.getConstant())
      return mlir::failure();
 
    cir::StoreOp initStoreOp =
        findDominatingInitOp(allocaOp, op, domInfo.getDominanceInfo());
    if (!initStoreOp)
      return mlir::failure();
    if (initStoreOp.getIsVolatile() || initStoreOp.getMemOrder()) {
      // We intentionally act conservatively here and we don't want to simplify
      // the load if the corresponding store is either volatile or atomic.
      return mlir::failure();
    }
 
    rewriter.replaceOp(op, initStoreOp.getValue());
    return mlir::success();
  }
 
private:
  RewriteDominanceInfo &domInfo;
};
 
/// Simplify suitable ternary operations into select operations.
///
/// For now we only simplify those ternary operations whose true and false
/// branches directly yield a value or a constant. That is, both of the true and
/// the false branch must either contain a cir.yield operation as the only
/// operation in the branch, or contain a cir.const operation followed by a
/// cir.yield operation that yields the constant value.
///
/// For example, we will simplify the following ternary operation:
///
///   %0 = ...
///   %1 = cir.ternary (%condition, true {
///     %2 = cir.const ...
///     cir.yield %2
///   } false {
///     cir.yield %0
///
/// into the following sequence of operations:
///
///   %1 = cir.const ...
///   %0 = cir.select if %condition then %1 else %2
struct SimplifyTernary final : public OpRewritePattern<TernaryOp> {
  using OpRewritePattern<TernaryOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(TernaryOp op,
                                PatternRewriter &rewriter) const override {
    if (op->getNumResults() != 1)
      return mlir::failure();
 
    if (!isSimpleTernaryBranch(op.getTrueRegion()) ||
        !isSimpleTernaryBranch(op.getFalseRegion()))
      return mlir::failure();
 
    cir::YieldOp trueBranchYieldOp =
        mlir::cast<cir::YieldOp>(op.getTrueRegion().front().getTerminator());
    cir::YieldOp falseBranchYieldOp =
        mlir::cast<cir::YieldOp>(op.getFalseRegion().front().getTerminator());
    mlir::Value trueValue = trueBranchYieldOp.getArgs()[0];
    mlir::Value falseValue = falseBranchYieldOp.getArgs()[0];
 
    rewriter.inlineBlockBefore(&op.getTrueRegion().front(), op);
    rewriter.inlineBlockBefore(&op.getFalseRegion().front(), op);
    rewriter.eraseOp(trueBranchYieldOp);
    rewriter.eraseOp(falseBranchYieldOp);
    rewriter.replaceOpWithNewOp<cir::SelectOp>(op, op.getCond(), trueValue,
                                               falseValue);
 
    return mlir::success();
  }
 
private:
  bool isSimpleTernaryBranch(mlir::Region &region) const {
    if (!region.hasOneBlock())
      return false;
 
    mlir::Block &onlyBlock = region.front();
    mlir::Block::OpListType &ops = onlyBlock.getOperations();
 
    // The region/block could only contain at most 2 operations.
    if (ops.size() > 2)
      return false;
 
    if (ops.size() == 1) {
      // The region/block only contain a cir.yield operation.
      return true;
    }
 
    // Check whether the region/block contains a cir.const followed by a
    // cir.yield that yields the value.
    auto yieldOp = mlir::cast<cir::YieldOp>(onlyBlock.getTerminator());
    auto yieldValueDefOp =
        yieldOp.getArgs()[0].getDefiningOp<cir::ConstantOp>();
    return yieldValueDefOp && yieldValueDefOp->getBlock() == &onlyBlock;
  }
};
 
/// Simplify select operations with boolean constants into simpler forms.
///
/// This pattern simplifies select operations where both true and false values
/// are boolean constants. Two specific cases are handled:
///
/// 1. When selecting between true and false based on a condition,
///    the operation simplifies to just the condition itself:
///
///    %0 = cir.select if %condition then true else false
///    ->
///    (replaced with %condition directly)
///
/// 2. When selecting between false and true based on a condition,
///    the operation simplifies to the logical negation of the condition:
///
///    %0 = cir.select if %condition then false else true
///    ->
///    %0 = cir.not %condition
struct SimplifySelect : public OpRewritePattern<SelectOp> {
  using OpRewritePattern<SelectOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(SelectOp op,
                                PatternRewriter &rewriter) const final {
    auto trueValueOp = op.getTrueValue().getDefiningOp<cir::ConstantOp>();
    auto falseValueOp = op.getFalseValue().getDefiningOp<cir::ConstantOp>();
    if (!trueValueOp || !falseValueOp)
      return mlir::failure();
 
    auto trueValue = trueValueOp.getValueAttr<cir::BoolAttr>();
    auto falseValue = falseValueOp.getValueAttr<cir::BoolAttr>();
    if (!trueValue || !falseValue)
      return mlir::failure();
 
    // cir.select if %0 then #true else #false -> %0
    if (trueValue.getValue() && !falseValue.getValue()) {
      rewriter.replaceAllUsesWith(op, op.getCondition());
      rewriter.eraseOp(op);
      return mlir::success();
    }
 
    // cir.select if %0 then #false else #true -> cir.not %0
    if (!trueValue.getValue() && falseValue.getValue()) {
      rewriter.replaceOpWithNewOp<cir::NotOp>(op, op.getCondition());
      return mlir::success();
    }
 
    return mlir::failure();
  }
};
 
/// Simplify `cir.switch` operations by folding cascading cases
/// into a single `cir.case` with the `anyof` kind.
///
/// This pattern identifies cascading cases within a `cir.switch` operation.
/// Cascading cases are defined as consecutive `cir.case` operations of kind
/// `equal`, each containing a single `cir.yield` operation in their body.
///
/// The pattern merges these cascading cases into a single `cir.case` operation
/// with kind `anyof`, aggregating all the case values.
///
/// The merging process continues until a `cir.case` with a different body
/// (e.g., containing `cir.break` or compound stmt) is encountered, which
/// breaks the chain.
///
/// Example:
///
/// Before:
///   cir.case equal, [#cir.int<0> : !s32i] {
///     cir.yield
///   }
///   cir.case equal, [#cir.int<1> : !s32i] {
///     cir.yield
///   }
///   cir.case equal, [#cir.int<2> : !s32i] {
///     cir.break
///   }
///
/// After applying SimplifySwitch:
///   cir.case anyof, [#cir.int<0> : !s32i, #cir.int<1> : !s32i, #cir.int<2> :
///   !s32i] {
///     cir.break
///   }
struct SimplifySwitch : public OpRewritePattern<SwitchOp> {
  using OpRewritePattern<SwitchOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(SwitchOp op,
                                PatternRewriter &rewriter) const override {
 
    LogicalResult changed = mlir::failure();
    SmallVector<CaseOp, 8> cases;
    SmallVector<CaseOp, 4> cascadingCases;
    SmallVector<mlir::Attribute, 4> cascadingCaseValues;
 
    op.collectCases(cases);
    if (cases.empty())
      return mlir::failure();
 
    auto resetMergeState = [&]() {
      cascadingCases.clear();
      cascadingCaseValues.clear();
    };
 
    auto flushMergedOps = [&]() {
      for (CaseOp &c : cascadingCases)
        rewriter.eraseOp(c);
      resetMergeState();
    };
 
    auto mergeCascadingInto = [&](CaseOp &target) {
      rewriter.modifyOpInPlace(target, [&]() {
        target.setValueAttr(rewriter.getArrayAttr(cascadingCaseValues));
        target.setKind(CaseOpKind::Anyof);
      });
      changed = mlir::success();
    };
 
    // Merge all pending cascading cases into the last one collected, which
    // survives as a distinct `cir.case`; the rest are erased since their
    // values have been folded into it.
    auto mergeLastCascadingAndFlush = [&]() {
      assert(!cir::MissingFeatures::foldRangeCase());
      CaseOp lastCascadingCase = cascadingCases.back();
      mergeCascadingInto(lastCascadingCase);
      cascadingCases.pop_back();
      flushMergedOps();
    };
 
    for (CaseOp c : cases) {
      // Cascading cases must be textually adjacent to the previously
      // collected cascading case. This is false when something is in the
      // way (e.g. a goto) or when the previous cascading case was found
      // nested inside a sibling case's body (e.g. a case label that falls
      // through into a compound statement) rather than next to `c`.
      bool isAdjacentToLastCascadingCase =
          !cascadingCases.empty() &&
          c->getPrevNode() == cascadingCases.back().getOperation();
 
      if (!cascadingCases.empty() && !isAdjacentToLastCascadingCase) {
        if (cascadingCases.size() > 1)
          mergeLastCascadingAndFlush();
        else
          resetMergeState();
      }
 
      cir::CaseOpKind kind = c.getKind();
      if (kind == cir::CaseOpKind::Equal &&
          isa<YieldOp>(c.getCaseRegion().front().front())) {
        // If the case contains only a YieldOp, collect it for cascading merge
        cascadingCases.push_back(c);
        cascadingCaseValues.push_back(c.getValue()[0]);
      } else if (kind == cir::CaseOpKind::Equal && !cascadingCases.empty()) {
        // merge previously collected cascading cases
        cascadingCaseValues.push_back(c.getValue()[0]);
        mergeCascadingInto(c);
        flushMergedOps();
      } else if (kind != cir::CaseOpKind::Equal && cascadingCases.size() > 1) {
        // If a Default, Anyof or Range case is found and there are previous
        // cascading cases, merge all of them into the last cascading case.
        // We don't currently fold case range statements with other case
        // statements.
        mergeLastCascadingAndFlush();
      } else {
        resetMergeState();
      }
    }
 
    // Edge case: all cases are simple cascading cases
    if (cascadingCases.size() == cases.size())
      mergeLastCascadingAndFlush();
 
    return changed;
  }
};
 
struct SimplifyVecSplat : public OpRewritePattern<VecSplatOp> {
  using OpRewritePattern<VecSplatOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(VecSplatOp op,
                                PatternRewriter &rewriter) const override {
    mlir::Value splatValue = op.getValue();
    auto constant = splatValue.getDefiningOp<cir::ConstantOp>();
    if (!constant)
      return mlir::failure();
 
    auto value = constant.getValue();
    if (!mlir::isa_and_nonnull<cir::IntAttr>(value) &&
        !mlir::isa_and_nonnull<cir::FPAttr>(value))
      return mlir::failure();
 
    // FIXME(CIR): We should consider making a matrix constant attribute so that
    // we can simplify it here too.
    assert(!MissingFeatures::matrixType());
    auto resultType = mlir::dyn_cast<cir::VectorType>(op.getResult().getType());
    if (!resultType || resultType.isScalable())
      return mlir::failure();
    SmallVector<mlir::Attribute, 16> elements(resultType.getFixedSize(), value);
    auto constVecAttr = cir::ConstVectorAttr::get(
        resultType, mlir::ArrayAttr::get(getContext(), elements));
 
    rewriter.replaceOpWithNewOp<cir::ConstantOp>(op, constVecAttr);
    return mlir::success();
  }
};
 
//===----------------------------------------------------------------------===//
// CIRSimplifyPass
//===----------------------------------------------------------------------===//
 
struct CIRSimplifyPass : public impl::CIRSimplifyBase<CIRSimplifyPass> {
  using CIRSimplifyBase::CIRSimplifyBase;
 
  void runOnOperation() override;
};
 
void populateMergeCleanupPatterns(RewritePatternSet &patterns,
                                  RewriteDominanceInfo &domInfo) {
  // clang-format off
  patterns.add<
    SimplifyTernary,
    SimplifySelect,
    SimplifySwitch,
    SimplifyVecSplat
  >(patterns.getContext());
  // clang-format on
  patterns.add<SimplifyConstantLoad>(patterns.getContext(), domInfo);
}
 
void CIRSimplifyPass::runOnOperation() {
  // Collect rewrite patterns. SimplifyConstantLoad queries dominance, so the
  // dominance info doubles as the driver's listener to stay coherent.
  RewriteDominanceInfo domInfo;
  RewritePatternSet patterns(&getContext());
  populateMergeCleanupPatterns(patterns, domInfo);
 
  // Collect operations to apply patterns.
  llvm::SmallVector<Operation *, 16> ops;
  getOperation()->walk([&](Operation *op) {
    if (isa<TernaryOp, SelectOp, SwitchOp, VecSplatOp, LoadOp>(op))
      ops.push_back(op);
  });
 
  // Apply patterns. Scope the rewrite to the body region, which encloses
  // the common ancestor the driver would otherwise search for.
  GreedyRewriteConfig config;
  config.setScope(&getOperation()->getRegion(0));
  config.setListener(&domInfo);
  if (applyOpPatternsGreedily(ops, std::move(patterns), config).failed())
    signalPassFailure();
}
 
} // namespace
 
std::unique_ptr<Pass> mlir::createCIRSimplifyPass() {
  return std::make_unique<CIRSimplifyPass>();
}