// Copyright 2021 The IREE Authors
//
// Licensed 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 "iree/compiler/Dialect/Util/Analysis/Constant/ConstExpr.h"
#include "iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
#include "iree/compiler/Utils/StringUtils.h"
#include "llvm/Support/Debug.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Iterators.h"
#include "mlir/IR/SymbolTable.h"

#define DEBUG_TYPE "iree-constexpr"

namespace mlir::iree_compiler::IREE::Util {

#define GEN_PASS_DEF_HOISTINTOGLOBALSPASS
#include "iree/compiler/Dialect/Util/Transforms/Passes.h.inc"

namespace {

static llvm::cl::opt<std::string> clPrintDotGraphToFile(
    "iree-util-hoist-into-globals-print-constexpr-dotgraph-to",
    llvm::cl::desc(
        "Prints a dot graph representing the const-expr analysis. The red "
        "nodes represent roots and the green nodes represent hoisted values."),
    llvm::cl::value_desc("filename"));

// Maps an original value in the program to the symbol name of a global.
using HoistedValueMap = llvm::DenseMap<Value, GlobalOp>;

// Returns true if |op| is a flow dispatch region/workgroups op, i.e. a region
// that is later outlined into an isolated executable module.
// The dialect/op names are checked by string since Flow depends on Util
// and we don't want to create a circular dependency.
static bool isDispatchContainer(Operation *op) {
  StringRef opName = op->getName().getStringRef();
  return opName == "flow.dispatch.region" ||
         opName == "flow.dispatch.workgroups";
}

static std::string getHoistedName(Type type) {
  std::string str;
  llvm::raw_string_ostream os(str);
  os << "__hoisted_";
  auto rankedTensorType = dyn_cast<RankedTensorType>(type);
  if (rankedTensorType && rankedTensorType.getEncoding()) {
    rankedTensorType.dropEncoding().print(os);
    os << "_encoded";
  } else {
    type.print(os);
  }
  str = sanitizeSymbolName(str);
  if (str.substr(str.size() - 1) == "_") {
    str = str.substr(0, str.size() - 1); // strip trailing _
  }
  return str;
}

// Hoist expressions into globals. It is not expected that such a greedy
// algorithm is great, but it is simple. Naive use of this algorithm very likely
// favors programs that consume more memory at runtime than is strictly
// necessary. Either this algorithm can be made smarter or a follow-on pass
// can sink globals into the program where it is profitable to reduce
// working set size.
class HoistIntoGlobalsPass
    : public impl::HoistIntoGlobalsPassBase<HoistIntoGlobalsPass> {
public:
  using Base::Base;

  void getDependentDialects(DialectRegistry &registry) const override {
    registerConstExprDependentDialects(registry);
    if (this->registerDependentDialectsFn) {
      (*registerDependentDialectsFn)(registry);
    }
  }

  HoistIntoGlobalsPass(const ExprHoistingOptions &options)
      : Base(),
        registerDependentDialectsFn(options.registerDependentDialectsFn) {
    this->maxSizeIncreaseThreshold.setValue(options.maxSizeIncreaseThreshold);
  }

  void runOnOperation() override {
    SymbolTable moduleSymbols(getOperation());
    const auto &constExprs = getAnalysis<ConstExprAnalysis>();
    ConstExprHoistingPolicy policy(constExprs, this->maxSizeIncreaseThreshold);
    policy.initialize();

    // Print analysis dot graph if requested.
    if (!clPrintDotGraphToFile.empty()) {
      std::error_code ec;
      llvm::raw_fd_ostream file(clPrintDotGraphToFile, ec);
      if (ec) {
        getOperation().emitError()
            << "failed to open file for printing dot graph: " << ec.message();
        return signalPassFailure();
      }
      policy.printDotGraph(file);
      file.close();
    }

    // Walk all operations in the program and hoist any escapes from
    // const-expr values into globals. Note that we must walk the const-exprs
    // in topological order so that corresponding initializers will be created
    // in order without depending on globals that have not been initialized
    // yet.
    OpBuilder builder(&getContext());
    for (auto funcOp : getOperation().getOps<FunctionOpInterface>()) {
      // Ignore initializers.
      if (isa<IREE::Util::InitializerOpInterface>(funcOp.getOperation())) {
        continue;
      }

      // Maps original values to newly materialized globals (per-function).
      HoistedValueMap hoistedMap;

      // Operation order for deterministic sorting (per-function).
      llvm::DenseMap<Operation *, unsigned> opOrder;
      unsigned orderIdx = 0;

      auto walkRes = funcOp.walk<WalkOrder::PreOrder>([&](Operation *iterOp) {
        // We only want to look at const-expr ops (non roots) since they may
        // have interesting escapes. Non const-expr ops are ignored here, but
        // we still fall through to the dispatch check below so that we stop
        // descending into dispatch bodies even when the dispatch op itself is
        // not a const-expr.
        auto *iterInfo = constExprs.lookup(iterOp);
        if (iterInfo) {
          // Record operation order for deterministic sorting. Since we walk in
          // PreOrder, producers are visited before their users.
          opOrder[iterOp] = orderIdx++;
          for (Value constExprResult : iterOp->getResults()) {
            auto *resultInfo = constExprs.lookup(constExprResult);
            assert(resultInfo && "must have const-expr info");
            if (policy.getDecision(resultInfo)->getOutcome() !=
                ConstExprHoistingPolicy::ENABLE_HOIST) {
              continue;
            }
            if (failed(hoistConstExpr(constExprResult, hoistedMap,
                                      moduleSymbols, constExprs, opOrder))) {
              return WalkResult::interrupt();
            }
          }
        }

        // FIXME: this is a temporary workaround for issue #24671
        // Never hoist values defined inside a dispatch region
        if (isDispatchContainer(iterOp)) {
          return WalkResult::skip();
        }
        return WalkResult::advance();
      });
      if (walkRes.wasInterrupted()) {
        return signalPassFailure();
      }

      // Apply RAUW cleanups for this function. We do this after cloning to
      // avoid invalidating the analysis during the walk.
      // Sort the hoisted values by program order for deterministic output.
      using HoistedValue = std::pair<Value, GlobalOp>;
      auto sortedHoisted = llvm::to_vector_of<HoistedValue>(hoistedMap);
      llvm::sort(sortedHoisted,
                 [&opOrder](const HoistedValue &lhs, const HoistedValue &rhs) {
                   return opOrder[lhs.first.getDefiningOp()] <
                          opOrder[rhs.first.getDefiningOp()];
                 });

      for (auto [originalValue, globalOp] : sortedHoisted) {
        builder.setInsertionPointAfterValue(originalValue);
        auto loadOp = globalOp.createLoadOp(globalOp->getLoc(), builder);
        if (!originalValue.getDefiningOp()
                 ->getParentOfType<IREE::Util::InitializerOpInterface>()) {
          loadOp.setGlobalImmutable(true);
        }
        Value loadedValue = loadOp.getLoadedGlobalValue();
        // Call user hook to cast back to the original type.
        if (auto hoistableType = dyn_cast<IREE::Util::HoistableTypeInterface>(
                originalValue.getType())) {
          loadedValue = hoistableType.decodeStorageType(
              builder, loadedValue.getLoc(), originalValue.getType(),
              loadedValue);
        }
        if (loadedValue.getType() != originalValue.getType()) {
          getOperation().emitError()
              << "Unresolved conflict between casted global of type "
              << loadedValue.getType() << " and original type "
              << originalValue.getType();
          return signalPassFailure();
        }
        originalValue.replaceAllUsesWith(loadedValue);
      }
    }
    cleanupDeadOps(constExprs);
  }

  Operation *getTopLevelOp(Operation *childOp) {
    auto *moduleBlock = getOperation().getBody();
    auto *op = childOp;
    while (op->getBlock() != moduleBlock) {
      op = op->getParentOp();
    }
    return op;
  }

  LogicalResult
  hoistConstExpr(Value originalValue, HoistedValueMap &hoistedMap,
                 SymbolTable &moduleSymbols,
                 const ConstExprAnalysis &constExprs,
                 const llvm::DenseMap<Operation *, unsigned> &opOrder) {
    IREE::Util::GlobalOp existingGlobal = hoistedMap.lookup(originalValue);
    if (existingGlobal) {
      return success();
    }

    // Gather any dialect attributes we may need to preserve.
    auto *topLevelOp = getTopLevelOp(originalValue.getDefiningOp());
    NamedAttrList dialectAttrs;
    IREE::Util::HoistableAttrInterface::gatherHoistableAttrs(topLevelOp,
                                                             dialectAttrs);

    // No existing mapping - create a new global.
    OpBuilder moduleBuilder(topLevelOp);
    auto initializerOp = IREE::Util::InitializerOp::create(
        moduleBuilder, originalValue.getLoc());
    initializerOp->setDialectAttrs(dialectAttrs);
    auto initializerBuilder =
        OpBuilder::atBlockEnd(initializerOp.addEntryBlock());
    moduleBuilder.setInsertionPoint(initializerOp);
    if (failed(cloneConstExprInto(initializerOp.getLoc(), moduleBuilder,
                                  initializerBuilder, originalValue,
                                  dialectAttrs, hoistedMap, moduleSymbols,
                                  constExprs, opOrder))) {
      return failure();
    }

    existingGlobal = hoistedMap.lookup(originalValue);
    (void)existingGlobal;
    assert(existingGlobal &&
           "hoisting const-expr should have mapped a global for the requested "
           "value");
    return success();
  }

  void
  cloneProducerTreeInto(OpBuilder &initializerBuilder,
                        const ConstExprAnalysis::ConstValueInfo *producerInfo,
                        HoistedValueMap &hoistedMap, IRMapping &cloneMapping,
                        const ConstExprAnalysis &constExprs,
                        const llvm::DenseMap<Operation *, unsigned> &opOrder) {
    if (cloneMapping.contains(producerInfo->constValue)) {
      return;
    }

    // We either have a global associated already or we need to traverse
    // down and materialize producers.
    IREE::Util::GlobalOp existingGlobal =
        hoistedMap.lookup(producerInfo->constValue);
    if (existingGlobal) {
      Value newGlobal =
          existingGlobal
              .createLoadOp(existingGlobal.getLoc(), initializerBuilder)
              .getLoadedGlobalValue();
      // Call user hook to cast back to the original type.
      if (auto hoistableType = dyn_cast<IREE::Util::HoistableTypeInterface>(
              producerInfo->constValue.getType())) {
        newGlobal = hoistableType.decodeStorageType(
            initializerBuilder, newGlobal.getLoc(),
            producerInfo->constValue.getType(), newGlobal);
      }
      cloneMapping.map(producerInfo->constValue, newGlobal);
      return;
    }

    // Materialize all producers recursively. Sort producers by their program
    // order for deterministic output.
    auto sortedProducers =
        llvm::to_vector_of<ConstExprAnalysis::ConstValueInfo *>(
            producerInfo->producers);
    llvm::sort(sortedProducers,
               [&opOrder](ConstExprAnalysis::ConstValueInfo *lhs,
                          ConstExprAnalysis::ConstValueInfo *rhs) {
                 return opOrder.lookup(lhs->constValue.getDefiningOp()) <
                        opOrder.lookup(rhs->constValue.getDefiningOp());
               });
    for (ConstExprAnalysis::ConstValueInfo *prodInfo : sortedProducers) {
      cloneProducerTreeInto(initializerBuilder, prodInfo, hoistedMap,
                            cloneMapping, constExprs, opOrder);
    }

    // And clone the requested op.
    Operation *sourceOp = producerInfo->constValue.getDefiningOp();
    assert(sourceOp && "must have defining op for const-expr values");
    LLVM_DEBUG({
      llvm::dbgs() << "[HoistIntoGlobals]    + clone op: ";
      sourceOp->print(llvm::dbgs(), constExprs.getAsmState());
      llvm::dbgs() << "\n";
    });
    Operation *clonedOp = sourceOp->clone(cloneMapping);
    initializerBuilder.insert(clonedOp);
  }

  // Clones the const expr tree rooted at `constExprValue` into the given
  // initializer, noting any new hoisted value mappings that result. At
  // a minimum, a mapping will be created for the requested value.
  LogicalResult
  cloneConstExprInto(Location loc, OpBuilder &moduleBuilder,
                     OpBuilder &initializerBuilder, Value constExprValue,
                     NamedAttrList dialectAttrs, HoistedValueMap &hoistedMap,
                     SymbolTable &moduleSymbols,
                     const ConstExprAnalysis &constExprs,
                     const llvm::DenseMap<Operation *, unsigned> &opOrder) {
    // Do a depth first traversal of the producers, emitting them in a valid
    // def-use order.
    Operation *rootOp = constExprValue.getDefiningOp();
    assert(rootOp && "const-expr value should have a defining op");
    auto *rootInfo = constExprs.lookup(rootOp);
    assert(rootInfo && "must have const-value-info for const-expr root op");

    // Clone the whole tree as needed.
    IRMapping cloneMapping;
    cloneProducerTreeInto(initializerBuilder, rootInfo, hoistedMap,
                          cloneMapping, constExprs, opOrder);

    // And for each result, create a global and store into it.
    for (Value origResult : rootOp->getResults()) {
      Value clonedResult = cloneMapping.lookup(origResult);
      Type globalType = origResult.getType();
      // If the original type is registered as hoistable, invoke the interface
      // functions for setting the preferred storage type.
      auto hoistableType =
          dyn_cast<IREE::Util::HoistableTypeInterface>(globalType);
      if (hoistableType) {
        // Allow the storage type of the global to differ from the local type.
        globalType = hoistableType.getPreferredStorageType();
      }
      auto globalOp = IREE::Util::GlobalOp::create(
          moduleBuilder, loc, getHoistedName(globalType), false, globalType);
      moduleSymbols.insert(globalOp);
      SymbolTable::setSymbolVisibility(globalOp,
                                       SymbolTable::Visibility::Private);
      globalOp->setDialectAttrs(dialectAttrs);

      // Save the mapping for the future.
      hoistedMap[origResult] = globalOp;

      // And store into it.
      LLVM_DEBUG({
        llvm::dbgs() << "[HoistIntoGlobals]    + create global @"
                     << globalOp.getSymName() << " = ";
        clonedResult.print(llvm::dbgs());
        llvm::dbgs() << "\n";
      });
      if (hoistableType) {
        // Allow casting to the global type if it differs from the local type.
        clonedResult = hoistableType.encodeStorageType(
            initializerBuilder, clonedResult.getLoc(), globalType,
            clonedResult);
      }
      if (clonedResult.getType() != globalType) {
        return globalOp.emitError()
               << "unresolved conflict between global of type " << globalType
               << " and stored type " << clonedResult.getType();
      }
      globalOp.createStoreOp(loc, clonedResult, initializerBuilder);
    }

    IREE::Util::ReturnOp::create(initializerBuilder, loc);
    return success();
  }

  void cleanupDeadOps(const ConstExprAnalysis &constExprs) {
    llvm::DenseSet<Operation *> allOps;
    constExprs.populateConstExprOperations(allOps);

    // Since we are mutating the const-expr ops, the ConstExprAnalysis will no
    // longer be valid after this point.
    for (auto funcOp : getOperation().getOps<FunctionOpInterface>()) {
      // Ignore initializers.
      if (isa<IREE::Util::InitializerOpInterface>(funcOp.getOperation())) {
        continue;
      }
      funcOp.walk<WalkOrder::PostOrder, ReverseIterator>(
          [&](Operation *iterOp) {
            if (allOps.contains(iterOp) && iterOp->use_empty()) {
              // Bingo.
              LLVM_DEBUG({
                llvm::dbgs() << "[HoistIntoGlobals] erase dead op: ";
                iterOp->print(llvm::dbgs(), constExprs.getAsmState());
                llvm::dbgs() << "\n";
              });
              allOps.erase(iterOp);
              iterOp->erase();
            }
            return WalkResult::advance();
          });
    }
  }

private:
  const std::optional<ExprHoistingOptions::RegisterDialectsFn>
      registerDependentDialectsFn = std::nullopt;
};

} // namespace

std::unique_ptr<Pass>
createHoistIntoGlobalsPass(const ExprHoistingOptions &options) {
  return std::make_unique<HoistIntoGlobalsPass>(options);
}

} // namespace mlir::iree_compiler::IREE::Util
