#include "Analysis/AMDGPUAllocation.h"
#include "Dialect/TritonAMDGPU/IR/Dialect.h"
#include "triton/Analysis/Allocation.h"
#include "triton/Analysis/Utility.h"
#include "triton/Dialect/Triton/IR/Dialect.h"
#include "triton/Dialect/Triton/IR/Utility.h"
#include "triton/Dialect/TritonGPU/IR/Dialect.h"
#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h"
#include "triton/Dialect/TritonInstrument/IR/ConSanConstants.h"

#include "third_party/amd/include/Dialect/TritonAMDGPU/Utility/CommonUtils.h"

namespace mlir::triton::AMD {

unsigned getConvertLayoutScratchInBytes(gpu::ConvertLayoutOp op,
                                        TargetInfoBase &targetInfo) {
  if (!cvtNeedsSharedMemory(op))
    return 0;
  auto srcTy = op.getSrc().getType();
  auto dstTy = op.getType();
  int numBanks = targetInfo.getSharedMemoryBanks();
  auto srcLayout = gpu::toLinearLayout(srcTy);
  auto dstLayout = gpu::toLinearLayout(dstTy);
  auto bitwidth = getBitwidth(srcTy);
  auto vecBitwidth =
      triton::gpu::getVecBitwidthLdSt(srcLayout, dstLayout, bitwidth);
  auto [dstTile, srcTile] = targetInfo.getSharedLdStTiles(vecBitwidth);
  unsigned elems = getNumScratchElemsSwizzledCvt(srcLayout, dstLayout, bitwidth,
                                                 numBanks, srcTile, dstTile);
  return elems * bitwidth / 8;
}

static unsigned getBufferAtomicScratchSizeInBytes(Operation *op) {
  Value result = op->getResult(0);
  if (result.use_empty())
    return 0;
  auto tensorTy = dyn_cast<RankedTensorType>(result.getType());
  if (!tensorTy)
    return 0;
  auto freeVariableMasks = gpu::toLinearLayout(tensorTy).getFreeVariableMasks();
  bool hasBroadcast = llvm::any_of(freeVariableMasks,
                                   [](auto mask) { return mask.second != 0; });
  if (!hasBroadcast)
    return 0;
  auto smemShape = convertType<unsigned>(gpu::getShapePerCTA(tensorTy));
  auto elems = getNumScratchElements(smemShape);
  if (elems == 0)
    return 0;
  auto elemTy = tensorTy.getElementType();
  return elems * std::max<int>(8, elemTy.getIntOrFloatBitWidth()) / 8;
}

unsigned AMDAllocationAnalysisScratchSizeFn(Operation *op,
                                            TargetInfoBase &targetInfo) {

  if (auto reduceOp = dyn_cast<ReduceOp>(op)) {
    ReduceOpHelper::GetNumScratchElemsFn AMDGetNumScratchElemsFn =
        [&targetInfo](const triton::LinearLayout &src,
                      const triton::LinearLayout &dst, unsigned bitwidth) {
          int numBanks = targetInfo.getSharedMemoryBanks();
          auto vecBitwidth =
              triton::gpu::getVecBitwidthLdSt(src, dst, bitwidth);
          auto [dstTile, srcTile] = targetInfo.getSharedLdStTiles(vecBitwidth);
          return getNumScratchElemsSwizzledCvt(src, dst, bitwidth, numBanks,
                                               srcTile, dstTile);
        };
    return ReduceOpHelper(reduceOp).getScratchSizeInBytes(
        AMDGetNumScratchElemsFn);
  }

  if (auto cvtLayout = dyn_cast<mlir::triton::gpu::ConvertLayoutOp>(op)) {
    return getConvertLayoutScratchInBytes(cvtLayout, targetInfo);
  }

  if (auto ws = dyn_cast<mlir::triton::gpu::WarpSpecializeOp>(op)) {
    uint64_t captureSize = 0;
    // Tightly pack the captures in memory.
    for (Type type : ws.getPartitionOp().getOperandTypes()) {
      if (auto descType = dyn_cast<TensorDescType>(type))
        captureSize +=
            mlir::triton::amdgpu::getTensorDescNumDwords(descType) * 4;
      else
        captureSize += mlir::triton::gpu::getSharedMemorySize(type);
    }
    // ConSan adds captures after allocation; reserve space pre-computed by the
    // common TritonInstrumentPrepareConSanCaptures pass.
    if (auto extra = ws->getAttrOfType<IntegerAttr>(
            mlir::triton::instrument::kConSanExtraCaptureBytesAttr))
      captureSize += extra.getInt();
    return captureSize;
  }

  if (isa<amdgpu::BufferAtomicCASOp, amdgpu::BufferAtomicRMWOp>(op))
    return getBufferAtomicScratchSizeInBytes(op);

  return defaultAllocationAnalysisScratchSizeFn(op);
}

} // namespace mlir::triton::AMD
