#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
#include "triton/Conversion/TritonGPUToLLVM/Utility.h"

namespace {

using namespace mlir;
using namespace mlir::triton;

// NOTE: [Additional Function Arguments]
// Triton patches additional arguments to the function signature to support
// (1) shared memory, (2) global scratch memory, and (3) profile scratch memory.
// To support use of shared memory and global scratch memory inside of a
// function, the caller allocates a single large block of the relevant memory
// and calls the function with these extra arguments at the end.
// Profile scratch memory is only used when the function is instrumented for
// profiling.
//
// For the kernel function itself, the shared memory base is a global symbol
// so no additional function argument is required but global scratch memory
// allocation is still passed in as the last argument. Though here the scratch
// memory is shared between all programs, so a linear offset based on the
// program id is required to get the local scratch base.
struct FuncOpConversion : public ConvertOpToLLVMPattern<triton::FuncOp> {
  FuncOpConversion(LLVMTypeConverter &converter,
                   const TargetInfoBase &targetInfo, PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit), targetInfo(targetInfo) {}

  // Map the MLIR attribute `tt.nv_tma_desc` to the appropriate LLVM and NVVM
  // attributes.
  static LogicalResult handleByvalTmaDescArgs(LLVM::LLVMFuncOp &llvmFuncOp) {
    const bool isKernel = triton::isKernel(llvmFuncOp);
    for (unsigned i = 0; i < llvmFuncOp.getNumArguments(); ++i) {
      const auto attrs = llvmFuncOp.getArgAttrDict(i);
      if (!attrs) {
        continue;
      }

      for (const auto &attr : attrs) {
        if (attr.getName() == "tt.nv_tma_desc") {
          const auto i32_type = IntegerType::get(llvmFuncOp.getContext(), 32);
          assert(attr.getValue() == mlir::IntegerAttr::get(i32_type, 1));
          if (!isKernel) {
            return llvmFuncOp.emitError(
                "tt.nv_tma_desc is not supported for device functions");
          }

          // See
          // https://github.com/google/jax/blob/main/jaxlib/mosaic/gpu/passes.cc
          const auto byteType =
              mlir::IntegerType::get(llvmFuncOp.getContext(), 8);
          const auto arrayType = mlir::LLVM::LLVMArrayType::get(
              llvmFuncOp.getContext(), byteType, 128);
          llvmFuncOp.setArgAttr(i, LLVM::LLVMDialect::getByValAttrName(),
                                mlir::TypeAttr::get(arrayType));
          llvmFuncOp.setArgAttr(i, NVVM::NVVMDialect::getGridConstantAttrName(),
                                mlir::UnitAttr::get(llvmFuncOp.getContext()));
          llvmFuncOp.setArgAttr(i, LLVM::LLVMDialect::getAlignAttrName(),
                                mlir::IntegerAttr::get(i32_type, 64));
        }
      }
    }
    return success();
  }

  static void attachKernelAttributes(LLVM::LLVMFuncOp newFuncOp,
                                     triton::FuncOp funcOp,
                                     ConversionPatternRewriter &rewriter) {
    newFuncOp->setAttr(
        NVVM::NVVMDialect::getKernelFuncAttrName(),
        rewriter.getIntegerAttr(type::u1Ty(rewriter.getContext()), 1));

    // Determine the actual number of required warps.
    int numWarps = triton::gpu::lookupNumWarps(funcOp);
    if (auto totalNumWarps = funcOp.getParentOp()->getAttrOfType<IntegerAttr>(
            "ttg.total-num-warps"))
      numWarps = totalNumWarps.getInt();

    int numCTAs = 1;
    if (auto module = funcOp->getParentOfType<ModuleOp>()) {
      if (auto moduleAttr =
              module->getAttrOfType<IntegerAttr>(triton::gpu::AttrNumCTAsName))
        numCTAs = moduleAttr.getInt();
    }

    // Set `nvvm.maxnreg` if it was specified on the module.
    if (Attribute maxnregAttr =
            funcOp.getParentOp()->getAttr(triton::gpu::AttrMaxRegistersName))
      newFuncOp->setAttr(NVVM::NVVMDialect::getMaxnregAttrName(), maxnregAttr);

    // Do we want to do this for nCTAs == 1 whenever sm >= 90?
    if (numCTAs > 1) {
      // Request a specific number of CTAs per cluster in the generated PTX.
      newFuncOp->setAttr(NVVM::NVVMDialect::getClusterDimAttrName(),
                         rewriter.getDenseI32ArrayAttr(numCTAs));
    }

    // Set an attribute for reqntidx, it could be used in latter LLVM codegen
    // for `nvvm.annotation` metadata.
    newFuncOp->setAttr(NVVM::NVVMDialect::getReqntidAttrName(),
                       rewriter.getDenseI32ArrayAttr(32 * numWarps));
  }

  LogicalResult
  matchAndRewrite(triton::FuncOp funcOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Prevent LLVM's inliner to inline this function
    auto amendedFuncOp = amendFuncOp(funcOp, rewriter, targetInfo);

    FailureOr<LLVM::LLVMFuncOp> maybeNewFuncOp =
        mlir::convertFuncOpToLLVMFuncOp(amendedFuncOp, rewriter,
                                        *getTypeConverter());
    if (failed(maybeNewFuncOp)) {
      return failure();
    }

    LLVM::LLVMFuncOp newFuncOp = *maybeNewFuncOp;
    handleArgPtrDatatype(funcOp, newFuncOp);

    auto ctx = funcOp->getContext();

    if (triton::isKernel(funcOp)) {
      // Set an attribute to indicate this function is a kernel entry.
      attachKernelAttributes(newFuncOp, funcOp, rewriter);
      newFuncOp.setLinkage(LLVM::Linkage::External);
    } else {
      // The noinline attribute will be used by the LLVM codegen to prevent
      // inlining.
      // https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp#L267
      newFuncOp.setPassthroughAttr(
          ArrayAttr::get(ctx, rewriter.getStringAttr("noinline")));
      newFuncOp.setLinkage(LLVM::Linkage::Internal);
      if (Attribute numWarps = funcOp->getAttr(triton::gpu::AttrNumWarpsName))
        newFuncOp->setAttr("ws_num_warps", numWarps);
    }

    rewriter.eraseOp(funcOp);
    rewriter.eraseOp(amendedFuncOp);

    // Add attributes for by-value TMA descriptor args (nvidia)
    return handleByvalTmaDescArgs(newFuncOp);
  }

private:
  const TargetInfoBase &targetInfo;
};

} // namespace

void mlir::triton::populateFuncOpConversionPattern(
    LLVMTypeConverter &typeConverter, RewritePatternSet &patterns,
    const TargetInfoBase &targetInfo, PatternBenefit benefit) {
  patterns.add<FuncOpConversion>(typeConverter, targetInfo, benefit);
}
