#include "PatternTritonGPUOpToLLVM.h"
#include "TargetInfo.h"
#include "Utility.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "triton/Dialect/Triton/IR/Dialect.h"

#include "BufferOpsEmitter.h"

using namespace triton::AMD;
using mlir::triton::amdgpu::ISAFamily;

namespace {

// Utility function to determine if a scalar/tensor value is zero
bool isZero(Value v) {
  if (auto constantOp = v.getDefiningOp<LLVM::ConstantOp>()) {
    if (auto attr = dyn_cast<IntegerAttr>(constantOp.getValue()))
      return attr.getValue().isZero();
    if (auto attr = dyn_cast<FloatAttr>(constantOp.getValue()))
      return attr.getValue().isZero();
    if (auto denseAttr =
            dyn_cast<DenseFPElementsAttr>(constantOp.getValueAttr()))
      return denseAttr.isSplat() && denseAttr.getSplatValue<APFloat>().isZero();
    if (auto denseAttr =
            dyn_cast<DenseIntElementsAttr>(constantOp.getValueAttr()))
      return denseAttr.isSplat() && denseAttr.getSplatValue<APInt>().isZero();
  }
  return false;
}
} // namespace

namespace mlir::LLVM::AMD {
BufferEmitter::BufferEmitter(RewriterBase &rw, Location loc, TargetInfo ti)
    : rewriter(rw), loc(loc), targetInfo(ti) {}

Value BufferEmitter::createResourceDescriptor(Value basePtr,
                                              Value blockStride) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  // 1. Create the resource descriptor
  // bits 0-11: dst sel, ignored by these intrinsics
  // bits 12-14: data format (ignored, must be nonzero, 7=float)
  // bits 15-18: data format (ignored, must be nonzero, 4=32bit)
  // bit 19: In nested heap (0 here)
  // bit 20: Behavior on unmap (0 means  "return 0 / ignore")
  // bits 21-22: Index stride for swizzles (N/A)
  // bit 23: Add thread ID (0)
  // bit 24: Reserved to 1 (RDNA) or 0 (CDNA)
  // bits 25-26: Reserved (0)
  // bit 27: Buffer is non-volatile (CDNA only)
  // bits 28-29: Out of bounds select (RDNA only)
  //             (0 = structured,
  //              1 = check index,
  //              2 = none,
  //              3 = either swizzles or testing against offset field)
  // bits 30-31: Type (must be 0)
  //
  // For GFX12+ (RDNA4, GFX1250): LLVM's lowerPointerAsRsrcIntrin()
  // (SIISelLowering.cpp) rebuilds the descriptor in v2i64 format (57-bit
  // base, 45-bit num_records) and shifts the flags operand left by 28 bits
  // into bits [127:124] of the descriptor. Therefore only flags bits [3:0]
  // survive, mapping to the hardware descriptor fields:
  //   bit 0 -> bit 124: swizzle_enable (0)
  //   bit 1 -> bit 125: OOB_select (0=structured, 1=check offset only)
  //   bits 2-3 -> bits 127:126: type (must be 0)
  // OOB_select=0 is correct for raw buffer ops with stride=0
  // (structured and unstructured modes are equivalent in this case).
  // The RDNA-style flags below have bits [3:0]=0, so they are effectively
  // ignored on GFX12+ but we include GFX1250 in the check for consistency.
  uint32_t flags = (7 << 12) | (4 << 15);
  if (llvm::is_contained({ISAFamily::RDNA2, ISAFamily::RDNA3, ISAFamily::RDNA4m,
                          ISAFamily::RDNA4, ISAFamily::GFX1250},
                         targetInfo.getISAFamily())) {
    flags |= (1 << 24);
    uint32_t oob = 3;
    flags |= (oob << 28);
  }

