#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h"
#include "mlir/Analysis/TopologicalSortUtils.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "triton/Analysis/AxisInfo.h"
#include "triton/Dialect/Triton/IR/Utility.h"
#include "triton/Dialect/TritonGPU/IR/Dialect.h"
#include "triton/Dialect/TritonGPU/Transforms/Passes.h"
#include "triton/Dialect/TritonGPU/Transforms/Schedule.h"
#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
#include "triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h"
#include "triton/Tools/LayoutUtils.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include <queue>

#define DEBUG_TYPE "triton-loop-pipeline"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")

using namespace mlir;
namespace tt = mlir::triton;
namespace ttg = mlir::triton::gpu;
namespace ttng = mlir::triton::nvidia_gpu;

//===----------------------------------------------------------------------===//
// Hoisting Utilities
//===----------------------------------------------------------------------===//

bool triton::isPureScalarOp(Operation *op) {
  auto isScalar = [](Type type) { return type.isIntOrIndexOrFloat(); };
  return isPure(op) && llvm::all_of(op->getOperandTypes(), isScalar) &&
         llvm::all_of(op->getResultTypes(), isScalar);
}

bool triton::getDominatingValueSetOpsToHoist(
    DominanceInfo &domInfo, Operation *refOp, ArrayRef<Value> valueSet,
    llvm::SetVector<Operation *> &toHoist,
    function_ref<bool(Operation *)> canHoist,
    function_ref<bool(BlockArgument)> canUseArg) {
  // The set of operations below `refOp` that are being checked if they can be
  // hoisted. This set prevents checking operations twice but also if the
  // computation can be hoisted, this becomes the set of operations to hoist.
  llvm::SetVector<Operation *> visited;

  // Climb the use-def chain breadth-first so that operations can be hoisted in
  // the reverse visitation order.
  std::queue<Value> queue;
  for (Value value : valueSet)
    queue.push(value);

  while (!queue.empty()) {
    Value value = queue.front();
    queue.pop();

    // If the value properly dominates the outer loop, then it must be invariant
    // to it.
    if (domInfo.properlyDominates(value, refOp))
      continue;
    // If the value is a block argument, check if it can be used.
    if (auto arg = dyn_cast<BlockArgument>(value)) {
      if (!canUseArg(arg))
        return false;
      continue;
    }

    Operation *op = value.getDefiningOp();
    // Check if the op was already visited.
    if (visited.contains(op))
      continue;
    // If the defining op cannot be hoisted, then the value cannot be made loop
    // invariant.
    if (!canHoist(op))
      return false;
    visited.insert(op);
    // Recurse on the operands of the op.
    for (Value operand : op->getOperands())
      queue.push(operand);
  }

  // The operations in `visited` must be hoisted. Note that operations are not
  // added to `toHoist` unless all of `values` can be hoisted. This is to avoid
  // hoisting operations for loops that don't end up getting fused if one of
  // their bounds operands cannot be hoisted.
  toHoist.insert(visited.begin(), visited.end());

  return true;
}

void triton::hoistOpsBefore(Operation *refOp,
                            const llvm::SetVector<Operation *> &toHoist) {
  return hoistOpsBefore(refOp->getBlock(), refOp->getIterator(), toHoist);
}
void triton::hoistOpsBefore(Block *block, Block::iterator it,
                            const llvm::SetVector<Operation *> &toHoist) {
  for (Operation *op : topologicalSort(toHoist)) {
    op->moveBefore(block, it);
  }
}

//===----------------------------------------------------------------------===//
// Sinking Utilities
//===----------------------------------------------------------------------===//

Value triton::sinkValueRedefinition(RewriterBase &rewriter, Value in, Value out,
                                    Block *block) {
  OpBuilder::InsertionGuard guard(rewriter);
  for (; block != in.getParentBlock();
       block = block->getParentOp()->getBlock()) {
    Operation *op = block->getParentOp();
    rewriter.setInsertionPoint(op);

    // `in` is live into the loop body. `out` becomes the live-out if the
    // loop executes at least once.
    if (auto forOp = dyn_cast<scf::ForOp>(op)) {
      forOp = addIterArgsToLoop(rewriter, forOp, in);
      appendToForOpYield(forOp, out);
      out = forOp.getResults().back();
      continue;
    }

    // `in` is live into both branches. `out` becomes the live-out if the
    // particular branch is taken.
    if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
      scf::IfOp newIfOp =
          replaceIfOpWithNewSignature(rewriter, ifOp, out.getType());
      scf::YieldOp taken = newIfOp.thenYield();
      scf::YieldOp other = newIfOp.elseYield();
      if (block == newIfOp.elseBlock())
        std::swap(taken, other);
      taken->insertOperands(taken.getNumOperands(), out);
      other->insertOperands(other.getNumOperands(), in);
      out = newIfOp.getResults().back();
      rewriter.eraseOp(ifOp);
      continue;
    }

    // TODO: Handle `scf.while`, etc.
    llvm::report_fatal_error("FIXME: sinking into unhandled control flow op: " +
                             op->getName().getStringRef());
  }

  return out;
}

