#include "PatternTritonGPUOpToLLVM.h"
#include "Utility.h"

#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
#include "triton/Dialect/TritonInstrument/IR/Dialect.h"

using namespace mlir;
using namespace mlir::triton;

using ::mlir::triton::gpu::getShapePerCTA;
using ::mlir::triton::gpu::isPermutationMatrixLayout;
using ::mlir::triton::gpu::NvidiaMmaEncodingAttr;
using ::mlir::triton::gpu::toLinearLayout;

LogicalResult convertMMA(triton::DotOp op, triton::DotOp::Adaptor adaptor,
                         const LLVMTypeConverter *typeConverter,
                         ConversionPatternRewriter &rewriter, bool isTuring);

LogicalResult convertMMA(triton::instrument::DotI8Op op,
                         triton::instrument::DotI8Op::Adaptor adaptor,
                         const LLVMTypeConverter *typeConverter,
                         ConversionPatternRewriter &rewriter, bool isTuring);

LogicalResult convertMMADotScaled(triton::DotScaledOp op,
                                  triton::DotScaledOp::Adaptor adaptor,
                                  const LLVMTypeConverter *typeConverter,
                                  ConversionPatternRewriter &rewriter);

LogicalResult convertWGMMA(triton::nvidia_gpu::WarpGroupDotOp op,
                           triton::nvidia_gpu::WarpGroupDotOp::Adaptor adaptor,
                           const LLVMTypeConverter *typeConverter,
                           ConversionPatternRewriter &rewriter, Value thread);

namespace {
struct ScaledDotOpConversion
    : public ConvertOpToLLVMPattern<triton::DotScaledOp> {
  using ConvertOpToLLVMPattern<triton::DotScaledOp>::ConvertOpToLLVMPattern;

  ScaledDotOpConversion(LLVMTypeConverter &converter, int,
                        PatternBenefit benefit)
      : ConvertOpToLLVMPattern<triton::DotScaledOp>(converter, benefit) {}

  LogicalResult
  matchAndRewrite(triton::DotScaledOp op, triton::DotScaledOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto rty = cast<RankedTensorType>(op.getResult().getType());
    if (!isPermutationMatrixLayout(
            toLinearLayout(rty.getShape(), rty.getEncoding())))
      return rewriter.notifyMatchFailure(
          op, "ScaledDotOp result encoding must have a permutation-matrix "
              "linear layout");
    return convertMMADotScaled(op, adaptor, getTypeConverter(), rewriter);
  }
};

struct DotOpConversion : public ConvertOpToLLVMPattern<triton::DotOp> {
  using ConvertOpToLLVMPattern<triton::DotOp>::ConvertOpToLLVMPattern;

  DotOpConversion(LLVMTypeConverter &converter, int, PatternBenefit benefit)
      : ConvertOpToLLVMPattern<triton::DotOp>(converter, benefit) {}

