#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
#include "triton/Conversion/TritonGPUToLLVM/Utility.h"

namespace {

using namespace mlir;
using namespace mlir::triton;

struct ReturnOpConversion : public ConvertOpToLLVMPattern<triton::ReturnOp> {
  using ConvertOpToLLVMPattern<triton::ReturnOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(triton::ReturnOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto funcOp = op->getParentOfType<LLVM::LLVMFuncOp>();
    auto loc = op.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    if (funcOp->hasAttr("nvvm.kernel")) {
      // A GPU kernel
      if (op.getNumOperands() > 0) {
        return rewriter.notifyMatchFailure(
            op, "Kernel functions do not support return with operands");
      }
      rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(op, TypeRange(), ValueRange(),
                                                  op->getAttrs());
    } else {
      // A device function
      LLVM::ReturnOp newOp;
      if (adaptor.getOperands().size() < 2) {
        // Single or no return value.
        newOp = LLVM::ReturnOp::create(rewriter, op.getLoc(),
                                       adaptor.getOperands());
      } else {
        // Pack the results into a struct.
        auto packedResultsTy = this->getTypeConverter()->packFunctionResults(
            funcOp.getResultTypes());
        Value packedResults =
            LLVM::UndefOp::create(rewriter, op.getLoc(), packedResultsTy);
        for (auto it : llvm::enumerate(adaptor.getOperands())) {
          packedResults = b.insert_val(packedResultsTy, packedResults,
                                       it.value(), it.index());
        }
        newOp = LLVM::ReturnOp::create(rewriter, op.getLoc(), packedResults);
      }
      newOp->setAttrs(op->getAttrs());
      rewriter.replaceOp(op, newOp->getResults());
    }
    return success();
  }
};

// CallOpInterfaceLowering is adapted from
// https://github.com/llvm/llvm-project/blob/fae656b2dd80246c3c6f01e9c77c49560368752c/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp#L485
struct CallOpConversion : public ConvertOpToLLVMPattern<triton::CallOp> {
  CallOpConversion(LLVMTypeConverter &converter,
                   const TargetInfoBase &targetInfo, PatternBenefit benefit)
      : ConvertOpToLLVMPattern<triton::CallOp>(converter, benefit),
        targetInfo(targetInfo) {}

  LogicalResult
  matchAndRewrite(triton::CallOp callOp,
                  typename triton::CallOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto promotedOperands = promoteOperands(callOp, adaptor, rewriter);
    auto newCallOp =
        convertCallOpToLLVMCallOp(callOp, promotedOperands, rewriter);
    if (!newCallOp)
      return failure();
    auto results = getCallOpResults(callOp, newCallOp, rewriter);
    rewriter.replaceOp(callOp, results);
    return success();
  }

private:
  SmallVector<Value, 4>
  promoteOperands(triton::CallOp callOp,
                  typename triton::CallOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const {
    // Get the last argument of the caller, which is the current stack pointer
    // of shared memory and append it to the operands of the callOp.
    auto loc = callOp.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    auto caller = callOp->getParentOfType<FunctionOpInterface>();
    auto promotedOperands = this->getTypeConverter()->promoteOperands(
        callOp.getLoc(), /*opOperands=*/callOp->getOperands(),
        adaptor.getOperands(), rewriter);
    if (!caller->hasAttr("allocation.offset") ||
        !callOp->hasAttr("allocation.offset")) {
      auto base = LLVM::getStackPointer(rewriter, caller);
      promotedOperands.push_back(base);
    } else {
      auto base = LLVM::getSharedMemoryBase(loc, rewriter, targetInfo, callOp);
      promotedOperands.push_back(base);
    }

    auto opOffsetAttr = callOp->getAttrOfType<mlir::IntegerAttr>(
        "ttg.global_scratch_memory_offset");
    Value globalOffsetVal;
    if (opOffsetAttr)
      globalOffsetVal = b.i32_val(opOffsetAttr.getValue().getZExtValue());
    auto profileOffsetAttr = callOp->getAttrOfType<mlir::IntegerAttr>(
        "ttg.profile_scratch_memory_offset");
    Value profileOffsetVal;
    if (profileOffsetAttr)
      profileOffsetVal = b.i32_val(profileOffsetAttr.getValue().getZExtValue());

    promotedOperands.push_back(LLVM::getGlobalScratchPtr(
        loc, rewriter, targetInfo, caller, globalOffsetVal));
    promotedOperands.push_back(LLVM::getProfileScratchPtr(
        loc, rewriter, targetInfo, caller, profileOffsetVal));
    return promotedOperands;
  }

  LLVM::CallOp
  convertCallOpToLLVMCallOp(triton::CallOp callOp,
                            ArrayRef<Value> promotedOperands,
                            ConversionPatternRewriter &rewriter) const {
    // Pack the result types into a struct.
    Type packedResult = nullptr;
    unsigned numResults = callOp.getNumResults();
    auto resultTypes = llvm::to_vector<4>(callOp.getResultTypes());

    if (numResults != 0) {
      if (!(packedResult =
                this->getTypeConverter()->packFunctionResults(resultTypes)))
        return nullptr;
    }
    auto newCallOp = LLVM::CallOp::create(rewriter, callOp.getLoc(),
                                          packedResult ? TypeRange(packedResult)
                                                       : TypeRange(),
                                          promotedOperands, callOp->getAttrs());
    newCallOp.getProperties().setOpBundleSizes(
        rewriter.getDenseI32ArrayAttr({}));
    newCallOp.getProperties().setOperandSegmentSizes(
        {static_cast<int>(promotedOperands.size()), 0});
    return newCallOp;
  }

  SmallVector<Value>
  getCallOpResults(triton::CallOp callOp, LLVM::CallOp newCallOp,
                   ConversionPatternRewriter &rewriter) const {
    auto numResults = callOp.getNumResults();
    SmallVector<Value> results;
    if (numResults < 2) {
      // If < 2 results, packing did not do anything and we can just return.
      results.append(newCallOp.result_begin(), newCallOp.result_end());
    } else {
      // Otherwise, it had been converted to an operation producing a structure.
      // Extract individual results from the structure and return them as list.
      results.reserve(numResults);
      for (unsigned i = 0; i < numResults; ++i) {
        results.push_back(LLVM::ExtractValueOp::create(
            rewriter, callOp.getLoc(), newCallOp->getResult(0), i));
      }
    }
    return results;
  }
  const TargetInfoBase &targetInfo;
};

} // namespace

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