//===----------------------------------------------------------------------===//
// Loop Pipelining Utilities
//===----------------------------------------------------------------------===//

bool mlir::triton::loopHasDistGreaterThanOne(scf::ForOp forOp) {
  return llvm::any_of(forOp.getBody()->getTerminator()->getOperands(),
                      [](Value operand) {
                        Operation *def = operand.getDefiningOp();
                        return !def;
                      });
}

bool mlir::triton::isOuterLoop(scf::ForOp forOp) {
  return llvm::any_of(forOp.getBody()->getOperations(), [](Operation &op) {
    return isa<scf::ForOp, scf::WhileOp>(op);
  });
}

// Function to mask operations during scheduling.
Operation *mlir::triton::predicateOp(RewriterBase &rewriter, Operation *op,
                                     Value pred) {
  OpBuilder::InsertionGuard guard(rewriter);
  if (mlir::isMemoryEffectFree(op))
    return op;
  if (isConstantIntValue(pred, 1))
    return op;
  if (isa<LLVM::AssumeOp, ttng::FenceAsyncSharedOp>(op))
    return op;
  if (isa<ttg::AsyncCommitGroupOp, ttg::AsyncWaitOp>(op))
    return op;
  if (op->hasTrait<OpTrait::LocalLoadTrait>())
    return op;
  if (isa<ttg::LocalStoreOp>(op))
    return op;
  if (isa<ttng::TMEMAllocOp, ttng::TMEMLoadOp>(op))
    return op;
  if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
    rewriter.setInsertionPoint(op);
    Value cnd = getPredMask(rewriter, ifOp.getCondition().getType(),
                            ifOp.getCondition(), pred);
    ifOp.getConditionMutable().assign(cnd);
    return op;
  }
  if (auto predicatedOp = dyn_cast<tt::PredicatedOpInterface>(op)) {
    rewriter.setInsertionPoint(op);
    Value mask =
        getPredMask(rewriter, predicatedOp.getPredicateOperandTypeLike(),
                    predicatedOp.getPredicateOperand(), pred);
    predicatedOp.setPredicateOperand(mask);
    return op;
  }
  if (!op->isRegistered()) {
    // Skip ops from unregistered dialects to make writing lit tests easier.
    return op;
  }

  op->emitOpError("pipeliner doesn't know how to predicate this op.");
  llvm::report_fatal_error("Fatal pipeliner error");
  return op;
}

Operation *mlir::triton::wrapInMaskOp(RewriterBase &rewriter, Operation *op,
                                      Value pred) {
  auto mask =
      ttg::MaskOp::create(rewriter, op->getLoc(), op->getResultTypes(), pred);
  rewriter.createBlock(&mask->getRegion(0));
  rewriter.setInsertionPointToStart(&mask->getRegion(0).front());
  auto newOp = rewriter.clone(*op);
  ttg::MaskReturnOp::create(rewriter, op->getLoc(), newOp->getResults());
  op->replaceAllUsesWith(mask->getResults());
  rewriter.eraseOp(op);
  return mask;
}

void mlir::triton::resolveMaskOp(ModuleOp moduleOp) {
  IRRewriter rewriter(moduleOp);

  // Canonicalize the IR to simplify the arithmetic ops defining the mask
  auto arithDialect =
      moduleOp.getContext()->getLoadedDialect<arith::ArithDialect>();
  RewritePatternSet patterns(moduleOp.getContext());
  arithDialect->getCanonicalizationPatterns(patterns);
  if (mlir::applyPatternsGreedily(moduleOp, std::move(patterns)).failed())
    return llvm::report_fatal_error("Failed to canonicalize the IR");

  SmallVector<ttg::MaskOp> maskOps;
  moduleOp->walk([&](ttg::MaskOp maskOp) { maskOps.push_back(maskOp); });
  for (auto maskOp : maskOps) {
    rewriter.setInsertionPoint(maskOp);
    while (&maskOp.getBody()->front() != maskOp.getBody()->getTerminator()) {
      Operation *op = &maskOp.getBody()->front();
      rewriter.moveOpBefore(op, maskOp);
      op = triton::predicateOp(rewriter, op, maskOp.getPred());
    }
    maskOp->replaceAllUsesWith(
        maskOp.getBody()->getTerminator()->getOperands());
    maskOp->erase();
  }
}

