#include "Utilities.h"
#include "triton/Dialect/TritonGPU/Transforms/Partition.h"
#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h"
#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"

using namespace mlir::triton;
using namespace mlir::triton::gpu;
using namespace mlir::triton::nvidia_gpu;

namespace mlir::triton::nvws {

Operation *createAlloc(OpBuilder &builder, Location loc,
                       MemDescType memDescType, Value src,
                       std::optional<int32_t> alignment) {
  if (isa<SharedMemorySpaceAttr>(memDescType.getMemorySpace())) {
    if (alignment)
      return LocalAllocOp::create(builder, loc, memDescType, src, *alignment);
    return LocalAllocOp::create(builder, loc, memDescType, src);
  } else {
    assert(isa<TensorMemorySpaceAttr>(memDescType.getMemorySpace()));
    return TMEMAllocOp::create(builder, loc, memDescType, src);
  }
}

ArefCreateOp createArefCreateOp(OpBuilder &builder, ArrayRef<Type> arefTypes,
                                ValueRange allocOps, Location loc) {
  auto ctx = builder.getContext();
  auto arefTy = ArefType::get(ctx, TypeArrayAttr::get(ctx, arefTypes));
  return ArefCreateOp::create(builder, loc, arefTy, allocOps);
}

int getArefDepth(MemDescType bufTy) {
  auto shape = bufTy.getShape();
  return isa<nvidia_gpu::TensorMemoryScalesEncodingAttr>(bufTy.getEncoding())
             ? 1
             : shape[0];
}

MemDescType getArefViewBufferType(MemDescType bufTy) {
  auto isScalesEnc =
      isa<nvidia_gpu::TensorMemoryScalesEncodingAttr>(bufTy.getEncoding());
  auto shape = bufTy.getShape();
  return gpu::MemDescType::get(isScalesEnc ? shape : shape.drop_front(),
                               bufTy.getElementType(), bufTy.getEncoding(),
                               bufTy.getMemorySpace(),
                               /*mutableMemory*/ true,
                               /*allocShape=*/bufTy.getAllocShape());
}

MemDescType getArefMultiBufferedType(MemDescType bufTy, int depth) {
  auto shape = bufTy.getShape();
  SmallVector<int64_t> bufferShape(shape.begin(), shape.end());
  if (!isa<nvidia_gpu::TensorMemoryScalesEncodingAttr>(bufTy.getEncoding()))
    bufferShape.insert(bufferShape.begin(), depth);
  return gpu::MemDescType::get(bufferShape, bufTy.getElementType(),
                               bufTy.getEncoding(), bufTy.getMemorySpace(),
                               /*mutableMemory*/ true);
}

LoopLikeOpInterface getOuterWSLoop(LoopLikeOpInterface loop) {
  auto wsLoop = loop;
  while (wsLoop && !wsLoop->hasAttr(triton::kWarpSpecializeAttrName)) {
    wsLoop = wsLoop->getParentOfType<LoopLikeOpInterface>();
  }
  return wsLoop;
}

} // namespace mlir::triton::nvws
