// Copyright 2020 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/Flow/IR/FlowDialect.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
#include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
#include "iree/compiler/Dialect/Util/Analysis/Explorer.h"
#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Pass/Pass.h"

namespace mlir::iree_compiler::IREE::Flow {

#define GEN_PASS_DEF_EXPORTBENCHMARKFUNCSPASS
#include "iree/compiler/Dialect/Flow/Transforms/Passes.h.inc"

// Creates a util.global with a primitive value of |type| initialized to zeros.
// Supports: ints, floats, vectors, and tensors.
//
// Example:
//  util.global @some_fn_arg0 = 4 : i32
//  util.global @some_fn_arg0 = dense<4> : tensor<4xi32>
static IREE::Util::GlobalOp
createPrimitiveDefaultGlobalOp(std::string name, Location loc, Type type,
                               SymbolTable &symbolTable,
                               OpBuilder &moduleBuilder) {
  // Get a zero-initialized constant attribute for the type, if supported.
  auto initialValue = moduleBuilder.getZeroAttr(type);
  if (!initialValue) {
    mlir::emitError(loc) << "unsupported function argument type: " << type;
    return {};
  }

  // Global with initializer; tensors will get turned into buffers eventually.
  auto globalOp =
      IREE::Util::GlobalOp::create(moduleBuilder, loc, name,
                                   /*isMutable=*/false, type, initialValue);
  globalOp.setPrivate();
  globalOp.setGlobalInliningPolicy(
      moduleBuilder.getAttr<IREE::Util::InlineNeverAttr>());
  symbolTable.insert(globalOp);
  return globalOp;
}

// Creates a util.global of the given |globalType| and initializes a buffer or
// buffer view as a zeroed |tensorType|.
static IREE::Util::GlobalOp
createBufferLikeGlobalOp(std::string name, Location loc, Type globalType,
                         TensorType tensorType, SymbolTable &symbolTable,
                         OpBuilder &moduleBuilder) {
  // Create !hal.buffer global for the storage buffer or buffer view.
  auto globalOp = IREE::Util::GlobalOp::create(moduleBuilder, loc, name,
                                               /*isMutable=*/false, globalType);
  globalOp.setPrivate();
  globalOp.setGlobalInliningPolicy(
      moduleBuilder.getAttr<IREE::Util::InlineNeverAttr>());
  symbolTable.insert(globalOp);

  // Create an initializer that allocates the buffer storage.
  // We do this by splatting and exporting to a buffer so that it looks like it
  // was created by the user.
  // Ensure initializer comes after the global.
  moduleBuilder.setInsertionPointAfter(globalOp);
  auto initializerOp = IREE::Util::InitializerOp::create(moduleBuilder, loc);
  auto initializerBuilder =
      OpBuilder::atBlockBegin(initializerOp.addEntryBlock());
  auto zeroAttr = moduleBuilder.getZeroAttr(tensorType.getElementType());
  auto zeroOp = arith::ConstantOp::create(initializerBuilder, loc, zeroAttr);
  // flow.tensor.splat 0
  auto splatOp =
      IREE::Flow::TensorSplatOp::create(initializerBuilder, loc, tensorType,
                                        zeroOp, /*result_dims=*/ValueRange{});
  // hal.tensor.export
  auto bufferExportOp = IREE::HAL::TensorExportOp::create(
      initializerBuilder, loc, globalOp.getType(), splatOp.getResult(),
      TypeAttr::get(splatOp.getType()), /*name=*/nullptr,
      /*affinity=*/nullptr);
  // util.optimization_barrier (try to prevent optimizations across the export)
  auto barrierOp = IREE::Util::OptimizationBarrierOp::create(
      initializerBuilder, loc, bufferExportOp.getTarget());
  // util.global.store
  globalOp.createStoreOp(loc, barrierOp.getResult(0), initializerBuilder);
  IREE::Util::ReturnOp::create(initializerBuilder, loc);

  return globalOp;
}

// Creates a util.global with a buffer view initialized to the same parameters
// as |arg| is used by. Fails if analysis cannot find a hal.tensor.import op.
//
// Example:
//  %0 = hal.tensor.import %arg : !hal.buffer_view -> tensor<4xi32>
// ->
//  util.global @some_fn_arg0 : !hal.buffer_view
//  util.initializer { ... }
static IREE::Util::GlobalOp
createImportBufferViewGlobalOp(std::string name, BlockArgument arg,
                               SymbolTable &symbolTable,
                               OpBuilder &moduleBuilder, Explorer &explorer) {
  auto loc = arg.getLoc();

  // Find a hal.tensor.import user.
  IREE::HAL::TensorImportOp importOp;
  if (explorer.walkTransitiveUsers(arg, [&](Operation *op) -> WalkResult {
        importOp = dyn_cast<IREE::HAL::TensorImportOp>(op);
        return importOp ? WalkResult::interrupt() : WalkResult::advance();
      }) == TraversalResult::INCOMPLETE) {
    // Analysis failed to find an import op. User needs to rework their program.
    mlir::emitError(loc) << "unsupported dynamic buffer view import on " << arg;
    return {};
  }

  // Extract the type, which must be a static tensor.
  auto targetType = importOp.getTarget().getType();
  auto tensorType = dyn_cast<RankedTensorType>(targetType);
  if (!tensorType || !tensorType.hasStaticShape()) {
    mlir::emitError(loc) << "unsupported buffer view import tensor type on "
                         << arg << " used as " << targetType;
    return {};
  }

  // Create the global and initialize it by allocating a zeroed buffer.
  return createBufferLikeGlobalOp(name, loc, arg.getType(), tensorType,
                                  symbolTable, moduleBuilder);
}

// Creates a util.global with a buffer initialized to the required storage
// capacity |arg| is used by. Fails if analysis cannot find a hal.tensor.export
// op.
//
// Example:
//  %1 = hal.tensor.export %0 into(%storage) : tensor<4xi32> -> !hal.buffer_view
// ->
//  util.global @some_fn_arg0 : !hal.buffer
//  util.initializer { ... }
static IREE::Util::GlobalOp createExportBufferGlobalOp(std::string name,
                                                       BlockArgument arg,
                                                       SymbolTable &symbolTable,
                                                       OpBuilder &moduleBuilder,
                                                       Explorer &explorer) {
  auto loc = arg.getLoc();

  // Find a hal.tensor.export or alias user and extract the encoding.
  Type sourceType;
  if (explorer.walkTransitiveUsers(arg, [&](Operation *op) -> WalkResult {
        if (auto aliasOp = dyn_cast<IREE::HAL::TensorAliasOp>(op)) {
          sourceType = aliasOp.getResult().getType();
          return WalkResult::interrupt();
        } else if (auto exportOp = dyn_cast<IREE::HAL::TensorExportOp>(op)) {
          sourceType = exportOp.getSourceEncoding();
          return WalkResult::interrupt();
        }
        return WalkResult::advance();
      }) == TraversalResult::INCOMPLETE) {
    // Analysis failed to find an export op. User needs to rework their program.
    mlir::emitError(loc) << "unsupported dynamic buffer view export on " << arg;
    return {};
  }

  // The type must be a static tensor for this pass to work.
  auto tensorType = dyn_cast<RankedTensorType>(sourceType);
  if (!tensorType || !tensorType.hasStaticShape()) {
    mlir::emitError(loc) << "unsupported buffer view export tensor type on "
                         << arg << " used as " << sourceType;
    return {};
  }

  // Create the global and initialize it by allocating a zeroed buffer.
  return createBufferLikeGlobalOp(name, loc, arg.getType(), tensorType,
                                  symbolTable, moduleBuilder);
}

static IREE::Util::GlobalOp createDummyInput(const std::string &namePrefix,
                                             BlockArgument arg,
                                             SymbolTable &symbolTable,
                                             OpBuilder &moduleBuilder,
                                             Explorer &explorer) {
  std::string name = namePrefix + "_arg" + std::to_string(arg.getArgNumber());
  auto globalOp =
      TypeSwitch<Type, IREE::Util::GlobalOp>(arg.getType())
          .Case([&](IREE::HAL::BufferViewType type) {
            return createImportBufferViewGlobalOp(name, arg, symbolTable,
                                                  moduleBuilder, explorer);
          })
          .Case([&](IREE::HAL::BufferType type) {
            return createExportBufferGlobalOp(name, arg, symbolTable,
                                              moduleBuilder, explorer);
          })
          .Default([&](Type type) {
            return createPrimitiveDefaultGlobalOp(name, arg.getLoc(), type,
                                                  symbolTable, moduleBuilder);
          });
  if (globalOp) {
    // Prevent globals from folding so that we have unique buffers for each arg.
    globalOp->setAttr("flow.unique_id", moduleBuilder.getStringAttr(name));
  }
  return globalOp;
}

static LogicalResult
createEntryPointBenchmarkFunc(mlir::ModuleOp moduleOp,
                              IREE::Util::FuncOp entryFuncOp,
                              Explorer &explorer) {
  auto symbolTable = explorer.getSymbolTables().getSymbolTable(moduleOp);
  OpBuilder moduleBuilder(moduleOp.getContext());
  moduleBuilder.setInsertionPointAfter(entryFuncOp);

  // We'll create all symbols based on this name prefix.
  auto funcName = std::string(entryFuncOp.getName()) + "_benchmark";

  // Create one dummy input variable per input. We may need to do some
  // analysis to find the actual type and initial value.
  SmallVector<IREE::Util::GlobalOp> dummyInputVariableOps;
  for (auto arg : entryFuncOp.getArguments()) {
    auto dummyVar =
        createDummyInput(funcName, arg, symbolTable, moduleBuilder, explorer);
    if (!dummyVar) {
      return failure();
    }
    dummyInputVariableOps.push_back(dummyVar);
  }

  // Create a `() -> ()` entry point op the benchmark tool can run.
  Location loc = entryFuncOp.getLoc();
  auto funcOp = IREE::Util::FuncOp::create(
      moduleBuilder, loc, funcName, moduleBuilder.getFunctionType({}, {}));
  funcOp.setPublic();
  funcOp->setAttr("iree.abi.stub", moduleBuilder.getUnitAttr());
  NamedAttribute reflectionAttr("iree.benchmark",
                                moduleBuilder.getStringAttr("entry"));
  funcOp->setAttr("iree.reflection",
                  moduleBuilder.getDictionaryAttr(reflectionAttr));
  Block *block = funcOp.addEntryBlock();

  // Call the existing function with dummy arguments.
  auto blockBuilder = OpBuilder::atBlockBegin(block);
  SmallVector<Value> args;
  for (int i = 0, e = entryFuncOp.getNumArguments(); i < e; ++i) {
    args.push_back(dummyInputVariableOps[i]
                       .createLoadOp(loc, blockBuilder)
                       .getLoadedGlobalValue());
  }
  auto callOp =
      IREE::Util::CallOp::create(blockBuilder, loc, entryFuncOp, args);

  // Sink all results with a barrier to ensure that DCE does not remove the
  // call.
  for (auto result : callOp.getResults()) {
    IREE::Util::OptimizationBarrierOp::create(blockBuilder, loc, result);
  }
  IREE::Util::ReturnOp::create(blockBuilder, loc);

  // Ensure the original function is not exported and not inlined.
  entryFuncOp->setAttr("noinline", moduleBuilder.getUnitAttr());
  entryFuncOp->removeAttr("iree.reflection");
  entryFuncOp.setPrivate();

  return success();
}

// Clones each exported functions (including those just created) with
// placeholder constant inputs instead of arguments and removes the exported
// attribute from the old functions.
// The input are provided using util.globals.
struct ExportBenchmarkFuncsPass
    : IREE::Flow::impl::ExportBenchmarkFuncsPassBase<ExportBenchmarkFuncsPass> {
  void runOnOperation() override {
    mlir::ModuleOp moduleOp = getOperation();

    // For analysis required on arguments and results.
    Explorer explorer(moduleOp, TraversalAction::SHALLOW);

    // Gather the functions we want to wrap for benchmarking and wrap them.
    // Since we are inserting new functions as part of this pass we must perform
    // the wrapping for only the inputs.
    SmallVector<IREE::Util::FuncOp> entryFuncOps;
    for (auto entryFuncOp : moduleOp.getOps<IREE::Util::FuncOp>()) {
      if (entryFuncOp.isPublic()) {
        entryFuncOps.push_back(entryFuncOp);
      }
    }
    for (auto entryFuncOp : entryFuncOps) {
      if (failed(
              createEntryPointBenchmarkFunc(moduleOp, entryFuncOp, explorer))) {
        signalPassFailure();
        return;
      }
    }
  }
};

} // namespace mlir::iree_compiler::IREE::Flow