// Return true if the given loop has the attribute
// `tt.disallow_acc_multi_buffer` set to true.
bool mlir::triton::getDisallowAccMultiBuffer(LoopLikeOpInterface loop) {
  return loop->hasAttr(mlir::triton::kDisallowAccMultiBufferAttrName);
}

std::pair<OpResult, int64_t>
mlir::triton::getDefinitionAndDistance(scf::ForOp forOp, Value value) {
  int64_t distance = 0;
  DenseSet<Value> seen;
  while (auto arg = dyn_cast<BlockArgument>(value)) {
    // Ignore implicit captures.
    if (arg.getOwner() != forOp.getBody())
      return {nullptr, 0};
    // Ignore induction variable.
    if (arg.getArgNumber() == 0)
      return {nullptr, 0};
    ++distance;
    value = forOp.getYieldedValues()[arg.getArgNumber() - 1];
    if (!seen.insert(value).second)
      return {nullptr, 0};
  }
  return {cast<OpResult>(value), distance};
}

std::pair<Operation *, int64_t>
mlir::triton::getDefiningOpAndDistance(scf::ForOp forOp, Value value) {
  auto [definition, distance] = getDefinitionAndDistance(forOp, value);
  return {definition ? definition.getDefiningOp() : nullptr, distance};
}

int mlir::triton::getCopyVecBytes(RankedTensorType registerTy,
                                  ttg::SharedEncodingTrait sharedEnc) {
  auto shape = registerTy.getShape();
  auto regLayout = triton::gpu::toLinearLayout(shape, registerTy.getEncoding());
  // FIXME: Here we should pass a MemDescType instead of a SharedEncodingTrait!!
  // This is currently broken for memdesc_subslice!
  auto sharedLayout = triton::gpu::toLinearLayout(shape, sharedEnc);
  auto regToSharedLayout = regLayout.invertAndCompose(sharedLayout);
  const int vecElems = regToSharedLayout.getNumConsecutiveInOut();
  return vecElems * registerTy.getElementTypeBitWidth() / 8;
}

bool mlir::triton::canBeConvertedToAsyncLoad(
    tt::LoadOp loadOp, tt::ModuleAxisInfoAnalysis &axisInfoAnalysis) {
  auto ptr = loadOp.getPtr();
  unsigned vec = axisInfoAnalysis.getContiguity(ptr);
  if (auto mask = loadOp.getMask())
    vec = std::min<unsigned>(vec, axisInfoAnalysis.getMaskAlignment(mask));

  auto tensorTy = dyn_cast<RankedTensorType>(ptr.getType());
  unsigned width = 0;
  if (tensorTy) {
    auto ty = cast<tt::PointerType>(tensorTy.getElementType()).getPointeeType();
    width = vec * ty.getIntOrFloatBitWidth();
  } else {
    width = cast<tt::PointerType>(ptr.getType())
                .getPointeeType()
                .getIntOrFloatBitWidth();
  }

  // We do not pipeline all loads for the following reasons:
  // 1. On nvidia GPUs, cp.async's cp-size can only be 4, 8, or 16.
  // 2. It's likely that pipling small loads won't offer much performance
  //    improvement and may even hurt performance by increasing register
  //    pressure.
  LDBG("Load " << *loadOp << " has width " << width);
  return width >= 32;
}

void mlir::triton::serializeLatencies(ModuleOp module,
                                      DenseMap<Operation *, int> &opLatency) {
  auto helper = TritonDialect::getLoaded(module)->getLatencyAttrHelper();
  auto builder = Builder(module);
  for (auto &[op, latency] : opLatency) {
    helper.setAttr(op, builder.getI32IntegerAttr(latency));
  }
}

void mlir::triton::serializeSelfLatencies(
    ModuleOp module, DenseMap<Operation *, int> &opSelfLatency) {
  auto helper = TritonDialect::getLoaded(module)->getSelfLatencyAttrHelper();
  auto builder = Builder(module);
  for (auto &[op, latency] : opSelfLatency) {
    helper.setAttr(op, builder.getI32IntegerAttr(latency));
  }
}

