mlir/lib/Analysis/DataFlow/ConstantPropagationAnalysis.cpp

//===- ConstantPropagationAnalysis.cpp - Constant propagation analysis ----===//
//
// 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 "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
#include <cassert>

#define DEBUG_TYPE "constant-propagation"

using namespace mlir;
using namespace mlir::dataflow;

//===----------------------------------------------------------------------===//
// ConstantValue
//===----------------------------------------------------------------------===//

void ConstantValue::print(raw_ostream &os) const {
  if (isUninitialized()) {
    os << "<UNINITIALIZED>";
    return;
  }
  if (getConstantValue() == nullptr) {
    os << "<UNKNOWN>";
    return;
  }
  return getConstantValue().print(os);
}

//===----------------------------------------------------------------------===//
// SparseConstantPropagation
//===----------------------------------------------------------------------===//

LogicalResult SparseConstantPropagation::visitOperation(
    Operation *op, ArrayRef<const Lattice<ConstantValue> *> operands,
    ArrayRef<Lattice<ConstantValue> *> results) {
  LDBG() << "SCP: Visiting operation: " << *op;

  // 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()) {
    setAllToEntryStates(results);
    return success();
  }

  SmallVector<Attribute, 8> constantOperands;
  constantOperands.reserve(op->getNumOperands());
  for (auto *operandLattice : operands) {
    if (operandLattice->getValue().isUninitialized())
      return success();
    constantOperands.push_back(operandLattice->getValue().getConstantValue());
  }

  // 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->getDiscardableAttrDictionary();
  Attribute originalProperties = op->getPropertiesAsAttribute();

  // Simulate the result of folding this operation to a constant.
  SmallVector<OpFoldResult, 8> foldResults;
  foldResults.reserve(op->getNumResults());
  LogicalResult folded = op->fold(constantOperands, foldResults);

  // `fold` can mutate the operation in place and still return an out-of-place
  // result, so the mutation must be reverted regardless of the outcome.
  // Only write the operands back if the fold changed them, as `setOperands`
  // relinks use-lists even for identical values.
  if (!llvm::equal(op->getOperands(), originalOperands))
    op->setOperands(originalOperands);
  op->setDiscardableAttrs(originalAttrs);
  if (originalProperties)
    (void)op->setPropertiesFromAttribute(originalProperties,
                                         /*emitError=*/nullptr);

  // If folding failed or was in-place, mark the results as overdefined. We
  // don't allow in-place folds here: the goal is simulated execution, not
  // general folding.
  if (failed(folded) || foldResults.empty()) {
    setAllToEntryStates(results);
    return success();
  }

  // Merge the fold results into the lattice for this operation.
  assert(foldResults.size() == op->getNumResults() && "invalid result size");
  for (const auto it : llvm::zip(results, foldResults)) {
    Lattice<ConstantValue> *lattice = std::get<0>(it);

    // Merge in the result of the fold, either a constant or a value.
    OpFoldResult foldResult = std::get<1>(it);
    if (Attribute attr = llvm::dyn_cast_if_present<Attribute>(foldResult)) {
      LDBG() << "Folded to constant: " << attr;
      propagateIfChanged(lattice,
                         lattice->join(ConstantValue(attr, op->getDialect())));
    } else {
      Value foldValue = cast<Value>(foldResult);
      LDBG() << "Folded to value: " << foldValue;
      // The folded value may not be an operand of `op`, so we need to use
      // `getLatticeElementFor` (and not `getLatticeElement`) so that
      // this operation is revisited if that value's lattice widens later.
      AbstractSparseForwardDataFlowAnalysis::join(
          lattice, *getLatticeElementFor(getProgramPointAfter(op), foldValue));
    }
  }
  return success();
}

void SparseConstantPropagation::setToEntryState(
    Lattice<ConstantValue> *lattice) {
  propagateIfChanged(lattice,
                     lattice->join(ConstantValue::getUnknownConstant()));
}