//===- PrintPass.cpp - Print the MLIR module                     ------------ //
//
// 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 print the MLIR module
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "src/enzyme_ad/jax/Passes/Passes.h"

#include "src/enzyme_ad/jax/Dialect/Ops.h"

#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "src/enzyme_ad/jax/Passes/Passes.h"
#include "stablehlo/dialect/StablehloOps.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
#include "llvm/Support/raw_ostream.h"

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"

#define DEBUG_TYPE "lower-kernel"

namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_LOWERKERNELPASS
#include "src/enzyme_ad/jax/Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir

using namespace mlir;
using namespace mlir::enzyme;
using namespace mlir::gpu;
using namespace enzyme;
using namespace mlir::enzymexla;
using namespace enzymexla;

using namespace stablehlo;

SymbolRefAttr SymRefAttrReplacingFunctionName(SymbolRefAttr origSymRef,
                                              StringRef newName) {
  auto newNameRef = FlatSymbolRefAttr::get(origSymRef.getContext(), newName);
  auto nestedRefsAttr = origSymRef.getNestedReferences();
  if (nestedRefsAttr.size() == 0) {
    return newNameRef;
  }

  auto rootRef = origSymRef.getRootReference();
  SmallVector<FlatSymbolRefAttr> nestedRefs;
  for (int i = 0; i < nestedRefsAttr.size() - 1; i++) {
    nestedRefs.push_back(nestedRefsAttr[i]);
  }
  nestedRefs.push_back(newNameRef);
  return SymbolRefAttr::get(origSymRef.getContext(), rootRef, nestedRefs);
}