DenseMap<Operation *, int> mlir::triton::deserializeLatencies(Operation *op) {
  DenseMap<Operation *, int> opLatency;
  auto latencyHelper = TritonDialect::getLoaded(op)->getLatencyAttrHelper();
  op->walk([&](Operation *op) {
    if (auto attr = latencyHelper.getAttr(op)) {
      opLatency[op] = attr.getInt();
      latencyHelper.removeAttr(op);
    }
  });
  return opLatency;
}

Value mlir::triton::createScalarAlloc(ImplicitLocOpBuilder &rewriter, Type type,
                                      unsigned numBuffers) {
  MLIRContext *ctx = rewriter.getContext();
  unsigned numCTAs = triton::gpu::TritonGPUDialect::getNumCTAs(
      rewriter.getBlock()->getParentOp()->getParentOfType<ModuleOp>());
  Attribute sharedMemorySpace =
      ttg::SharedMemorySpaceAttr::get(rewriter.getContext());
  auto barrierCGALayout = ttg::CGAEncodingAttr::get1DLayout(ctx, numCTAs);
  auto barrierEncoding =
      ttg::SwizzledSharedEncodingAttr::get(ctx, 1, 1, 1, {0}, barrierCGALayout);
  ttg::MemDescType memDescType = ttg::MemDescType::get(
      {numBuffers, numCTAs}, type, barrierEncoding, sharedMemorySpace,
      /*mutableMemory=*/true);
  return ttg::LocalAllocOp::create(rewriter, memDescType, Value());
}

// Create an allocation and init the mbarriers.
Value mlir::triton::createBarrierAlloc(Operation *op, int numBarriers,
                                       int arriveCount) {
  ImplicitLocOpBuilder rewriter(op->getLoc(), op);

  Value barrierAlloc =
      createScalarAlloc(rewriter, rewriter.getI64Type(), numBarriers);
  for (unsigned i = 0; i < numBarriers; i++) {
    Value barrierView = createSingleBufferView(rewriter, barrierAlloc, i);
    ttng::InitBarrierOp::create(rewriter, barrierView, arriveCount);
  }
  // Invalidate and deallocate the barriers.
  rewriter.setInsertionPointAfter(op);
  for (unsigned i = 0; i < numBarriers; i++) {
    Value barrierView = createSingleBufferView(rewriter, barrierAlloc, i);
    ttng::InvalBarrierOp::create(rewriter, barrierView);
  }
  ttg::LocalDeallocOp::create(rewriter, barrierAlloc);
  return barrierAlloc;
}

Value mlir::triton::createAlloc(Operation *insertBefore, RankedTensorType ty,
                                Location loc,
                                gpu::SharedEncodingTrait sharedEnc,
                                unsigned distance) {
  OpBuilder builder(insertBefore);
  Attribute sharedMemorySpace =
      ttg::SharedMemorySpaceAttr::get(insertBefore->getContext());
  SmallVector<int64_t> bufferShape(ty.getShape().begin(), ty.getShape().end());
  bufferShape.insert(bufferShape.begin(), distance);
  Type memdescType = ttg::MemDescType::get(bufferShape, ty.getElementType(),
                                           sharedEnc, sharedMemorySpace,
                                           /*mutableMemory=*/true);
  Value alloc = ttg::LocalAllocOp::create(builder, loc, memdescType);

  builder.setInsertionPointAfter(insertBefore);
  ttg::LocalDeallocOp::create(builder, insertBefore->getLoc(), alloc);
  return alloc;
}

bool mlir::triton::canPipelineTMALoad(Operation *op) {
  auto tensorTy = cast<RankedTensorType>(op->getResultTypes()[0]);
  auto sharedEncoding = getSharedEncoding(op);
  int64_t stageSizeInBits =
      ttg::getAllocationElems(sharedEncoding, tensorTy.getShape()) *
      tensorTy.getElementTypeBitWidth();
  return stageSizeInBits % (ttng::TMA_ALIGN * 8) == 0;
}

bool mlir::triton::canBeAsyncLoad(Operation *op) {
  if (mlir::triton::isTMALoad(op)) {
    return canPipelineTMALoad(op);
  }
  assert(isa<tt::LoadOp>(op));
  ttg::SharedEncodingTrait sharedEncoding = mlir::triton::getSharedEncoding(op);
  // Do not create async loads for small loads (cp.async requires at least 4
  // bytes)
  int copyVecBytes = mlir::triton::getCopyVecBytes(
      cast<RankedTensorType>(op->getResultTypes()[0]), sharedEncoding);
  if (copyVecBytes >= 4) {
    return true;
  }
  return false;
}

