//===- EnzymeBatchPass.cpp - Replace calls with their batched versions
//------------ //
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements a pass to lower gpu kernels in NVVM/gpu dialects into
// a generic parallel for representation
//===----------------------------------------------------------------------===//

#include "Passes/EnzymeBatchPass.h"
#include "Dialect/Ops.h"
#include "Interfaces/GradientUtilsReverse.h"
#include "PassDetails.h"
#include "Passes/Passes.h"

#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/Interfaces/FunctionInterfaces.h"

#define DEBUG_TYPE "enzyme-batch"

using namespace mlir;
using namespace mlir::enzyme;
using namespace enzyme;

namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_BATCHPASS
#include "Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir

namespace mlir {
namespace enzyme {
namespace batchutils {

mlir::TensorType applyBatchSizes(mlir::Type Ty,
                                 llvm::ArrayRef<int64_t> batchSizes) {
  auto T = dyn_cast<TensorType>(Ty);
  if (!T) {
    return RankedTensorType::get(batchSizes, Ty);
  }

  SmallVector<int64_t> shape(batchSizes.begin(), batchSizes.end());
  shape.append(T.getShape().begin(), T.getShape().end());
  auto T2 = T.clone(shape);
  return T2;
}

LogicalResult handleCallOp(
    func::CallOp callOp, OpBuilder &builder, IRMapping &mapper,
    llvm::ArrayRef<int64_t> batchSizes,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache) {
  // Get the called function
  auto moduleOp = callOp->getParentOfType<ModuleOp>();
  auto calledFunc =
      dyn_cast<FunctionOpInterface>(moduleOp.lookupSymbol(callOp.getCallee()));
  if (!calledFunc)
    return failure();

  // Create cache key for this function and batch size combination
  BatchCacheKey key{calledFunc,
                    SmallVector<int64_t>(batchSizes.begin(), batchSizes.end())};

  // Look up or create batched version of the called function
  FunctionOpInterface batchedFunc;
  auto it = batchedFunctionCache.find(key);
  if (it != batchedFunctionCache.end()) {
    batchedFunc = it->second;
  } else {
    std::string fnName = "batched_" + calledFunc.getName().str();
    batchedFunc = batchCloneFunction(builder, calledFunc, fnName, batchSizes,
                                     batchedFunctionCache);
    if (!batchedFunc)
      return failure();
    batchedFunctionCache[key] = batchedFunc;
  }

  // Create new call operation to the batched function
  SmallVector<Value> newOperands;
  for (auto operand : callOp->getOperands())
    newOperands.push_back(mapper.lookup(operand));

  auto newCall =
      func::CallOp::create(builder, callOp.getLoc(), batchedFunc.getName(),
                           batchedFunc.getResultTypes(), newOperands);

  // Map the results
  for (auto [oldResult, newResult] :
       llvm::zip(callOp.getResults(), newCall.getResults()))
    mapper.map(oldResult, newResult);

  return success();
}

void batchCloneBlock(
    OpBuilder &builder, Block *blk, IRMapping &mapper,
    llvm::ArrayRef<int64_t> batchSizes,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache,
    bool withoutTerminator) {
  for (auto &src : *blk) {
    if (auto callOp = dyn_cast<func::CallOp>(&src)) {
      if (succeeded(handleCallOp(callOp, builder, mapper, batchSizes,
                                 batchedFunctionCache)))
        continue;
    }

    if (auto ifaceOp = dyn_cast<BatchOpInterface>(&src)) {
      auto res = ifaceOp.createBatch(builder, mapper, batchSizes);
      if (res.succeeded())
        continue;
    }

    SmallVector<Value, 8> operands;
    SmallVector<Block *, 2> successors;

    // Remap the operands.
    operands.reserve(src.getNumOperands());
    for (auto opValue : src.getOperands())
      operands.push_back(mapper.lookup(opValue));

    if (withoutTerminator && src.hasTrait<OpTrait::IsTerminator>()) {
      // map the operands and the results
      for (unsigned i = 0, e = src.getNumResults(); i != e; ++i)
        mapper.map(src.getResult(i), operands[i]);
      continue;
    }

    // Remap the successors.
    successors.reserve(src.getNumSuccessors());
    for (Block *successor : src.getSuccessors())
      successors.push_back(mapper.lookup(successor));

    SmallVector<Type> resultTypes(src.getResultTypes().begin(),
                                  src.getResultTypes().end());
    for (auto &Ty : resultTypes) {
      Ty = applyBatchSizes(Ty, batchSizes);
    }

    Operation *newOp = Operation::create(
        src.getLoc(), src.getName(), resultTypes, operands, src.getAttrs(),
        OpaqueProperties(nullptr), successors, src.getNumRegions());

    // Clone the regions.
    for (auto &&[oldReg, newReg] :
         llvm::zip(src.getRegions(), newOp->getRegions())) {
      batchCloneRegion(builder, &oldReg, &newReg, mapper, batchSizes,
                       batchedFunctionCache);
    }

    // Remember the mapping of any results.
    for (unsigned i = 0, e = src.getNumResults(); i != e; ++i)
      mapper.map(src.getResult(i), newOp->getResult(i));

    builder.insert(newOp);
  }
}

void batchCloneRegion(
    OpBuilder &builder, Region *src, Region *dest, IRMapping &mapper,
    llvm::ArrayRef<int64_t> batchSizes,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache) {
  // For each block in src, generate a corresponding block in the dest region.
  for (auto &blk : *src) {
    auto newBlk = new Block();
    dest->push_back(newBlk);

    mapper.map(&blk, newBlk);

    for (auto arg : blk.getArguments()) {
      Value newArg = newBlk->addArgument(
          applyBatchSizes(arg.getType(), batchSizes), arg.getLoc());
      mapper.map(arg, newArg);
    }
  }

  for (auto &&[blk, newBlk] : llvm::zip(*src, *dest)) {
    IRRewriter::InsertionGuard insertGuard(builder);
    builder.setInsertionPointToEnd(&newBlk);
    batchCloneBlock(builder, &blk, mapper, batchSizes, batchedFunctionCache,
                    false);
  }
}

FunctionOpInterface batchCloneFunction(
    OpBuilder &builder, FunctionOpInterface F, Twine name,
    llvm::ArrayRef<int64_t> batchSizes,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache) {
  assert(!F.getFunctionBody().empty());

  auto FTy = cast<FunctionType>(F.getFunctionType());

  llvm::SmallVector<mlir::Type> RetTypes;
  RetTypes.reserve(FTy.getNumResults());

  for (auto Ty : FTy.getResults()) {
    RetTypes.push_back(applyBatchSizes(Ty, batchSizes));
  }

  SmallVector<mlir::Type, 4> ArgTypes;
  ArgTypes.reserve(FTy.getNumInputs());

  for (auto Ty : FTy.getInputs()) {
    ArgTypes.push_back(applyBatchSizes(Ty, batchSizes));
  }

  FunctionType newFTy = builder.getFunctionType(ArgTypes, RetTypes);

  auto NewF = cast<FunctionOpInterface>(F->cloneWithoutRegions());
  SymbolTable::setSymbolName(NewF, name.str());
  NewF.setType(newFTy);

  Operation *parent = F->getParentWithTrait<OpTrait::SymbolTable>();
  SymbolTable table(parent);
  table.insert(NewF);
  SymbolTable::setSymbolVisibility(NewF, SymbolTable::Visibility::Private);

  // Add the function to the cache BEFORE processing its body to support
  // recursion.
  BatchCacheKey key{F,
                    SmallVector<int64_t>(batchSizes.begin(), batchSizes.end())};
  batchedFunctionCache[key] = NewF;

  auto &origReg = F.getFunctionBody();
  auto &newReg = NewF.getFunctionBody();

  IRMapping mapper;
  batchCloneRegion(builder, &origReg, &newReg, mapper, batchSizes,
                   batchedFunctionCache);

  return NewF;
}

template <typename T>
FunctionOpInterface batchOperationWithoutInsertingCallOp(
    OpBuilder &builder, T CI, FunctionOpInterface fn,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache) {
  enzyme::batchutils::BatchCacheKey key{
      fn, SmallVector<int64_t>(CI.getBatchShape().begin(),
                               CI.getBatchShape().end())};

  // Check if we already have a batched version
  auto it = batchedFunctionCache.find(key);
  FunctionOpInterface newFunc;

  if (it != batchedFunctionCache.end()) {
    return it->second;
  } else {
    // Create new batched function and store in cache
    std::string newFnName = "batched_" + fn.getName().str();
    newFunc = batchCloneFunction(builder, fn, newFnName, CI.getBatchShape(),
                                 batchedFunctionCache);
    return newFunc;
  }
}

} // namespace batchutils
} // namespace enzyme
} // namespace mlir