  Value stride = b.int_val(16, 0);
  if (llvm::is_contained({ISAFamily::CDNA3, ISAFamily::CDNA4},
                         targetInfo.getISAFamily())) {
#if 0
    // Turn off cache-swizzling for the time being while we are figuring out
    // how to safely use it.
    if (blockStride) {
      Value enableSwizzle = b.int_val(16, 16384);
      Value mask14b = b.int_val(16, 16383);
      // Cache swizzle supports only upto 8k stride. Also simply swizzling the
      // largest available stride (8k) doesn't help those unsupported large
      // stride. Especially better to avoid using the stride which is 2^N when
      // N>13, e.g. by add padding to the buffer.
      Value stride16b =
          LLVM::TruncOp::create(rewriter, loc, i16_ty, blockStride);
      Value strideSat = LLVM::AndOp::create(rewriter, loc, stride16b, mask14b);
      // stride[13:0] = swizzling stride
      // stride[14] = swizzle enabling bit
      stride = LLVM::OrOp::create(rewriter, loc, enableSwizzle, strideSat);
    }
#endif
  }

  Value flagsConst = b.int_val(32, flags);
  Type rsrcType = LLVM::LLVMPointerType::get(rewriter.getContext(), 8);
  Value numRecordsByte = b.int_val(64, std::numeric_limits<int>::max() - 1);

  Value resource = rewriter.createOrFold<ROCDL::MakeBufferRsrcOp>(
      loc, rsrcType, basePtr, stride, numRecordsByte, flagsConst);
  return resource;
}

Value BufferEmitter::emitLoad(Type type, Value rsrcDesc, Value offset,
                              Value pred, Value falseVal,
                              triton::CacheModifier cm) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  SmallVector<Value, 6> args;
  int32_t aux = 0;
  fillCommonArgs(type, rsrcDesc, offset, pred, cm, /*isBufferLoad=*/true, args,
                 aux);
  Type bufferType = getBufferOpType(type, false);
  Value data = ROCDL::RawPtrBufferLoadOp::create(
      rewriter, loc, bufferType, args[0], args[1], args[2],
      rewriter.getI32IntegerAttr(aux), /*alias_scopes=*/nullptr,
      /*noalias_scopes=*/nullptr, /*tbaa=*/nullptr);
  data = b.bitcast(data, type);
  if (!isZero(falseVal))
    data = b.select(pred, data, falseVal);
  return data;
}

ROCDL::RawPtrBufferLoadAsyncLdsOp
BufferEmitter::emitLoadToLds(Type type, Value byteWidth, Value rsrcDesc,
                             Value offset, Value dst, Value pred,
                             triton::CacheModifier cm) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  SmallVector<Value, 6> commonArgs;
  int32_t aux = 0;
  fillCommonArgs(type, rsrcDesc, offset, pred, cm, /*isBufferLoad=*/true,
                 commonArgs, aux);

  // buffer_load_to_lds is only supported on gfx942/gfx950 which always use
  // asyncmark. Emit the async intrinsic so LLVM's SIInsertWaitcnts tracks
  // these operations via asyncmark/wait_asyncmark.
  return ROCDL::RawPtrBufferLoadAsyncLdsOp::create(
      rewriter, loc,
      commonArgs[0], // ArgIndex 0: rsrc
      dst,           // ArgIndex 1: LDS base ptr
      byteWidth,     // ArgIndex 2: data byte size (immarg)
      commonArgs[1], // ArgIndex 3: voffset (per-lane VGPR)
      commonArgs[2], // ArgIndex 4: soffset (always 0 here)
      b.i32_val(0),  // ArgIndex 5: imm offset (immarg, always 0 here)
      rewriter.getI32IntegerAttr(aux), /*alias_scopes=*/nullptr,
      /*noalias_scopes=*/nullptr, /*tbaa=*/nullptr);
}

Value BufferEmitter::emitAtomicCAS(Type type, Value rsrcDesc, Value offset,
                                   Value casCmpVal, Value casStoreVal,
                                   Value pred, bool hasUsers) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  VectorType storeVecTy = cast<VectorType>(casStoreVal.getType());
  VectorType cmpVecTy = cast<VectorType>(casCmpVal.getType());
  Type bufferType = getBufferOpType(type, true);
  if (storeVecTy != bufferType)
    casStoreVal = b.bitcast(casStoreVal, bufferType);
  if (cmpVecTy != bufferType)
    casCmpVal = b.bitcast(casCmpVal, bufferType);
  // Note: rocdl.raw.ptr.buffer.atomic.cmpswap expects
  // val to be before cmp in the arg list. This is
  // the opposite of the order in tl.atomic_cmpxchg
  // and amdg.buffer_atomic_cas
  SmallVector<Value, 6> args{casStoreVal, casCmpVal};
  int32_t aux = 0;
  fillCommonArgsAtomics(type, rsrcDesc, offset, pred, hasUsers, aux, args);

  Value data = ROCDL::RawPtrBufferAtomicCmpSwap::create(
      rewriter, loc, bufferType, args[0], args[1], args[2], args[3], args[4],
      rewriter.getI32IntegerAttr(aux), /*alias_scopes=*/nullptr,
      /*noalias_scopes=*/nullptr, /*tbaa=*/nullptr);
  data = b.bitcast(data, type);
  return data;
}