void mlir::triton::combineRedundantWaitOps(
    llvm::SmallSetVector<Operation *, 8> &waitOps,
    llvm::function_ref<bool(Operation *)> isCounterBarrier,
    llvm::function_ref<Operation *(OpBuilder &, Location, ValueRange, unsigned)>
        createWait) {
  llvm::MapVector<Operation *, Operation *> toDelete;
  for (Operation *waitOp : waitOps) {
    if (toDelete.count(waitOp))
      continue;
    StringRef waitName = waitOp->getName().getStringRef();
    auto getNum = [](Operation *op) {
      return static_cast<unsigned>(
          op->getAttrOfType<IntegerAttr>("num").getInt());
    };
    SmallVector<Operation *> waitGroup = {waitOp};
    SmallVector<Value> depTokens(waitOp->getOperands().begin(),
                                 waitOp->getOperands().end());
    unsigned minWaitNumber = getNum(waitOp);
    Operation *next = waitOp->getNextNode();
    // Stop at the end of the block, at any branching op (forOp, ifOp, whileOp),
    // or at any caller-declared counter barrier.
    while (next && !isa<RegionBranchOpInterface>(next) &&
           !isCounterBarrier(next)) {
      if (next->getName().getStringRef() == waitName) {
        waitGroup.push_back(next);
        minWaitNumber = std::min(minWaitNumber, getNum(next));
        depTokens.append(next->getOperands().begin(),
                         next->getOperands().end());
      }
      next = next->getNextNode();
    }
    if (waitGroup.size() == 1)
      continue;
    OpBuilder builder(waitGroup.front());
    Operation *newWaitOp =
        createWait(builder, waitOp->getLoc(), depTokens, minWaitNumber);
    for (Operation *wo : waitGroup) {
      toDelete[wo] = newWaitOp;
    }
  }
  for (auto entry : toDelete) {
    entry.first->replaceAllUsesWith(entry.second);
    entry.first->erase();
  }
}

ttg::MemDescType mlir::triton::getBufferViewType(ttg::MemDescType allocTy,
                                                 bool mutableMemory) {
  return ttg::MemDescType::get(allocTy.getShape().drop_front(),
                               allocTy.getElementType(), allocTy.getEncoding(),
                               allocTy.getMemorySpace(), mutableMemory,
                               /*allocShape=*/allocTy.getAllocShape());
}

ttg::MemDescType
mlir::triton::getMultiBufferedType(ttg::MemDescType memDescType,
                                   int32_t depth) {
  auto shape = memDescType.getShape();
  SmallVector<int64_t> bufferShape(shape.begin(), shape.end());
  bufferShape.insert(bufferShape.begin(), depth);
  return ttg::MemDescType::get(
      bufferShape, memDescType.getElementType(), memDescType.getEncoding(),
      memDescType.getMemorySpace(), /*mutableMemory*/ true);
}

ttg::SharedEncodingTrait mlir::triton::getSharedEncoding(RankedTensorType ty) {
  auto cgaLayout = ttg::getCGALayout(ty.getEncoding());
  auto order = ttg::getOrder(ty);
  // Use generic layout. This won't be optimal for 2D tensors.
  return ttg::SwizzledSharedEncodingAttr::get(ty.getContext(), 1, 1, 1, order,
                                              cgaLayout);
}

ttg::SharedEncodingTrait mlir::triton::getSharedEncoding(Operation *op) {
  // Try to use local alloc encoding if possible.
  ttg::SharedEncodingTrait localAllocEnc;
  if (llvm::any_of(op->getUsers(), [&](Operation *user) {
        return isa<ttg::LocalAllocOp>(user);
      })) {
    for (auto user : op->getUsers()) {
      auto localAlloc = dyn_cast<ttg::LocalAllocOp>(user);
      if (!localAlloc)
        continue;
      auto enc = mlir::cast<ttg::SharedEncodingTrait>(
          localAlloc.getType().getEncoding());
      if (!localAllocEnc) {
        localAllocEnc = enc;
      }
      if (enc != localAllocEnc) {
        // Some users have different encoding than others.
        // Use one of the encodings, and warn about the performance issue.
        op->emitRemark()
            << "Pipelining load with different use encodings. This will lead "
               "to layout conversions and performance degradation.";
        continue;
      }
    }
  }

  auto ty = cast<RankedTensorType>(op->getResultTypes()[0]);
  auto cgaLayout = ttg::getCGALayout(ty.getEncoding());
  auto order = ttg::getOrder(ty);
  if (auto load = dyn_cast<tt::DescriptorLoadLikeOpInterface>(op)) {
    // TMA encoding is set on the descriptor type
    return ttng::getEncodingFromDescriptor(op, ty, load.getDesc());
  }

  if (localAllocEnc)
    return localAllocEnc;

  // Try to use dot encoding if possible.
  bool incompatible = false;
  localAllocEnc =
      getSharedEncIfAllUsersAreDotEnc(op->getResult(0), incompatible)
          .value_or(nullptr);

  if (localAllocEnc)
    return localAllocEnc;

  // Use generic layout. This won't be optimal for 2D tensors.
  return ttg::SwizzledSharedEncodingAttr::get(ty.getContext(), 1, 1, 1, order,
                                              cgaLayout);
}

