//===-- CUFGPUToLLVMConversion.cpp ----------------------------------------===//
//
// 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 "flang/Optimizer/Transforms/CUFGPUToLLVMConversion.h"
#include "flang/Optimizer/Builder/CUFCommon.h"
#include "flang/Optimizer/CodeGen/TypeConverter.h"
#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
#include "flang/Optimizer/Support/DataLayout.h"
#include "flang/Runtime/CUDA/common.h"
#include "flang/Support/Fortran.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/Support/FormatVariadic.h"

namespace fir {
#define GEN_PASS_DEF_CUFGPUTOLLVMCONVERSION
#include "flang/Optimizer/Transforms/Passes.h.inc"
} // namespace fir

using namespace fir;
using namespace mlir;
using namespace Fortran::runtime;

namespace {

static mlir::Value createKernelArgArray(mlir::Location loc,
                                        mlir::ValueRange operands,
                                        mlir::PatternRewriter &rewriter) {

  auto *ctx = rewriter.getContext();
  llvm::SmallVector<mlir::Type> structTypes(operands.size(), nullptr);

  for (auto [i, arg] : llvm::enumerate(operands))
    structTypes[i] = arg.getType();

  auto structTy = mlir::LLVM::LLVMStructType::getLiteral(ctx, structTypes);
  auto ptrTy = mlir::LLVM::LLVMPointerType::get(rewriter.getContext());
  mlir::Type i32Ty = rewriter.getI32Type();
  auto zero = rewriter.create<mlir::LLVM::ConstantOp>(
      loc, i32Ty, rewriter.getIntegerAttr(i32Ty, 0));
  auto one = rewriter.create<mlir::LLVM::ConstantOp>(
      loc, i32Ty, rewriter.getIntegerAttr(i32Ty, 1));
  mlir::Value argStruct =
      rewriter.create<mlir::LLVM::AllocaOp>(loc, ptrTy, structTy, one);
  auto size = rewriter.create<mlir::LLVM::ConstantOp>(
      loc, i32Ty, rewriter.getIntegerAttr(i32Ty, structTypes.size()));
  mlir::Value argArray =
      rewriter.create<mlir::LLVM::AllocaOp>(loc, ptrTy, ptrTy, size);

  for (auto [i, arg] : llvm::enumerate(operands)) {
    auto indice = rewriter.create<mlir::LLVM::ConstantOp>(
        loc, i32Ty, rewriter.getIntegerAttr(i32Ty, i));
    mlir::Value structMember = rewriter.create<LLVM::GEPOp>(
        loc, ptrTy, structTy, argStruct,
        mlir::ArrayRef<mlir::Value>({zero, indice}));
    rewriter.create<LLVM::StoreOp>(loc, arg, structMember);
    mlir::Value arrayMember = rewriter.create<LLVM::GEPOp>(
        loc, ptrTy, ptrTy, argArray, mlir::ArrayRef<mlir::Value>({indice}));
    rewriter.create<LLVM::StoreOp>(loc, structMember, arrayMember);
  }
  return argArray;
}

struct GPULaunchKernelConversion
    : public mlir::ConvertOpToLLVMPattern<mlir::gpu::LaunchFuncOp> {
  explicit GPULaunchKernelConversion(
      const fir::LLVMTypeConverter &typeConverter, mlir::PatternBenefit benefit)
      : mlir::ConvertOpToLLVMPattern<mlir::gpu::LaunchFuncOp>(typeConverter,
                                                              benefit) {}

  using OpAdaptor = typename mlir::gpu::LaunchFuncOp::Adaptor;

  mlir::LogicalResult
  matchAndRewrite(mlir::gpu::LaunchFuncOp op, OpAdaptor adaptor,
                  mlir::ConversionPatternRewriter &rewriter) const override {
    mlir::Location loc = op.getLoc();
    auto *ctx = rewriter.getContext();
    mlir::ModuleOp mod = op->getParentOfType<mlir::ModuleOp>();
    mlir::Value dynamicMemorySize = op.getDynamicSharedMemorySize();
    mlir::Type i32Ty = rewriter.getI32Type();
    if (!dynamicMemorySize)
      dynamicMemorySize = rewriter.create<mlir::LLVM::ConstantOp>(
          loc, i32Ty, rewriter.getIntegerAttr(i32Ty, 0));

    mlir::Value kernelArgs =
        createKernelArgArray(loc, adaptor.getKernelOperands(), rewriter);

    auto ptrTy = mlir::LLVM::LLVMPointerType::get(rewriter.getContext());
    auto kernel = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(op.getKernelName());
    mlir::Value kernelPtr;
    if (!kernel) {
      auto funcOp = mod.lookupSymbol<mlir::func::FuncOp>(op.getKernelName());
      if (!funcOp)
        return mlir::failure();
      kernelPtr =
          rewriter.create<LLVM::AddressOfOp>(loc, ptrTy, funcOp.getName());
    } else {
      kernelPtr =
          rewriter.create<LLVM::AddressOfOp>(loc, ptrTy, kernel.getName());
    }

    auto llvmIntPtrType = mlir::IntegerType::get(
        ctx, this->getTypeConverter()->getPointerBitwidth(0));
    auto voidTy = mlir::LLVM::LLVMVoidType::get(ctx);

    mlir::Value nullPtr = rewriter.create<LLVM::ZeroOp>(loc, ptrTy);

    if (op.hasClusterSize()) {
      auto funcOp = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(
          RTNAME_STRING(CUFLaunchClusterKernel));
      auto funcTy = mlir::LLVM::LLVMFunctionType::get(
          voidTy,
          {ptrTy, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,
           llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,
           llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, i32Ty, ptrTy, ptrTy},
          /*isVarArg=*/false);
      auto cufLaunchClusterKernel = mlir::SymbolRefAttr::get(
          mod.getContext(), RTNAME_STRING(CUFLaunchClusterKernel));
      if (!funcOp) {
        mlir::OpBuilder::InsertionGuard insertGuard(rewriter);
        rewriter.setInsertionPointToStart(mod.getBody());
        auto launchKernelFuncOp = rewriter.create<mlir::LLVM::LLVMFuncOp>(
            loc, RTNAME_STRING(CUFLaunchClusterKernel), funcTy);
        launchKernelFuncOp.setVisibility(
            mlir::SymbolTable::Visibility::Private);
      }
      mlir::Value stream = adaptor.getAsyncObject();
      if (!stream)
        stream = rewriter.create<mlir::LLVM::ConstantOp>(
            loc, llvmIntPtrType, rewriter.getIntegerAttr(llvmIntPtrType, -1));
      rewriter.replaceOpWithNewOp<mlir::LLVM::CallOp>(
          op, funcTy, cufLaunchClusterKernel,
          mlir::ValueRange{kernelPtr, adaptor.getClusterSizeX(),
                           adaptor.getClusterSizeY(), adaptor.getClusterSizeZ(),
                           adaptor.getGridSizeX(), adaptor.getGridSizeY(),
                           adaptor.getGridSizeZ(), adaptor.getBlockSizeX(),
                           adaptor.getBlockSizeY(), adaptor.getBlockSizeZ(),
                           stream, dynamicMemorySize, kernelArgs, nullPtr});
    } else {
      auto procAttr =
          op->getAttrOfType<cuf::ProcAttributeAttr>(cuf::getProcAttrName());
      bool isGridGlobal =
          procAttr && procAttr.getValue() == cuf::ProcAttribute::GridGlobal;
      llvm::StringRef fctName = isGridGlobal
                                    ? RTNAME_STRING(CUFLaunchCooperativeKernel)
                                    : RTNAME_STRING(CUFLaunchKernel);
      auto funcOp = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(fctName);
      auto funcTy = mlir::LLVM::LLVMFunctionType::get(
          voidTy,
          {ptrTy, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,
           llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,
           i32Ty, ptrTy, ptrTy},
          /*isVarArg=*/false);
      auto cufLaunchKernel =
          mlir::SymbolRefAttr::get(mod.getContext(), fctName);
      if (!funcOp) {
        mlir::OpBuilder::InsertionGuard insertGuard(rewriter);
        rewriter.setInsertionPointToStart(mod.getBody());
        auto launchKernelFuncOp =
            rewriter.create<mlir::LLVM::LLVMFuncOp>(loc, fctName, funcTy);
        launchKernelFuncOp.setVisibility(
            mlir::SymbolTable::Visibility::Private);
      }

      mlir::Value stream = adaptor.getAsyncObject();
      if (!stream)
        stream = rewriter.create<mlir::LLVM::ConstantOp>(
            loc, llvmIntPtrType, rewriter.getIntegerAttr(llvmIntPtrType, -1));

      rewriter.replaceOpWithNewOp<mlir::LLVM::CallOp>(
          op, funcTy, cufLaunchKernel,
          mlir::ValueRange{kernelPtr, adaptor.getGridSizeX(),
                           adaptor.getGridSizeY(), adaptor.getGridSizeZ(),
                           adaptor.getBlockSizeX(), adaptor.getBlockSizeY(),
                           adaptor.getBlockSizeZ(), stream, dynamicMemorySize,
                           kernelArgs, nullPtr});
    }

    return mlir::success();
  }
};

static std::string getFuncName(cuf::SharedMemoryOp op) {
  if (auto gpuFuncOp = op->getParentOfType<mlir::gpu::GPUFuncOp>())
    return gpuFuncOp.getName().str();
  if (auto funcOp = op->getParentOfType<mlir::func::FuncOp>())
    return funcOp.getName().str();
  if (auto llvmFuncOp = op->getParentOfType<mlir::LLVM::LLVMFuncOp>())
    return llvmFuncOp.getSymName().str();
  return "";
}

static mlir::Value createAddressOfOp(mlir::ConversionPatternRewriter &rewriter,
                                     mlir::Location loc,
                                     gpu::GPUModuleOp gpuMod,
                                     std::string &sharedGlobalName) {
  auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(
      rewriter.getContext(), mlir::NVVM::NVVMMemorySpace::kSharedMemorySpace);
  if (auto g = gpuMod.lookupSymbol<fir::GlobalOp>(sharedGlobalName))
    return rewriter.create<mlir::LLVM::AddressOfOp>(loc, llvmPtrTy,
                                                    g.getSymName());
  if (auto g = gpuMod.lookupSymbol<mlir::LLVM::GlobalOp>(sharedGlobalName))
    return rewriter.create<mlir::LLVM::AddressOfOp>(loc, llvmPtrTy,
                                                    g.getSymName());
  return {};
}

struct CUFSharedMemoryOpConversion
    : public mlir::ConvertOpToLLVMPattern<cuf::SharedMemoryOp> {
  explicit CUFSharedMemoryOpConversion(
      const fir::LLVMTypeConverter &typeConverter, mlir::PatternBenefit benefit)
      : mlir::ConvertOpToLLVMPattern<cuf::SharedMemoryOp>(typeConverter,
                                                          benefit) {}
  using OpAdaptor = typename cuf::SharedMemoryOp::Adaptor;

  mlir::LogicalResult
  matchAndRewrite(cuf::SharedMemoryOp op, OpAdaptor adaptor,
                  mlir::ConversionPatternRewriter &rewriter) const override {
    mlir::Location loc = op->getLoc();
    if (!op.getOffset())
      mlir::emitError(loc,
                      "cuf.shared_memory must have an offset for code gen");

    auto gpuMod = op->getParentOfType<gpu::GPUModuleOp>();
    std::string sharedGlobalName =
        (getFuncName(op) + llvm::Twine(cudaSharedMemSuffix)).str();
    mlir::Value sharedGlobalAddr =
        createAddressOfOp(rewriter, loc, gpuMod, sharedGlobalName);

    if (!sharedGlobalAddr)
      mlir::emitError(loc, "Could not find the shared global operation\n");

    auto castPtr = rewriter.create<mlir::LLVM::AddrSpaceCastOp>(
        loc, mlir::LLVM::LLVMPointerType::get(rewriter.getContext()),
        sharedGlobalAddr);
    mlir::Type baseType = castPtr->getResultTypes().front();
    llvm::SmallVector<mlir::LLVM::GEPArg> gepArgs = {op.getOffset()};
    mlir::Value shmemPtr = rewriter.create<mlir::LLVM::GEPOp>(
        loc, baseType, rewriter.getI8Type(), castPtr, gepArgs);
    rewriter.replaceOp(op, {shmemPtr});
    return mlir::success();
  }
};

class CUFGPUToLLVMConversion
    : public fir::impl::CUFGPUToLLVMConversionBase<CUFGPUToLLVMConversion> {
public:
  void runOnOperation() override {
    auto *ctx = &getContext();
    mlir::RewritePatternSet patterns(ctx);
    mlir::ConversionTarget target(*ctx);

    mlir::Operation *op = getOperation();
    mlir::ModuleOp module = mlir::dyn_cast<mlir::ModuleOp>(op);
    if (!module)
      return signalPassFailure();

    std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(
        module, /*allowDefaultLayout=*/false);
    fir::LLVMTypeConverter typeConverter(module, /*applyTBAA=*/false,
                                         /*forceUnifiedTBAATree=*/false, *dl);
    cuf::populateCUFGPUToLLVMConversionPatterns(typeConverter, patterns);
    target.addIllegalOp<mlir::gpu::LaunchFuncOp>();
    target.addIllegalOp<cuf::SharedMemoryOp>();
    target.addLegalDialect<mlir::LLVM::LLVMDialect>();
    if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,
                                                  std::move(patterns)))) {
      mlir::emitError(mlir::UnknownLoc::get(ctx),
                      "error in CUF GPU op conversion\n");
      signalPassFailure();
    }
  }
};
} // namespace

void cuf::populateCUFGPUToLLVMConversionPatterns(
    const fir::LLVMTypeConverter &converter, mlir::RewritePatternSet &patterns,
    mlir::PatternBenefit benefit) {
  patterns.add<CUFSharedMemoryOpConversion, GPULaunchKernelConversion>(
      converter, benefit);
}