Value BufferEmitter::emitAtomicRMW(RMWOp rmwType, Type type, Value rsrcDesc,
                                   Value offset, Value data, Value pred,
                                   bool hasUsers) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  VectorType vecTy = cast<VectorType>(data.getType());
  Type bufferType = getBufferOpType(type, true);
  if (vecTy != bufferType)
    data = b.bitcast(data, bufferType);

  SmallVector<Value, 6> args{data};
  int32_t aux = 0;
  fillCommonArgsAtomics(type, rsrcDesc, offset, pred, hasUsers, aux, args);

  // TODO:
  //   The ops in ROCDL (e.g., RawPtrBufferAtomicFaddOp) have no return value,
  //   but they lower to instrinsics that can return values. This causes the
  //   LLVM verifier to fail. When this is fixed, the ROCDL ops should be used
  //   here.
  auto rmwOpStr = stringifyRMWOp(rmwType).str();
  if (rmwType == RMWOp::XCHG) {
    // RMWOp::XCHG stringifies to "exch", but the AMDGPU buffer-atomic
    // intrinsic uses the "swap" suffix.
    rmwOpStr = "swap";
  } else if (rmwType == RMWOp::MAX || rmwType == RMWOp::MIN) {
    // RMWOp::MAX / MIN stringify to "max" / "min", which are not real AMDGPU
    // buffer-atomic intrinsic suffixes. The valid suffixes are
    // .{s,u,f}{max,min}. RMWOp::UMAX and RMWOp::UMIN already stringify to
    // "umax" / "umin" and need no override.
    StringRef prefix = isa<FloatType>(getElementTypeOrSelf(type)) ? "f" : "s";
    rmwOpStr = (prefix + rmwOpStr).str();
  }
  auto instrinsic = "llvm.amdgcn.raw.ptr.buffer.atomic." + rmwOpStr;
  SmallVector<Value, 6> intrinsicArgs = args;
  intrinsicArgs.push_back(b.i32_val(aux));
  auto bufferAtomicRMW = LLVM::createLLVMIntrinsicCallOp(
      rewriter, loc, instrinsic, bufferType, intrinsicArgs);

  return b.bitcast(bufferAtomicRMW.getResult(0), type);
}

void BufferEmitter::emitStore(Value rsrcDesc, Value offset, Value data,
                              Value pred, triton::CacheModifier cm) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  VectorType vecTy = cast<VectorType>(data.getType());
  Type bufferType = getBufferOpType(vecTy, false);
  if (vecTy != bufferType)
    data = b.bitcast(data, bufferType);
  SmallVector<Value, 6> args{data};
  int32_t aux = 0;
  fillCommonArgs(vecTy, rsrcDesc, offset, pred, cm, /*isBufferLoad=*/false,
                 args, aux);
  ROCDL::RawPtrBufferStoreOp::create(
      rewriter, loc, TypeRange{}, args[0], args[1], args[2], args[3],
      rewriter.getI32IntegerAttr(aux), /*alias_scopes=*/nullptr,
      /*noalias_scopes=*/nullptr, /*tbaa=*/nullptr);
}