int mlir::triton::getNumStagesOrDefault(scf::ForOp forOp,
                                        int defaultNumStages) {
  // Use the attribute attached to the loop if it exists otherwise use the
  // global control.
  auto helper = TritonDialect::getLoaded(forOp)->getNumStagesAttrHelper();
  if (auto attr = helper.getAttr(forOp))
    return attr.getInt();
  return defaultNumStages;
}

TypedValue<ttg::MemDescType>
triton::createSingleBufferView(OpBuilder &builder, Value alloc, Value idx) {
  assert(isa<ttg::MemDescType>(alloc.getType()) && "Expected MemDescType");
  auto allocDescType = cast<ttg::MemDescType>(alloc.getType());
  SmallVector<int64_t> shape;
  assert(allocDescType.getShape().size() > 1 &&
         "Expected multi-dimensional memdesc (e.g., Nx...) for subview");
  shape.insert(shape.end(), allocDescType.getShape().begin() + 1,
               allocDescType.getShape().end());
  auto viewDescType = ttg::MemDescType::get(
      shape, allocDescType.getElementType(), allocDescType.getEncoding(),
      allocDescType.getMemorySpace(), allocDescType.getMutableMemory());
  return ttg::MemDescIndexOp::create(builder, alloc.getLoc(), viewDescType,
                                     alloc, idx);
}

TypedValue<ttg::MemDescType>
triton::createSingleBufferView(OpBuilder &builder, Value alloc, int idx) {
  Value idxVal = arith::ConstantIntOp::create(builder, alloc.getLoc(), idx, 32);
  return createSingleBufferView(builder, alloc, idxVal);
}

Value triton::createIncrementModulo(OpBuilder &builder, Location loc,
                                    Value counter, Value modulus, Value zero,
                                    Value one, Value *outWrapCond) {
  Value addOne = arith::AddIOp::create(builder, loc, counter, one);
  Value outOfRangeCond = arith::CmpIOp::create(
      builder, loc, arith::CmpIPredicate::sge, addOne, modulus);
  if (outWrapCond)
    *outWrapCond = outOfRangeCond;
  return arith::SelectOp::create(builder, loc, outOfRangeCond, zero, addOne);
}

/////////////////////////////
// LOWER TMA DESCRIPTORS
/////////////////////////////

static void
allocTMABuffers(LoopLikeOpInterface loop,
                llvm::MapVector<Operation *, Value> &tmaBufferMapping,
                int maxStage) {
  IRRewriter rewriter(loop);
  Block *body = &loop.getLoopRegions().back()->front();

  // Create a multi-buffered allocation for each MakeTensorDescOp call in the
  // loop
  body->walk([&](tt::MakeTensorDescOp op) {
    // TODO peter: walk to loop yield to find the init value if this is a
    // loop-carried value. That would save us from allocating another buffer
    // just for the init value
    auto loc = op.getLoc();
    Value alloc = triton::gpu::GlobalScratchAllocOp::create(
        rewriter, loc, triton::getPointerType(rewriter.getI8Type()),
        maxStage * ttng::TMA_SIZE_BYTES, ttng::TMA_ALIGN, UnitAttr());
    tmaBufferMapping[op.getOperation()] = alloc;
  });
}

static Value subviewTMADescriptor(OpBuilder &builder, Location loc, Value alloc,
                                  Value counter) {
  Value tmaSizeVal =
      arith::ConstantIntOp::create(builder, loc, ttng::TMA_SIZE_BYTES, 32);
  Value offset = arith::MulIOp::create(builder, loc, tmaSizeVal, counter);
  return triton::AddPtrOp::create(builder, loc, alloc.getType(), alloc, offset);
}