bool CompileGPUKernel(SymbolTableCollection &symbolTable, mlir::Location loc,
                      FunctionOpInterface op, size_t gridx, size_t gridy,
                      size_t gridz, size_t blockx, size_t blocky, size_t blockz,
                      size_t shmem, size_t clusterx, size_t clustery,
                      size_t clusterz, enzymexla::KernelCallOp kcall) {

  OpBuilder builder(op);

  FunctionType gpuTy0 = dyn_cast<FunctionType>(op.getFunctionType());
  if (!gpuTy0) {
    if (auto lty = dyn_cast<LLVM::LLVMFunctionType>(op.getFunctionType())) {
      gpuTy0 = builder.getFunctionType(lty.getParams(), {});
    } else {
      op.emitError(
          "Require target operand to have functiontype or llvmfunctiontype");
      return false;
    }
  }
  SmallVector<Type, 1> newParams;
  for (Type p : gpuTy0.getInputs()) {
    if (auto AT = dyn_cast<LLVM::LLVMArrayType>(p)) {
      p = AT.getElementType();
    }
    newParams.push_back(p);
  }
  FunctionType gpuTy = builder.getFunctionType(newParams, {});

  std::string legalName = op.getName().str();
  std::replace(legalName.begin(), legalName.end(), '#', '_');
  std::string gpumodname = "gpumod_" + legalName;

  gpu::GPUModuleOp gpumod = nullptr;
  gpu::GPUFuncOp gpufunc = nullptr;

  for (auto &gop : *op.getOperation()->getBlock()) {
    auto gmod = dyn_cast<gpu::GPUModuleOp>(gop);
    if (!gmod)
      continue;
    if (gmod.getName() == gpumodname) {
      gpumod = gmod;
      break;
    }
  }

  if (gpumod) {
    for (auto &gop : *(gpumod.getBody())) {
      auto gfunc = dyn_cast<gpu::GPUFuncOp>(gop);
      if (!gfunc)
        continue;
      if (gfunc.getName() == legalName) {
        gpufunc = gfunc;
        break;
      }
    }
    assert(gpufunc);
  } else {
    gpumod = gpu::GPUModuleOp::create(builder, loc, gpumodname);
    builder.setInsertionPointToStart(&gpumod.getBodyRegion().front());

    gpufunc = gpu::GPUFuncOp::create(builder, loc, legalName, gpuTy);
    {
      auto entry = &gpufunc.getBody().front();
      builder.setInsertionPointToEnd(entry);
      IRMapping map;
      for (auto &&[oldarg, newarg] :
           zip(op.getArguments(), gpufunc.getArguments())) {
        Value newval = newarg;

        if (auto AT = dyn_cast<LLVM::LLVMArrayType>(oldarg.getType())) {
          auto ud =
              LLVM::UndefOp::create(builder, newarg.getLoc(), oldarg.getType());
          int64_t c0[1] = {0};
          newval = LLVM::InsertValueOp::create(
              builder, newarg.getLoc(), oldarg.getType(), ud, newval, c0);
        }

        map.map(oldarg, newval);
      }

      op.getFunctionBody().cloneInto(&gpufunc.getBody(), map);
      gpufunc.setKernelAttr(builder.getUnitAttr());

      auto second = entry->getNextNode();
      entry->getOperations().splice(entry->getOperations().end(),
                                    second->getOperations());

      second->erase();

      gpufunc->walk([](LLVM::ReturnOp op) {
        OpBuilder rewriter(op);
        gpu::ReturnOp::create(rewriter, op.getLoc());
        op.erase();
      });

      gpufunc->walk([](LLVM::UnreachableOp op) {
        OpBuilder rewriter(op);
        gpu::ReturnOp::create(rewriter, op.getLoc());
        op.erase();
      });

      gpufunc->walk([](func::ReturnOp op) {
        OpBuilder rewriter(op);
        gpu::ReturnOp::create(rewriter, op.getLoc());
        op.erase();
      });
    }
    SmallVector<Operation *> tocopy;
    op->walk([&](CallOpInterface cop) {
      if (auto op2 = cop.resolveCallable())
        tocopy.push_back(op2);
    });
    op->walk([&](LLVM::AddressOfOp cop) {
      if (auto op2 = cop.getGlobal(symbolTable))
        tocopy.push_back(op2);
      else if (auto op2 = cop.getFunction(symbolTable))
        tocopy.push_back(op2);
    });
    SmallPtrSet<Operation *, 1> done;

    builder.setInsertionPointToStart(&gpumod.getBodyRegion().front());
    while (tocopy.size()) {
      auto cur = tocopy.pop_back_val();
      if (done.count(cur))
        continue;
      done.insert(cur);
      builder.clone(*cur);
      cur->walk([&](CallOpInterface cop) {
        if (auto op2 = cop.resolveCallable())
          tocopy.push_back(op2);
      });
      cur->walk([&](LLVM::AddressOfOp cop) {
        if (auto op2 = cop.getGlobal(symbolTable))
          tocopy.push_back(op2);
        else if (auto op2 = cop.getFunction(symbolTable))
          tocopy.push_back(op2);
      });
    }
  }

  builder.setInsertionPoint(op);

  static size_t callid = 0;
  callid++;
  auto callName = (op.getName() + "$call$" + std::to_string(callid)).str();
  auto func = func::FuncOp::create(builder, loc, callName, gpuTy);
  func->setDiscardableAttrs(op->getDiscardableAttrDictionary());
  func.setVisibility(SymbolTable::Visibility::Private);

  auto &entryBlock = *func.addEntryBlock();
  builder.setInsertionPointToStart(&entryBlock);

  auto idx = builder.getIntegerType(64);
  auto i32 = builder.getIntegerType(32);
  gpu::KernelDim3 gridSize{
      arith::ConstantIntOp::create(builder, loc, idx, gridx),
      arith::ConstantIntOp::create(builder, loc, idx, gridy),
      arith::ConstantIntOp::create(builder, loc, idx, gridz),
  };

  gpu::KernelDim3 blockSize{
      arith::ConstantIntOp::create(builder, loc, idx, blockx),
      arith::ConstantIntOp::create(builder, loc, idx, blocky),
      arith::ConstantIntOp::create(builder, loc, idx, blockz),
  };

  auto dynshmem = arith::ConstantIntOp::create(builder, loc, i32, shmem);

  Value stream = enzymexla::GetStreamOp::create(
      builder, loc, gpu::AsyncTokenType::get(kcall.getContext()));

  if (clusterx != 1 || clustery != 1 || clusterz != 1) {
    gpu::KernelDim3 clusterSize{
        arith::ConstantIntOp::create(builder, loc, idx, clusterx),
        arith::ConstantIntOp::create(builder, loc, idx, clustery),
        arith::ConstantIntOp::create(builder, loc, idx, clusterz),
    };

    gpu::LaunchFuncOp::create(builder, loc, gpufunc, gridSize, blockSize,
                              dynshmem, entryBlock.getArguments(),
                              stream.getType(), ValueRange(stream),
                              clusterSize);
  } else {
    gpu::LaunchFuncOp::create(builder, loc, gpufunc, gridSize, blockSize,
                              dynshmem, entryBlock.getArguments(),
                              stream.getType(), ValueRange(stream));
  }

  mlir::func::ReturnOp::create(builder, loc);

  if (!op->getParentOp()->hasAttr(
          gpu::GPUDialect::getContainerModuleAttrName()))
    op->getParentOp()->setAttr(gpu::GPUDialect::getContainerModuleAttrName(),
                               UnitAttr::get(kcall.getContext()));

  OpBuilder rewriter(kcall);
  auto replacement = enzymexla::JITCallOp::create(
      rewriter, kcall.getLoc(), kcall.getResultTypes(),
      SymRefAttrReplacingFunctionName(kcall.getFn(), callName),
      kcall.getInputs(), kcall.getBackendConfigAttr(),
      kcall.getOperandLayoutsAttr(), kcall.getResultLayoutsAttr(),
      kcall.getArgAttrsAttr(), kcall.getResAttrsAttr(),
      kcall.getOutputOperandAliasesAttr(), kcall.getXlaSideEffectFreeAttr());
  kcall.replaceAllUsesWith(replacement);
  kcall.erase();
  return true;
};