namespace {

struct BatchPass : public enzyme::impl::BatchPassBase<BatchPass> {
  void runOnOperation() override;

  // Cache mapping original function and batch sizes to batched function
  std::map<enzyme::batchutils::BatchCacheKey, FunctionOpInterface>
      batchedFunctionCache;

  void lowerEnzymeBatchCalls(SymbolTableCollection &symbolTable,
                             FunctionOpInterface op) {
    {
      SmallVector<Operation *> toLower;
      op->walk([&](enzyme::BatchOp dop) {
        auto *symbolOp =
            symbolTable.lookupNearestSymbolFrom(dop, dop.getFnAttr());
        auto callableOp = cast<FunctionOpInterface>(symbolOp);

        lowerEnzymeBatchCalls(symbolTable, callableOp);
        toLower.push_back(dop);
      });

      OpBuilder builder(op);

      for (auto T : toLower) {
        if (auto F = dyn_cast<enzyme::BatchOp>(T)) {
          auto res = enzyme::batchutils::batchOperation(symbolTable, builder, F,
                                                        batchedFunctionCache);
          if (!res.succeeded()) {
            signalPassFailure();
            return;
          }
        } else {
          llvm_unreachable("Illegal type");
        }
      }
    };
  };
};

} // end anonymous namespace

void BatchPass::runOnOperation() {
  SymbolTableCollection symbolTable;
  symbolTable.getSymbolTable(getOperation());
  getOperation()->walk(
      [&](FunctionOpInterface op) { lowerEnzymeBatchCalls(symbolTable, op); });
}