static LogicalResult rewriteTMABufferUpdates(
    LoopLikeOpInterface loop,
    const llvm::MapVector<Operation *, Value> &tmaBufferMapping,
    ArrayRef<BlockArgument> tmaCounters, int numBuffers, Value one, Value zero,
    MutableArrayRef<OpOperand> tmaCounterYields,
    triton::CoarseSchedule &schedule) {
  assert(tmaBufferMapping.size() == tmaCounters.size());

  auto auxBuilder = mlir::OpBuilder(loop);
  Value numBuffersVal =
      arith::ConstantIntOp::create(auxBuilder, loop.getLoc(), numBuffers, 32);

  for (auto [iOp, pair] : llvm::enumerate(tmaBufferMapping)) {
    auto &[op, alloc] = pair;

    // Rewriter MakeTensorDescOp as writing a TMA descriptor
    auto makeDescOp = cast<tt::MakeTensorDescOp>(op);

    triton::OpBuilderForStage builder(makeDescOp.getLoc(), makeDescOp,
                                      schedule);

    BlockArgument counter = tmaCounters[iOp];
    Value nextBuf =
        subviewTMADescriptor(builder, builder.getLoc(), alloc, counter);
    if (failed(ttng::createTMADesc(nextBuf, makeDescOp, builder))) {
      return failure();
    }
    ttng::TensormapFenceproxyAcquireOp::create(builder, nextBuf);
    Value nextDesc = ttng::ReinterpretTensorDescOp::create(
        builder, makeDescOp.getType(), nextBuf);

    makeDescOp.getResult().replaceAllUsesWith(nextDesc);

    // Increment the buffer index counter
    Value nextCounter = createIncrementModulo(
        builder, builder.getLoc(), counter, numBuffersVal, zero, one);

    // If we are in a (potentially nested) if region, propagate the counter
    // up to the main loop body scope
    IRRewriter rewriter(loop);
    nextCounter = triton::sinkValueRedefinition(rewriter, counter, nextCounter,
                                                op->getBlock());

    // Finally, rewrite the loop level yield
    tmaCounterYields[iOp].set(nextCounter);
    makeDescOp.erase();
  }
  return success();
}

LoopLikeOpInterface triton::lowerTMADescriptors(LoopLikeOpInterface loop,
                                                CoarseSchedule &schedule) {
  llvm::MapVector<Operation *, Value> tmaBufferMapping;
  int maxStage = schedule.getNumStages() - 1;
  Block *body = &loop.getLoopRegions().back()->front();
  for (Operation &op : body->without_terminator()) {
    if (isa<ttng::WarpGroupDotOp>(op)) {
      // Hopper only: Add one more buffer slice if there is a WarpGroupDotOp,
      // as if it will be pipelined, we will effectively make the pipeline
      // one stage longer.
      maxStage += 1;
      break;
    }
  }
  allocTMABuffers(loop, tmaBufferMapping, maxStage);
  if (tmaBufferMapping.empty())
    return loop;

  IRRewriter builder(loop);
  Location loc = loop.getLoc();
  Value zero = arith::ConstantIntOp::create(builder, loc, 0, 32);
  Value one = arith::ConstantIntOp::create(builder, loc, 1, 32);
  SmallVector<Value> newOperands;
  // Create one counter per TMA buffer. This allows the descriptors to be
  // updated independently without needing to write duplicate of existing tma
  // descriptors.
  for (int i = 0; i < tmaBufferMapping.size(); ++i) {
    newOperands.push_back(zero);
  }

  if (auto forOp = dyn_cast<scf::ForOp>(loop.getOperation())) {
    forOp = addIterArgsToLoop(builder, forOp, newOperands);
    auto yield = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
    for (Value operand : newOperands)
      yield.getResultsMutable().append(operand);
    loop = forOp;
  } else {
    auto whileOp = cast<scf::WhileOp>(loop.getOperation());
    auto numBeforeArgs = whileOp.getBeforeArguments().size();
    auto numAfterArgs = whileOp.getAfterArguments().size();
    SmallVector<Type> resultTypes(newOperands.size(), builder.getI32Type());
    auto oldWhile = whileOp;
    whileOp = replaceWhileOpWithNewSignature(builder, whileOp, newOperands,
                                             resultTypes);
    oldWhile.erase();
    for (int i = 0; i < newOperands.size(); ++i) {
      whileOp.getConditionOp().getArgsMutable().append(
          whileOp.getBeforeArguments()[numBeforeArgs + i]);
      whileOp.getYieldOp().getResultsMutable().append(
          whileOp.getAfterArguments()[numAfterArgs + i]);
    }
    loop = whileOp;
  }

  body = &loop.getLoopRegions().back()->front();
  auto tmaCounters = body->getArguments().take_back(tmaBufferMapping.size());
  auto tmaCounterYields =
      loop.getYieldedValuesMutable()->take_back(tmaBufferMapping.size());
  if (failed(rewriteTMABufferUpdates(loop, tmaBufferMapping, tmaCounters,
                                     maxStage, one, zero, tmaCounterYields,
                                     schedule))) {
    llvm_unreachable("Failed to rewrite TMA ops");
  }
  return loop;
}