  LogicalResult
  matchAndRewrite(triton::DotOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // D = A * B + C
    auto dType = op.getResult().getType();
    auto dEncoding = dType.getEncoding();

    if (!isPermutationMatrixLayout(toLinearLayout(dType.getShape(), dEncoding)))
      return rewriter.notifyMatchFailure(
          op,
          "DotOp result encoding must have a permutation-matrix linear layout");

    NvidiaMmaEncodingAttr mmaLayout =
        dyn_cast<NvidiaMmaEncodingAttr>(dEncoding);
    if (mmaLayout) {
      if (mmaLayout.getVersionMajor() == 2) {
        return convertMMA(op, adaptor, getTypeConverter(), rewriter,
                          mmaLayout.isTuring());
      }

      llvm::report_fatal_error(
          "Unsupported MMA kind found when converting DotOp to LLVM.");
    }

    if (isa<BlockedEncodingAttr>(dEncoding))
      return convertFMADot(op, adaptor, getTypeConverter(), rewriter);

    llvm::report_fatal_error(
        "Unsupported DotOp found when converting TritonGPU to LLVM.");
  }
};

struct DotI8OpConversion
    : public ConvertOpToLLVMPattern<triton::instrument::DotI8Op> {
  using ConvertOpToLLVMPattern<
      triton::instrument::DotI8Op>::ConvertOpToLLVMPattern;

  DotI8OpConversion(LLVMTypeConverter &converter, int, PatternBenefit benefit)
      : ConvertOpToLLVMPattern<triton::instrument::DotI8Op>(converter,
                                                            benefit) {}

  LogicalResult
  matchAndRewrite(triton::instrument::DotI8Op op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto dType = op.getD().getType();
    auto dEncoding = dType.getEncoding();
    if (!isPermutationMatrixLayout(toLinearLayout(dType.getShape(), dEncoding)))
      return rewriter.notifyMatchFailure(
          op, "DotI8Op result encoding must have a permutation-matrix linear "
              "layout");

    auto mmaLayout = dyn_cast<NvidiaMmaEncodingAttr>(dEncoding);
    if (!mmaLayout || mmaLayout.getVersionMajor() != 2)
      return rewriter.notifyMatchFailure(op,
                                         "DotI8Op requires an MMAv2 layout");
    return convertMMA(op, adaptor, getTypeConverter(), rewriter,
                      mmaLayout.isTuring());
  }
};

struct WarpGroupDotOpConversion
    : public ConvertOpToLLVMPattern<triton::nvidia_gpu::WarpGroupDotOp> {
  using ConvertOpToLLVMPattern<
      triton::nvidia_gpu::WarpGroupDotOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(triton::nvidia_gpu::WarpGroupDotOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    return convertWGMMA(op, adaptor, getTypeConverter(), rewriter,
                        getThreadId(rewriter, loc));
  }
};

struct WarpGroupDotWaitOpConversion
    : public ConvertOpToLLVMPattern<triton::nvidia_gpu::WarpGroupDotWaitOp> {
  using ConvertOpToLLVMPattern<
      triton::nvidia_gpu::WarpGroupDotWaitOp>::ConvertOpToLLVMPattern;

  static FailureOr<Value> packInputs(Location loc, RewriterBase &rewriter,
                                     ValueRange inputs) {
    if (inputs.size() == 1)
      return inputs.front();

    SmallVector<Type> types;
    // Pack the inputs into a single struct.
    for (Type type : inputs.getTypes()) {
      auto structType = dyn_cast<LLVM::LLVMStructType>(type);
      if (!structType)
        return failure();
      llvm::append_range(types, structType.getBody());
    }
    auto packedType =
        LLVM::LLVMStructType::getLiteral(rewriter.getContext(), types);
    Value packed = LLVM::UndefOp::create(rewriter, loc, packedType);
    unsigned outputStructIndex = 0;
    for (Value input : inputs) {
      for (auto [i, type] : llvm::enumerate(
               cast<LLVM::LLVMStructType>(input.getType()).getBody())) {
        Value value = LLVM::ExtractValueOp::create(rewriter, loc, input, i);
        packed = LLVM::InsertValueOp::create(rewriter, loc, packedType, packed,
                                             value, outputStructIndex++);
      }
    }
    return packed;
  }

  static SmallVector<Value> unpackOutput(Location loc, RewriterBase &rewriter,
                                         Value packedOutput, TypeRange types) {
    SmallVector<Value> outputs;
    if (types.size() == 1) {
      outputs.push_back(packedOutput);
      return outputs;
    }

    // Unpack the output into the original struct types.
    unsigned outputStructIndex = 0;
    for (Type type : types) {
      auto structType = cast<LLVM::LLVMStructType>(type);
      Value unpacked = LLVM::UndefOp::create(rewriter, loc, structType);
      for (auto [i, type] : llvm::enumerate(structType.getBody())) {
        Value value = LLVM::ExtractValueOp::create(rewriter, loc, packedOutput,
                                                   outputStructIndex++);
        unpacked = LLVM::InsertValueOp::create(rewriter, loc, structType,
                                               unpacked, value, i);
      }
      outputs.push_back(unpacked);
    }
    return outputs;
  }

  LogicalResult
  matchAndRewrite(triton::nvidia_gpu::WarpGroupDotWaitOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    ValueRange inputs = adaptor.getInputs();
    auto packed = packInputs(loc, rewriter, inputs);
    if (failed(packed))
      return failure();

    auto wait = triton::nvgpu::WGMMAWaitGroupOp::create(rewriter, loc, *packed,
                                                        op.getPendings());
    auto outputs =
        unpackOutput(loc, rewriter, wait.getResult(), inputs.getTypes());
    bool synchronizeWarpGroups =
        !op.getWarpGroupLocal() && triton::gpu::lookupNumWarps(op) > 4;
    rewriter.replaceOp(op, outputs);

    // WGMMA waits only synchronize the issuing warp group. Synchronize all
    // warp groups before releasing shared-memory dependencies.
    if (synchronizeWarpGroups) {
      triton::gpu::BarrierOp::create(rewriter, loc,
                                     triton::gpu::AddrSpace::Local);
    }
    return success();
  }
};
} // namespace

void mlir::triton::NVIDIA::populateDotOpToLLVMPatterns(
    LLVMTypeConverter &typeConverter, RewritePatternSet &patterns,
    int computeCapability, PatternBenefit benefit) {
  patterns.add<DotOpConversion>(typeConverter, computeCapability, benefit);
  patterns.add<DotI8OpConversion>(typeConverter, computeCapability, benefit);
  patterns.add<WarpGroupDotOpConversion>(typeConverter, benefit);
  patterns.add<WarpGroupDotWaitOpConversion>(typeConverter, benefit);
  patterns.add<ScaledDotOpConversion>(typeConverter, computeCapability,
                                      benefit);
}
