blob: 3d725d77db8ec4d59faed9ab3cb7b2ae09714f33 [file] [log] [blame]
//===- SCCP.cpp - Sparse Conditional Constant Propagation -----------------===//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// This transformation pass performs a sparse conditional constant propagation
// in MLIR. It identifies values known to be constant, propagates that
// information throughout the IR, and replaces them. This is done with an
// optimistic dataflow analysis that assumes that all values are constant until
// proven otherwise.
#include "PassDetail.h"
#include "mlir/Analysis/DataFlowAnalysis.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Dialect.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/Passes.h"
using namespace mlir;
// SCCP Analysis
namespace {
struct SCCPLatticeValue {
SCCPLatticeValue(Attribute constant = {}, Dialect *dialect = nullptr)
: constant(constant), constantDialect(dialect) {}
/// The pessimistic state of SCCP is non-constant.
static SCCPLatticeValue getPessimisticValueState(MLIRContext *context) {
return SCCPLatticeValue();
static SCCPLatticeValue getPessimisticValueState(Value value) {
return SCCPLatticeValue();
/// Equivalence for SCCP only accounts for the constant, not the originating
/// dialect.
bool operator==(const SCCPLatticeValue &rhs) const {
return constant == rhs.constant;
/// To join the state of two values, we simply check for equivalence.
static SCCPLatticeValue join(const SCCPLatticeValue &lhs,
const SCCPLatticeValue &rhs) {
return lhs == rhs ? lhs : SCCPLatticeValue();
/// The constant attribute value.
Attribute constant;
/// The dialect the constant originated from. This is not used as part of the
/// key, and is only needed to materialize the held constant if necessary.
Dialect *constantDialect;
struct SCCPAnalysis : public ForwardDataFlowAnalysis<SCCPLatticeValue> {
using ForwardDataFlowAnalysis<SCCPLatticeValue>::ForwardDataFlowAnalysis;
~SCCPAnalysis() override = default;
visitOperation(Operation *op,
ArrayRef<LatticeElement<SCCPLatticeValue> *> operands) final {
// Don't try to simulate the results of a region operation as we can't
// guarantee that folding will be out-of-place. We don't allow in-place
// folds as the desire here is for simulated execution, and not general
// folding.
if (op->getNumRegions())
return markAllPessimisticFixpoint(op->getResults());
SmallVector<Attribute> constantOperands(
llvm::map_range(operands, [](LatticeElement<SCCPLatticeValue> *value) {
return value->getValue().constant;
// Save the original operands and attributes just in case the operation
// folds in-place. The constant passed in may not correspond to the real
// runtime value, so in-place updates are not allowed.
SmallVector<Value, 8> originalOperands(op->getOperands());
DictionaryAttr originalAttrs = op->getAttrDictionary();
// Simulate the result of folding this operation to a constant. If folding
// fails or was an in-place fold, mark the results as overdefined.
SmallVector<OpFoldResult, 8> foldResults;
if (failed(op->fold(constantOperands, foldResults)))
return markAllPessimisticFixpoint(op->getResults());
// If the folding was in-place, mark the results as overdefined and reset
// the operation. We don't allow in-place folds as the desire here is for
// simulated execution, and not general folding.
if (foldResults.empty()) {
return markAllPessimisticFixpoint(op->getResults());
// Merge the fold results into the lattice for this operation.
assert(foldResults.size() == op->getNumResults() && "invalid result size");
Dialect *dialect = op->getDialect();
ChangeResult result = ChangeResult::NoChange;
for (unsigned i = 0, e = foldResults.size(); i != e; ++i) {
LatticeElement<SCCPLatticeValue> &lattice =
// Merge in the result of the fold, either a constant or a value.
OpFoldResult foldResult = foldResults[i];
if (Attribute attr = foldResult.dyn_cast<Attribute>())
result |= lattice.join(SCCPLatticeValue(attr, dialect));
result |= lattice.join(getLatticeElement(foldResult.get<Value>()));
return result;
/// Implementation of `getSuccessorsForOperands` that uses constant operands
/// to potentially remove dead successors.
LogicalResult getSuccessorsForOperands(
BranchOpInterface branch,
ArrayRef<LatticeElement<SCCPLatticeValue> *> operands,
SmallVectorImpl<Block *> &successors) final {
SmallVector<Attribute> constantOperands(
llvm::map_range(operands, [](LatticeElement<SCCPLatticeValue> *value) {
return value->getValue().constant;
if (Block *singleSucc = branch.getSuccessorForOperands(constantOperands)) {
return success();
return failure();
/// Implementation of `getSuccessorsForOperands` that uses constant operands
/// to potentially remove dead region successors.
void getSuccessorsForOperands(
RegionBranchOpInterface branch, Optional<unsigned> sourceIndex,
ArrayRef<LatticeElement<SCCPLatticeValue> *> operands,
SmallVectorImpl<RegionSuccessor> &successors) final {
SmallVector<Attribute> constantOperands(
llvm::map_range(operands, [](LatticeElement<SCCPLatticeValue> *value) {
return value->getValue().constant;
branch.getSuccessorRegions(sourceIndex, constantOperands, successors);
} // namespace
// SCCP Rewrites
/// Replace the given value with a constant if the corresponding lattice
/// represents a constant. Returns success if the value was replaced, failure
/// otherwise.
static LogicalResult replaceWithConstant(SCCPAnalysis &analysis,
OpBuilder &builder,
OperationFolder &folder, Value value) {
LatticeElement<SCCPLatticeValue> *lattice =
if (!lattice)
return failure();
SCCPLatticeValue &latticeValue = lattice->getValue();
if (!latticeValue.constant)
return failure();
// Attempt to materialize a constant for the given value.
Dialect *dialect = latticeValue.constantDialect;
Value constant = folder.getOrCreateConstant(
builder, dialect, latticeValue.constant, value.getType(), value.getLoc());
if (!constant)
return failure();
return success();
/// Rewrite the given regions using the computing analysis. This replaces the
/// uses of all values that have been computed to be constant, and erases as
/// many newly dead operations.
static void rewrite(SCCPAnalysis &analysis, MLIRContext *context,
MutableArrayRef<Region> initialRegions) {
SmallVector<Block *> worklist;
auto addToWorklist = [&](MutableArrayRef<Region> regions) {
for (Region &region : regions)
for (Block &block : llvm::reverse(region))
// An operation folder used to create and unique constants.
OperationFolder folder(context);
OpBuilder builder(context);
while (!worklist.empty()) {
Block *block = worklist.pop_back_val();
for (Operation &op : llvm::make_early_inc_range(*block)) {
// Replace any result with constants.
bool replacedAll = op.getNumResults() != 0;
for (Value res : op.getResults())
replacedAll &=
succeeded(replaceWithConstant(analysis, builder, folder, res));
// If all of the results of the operation were replaced, try to erase
// the operation completely.
if (replacedAll && wouldOpBeTriviallyDead(&op)) {
assert(op.use_empty() && "expected all uses to be replaced");
// Add any the regions of this operation to the worklist.
// Replace any block arguments with constants.
for (BlockArgument arg : block->getArguments())
(void)replaceWithConstant(analysis, builder, folder, arg);
// SCCP Pass
namespace {
struct SCCP : public SCCPBase<SCCP> {
void runOnOperation() override;
} // end anonymous namespace
void SCCP::runOnOperation() {
Operation *op = getOperation();
SCCPAnalysis analysis(op->getContext());;
rewrite(analysis, op->getContext(), op->getRegions());
std::unique_ptr<Pass> mlir::createSCCPPass() {
return std::make_unique<SCCP>();