DenseSet<Operation *>
triton::getTopLevelUsersInLoop(Operation *op, scf::ForOp forOp,
                               std::function<bool(Operation *)> filter) {
  DenseSet<Operation *> topLevelUsers;
  SmallVector<OpOperand *> q;
  for (auto &use : op->getUses())
    q.push_back(&use);
  while (!q.empty()) {
    auto use = q.pop_back_val();
    auto yieldOp = dyn_cast<scf::YieldOp>(use->getOwner());
    if (yieldOp && yieldOp->getParentOp() == forOp) {
      for (auto &use :
           forOp.getRegionIterArgs()[use->getOperandNumber()].getUses())
        q.push_back(&use);
      continue;
    }
    // Don't count view operations as uses. Follow them through to their
    // users.
    if (use->getOwner()->hasTrait<OpTrait::MemDescViewTrait>()) {
      for (auto &use : use->getOwner()->getUses())
        q.push_back(&use);
      continue;
    }
    if (filter && !filter(use->getOwner()))
      continue;
    Operation *topLevelUser =
        forOp.getBody()->findAncestorOpInBlock(*use->getOwner());
    topLevelUsers.insert(topLevelUser);
  }
  return topLevelUsers;
}

// Helper function that finds an operation based on a comparison predicate
static Operation *getUseOfPipelinedOp(
    ArrayRef<Operation *> ops, scf::ForOp forOp,
    triton::CoarseSchedule &schedule,
    std::function<bool(Operation *)> filterUse,
    std::function<bool(Operation *, Operation *)> shouldPrefer) {
  DenseSet<Operation *> topLevelUsers;
  Operation *selectedUser = nullptr;
  for (Operation *op : ops) {
    auto users = triton::getTopLevelUsersInLoop(op, forOp, filterUse);
    topLevelUsers.insert(users.begin(), users.end());
  }
  for (Operation *topLevelUser : topLevelUsers) {
    assert(schedule.count(topLevelUser) && "op user not found in the schedule");
    if (!selectedUser || shouldPrefer(topLevelUser, selectedUser)) {
      selectedUser = topLevelUser;
    }
  }
  return selectedUser;
}

Operation *
triton::getFirstUseOfPipelinedOp(ArrayRef<Operation *> ops, scf::ForOp forOp,
                                 triton::CoarseSchedule &schedule,
                                 std::function<bool(Operation *)> filterUse) {
  return getUseOfPipelinedOp(
      ops, forOp, schedule, filterUse,
      [&](Operation *candidate, Operation *current) {
        auto [candidateStage, candidateCluster] = schedule[candidate];
        auto [currentStage, currentCluster] = schedule[current];

        return candidateStage < currentStage ||
               (candidateStage == currentStage &&
                schedule.clusters.isBefore(candidateCluster, currentCluster)) ||
               (candidateStage == currentStage &&
                candidateCluster == currentCluster &&
                candidate->isBeforeInBlock(current));
      });
}

Operation *
triton::getLastUseOfPipelinedOp(ArrayRef<Operation *> ops, scf::ForOp forOp,
                                triton::CoarseSchedule &schedule,
                                std::function<bool(Operation *)> filterUse) {
  return getUseOfPipelinedOp(
      ops, forOp, schedule, filterUse,
      [&](Operation *candidate, Operation *current) {
        auto [candidateStage, candidateCluster] = schedule[candidate];
        auto [currentStage, currentCluster] = schedule[current];

        return candidateStage > currentStage ||
               (candidateStage == currentStage &&
                schedule.clusters.isBefore(currentCluster, candidateCluster)) ||
               (candidateStage == currentStage &&
                candidateCluster == currentCluster &&
                current->isBeforeInBlock(candidate));
      });
}

void triton::removePipeliningAttributes(ModuleOp moduleOp) {
  moduleOp->walk([&](Operation *op) {
    op->removeAttr(mlir::triton::kLoopStageAttrName);
    op->removeAttr(mlir::triton::kLoopClusterAttrName);
    op->removeAttr(mlir::triton::kScheduledMaxStageAttrName);
  });
}