Type BufferEmitter::getBufferOpType(Type type, bool atomicsOp) {
  int64_t vecSize = 1;
  Type elementType = type;
  if (auto vecType = dyn_cast<VectorType>(type)) {
    vecSize = vecType.getNumElements();
    elementType = vecType.getElementType();
  }

  const int valueElemNBits = std::max(8u, elementType.getIntOrFloatBitWidth());
  const size_t totalWidthBits = valueElemNBits * vecSize;

  Type bufferElementType = elementType;
  // We don't want to cast from bf16 if we are emitting buffer atomics
  if (elementType.isBF16() && !atomicsOp) {
    bufferElementType = rewriter.getI16Type();
  }

  // If we are dealing with a subword type (e.g., i8 or f16) but we
  // still need multiple words, then pack the subwords into 32bit integers
  // and update the vector length and the type
  // We never need to pack for buffer atomics because we ensure
  // 1) We can always emit a 32-bit / 64-bit atomics op
  // 2) For tensors of 16-bit values that the values are contiguous
  int64_t bufferVecSize = vecSize;
  if (valueElemNBits < 32 && !atomicsOp) {
    if (totalWidthBits > 32) {
      bufferElementType = rewriter.getI32Type();
      bufferVecSize = totalWidthBits / 32;
    } else {
      bufferElementType = rewriter.getIntegerType(totalWidthBits);
      bufferVecSize = 1;
    }
  }

  // This is the buffer type that the buffer operation will use. It
  // will be bitcast-able to the original type. So if the types
  // ended up different, we simply have to emit a `bitcastOp` to convert
  Type bufferType = type;
  if (bufferVecSize != vecSize || bufferElementType != elementType)
    bufferType = VectorType::get(bufferVecSize, bufferElementType);
  if (bufferVecSize == 1)
    bufferType = getElementTypeOrSelf(bufferType);

  return bufferType;
}

void BufferEmitter::fillCommonArgs(Type type, Value rsrcDesc,
                                   Value vOffsetElems, Value pred,
                                   triton::CacheModifier cm, bool isBufferLoad,
                                   SmallVector<Value> &args, int32_t &aux) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  // 1. Create the (masked) offset
  Type elementType = getElementTypeOrSelf(type);
  const int valueElemNBits = std::max(8u, elementType.getIntOrFloatBitWidth());
  const int elementByteWidth = valueElemNBits / 8;
  // Please note: the index passed is not in bytes, but in number of elements
  // In order to pass the index to the buffer operation, we need to convert in
  // bytes (i.e., we need to multiply by `elementByteWidth`)
  Value vOffsetOutOfBunds = b.int_val(
      32, static_cast<int>(std::numeric_limits<int>::max() + int64_t(1)));
  Value vOffsetBytes = b.mul(b.int_val(32, elementByteWidth), vOffsetElems);
  Value maskedOffsetBytes = b.select(pred, vOffsetBytes, vOffsetOutOfBunds);

  // 2. Set the sgprOffset to 0
  Value sgprOffset = b.int_val(32, 0);

  aux = getCtrlBitsForCacheModifierOnTarget(cm, isBufferLoad, targetInfo);

  // 4. Add the arguments
  args.push_back(rsrcDesc);
  args.push_back(maskedOffsetBytes);
  args.push_back(sgprOffset);
}

void BufferEmitter::fillCommonArgsAtomics(Type type, Value rsrcDesc,
                                          Value vOffsetElems, Value pred,
                                          bool hasUsers, int32_t &aux,
                                          SmallVector<Value> &args) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  // 1. Create the (masked) offset
  Type elementType = getElementTypeOrSelf(type);
  const int valueElemNBits = std::max(8u, elementType.getIntOrFloatBitWidth());
  const int elementByteWidth = valueElemNBits / 8;
  // Please note: the index passed is not in bytes, but in number of elements
  // In order to pass the index to the buffer operation, we need to convert in
  // bytes (i.e., we need to multiply by `elementByteWidth`)
  Value vOffsetOutOfBunds = b.int_val(
      32, static_cast<int>(std::numeric_limits<int>::max() + int64_t(1)));
  Value vOffsetBytes = b.mul(b.int_val(32, elementByteWidth), vOffsetElems);
  Value maskedOffsetBytes = b.select(pred, vOffsetBytes, vOffsetOutOfBunds);

  // 2. Set the sgprOffset to 0
  Value sgprOffset = b.int_val(32, 0);

  aux = targetInfo.getBufferAtomicCachePolicy(hasUsers);

  // 4. Add the arguments
  args.push_back(rsrcDesc);
  args.push_back(maskedOffsetBytes);
  args.push_back(sgprOffset);
}

} // namespace mlir::LLVM::AMD