bool CompileCPUKernel(SymbolTableCollection &symbolTable, mlir::Location loc,
                      FunctionOpInterface op, size_t gridx, size_t gridy,
                      size_t gridz, size_t blockx, size_t blocky, size_t blockz,
                      size_t shmem, size_t clusterx, size_t clustery,
                      size_t clusterz, enzymexla::KernelCallOp kcall) {
  OpBuilder builder(op);

  if (clusterx != 1 || clustery != 1 || clusterz != 1) {
    op.emitError("CPU kernels do not support cluster");
    return false;
  }

  FunctionType gpuTy0 = dyn_cast<FunctionType>(op.getFunctionType());
  if (!gpuTy0) {
    if (auto lty = dyn_cast<LLVM::LLVMFunctionType>(op.getFunctionType())) {
      gpuTy0 = builder.getFunctionType(lty.getParams(), {});
    } else {
      op.emitError(
          "Require target operand to have functiontype or llvmfunctiontype");
      return false;
    }
  }
  SmallVector<Type, 1> newParams;
  for (Type p : gpuTy0.getInputs()) {
    if (auto AT = dyn_cast<LLVM::LLVMArrayType>(p)) {
      p = AT.getElementType();
    }
    newParams.push_back(p);
  }

  static int id = 0;
  auto callName = (op.getName() + "$" + "par" + std::to_string(id)).str();
  id++;
  auto func = func::FuncOp::create(builder, loc, callName, gpuTy0);
  func->setDiscardableAttrs(op->getDiscardableAttrDictionary());
  func.setVisibility(SymbolTable::Visibility::Private);
  auto &entryBlock = *func.addEntryBlock();
  builder.setInsertionPointToStart(&entryBlock);

  SmallVector<mlir::Value> inits;
  SmallVector<mlir::Value> finals;
  SmallVector<mlir::Value> incs;
  for (auto val : {gridx, gridy, gridz, blockx, blocky, blockz}) {
    inits.push_back(arith::ConstantIndexOp::create(builder, loc, 0));
    incs.push_back(arith::ConstantIndexOp::create(builder, loc, 1));
    finals.push_back(arith::ConstantIndexOp::create(builder, loc, val));
  }

  IRMapping map;
  map.map(op.getArguments(), entryBlock.getArguments());

  auto context = loc.getContext();
  SmallVector<AffineMap> idMaps, zeroMaps;
  auto zeroMap = AffineMap::getConstantMap(0, context);
  zeroMaps.insert(zeroMaps.begin(), 6, zeroMap);
  for (unsigned i = 0; i < 6; i++) {
    auto idMap = AffineMap::get(0, 6, getAffineSymbolExpr(i, context));
    idMaps.push_back(idMap);
  }

  SmallVector<int64_t> steps(6, 1);
  auto par = affine::AffineParallelOp::create(
      builder, loc, TypeRange(), ArrayRef<arith::AtomicRMWKind>(), zeroMaps,
      ValueRange(), idMaps, finals, steps);

  mlir::func::ReturnOp::create(builder, loc);

  builder.setInsertionPointToStart(&par.getRegion().front());
  auto executeRegion =
      scf::ExecuteRegionOp::create(builder, loc, ArrayRef<mlir::Type>());

  op.getFunctionBody().cloneInto(&executeRegion.getRegion(), map);

  executeRegion->walk([](LLVM::ReturnOp op) {
    OpBuilder rewriter(op);
    scf::YieldOp::create(rewriter, op.getLoc());
    op.erase();
  });

  executeRegion->walk([](func::ReturnOp op) {
    OpBuilder rewriter(op);
    scf::YieldOp::create(rewriter, op.getLoc());
    op.erase();
  });

  executeRegion->walk([](LLVM::UnreachableOp op) {
    OpBuilder rewriter(op);
    scf::YieldOp::create(rewriter, op.getLoc());
    op.erase();
  });

  // block idx
  executeRegion->walk([&](NVVM::BlockIdXOp idxOp) {
    OpBuilder rewriter(idxOp);
    auto rep = arith::IndexCastUIOp::create(rewriter, op.getLoc(),
                                            idxOp.getType(), par.getIVs()[0]);
    idxOp.replaceAllUsesWith(rep.getResult());
    idxOp.erase();
  });
  executeRegion->walk([&](NVVM::BlockIdYOp idxOp) {
    OpBuilder rewriter(idxOp);
    auto rep = arith::IndexCastUIOp::create(rewriter, op.getLoc(),
                                            idxOp.getType(), par.getIVs()[1]);
    idxOp.replaceAllUsesWith(rep.getResult());
    idxOp.erase();
  });
  executeRegion->walk([&](NVVM::BlockIdZOp idxOp) {
    OpBuilder rewriter(idxOp);
    auto rep = arith::IndexCastUIOp::create(rewriter, op.getLoc(),
                                            idxOp.getType(), par.getIVs()[2]);
    idxOp.replaceAllUsesWith(rep.getResult());
    idxOp.erase();
  });
  executeRegion->walk([&](gpu::BlockIdOp idxOp) {
    Value val = nullptr;
    if (idxOp.getDimension() == gpu::Dimension::x)
      val = par.getIVs()[0];
    else if (idxOp.getDimension() == gpu::Dimension::y)
      val = par.getIVs()[1];
    else if (idxOp.getDimension() == gpu::Dimension::z)
      val = par.getIVs()[2];
    else
      llvm_unreachable("illegal dimension");
    idxOp.replaceAllUsesWith(val);
    idxOp.erase();
  });

  // thread idx
  executeRegion->walk([&](NVVM::ThreadIdXOp idxOp) {
    OpBuilder rewriter(idxOp);
    auto rep = arith::IndexCastUIOp::create(rewriter, op.getLoc(),
                                            idxOp.getType(), par.getIVs()[3]);
    idxOp.replaceAllUsesWith(rep.getResult());
    idxOp.erase();
  });
  executeRegion->walk([&](NVVM::ThreadIdYOp idxOp) {
    OpBuilder rewriter(idxOp);
    auto rep = arith::IndexCastUIOp::create(rewriter, op.getLoc(),
                                            idxOp.getType(), par.getIVs()[4]);
    idxOp.replaceAllUsesWith(rep.getResult());
    idxOp.erase();
  });
  executeRegion->walk([&](NVVM::ThreadIdZOp idxOp) {
    OpBuilder rewriter(idxOp);
    auto rep = arith::IndexCastUIOp::create(rewriter, op.getLoc(),
                                            idxOp.getType(), par.getIVs()[5]);
    idxOp.replaceAllUsesWith(rep.getResult());
    idxOp.erase();
  });
  executeRegion->walk([&](gpu::ThreadIdOp idxOp) {
    Value val = nullptr;
    if (idxOp.getDimension() == gpu::Dimension::x)
      val = par.getIVs()[3];
    else if (idxOp.getDimension() == gpu::Dimension::y)
      val = par.getIVs()[4];
    else if (idxOp.getDimension() == gpu::Dimension::z)
      val = par.getIVs()[5];
    else
      llvm_unreachable("illegal dimension");
    idxOp.replaceAllUsesWith(val);
    idxOp.erase();
  });

  OpBuilder rewriter(kcall);
  auto replacement = enzymexla::JITCallOp::create(
      rewriter, kcall.getLoc(), kcall.getResultTypes(),
      SymRefAttrReplacingFunctionName(kcall.getFn(), callName),
      kcall.getInputs(), kcall.getBackendConfigAttr(),
      kcall.getOperandLayoutsAttr(), kcall.getResultLayoutsAttr(),
      kcall.getArgAttrsAttr(), kcall.getResAttrsAttr(),
      kcall.getOutputOperandAliasesAttr(), kcall.getXlaSideEffectFreeAttr());
  kcall.replaceAllUsesWith(replacement);
  kcall.erase();
  return true;
};

namespace {

struct LowerKernelPass
    : public mlir::enzyme::impl::LowerKernelPassBase<LowerKernelPass> {
  using LowerKernelPassBase::LowerKernelPassBase;

  void runOnOperation() override {
    SymbolTableCollection symbolTable;
    symbolTable.getSymbolTable(getOperation());

    getOperation()->walk([&](KernelCallOp op) {
      size_t data[11];

      auto *symbolOp = symbolTable.lookupNearestSymbolFrom(op, op.getFnAttr());
      auto fn = cast<FunctionOpInterface>(symbolOp);
      if (fn.getArguments().size() != op.getInputs().size()) {
        op->emitError() << "Kernel_call had " << op.getInputs().size()
                        << " whereas called kernel requires "
                        << fn.getArguments().size() << "\n";
        return;
      }

      Value vals[] = {op.getGridx(),  op.getGridy(),  op.getGridz(),
                      op.getBlockx(), op.getBlocky(), op.getBlockz(),
                      op.getShmem()};
      for (auto en : llvm::enumerate(vals)) {
        auto val = getConstantValue(op, en.value(), 0);
        if (val == -1)
          return;
        data[1 + en.index()] = val;
      }

      // Resolve the optional cluster dimensions
      Value clusterVals[] = {op.getClusterx(), op.getClustery(),
                             op.getClusterz()};
      for (auto en : llvm::enumerate(clusterVals)) {
        auto val = getConstantValue(op, en.value(), 1);
        if (val == -1)
          return;
        data[8 + en.index()] = val;
      }

      // Compiled kernel goes here once ready
      if (backend == "cuda") {
        CompileGPUKernel(symbolTable, op.getLoc(), fn, data[1], data[2],
                         data[3], data[4], data[5], data[6], data[7], data[8],
                         data[9], data[10], op);
      } else if (backend == "cpu") {
        CompileCPUKernel(symbolTable, op.getLoc(), fn, data[1], data[2],
                         data[3], data[4], data[5], data[6], data[7], data[8],
                         data[9], data[10], op);
      } else {
        op->emitError() << "Cannot lower kernel to unknown backend \""
                        << backend << "\"";
      }
    });
  }

private:
  size_t getConstantValue(KernelCallOp op, Value v, size_t defaultValue) {
    if (!v)
      return defaultValue;

    DenseIntElementsAttr stepAttr;
    if (!matchPattern(v, m_Constant(&stepAttr))) {
      op->emitError() << "Cannot lower kernel with a value which "
                         "is not a constant integer tensor";
      return -1;
    }
    if (stepAttr.size() != 1) {
      op->emitError() << "Cannot lower kernel with a value which "
                         "is not a constant integer tensor of size 1";
      return -1;
    }
    return (*stepAttr.begin()).getZExtValue();
  }
};

} // end anonymous namespace
