#include "Dialect/TritonAMDGPU/IR/TargetFeatures.h"
#include "TritonAMDGPUTransforms/Passes.h"
#include "Utility.h"
#include "amd/lib/TritonAMDGPUToLLVM/AsyncUtility.h"
#include "amd/lib/TritonAMDGPUToLLVM/TargetInfo.h"
#include "amd/lib/TritonAMDGPUTransforms/PipelineUtility.h"
#include "triton/Dialect/Triton/IR/Utility.h"
#include "triton/Dialect/TritonGPU/IR/CGAEncodingAttr.h"
#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h"
#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
#include "llvm/Support/Debug.h"
#include <variant>

#undef DEBUG_TYPE
#define DEBUG_TYPE "tritonamdgpu-pipeline-lower-loops"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")

namespace tt = mlir::triton;
namespace ttg = mlir::triton::gpu;
using mlir::triton::amdgpu::TargetFeatures;

//===----------------------------------------------------------------------===//
// This file will conditionally allocate lds memory, create local/async load
// operations, and create schedule for these operations. After lowerLoops,
// schedule will be passed to expandLoops and eventually to PipelineExpander.
//===----------------------------------------------------------------------===//

using mlir::triton::AMD::AttrBypassLDS;

namespace mlir {
struct StreamCopyChainOps {
  tt::LoadOp loadOp;
  ttg::MemDescIndexOp subviewOp;
  ttg::LocalStoreOp localStoreOp;
  ttg::LocalLoadOp maybeLocalLoadOp;
};

struct AsyncCopyChainOps {
  ttg::AsyncCopyGlobalToLocalOp copyOp;
  ttg::AsyncCommitGroupOp commitOp;
  ttg::AsyncWaitOp waitOp;
  ttg::LocalLoadOp maybeLocalLoadOp;
};

// Common shape for an async TDM chain (copy or gather). The first op (copy or
// gather) is type-erased because every consumer below only needs Operation *.
struct TDMChainOps {
  Operation *tdmOp;
  triton::amdgpu::AsyncTDMWait waitOp;
  ttg::LocalLoadOp maybeLocalLoadOp;
};

using StreamOpVariant =
    std::variant<StreamCopyChainOps, AsyncCopyChainOps, TDMChainOps>;
using LoadToStreamOpMap = llvm::MapVector<Operation *, StreamOpVariant>;

// Shared skeleton for building a TDM chain:
//   <buffer view> = createSingleBufferView(...)
//   <tdmOp>       = buildTDMOp(builder, view, pred)   <- caller-supplied
//   <waitOp>      = amdgpu::AsyncTDMWait(tdmOp)
//   replace uses of original load with a local_load after waitOp
TDMChainOps createTDMAsync(
    Operation *origLoad, Value alloc, Value extractIdx,
    function_ref<Operation *(OpBuilder &, Location, Value, Value)> buildTDMOp) {
  OpBuilder builder(origLoad);
  Location loc = origLoad->getLoc();

  Value pred = arith::ConstantIntOp::create(builder, loc, 1, 32);

  auto viewLoad = triton::createSingleBufferView(builder, alloc, extractIdx)
                      .getDefiningOp<ttg::MemDescIndexOp>();

  Operation *tdmOp = buildTDMOp(builder, loc, viewLoad, pred);

  auto waitOp = triton::amdgpu::AsyncTDMWait::create(builder, loc,
                                                     tdmOp->getResult(0), 0);

  auto maybeSharedLoad = tt::replaceUsesWithLocalLoad(
      builder, origLoad->getResult(0), viewLoad, waitOp);

  return {tdmOp, waitOp, maybeSharedLoad};
}

TDMChainOps createTDMAsyncCopy(tt::DescriptorLoadOp loadOp, Value alloc,
                               Value extractIdx) {
  return createTDMAsync(
      loadOp, alloc, extractIdx,
      [&](OpBuilder &builder, Location loc, Value view, Value pred) {
        Value desc = createUpdateTDMDescriptorOp(builder, loc, loadOp.getDesc(),
                                                 loadOp.getIndices(),
                                                 /*pred=*/pred);
        return triton::amdgpu::AsyncTDMCopyGlobalToLocalOp::create(builder, loc,
                                                                   desc, view);
      });
}

TDMChainOps createTDMAsyncGather(tt::DescriptorGatherOp gatherOp, Value alloc,
                                 Value extractIdx) {
  return createTDMAsync(
      gatherOp, alloc, extractIdx,
      [&](OpBuilder &builder, Location loc, Value view, Value pred) {
        // Convert to the TDM index layout for better gather performance;
        // insert a layout convert if the indices don't already have it.
        auto indices = gatherOp.getXOffsets();
        auto indicesType = cast<RankedTensorType>(indices.getType());
        auto idxEnc = getTDMGatherScatterIndexEncoding(gatherOp, indicesType);
        if (indicesType.getEncoding() != idxEnc) {
          auto newIdxType = RankedTensorType::get(
              indicesType.getShape(), indicesType.getElementType(), idxEnc);
          indices =
              ttg::ConvertLayoutOp::create(builder, loc, newIdxType, indices);
        }
        Value zero = arith::ConstantIntOp::create(builder, loc, 0, 32);
        Value desc = createUpdateTDMDescriptorOp(
            builder, loc, gatherOp.getDesc(), {zero, gatherOp.getYOffset()},
            /*pred=*/pred);
        return triton::amdgpu::AsyncTDMGatherOp::create(builder, loc, desc,
                                                        indices, view);
      });
}

namespace {

bool canBeConvertedToAsyncLoad(unsigned numBuffers, tt::LoadOp loadOp,
                               ttg::SharedEncodingTrait sharedEnc,
                               tt::ModuleAxisInfoAnalysis &axisInfoAnalysis,
                               const tt::AMD::TargetInfo &targetInfo);

AsyncCopyChainOps createAsyncCopy(tt::LoadOp loadOp, Value alloc,
                                  Value extractIdx, int contiguity) {
  OpBuilder builder(loadOp);
  Location loc = loadOp.getLoc();

  // Extract local subview from shared allocation
  auto viewLoad = triton::createSingleBufferView(builder, alloc, extractIdx)
                      .getDefiningOp<ttg::MemDescIndexOp>();

  auto copyOp = ttg::AsyncCopyGlobalToLocalOp::create(
      builder, loc, loadOp.getPtr(), viewLoad, loadOp.getMask(),
      loadOp.getOther(), loadOp.getCache(), loadOp.getEvict(),
      loadOp.getIsVolatile(), contiguity);
  auto commitOp =
      ttg::AsyncCommitGroupOp::create(builder, loc, copyOp->getResult(0));
  ttg::AsyncWaitOp waitOp =
      ttg::AsyncWaitOp::create(builder, loc, commitOp->getResult(0), 0);

  auto maybeSharedLoad = tt::replaceUsesWithLocalLoad(
      builder, loadOp->getResult(0), viewLoad, waitOp);

  return {copyOp, commitOp, waitOp, maybeSharedLoad};
}

void scheduleLocalLoad(ttg::LocalLoadOp localLoadOp,
                       tt::CoarseSchedule &schedule, int stage,
                       const tt::CoarseSchedule::Cluster &cluster) {
  schedule.insert(localLoadOp, stage, cluster);
  // If its only user is a ConvertLayout, we place it into the same stage so
  // it can be folded by a later pass
  if (localLoadOp->hasOneUse()) {
    auto cvt = *localLoadOp->getUsers().begin();
    if (isa<ttg::ConvertLayoutOp>(cvt)) {
      schedule.insert(cvt, stage, cluster);
    }
  }
}

StreamCopyChainOps createStreamCopy(tt::LoadOp loadOp, Value alloc,
                                    Value extractIdx) {
  OpBuilder builder(loadOp);
  Location loc = loadOp.getLoc();

  // Extract local subview from shared allocation
  auto viewLoad = triton::createSingleBufferView(builder, alloc, extractIdx)
                      .getDefiningOp<ttg::MemDescIndexOp>();

  tt::LoadOp newLoadOp = cast<tt::LoadOp>(builder.clone(*loadOp));
  auto storeOp = ttg::LocalStoreOp::create(builder, loc, newLoadOp, viewLoad);
  auto maybeLocalLoad =
      tt::replaceUsesWithLocalLoad(builder, loadOp->getResult(0), viewLoad);

  return {newLoadOp, viewLoad, storeOp, maybeLocalLoad};
}

// Adapted from
// lib/Dialect/TritonGPU/Transforms/Utility.cpp::getSharedEncIfAllUsersAreDotEnc
// to support AMDMfmaEncodingAttr.
// TODO(max): figure out how to refactor to use upstream
//
// If all the transitive uses of the given value have are used by a convert to
// the same dot operand encoding, return true and get the shared encoding that
// needs to be used to be compatible with users' layouts.
std::optional<ttg::SharedEncodingTrait> getSharedEncIfAllUsersAreDotEnc(
    Operation *loadOp, tt::ModuleAxisInfoAnalysis &axisInfoAnalysis,
    const TargetFeatures &targetFeatures, const tt::AMD::TargetInfo &targetInfo,
    bool useAsyncCopy) {
  assert(loadOp);
  Value loadedValue = loadOp->getResult(0);
  llvm::SmallVector<ttg::SharedEncodingTrait> sharedEncs;
  for (Operation *user : loadedValue.getUsers()) {
    LDBG(" getSharedEncIfAllUsersAreDotEnc current user: " << *user);
    if (user->getNumResults() != 1)
      return std::nullopt;

    ttg::SharedEncodingTrait tempAttr;
    Value userResult = user->getResult(0);
    Type userResType = userResult.getType();
    if (auto memDesc = dyn_cast<ttg::MemDescType>(userResType)) {
      // First time we find a shared encoding in the chain, save it and try to
      // use it if it is compatible with the other users.
      tempAttr = cast<ttg::SharedEncodingTrait>(memDesc.getEncoding());
      // If the immediate user is ttg::LocalAllocOp, likely it's created in
      // TritonAMDGPUOptimizeDotOperands. We should just respect it.
      if (!isa<ttg::LocalAllocOp>(user) &&
          !getSharedEncIfAllUsersAreDotEnc(user, axisInfoAnalysis,
                                           targetFeatures, targetInfo,
                                           useAsyncCopy)) {
        return std::nullopt;
      }
      LDBG("Deduced shared encoding candidate from memDesc: " << tempAttr);
      sharedEncs.push_back(tempAttr);
    } else {
      if (!(isa<ttg::ConvertLayoutOp>(user) ||
            user->hasTrait<OpTrait::LocalLoadTrait>()))
        return std::nullopt;

      auto srcTy = cast<ttg::TensorOrMemDesc>(loadedValue.getType());
      auto cgaLayout = ttg::getCGALayout(srcTy.getEncoding());

      auto order = getOrderForMemory(srcTy);
      if (useAsyncCopy && !targetInfo.supportsDirectToLdsScatter()) {
        // For architectures that don't support scattering into LDS we must
        // ensure that each warp writes a contiguous memory chunk. This requires
        // the shared memory order to follow the thread order, while preserving
        // the fastest dimension from the memory order if it's contiguous > 1 to
        // keep vectorization.
        auto srcTensorTy = cast<RankedTensorType>(srcTy);
        auto regOrder = triton::gpu::getOrder(srcTensorTy);
        auto threadOrder = triton::gpu::getThreadOrder(srcTensorTy);

        SetVector<unsigned> orderSet;

        auto regContig =
            triton::gpu::getContigPerThread(srcTensorTy)[regOrder[0]];
        unsigned elemBitWidth = srcTy.getElementType().getIntOrFloatBitWidth();
        unsigned finalRegContig =
            fitToValidDirectToLdsVecSize(regContig, elemBitWidth, targetInfo);
        // We only apply the order change if we find a vector size which works
        // for async copy.
        if (finalRegContig > 0) {
          // Preserve the fastest reg dim for contig > 1 to keep vectorization.
          if (finalRegContig > 1)
            orderSet.insert(order[0]);
          orderSet.insert(threadOrder.begin(), threadOrder.end());
          order = orderSet.takeVector();
        }
      }

      unsigned bitWidth = srcTy.getElementType().getIntOrFloatBitWidth();
      SmallVector<unsigned> sharedOrder;
      int rank = order.size();
      // TODO rework this when shared -> dotOperand conversions support
      // arbitrary shared memory ordering
      if (rank == 3) {
        // Move the batch dimension (dim #0) to be the last so that it will be
        // the slowest varying dimension.
        for (unsigned i = 0; i < rank; ++i)
          if (order[i] != 0)
            sharedOrder.emplace_back(order[i]);
        sharedOrder.emplace_back(0);
      } else {
        sharedOrder = order;
      }

      auto userResEnc = cast<ttg::TensorOrMemDesc>(userResType).getEncoding();
      if (auto dotOpEnc = dyn_cast<ttg::DotOperandEncodingAttr>(userResEnc)) {
        // Determine if we can use padded layouts and fallback to swizzled
        // layouts if not
        bool canUseAsyncCopy = false;
        if (useAsyncCopy && isa<tt::LoadOp>(loadOp)) {
          // We pass numBuffers=2 because we assume the schedule will not
          // determine a single buffer (which does not work with AsyncCopy)
          canUseAsyncCopy = canBeConvertedToAsyncLoad(
              2, cast<tt::LoadOp>(loadOp), {}, axisInfoAnalysis, targetInfo);
        }
        tempAttr = composePaddedLayout(targetFeatures, dotOpEnc.getOpIdx(),
                                       dotOpEnc.getKWidth(), srcTy, sharedOrder,
                                       dotOpEnc, canUseAsyncCopy);
        if (!tempAttr) {
          tempAttr = ttg::SwizzledSharedEncodingAttr::get(
              loadedValue.getContext(), dotOpEnc, srcTy.getShape(), sharedOrder,
              cgaLayout, bitWidth, /*needTrans=*/false);
        }
        LDBG("Deduced shared encoding candidate from dot layout: " << tempAttr);
        sharedEncs.push_back(tempAttr);
      } else if (auto llEnc = dyn_cast<ttg::LinearEncodingAttr>(userResEnc)) {
        // We use linear layout directly for scaled dot fp8 operands. For such
        // cases, we need to look further down the def-use chain to find the dot
        // op for the mfma layout to deduce operand index and other information.
        unsigned opIdx;
        unsigned vecSize;
        if (auto mfmaEnc = getDotEncoding<ttg::AMDMfmaEncodingAttr>(
                userResult, &opIdx, &vecSize)) {
          LDBG("deduced opIdx: " << opIdx << "; deduced vecSize: " << vecSize);
          tempAttr = mfmaEnc.composeSharedLayoutForOperand(
              cgaLayout, opIdx, srcTy.getShape(), order, vecSize, bitWidth,
              /*needTrans=*/false);
          LDBG("Deduced shared encoding candidate from mfma layout: "
               << tempAttr);
          sharedEncs.push_back(tempAttr);
        } else if (auto dotEnc = getDotEncoding<ttg::AMDWmmaEncodingAttr>(
                       userResult, &opIdx, &vecSize)) {
          LDBG("deduced opIdx: " << opIdx << "; deduced vecSize: " << vecSize);
          tempAttr = dotEnc.composeSharedLayoutForOperand(
              cgaLayout, opIdx, srcTy.getShape(), order, vecSize, bitWidth,
              /*needTrans=*/false);
          LDBG("Deduced shared encoding candidate from wmma layout: "
               << tempAttr);
          sharedEncs.push_back(tempAttr);
        }
      }
    }
  }

  auto equalSharedEncIgnoreVec = [](ttg::SwizzledSharedEncodingAttr a,
                                    ttg::SwizzledSharedEncodingAttr b) {
    if (!a || !b)
      return false;
    return (a.getPerPhase() == b.getPerPhase() &&
            a.getMaxPhase() == b.getMaxPhase() &&
            a.getOrder() == b.getOrder() &&
            a.getCGALayout() == b.getCGALayout());
  };
  if (sharedEncs.empty() || !sharedEncs.front())
    return std::nullopt;
  auto maxVecSharedEnc = sharedEncs.front();

  // TODO add support for padded layouts. Right now they will use a separate
  // allocation
  for (auto sharedEnc : llvm::drop_begin(sharedEncs, 1)) {
    auto maybeSwizzShared =
        dyn_cast<ttg::SwizzledSharedEncodingAttr>(sharedEnc);
    auto maybeSwizzMaxVec =
        dyn_cast<ttg::SwizzledSharedEncodingAttr>(maxVecSharedEnc);

    if (!equalSharedEncIgnoreVec(maybeSwizzShared, maybeSwizzMaxVec)) {
      LDBG("Incompatible shared encodings");
      return std::nullopt;
    }
    if (maybeSwizzShared.getVec() > maybeSwizzMaxVec.getVec()) {
      maxVecSharedEnc = sharedEnc;
    }
  }

  LDBG("Deduced shared encoding: " << maxVecSharedEnc);

  return maxVecSharedEnc;
}

// On targets without direct-to-LDS scatter support, a direct-to-LDS copy
// requires enough data per CTA for each lane to write at least 32 bits into
// LDS. The source layout itself may still be suboptimal here; CoalesceAsyncCopy
// can rewrite it before the final direct-to-LDS lowering checks coalescing
// legality.
bool hasEnoughCTABytesForDirectToLds(tt::LoadOp loadOp,
                                     const tt::AMD::TargetInfo &targetInfo) {
  if (targetInfo.supportsDirectToLdsScatter())
    return true;

  auto srcTy = dyn_cast<RankedTensorType>(loadOp.getResult().getType());
  if (!srcTy)
    return true;

  int64_t ctaTileBits = mlir::product(ttg::getShapePerCTA(srcTy)) *
                        srcTy.getElementTypeBitWidth();
  return ctaTileBits >= 32 * targetInfo.getWarpSize();
}

bool canBeConvertedToAsyncLoad(unsigned numBuffers, tt::LoadOp loadOp,
                               ttg::SharedEncodingTrait sharedEnc,
                               tt::ModuleAxisInfoAnalysis &axisInfoAnalysis,
                               const tt::AMD::TargetInfo &targetInfo) {
  // If we have a single buffer we would require another barrier after the
  // local_reads so instead we fall back to pipeline with registers
  // Removing this check will create incorrect IR, see
  // MembarUtility.h:membarFilter
  if (numBuffers <= 1)
    return false;

  // Checks whether the global pointer's contiguity and mask alignment allows
  // for at least 32 bit wide loads.
  if (!triton::canBeConvertedToAsyncLoad(loadOp, axisInfoAnalysis))
    return false;

  using tt::amdgpu::ISAFamily;
  bool hasDirectToLdsPath = llvm::is_contained(
      {ISAFamily::CDNA3, ISAFamily::CDNA4, ISAFamily::GFX1250},
      targetInfo.getISAFamily());

  // On targets without scatter support, reject CTA tiles that are too small to
  // provide one 32-bit direct-to-LDS write per lane. Layout-specific coalescing
  // is checked later, after CoalesceAsyncCopy has a chance to rewrite the src.
  if (hasDirectToLdsPath &&
      !hasEnoughCTABytesForDirectToLds(loadOp, targetInfo))
    return false;

  if (sharedEnc && hasDirectToLdsPath) {
    // Compute the final vecSize we can use for the combination of
    // sourceEncoding and sharedEncoding. We can only use AsyncCopy if the
    // target supports the requested or a smaller vecSize because we cannot
    // stride when loading directly to lds on GFX9
    auto srcTy = cast<RankedTensorType>(loadOp.getPtr().getType());
    auto regLayout = triton::gpu::toLinearLayout(srcTy);
    // It's the allocation so we trim the multibuffer dimension
    auto srcShape = srcTy.getShape();
    triton::LinearLayout sharedLayout;
    auto paddedEnc = dyn_cast<triton::gpu::PaddedSharedEncodingAttr>(sharedEnc);
    if (paddedEnc) {
      sharedLayout = paddedEnc.getLinearComponent();
    } else {
      sharedLayout = triton::gpu::toLinearLayout(srcShape, sharedEnc);
    }
    auto regToSharedLayout = regLayout.invertAndCompose(sharedLayout);

    unsigned elemBitWidth = tt::getPointeeBitWidth(srcTy);
    unsigned vecSize = regToSharedLayout.getNumConsecutiveInOut();
    if (paddedEnc)
      vecSize = std::min(vecSize, paddedEnc.getMinInterval());

    if (fitToValidDirectToLdsVecSize(vecSize, elemBitWidth, targetInfo) == 0)
      return false;
  }

  return true;
}

// Convert load ops into shared memory allocation loads and apply
// multi-buffering based on the required number of buffers.
LoadToStreamOpMap
createStreamOps(const LoadToInfoMap &loadToInfo, scf::ForOp &forOp,
                const int &numBuffers, bool useAsyncCopy,
                tt::ModuleAxisInfoAnalysis &axisInfoAnalysis) {
  IRRewriter builder(forOp);
  Location loc = forOp.getLoc();
  Value minusOne = arith::ConstantIntOp::create(builder, loc, -1, 32);
  Value zero = arith::ConstantIntOp::create(builder, loc, 0, 32);
  Value one = arith::ConstantIntOp::create(builder, loc, 1, 32);
  Value extractIdx = minusOne;
  Value numBuffersVal =
      arith::ConstantIntOp::create(builder, loc, numBuffers, 32);

  unsigned newOperandIndex = forOp.getBody()->getNumArguments();
  // Patch the loop to add the new loop carried dependency.
  forOp = addIterArgsToLoop(builder, forOp, {extractIdx});

  // Create one counter for the extract indices to avoid creating long
  // live range.
  extractIdx = forOp.getBody()->getArgument(newOperandIndex);

  builder.setInsertionPoint(forOp.getBody(), forOp.getBody()->begin());
  extractIdx = arith::AddIOp::create(builder, loc, extractIdx, one);
  Value cndExt = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::slt,
                                       extractIdx, numBuffersVal);
  extractIdx = arith::SelectOp::create(builder, loc, cndExt, extractIdx, zero);

  // Patch the yield with the updated counter.
  appendToForOpYield(forOp, {extractIdx});

  LoadToStreamOpMap loadToStreamOp;
  for (auto &[op, info] : loadToInfo) {
    if (!info.sharedEncoding)
      continue;

    auto loadOp = dyn_cast<tt::LoadOp>(op);
    auto descLoadOp = dyn_cast<tt::DescriptorLoadOp>(op);
    auto gatherOp = dyn_cast<tt::DescriptorGatherOp>(op);
    if (!loadOp && !descLoadOp && !gatherOp)
      continue;

    // Create an allocation that can hold distance number of loadOp shapes.
    auto ty = cast<RankedTensorType>(op->getResultTypes()[0]);
    Value alloc = triton::createAlloc(forOp, ty, op->getLoc(),
                                      info.sharedEncoding, numBuffers);
    assert(alloc && "Failed to create alloc for the async load.");
    triton::AMD::TargetInfo targetInfo(
        getAMDArch(op->getParentOfType<ModuleOp>()));

    // Replace the old load with multi-buffered loads
    if (descLoadOp) {
      loadToStreamOp[op] = createTDMAsyncCopy(descLoadOp, alloc, extractIdx);
      continue;
    }
    if (gatherOp) {
      loadToStreamOp[op] = createTDMAsyncGather(gatherOp, alloc, extractIdx);
      continue;
    }

    if (useAsyncCopy &&
        canBeConvertedToAsyncLoad(numBuffers, loadOp, info.sharedEncoding,
                                  axisInfoAnalysis, targetInfo)) {
      unsigned vec = axisInfoAnalysis.getContiguity(loadOp.getPtr());
      if (auto mask = loadOp.getMask())
        vec = std::min<unsigned>(vec, axisInfoAnalysis.getMaskAlignment(mask));
      loadToStreamOp[loadOp] = createAsyncCopy(loadOp, alloc, extractIdx, vec);
      continue;
    }

    loadToStreamOp[loadOp] = createStreamCopy(loadOp, alloc, extractIdx);
  }

  return loadToStreamOp;
}

} // namespace

namespace SingleDotSchedule {
namespace {
using namespace mlir::SingleDotSchedule;
using ClusterMap = DenseMap<tt::CoarseSchedule::ClusterHash, int>;

ClusterMap createClusterMap(tt::CoarseSchedule &schedule) {
  DenseMap<tt::CoarseSchedule::ClusterHash, int> clusterMap;
  for (auto &[op, stageAndCluster] : schedule.opToStageAndCluster) {
    auto [stage, cluster] = stageAndCluster;
    tt::CoarseSchedule::ClusterHash clusterHash =
        tt::CoarseSchedule::hashCluster(cluster);
    clusterMap[clusterHash] = *cluster;
  }

  return clusterMap;
}

// Remap global and compute clusters to the right place
void remapClusters(tt::CoarseSchedule &schedule, ClusterMap clusterMap,
                   Clusters &clusters) {
  for (auto &[op, stageAndCluster] : schedule.opToStageAndCluster) {
    tt::CoarseSchedule::ClusterHash clusterHash =
        tt::CoarseSchedule::hashCluster(stageAndCluster.second);
    int oldClusterId = clusterMap[clusterHash];
    if (oldClusterId == 0) {
      stageAndCluster.second = clusters[SCHED_GLOBAL_LOAD];
    } else {
      assert(oldClusterId == 1);
      stageAndCluster.second = clusters[SCHED_COMPUTE];
    }
  }
}

// Init Schedule Config based on settings and loop characteristics.
// Create clusters in order of ops in loop. This can interleave ops
// from different stages in the same cluster to achieve better backend
// scheduling.
//   WARNING: Changing the order of schedule.clusters.newAtBack() calls
//            can cause invalid schedules to be produced.
LogicalResult initSchedule(int maxDist, Stages &stages, int numStages,
                           int &numBuffers, bool useAsyncCopy, bool hasTDMLoad,
                           bool hasScaledDot, bool waitAtTail,
                           Clusters &clusters, tt::CoarseSchedule &schedule) {
  LDBG("Init SingleDotSchedule");
  int lastStage = numStages - 1;
  stages[SCHED_GLOBAL_LOAD] = 0;
  stages[SCHED_LOCAL_STORE] = 0;
  stages[SCHED_LOCAL_LOAD] = lastStage;
  stages[SCHED_COMPUTE] = lastStage;

  // LDS prefetch of dot operands: schedule ttg.local_load one stage before
  // the tt.dot that consumes it, so the LDS read of the next K-tile overlaps
  // the current tile's matrix op. for num_stages<3 there is no room to separate
  // them.
  // We take a conservative route and only enable LDS prefetch when the dot is
  // not scaled until more testing is done.
  // waitAtTail (pingpong) requires very specific scheduling, which is
  // incompatible with the prefetch cluster ordering below, so disable prefetch
  // in that case.
  bool ldsPrefetch =
      hasTDMLoad && !hasScaledDot && numStages >= 3 && !waitAtTail;
  if (ldsPrefetch) {
    stages[SCHED_LOCAL_LOAD] = lastStage - 1;
  }

  stages[SCHED_ASYNC_WAIT] = stages[SCHED_LOCAL_LOAD];

  bool pairedGlobalLoadLocalStore = stages[SCHED_LOCAL_STORE] == 0;
  stages[SCHED_LOCAL_STORE] += maxDist;
  if (waitAtTail) {
    stages[SCHED_ASYNC_WAIT] = std::max(0, stages[SCHED_LOCAL_LOAD] - 1);
  }

  LDBG(
      "Stage schedule:" << "  GLOBAL_LOAD stage = " << stages[SCHED_GLOBAL_LOAD]
                        << ", LOCAL_STORE stage = " << stages[SCHED_LOCAL_STORE]
                        << ", LOCAL_LOAD stage = " << stages[SCHED_LOCAL_LOAD]
                        << ", COMPUTE stage = " << stages[SCHED_COMPUTE]
                        << ", ASYNC_WAIT stage = " << stages[SCHED_ASYNC_WAIT]
                        << "; total = " << numStages);

  if (stages[SCHED_LOCAL_STORE] >= numStages ||
      stages[SCHED_LOCAL_STORE] > stages[SCHED_LOCAL_LOAD]) {
    LDBG("Invalid stage schedule");
    return failure();
  }

  // Calculate the number of buffers needed for each load.
  // TODO: Use the precise number of buffers needed by the particular load.
  numBuffers =
      std::max(1, stages[SCHED_LOCAL_LOAD] - stages[SCHED_LOCAL_STORE]);
  // If we use AsyncCopy or TDM loads, we need one more buffer since we are not
  // using a register buffer
  if (useAsyncCopy || hasTDMLoad) {
    numBuffers += 1;
  }

  LDBG("deduced max shared memory buffer number = " << numBuffers);

  // We place async wait as the first cluster because we want to have it being
  // the first in the main loop after pipelining.
  // In case we use async_copy with pingpong, we need to place async_wait at
  // the end of the previous iteration, so it can guarantee the correct
  // dependency when warp0 and warp1 are pipelined.
  int asyncWaitCluster = waitAtTail ? 4 : 0;
  // If tt.load and ttg.local_store are in the same stage
  //   spread them apart to allow overlap with compute
  // else
  //   Initiate ttg.local_store before tt.load
  int globalLoadCluster = 1;
  int localStoreCluster = 3;
  if (!pairedGlobalLoadLocalStore) {
    globalLoadCluster = 3;
    localStoreCluster = 2;
  }

  // If ttg.local_load and ttg.local_store are in the same stage
  //   spread them apart to allow overlap with compute
  // else if they share the buffer
  //   ttg.local_load must come first
  // else
  //   schedule ttg.local_load in the middle
  int localLoadCluster = globalLoadCluster;
  if (stages[SCHED_LOCAL_LOAD] == stages[SCHED_LOCAL_STORE]) {
    localLoadCluster = std::max(3, localStoreCluster + 1);
  } else if (numBuffers == 1 && localLoadCluster >= localStoreCluster) {
    // For 1 buffer, ttg.local_load must occur before ttg.local_store
    localLoadCluster = localStoreCluster - 1;
  }

  // Schedule compute with ttg.local_load if paired
  // otherwise, schedule in the middle
  int computeCluster = 2;
  if (stages[SCHED_LOCAL_LOAD] == stages[SCHED_COMPUTE]) {
    computeCluster = localLoadCluster;
  }

  // ldsPrefetch ordering: lead the loop body with the tt.dot (it consumes
  // operands pre-loaded into registers the previous iteration, so it has no
  // intra-iteration dependency), then the async_wait, the next-tile TDM copy,
  // and finally the ttg.local_load that feeds the next iteration.
  if (ldsPrefetch) {
    computeCluster = 0;
    asyncWaitCluster = 1;
    globalLoadCluster = 2;
    localLoadCluster = 3;
    localStoreCluster = 4;
  }

  // Create a hash map to associate cluster hash in old schedule with its
  // clusterID
  ClusterMap clusterMap = createClusterMap(schedule);

  // Make assignments
  Clusters clusterVec;
  schedule.clusters.clear();
  std::generate(clusterVec.begin(), clusterVec.end(),
                [&]() { return schedule.clusters.newAtBack(); });

  clusters[SCHED_GLOBAL_LOAD] = clusterVec[globalLoadCluster];
  clusters[SCHED_LOCAL_STORE] = clusterVec[localStoreCluster];
  clusters[SCHED_LOCAL_LOAD] = clusterVec[localLoadCluster];
  clusters[SCHED_COMPUTE] = clusterVec[computeCluster];
  clusters[SCHED_ASYNC_WAIT] = clusterVec[asyncWaitCluster];

  remapClusters(schedule, clusterMap, clusters);

  LDBG("Cluster schedule:" << "  GLOBAL_LOAD cluster = " << globalLoadCluster
                           << ", LOCAL_STORE cluster = " << localStoreCluster
                           << ", LOCAL_LOAD cluster = " << localLoadCluster
                           << ", COMPUTE cluster = " << computeCluster
                           << ", ASYNC_WAIT cluster = " << asyncWaitCluster
                           << "; total = " << SCHED_SIZE);

  return success();
}

void scheduleTDMOps(Operation *tdmOp, Operation *waitOp,
                    ttg::LocalLoadOp maybeLocalLoadOp, Operation *origLoadOp,
                    tt::CoarseSchedule &schedule, const Stages &stages,
                    const Clusters &clusters) {
  auto [loadStage, loadCluster] = schedule[origLoadOp];
  schedule.insert(tdmOp, loadStage, loadCluster);
  // If the LocalLoads are scheduled to a later stage than AsyncCopy we need to
  // place the AsyncCopy prefetches after the AsyncWaits which create a barrier
  // to ensure all warps are finished reading the shared buffer we will write
  // into. This is done by scheduling AsyncWait as the first cluster.
  // If AsyncCopy and LocalLoads are in the same stage we do not assign a
  // schdule so they are placed before the LocalLoads
  if (loadStage != stages[SCHED_LOCAL_LOAD])
    schedule.insert(waitOp, stages[SCHED_ASYNC_WAIT],
                    clusters[SCHED_ASYNC_WAIT]);

  if (maybeLocalLoadOp && stages[SCHED_LOCAL_LOAD] != stages[SCHED_COMPUTE]) {
    scheduleLocalLoad(maybeLocalLoadOp, schedule, stages[SCHED_LOCAL_LOAD],
                      clusters[SCHED_LOCAL_LOAD]);
  }
}

void scheduleAsyncCopy(const AsyncCopyChainOps &asyncOps, tt::LoadOp loadOp,
                       tt::CoarseSchedule &schedule, const Stages &stages,
                       const Clusters &clusters) {
  auto [copyOp, commitOp, waitOp, maybeLocalLoadOp] = asyncOps;
  auto [loadStage, loadCluster] = schedule[loadOp];
  schedule.insert(copyOp, loadStage, loadCluster);
  // Place ttg.async_commit_group op following AsyncCopyGlobalToLocal so the
  // later UpdateAsyncWaitCount pass can deduce better waitcnts
  schedule.insert(commitOp, loadStage, loadCluster);
  // If the LocalLoads are scheduled to a later stage than AsyncCopy we need to
  // place the AsyncCopy prefetches after the AsyncWaits which create a barrier
  // to ensure all warps are finished reading the shared buffer we will write
  // into. This is done by scheduling AsyncWait as the first cluster.
  // If AsyncCopy and LocalLoads are in the same stage we do not assign a
  // schdule so they are placed before the LocalLoads
  if (loadStage != stages[SCHED_LOCAL_LOAD])
    schedule.insert(waitOp, stages[SCHED_ASYNC_WAIT],
                    clusters[SCHED_ASYNC_WAIT]);

  if (maybeLocalLoadOp && stages[SCHED_LOCAL_LOAD] != stages[SCHED_COMPUTE]) {
    scheduleLocalLoad(maybeLocalLoadOp, schedule, stages[SCHED_LOCAL_LOAD],
                      clusters[SCHED_LOCAL_LOAD]);
  }
}

void scheduleStreamCopy(const StreamCopyChainOps &streamOps,
                        tt::LoadOp oldLoadOp, tt::CoarseSchedule &schedule,
                        const Stages &stages, const Clusters &clusters) {
  auto [newLoadOp, subviewOp, localStoreOp, maybeLocalLoadOp] = streamOps;
  auto [loadStage, loadCluster] = schedule[oldLoadOp];

  schedule.insert(newLoadOp, loadStage, loadCluster);
  schedule.insert(subviewOp, stages[SCHED_LOCAL_STORE],
                  clusters[SCHED_LOCAL_STORE]);
  schedule.insert(localStoreOp, stages[SCHED_LOCAL_STORE],
                  clusters[SCHED_LOCAL_STORE]);
  if (maybeLocalLoadOp && stages[SCHED_LOCAL_LOAD] != stages[SCHED_COMPUTE]) {
    scheduleLocalLoad(maybeLocalLoadOp, schedule, stages[SCHED_LOCAL_LOAD],
                      clusters[SCHED_LOCAL_LOAD]);
  }
}

void scheduleStreamOps(const LoadToStreamOpMap &loadToStreamOp,
                       tt::CoarseSchedule &schedule, const Stages &stages,
                       const Clusters &clusters) {
  for (auto [l, streamOps] : loadToStreamOp) {
    if (auto asyncOps = std::get_if<TDMChainOps>(&streamOps)) {
      auto [tdmOp, waitOp, localLoad] = *asyncOps;
      scheduleTDMOps(tdmOp, waitOp, localLoad, l, schedule, stages, clusters);
    } else if (auto asyncOps = std::get_if<AsyncCopyChainOps>(&streamOps)) {
      auto loadOp = cast<tt::LoadOp>(l);
      scheduleAsyncCopy(*asyncOps, loadOp, schedule, stages, clusters);
    } else if (auto sOps = std::get_if<StreamCopyChainOps>(&streamOps)) {
      auto loadOp = cast<tt::LoadOp>(l);
      scheduleStreamCopy(*sOps, loadOp, schedule, stages, clusters);
    }
  }
}

void updateSchedule(scf::ForOp &forOp, const LoadToInfoMap &loadToInfo,
                    tt::CoarseSchedule &schedule,
                    triton::AMD::ModuleAxisInfoAnalysis &axisInfoAnalysis,
                    bool useAsyncCopy, bool hasTDMLoad, bool waitAtTail) {
  LDBG("SingleDotSchedule::updateSchedule");
  Stages stages;
  Clusters clusters;

  int maxDist = 0;
  for (auto &[l, info] : loadToInfo) {
    maxDist = std::max(maxDist, info.distToUse);
  }

  bool hasScaledDot = false;
  forOp.walk([&](tt::DotScaledOp) { hasScaledDot = true; });

  int numBuffers = 1;
  if (failed(initSchedule(maxDist, stages, schedule.getNumStages(), numBuffers,
                          useAsyncCopy, hasTDMLoad, hasScaledDot, waitAtTail,
                          clusters, schedule)))
    return;

  // Convert the loads into shared memory allocations and loads from them.
  auto loadToStreamOps = createStreamOps(loadToInfo, forOp, numBuffers,
                                         useAsyncCopy, axisInfoAnalysis);

  scheduleStreamOps(loadToStreamOps, schedule, stages, clusters);
  dumpScheduleDebug(schedule, DEBUG_TYPE, "Coarse schedule stream ops:");

  for (auto [l, _] : loadToInfo) {
    if (isa<tt::DescriptorLoadOp, tt::DescriptorGatherOp>(l)) {
      schedule.erase(l);
      l->erase();
    }
  }

  scheduleDependencies(forOp, schedule);
  dumpScheduleDebug(schedule, DEBUG_TYPE, "Coarse schedule with dependencies:");
  ttg::scheduleDistanceOneDependencies(forOp, schedule);
  dumpScheduleDebug(schedule, DEBUG_TYPE, "Coarse schedule with dist 1:");
  tt::CoarseSchedule::Cluster computeCluster = clusters[SCHED_COMPUTE];
  ttg::scheduleRemainingToLastStage(forOp, schedule, computeCluster);
  dumpScheduleDebug(schedule, DEBUG_TYPE, "Final coarse schedule:");
}
} // namespace
} // namespace SingleDotSchedule

namespace ChainedDotSchedule {
namespace {
using namespace mlir::ChainedDotSchedule;

void scheduleAsyncCopy(const AsyncCopyChainOps &asyncOps, tt::LoadOp loadOp,
                       tt::CoarseSchedule &schedule,
                       const ChainedDotClusters &clusters) {
  auto [loadStage, loadCluster] = schedule[loadOp];
  auto [copyOp, commitOp, waitOp, maybeLocalLoadOp] = asyncOps;

  schedule.insert(copyOp, loadStage, loadCluster);
  // Place ttg.async_commit_group op following AsyncCopyGlobalToLocal so the
  // later UpdateAsyncWaitCount pass can deduce better waitcnts
  schedule.insert(commitOp, loadStage, loadCluster);

  if (loadStage == STAGE_GLOBAL_LOAD_1) {
    schedule.insert(waitOp, STAGE_LOCAL_LOAD_1, clusters[CLUSTER_ASYNC_WAIT_1]);
    if (maybeLocalLoadOp)
      scheduleLocalLoad(maybeLocalLoadOp, schedule, STAGE_LOCAL_LOAD_1,
                        clusters[CLUSTER_LOCAL_LOAD_1]);
  } else {
    schedule.insert(waitOp, STAGE_LOCAL_LOAD_2, clusters[CLUSTER_ASYNC_WAIT_2]);
    if (maybeLocalLoadOp)
      scheduleLocalLoad(maybeLocalLoadOp, schedule, STAGE_LOCAL_LOAD_2,
                        clusters[CLUSTER_LOCAL_LOAD_2]);
  }
}

void scheduleStreamCopy(const StreamCopyChainOps &streamOps, tt::LoadOp loadOp,
                        tt::CoarseSchedule &schedule,
                        const ChainedDotClusters &clusters) {
  auto [loadStage, loadCluster] = schedule[loadOp];
  auto [copyOp, subviewOp, localStoreOp, maybeLocalLoadOp] = streamOps;
  schedule.insert(copyOp, loadStage, loadCluster);

  if (loadStage == STAGE_GLOBAL_LOAD_1) {
    schedule.insert(subviewOp, STAGE_LOCAL_WRITE_1,
                    clusters[CLUSTER_LOCAL_WRITE_1]);
    schedule.insert(localStoreOp, STAGE_LOCAL_WRITE_1,
                    clusters[CLUSTER_LOCAL_WRITE_1]);

    if (maybeLocalLoadOp)
      schedule.insert(maybeLocalLoadOp, STAGE_LOCAL_LOAD_1,
                      clusters[CLUSTER_LOCAL_LOAD_1]);
  } else {
    schedule.insert(subviewOp, STAGE_LOCAL_WRITE_2,
                    clusters[CLUSTER_LOCAL_WRITE_2]);
    schedule.insert(localStoreOp, STAGE_LOCAL_WRITE_2,
                    clusters[CLUSTER_LOCAL_WRITE_2]);
    if (maybeLocalLoadOp)
      schedule.insert(maybeLocalLoadOp, STAGE_LOCAL_LOAD_2,
                      clusters[CLUSTER_LOCAL_LOAD_2]);
  }

  if (maybeLocalLoadOp) {
    if (auto cvt = dyn_cast<ttg::ConvertLayoutOp>(
            *maybeLocalLoadOp->getUsers().begin())) {
      auto [localLoadStage, localLoadCluster] = schedule[maybeLocalLoadOp];
      schedule.insert(cvt, localLoadStage, localLoadCluster);
    }
  }
}

void scheduleStreamOps(const LoadToStreamOpMap &loadToStreamOp,
                       tt::CoarseSchedule &schedule,
                       const ChainedDotClusters &clusters) {
  for (auto [l, streamOps] : loadToStreamOp) {
    auto loadOp = dyn_cast<tt::LoadOp>(l);
    if (!loadOp)
      continue;

    if (auto asyncOps = std::get_if<AsyncCopyChainOps>(&streamOps)) {
      scheduleAsyncCopy(*asyncOps, loadOp, schedule, clusters);
    } else if (auto sOps = std::get_if<StreamCopyChainOps>(&streamOps)) {
      scheduleStreamCopy(*sOps, loadOp, schedule, clusters);
    }
  }
}

void updateSchedule(scf::ForOp &forOp, const LoadToInfoMap &loadToInfo,
                    tt::CoarseSchedule &schedule,
                    triton::AMD::ModuleAxisInfoAnalysis &axisInfoAnalysis,
                    bool useAsyncCopy) {
  LDBG("ChainedDotSchedule::updateSchedule");
  ChainedDotClusters clusters;
  int cnt = clusters.size() - schedule.clusters.size();
  for (int i = 0; i < cnt; i++) {
    schedule.clusters.newAtBack();
  }
  auto it = schedule.clusters.begin();
  for (int i = 0; i < clusters.size(); i++, it++) {
    clusters[i] = it;
  }

  // Convert the loads into shared memory allocations and loads from them.
  // TODO support different numBuffers
  int numBuffers = useAsyncCopy ? 2 : 1;
  auto loadToStreamOps = createStreamOps(loadToInfo, forOp, numBuffers,
                                         useAsyncCopy, axisInfoAnalysis);
  scheduleStreamOps(loadToStreamOps, schedule, clusters);

  // Now that the original load is replaced by a sequence of new Ops. It become
  // dead, and hence should be removed from the schedule and can be erased.
  for (auto [l, _] : loadToStreamOps) {
    schedule.erase(l);
    l->erase();
  }

  scheduleDependencies(forOp, schedule);
  dumpScheduleDebug(schedule, DEBUG_TYPE, "Coarse schedule with dependencies:");

  triton::gpu::scheduleDistanceOneDependencies(forOp, schedule);
  dumpScheduleDebug(schedule, DEBUG_TYPE, "Coarse schedule with dist 1:");

  tt::CoarseSchedule::Cluster lastCluster = clusters.back();
  triton::gpu::scheduleRemainingToLastStage(forOp, schedule, lastCluster);
  dumpScheduleDebug(schedule, DEBUG_TYPE, "Final coarse schedule:");
}
} // namespace
} // namespace ChainedDotSchedule

static void lowerLoop(scf::ForOp forOp,
                      triton::AMD::ModuleAxisInfoAnalysis &axisInfoAnalysis,
                      bool useAsyncCopy, bool usePingpong) {
  tt::CoarseSchedule schedule;
  if (failed(schedule.deSerialize(forOp, /*normalizeClusterId=*/false))) {
    return;
  }

  dumpScheduleDebug(schedule, DEBUG_TYPE, "[lowerLoops]deserialized schedule:");

  int numStages = schedule.getNumStages();

  // i.e., we can still disable `waitAtTail` by explicitly disabling
  // pingpong, which is the only use case of this scheduling variant.
  bool waitAtTail = usePingpong && (numStages == 3) && useAsyncCopy;

  llvm::MapVector<Operation *, std::pair<int, Operation *>> loadOpToIndLevel =
      getIndirectLevel(axisInfoAnalysis, forOp, numStages);

  TargetFeatures targetFeatures(getAMDArch(forOp->getParentOfType<ModuleOp>()));
  triton::AMD::TargetInfo targetInfo(
      getAMDArch(forOp->getParentOfType<ModuleOp>()));

  bool hasTDMLoad = false;
  LoadToInfoMap loadToInfo;
  for (const auto &[load, info] : loadOpToIndLevel) {
    auto [distance, use] = info;
    if (load->hasAttrOfType<BoolAttr>(AttrBypassLDS)) {
      load->removeAttr(AttrBypassLDS);
      loadToInfo[load] = {nullptr, distance, use};
    } else {
      if (auto descLoadLike =
              dyn_cast<tt::DescriptorLoadLikeOpInterface>(load)) {
        hasTDMLoad = true;
        auto resultType = cast<RankedTensorType>(load->getResult(0).getType());
        auto paddedEncoding =
            getEncodingFromDescriptor(load, resultType, descLoadLike.getDesc());
        LoadInfo ldInfo;
        ldInfo.sharedEncoding = paddedEncoding;
        ldInfo.distToUse = distance;
        ldInfo.use = use;
        loadToInfo[load] = ldInfo;
      } else {
        LDBG("Deduce shared encoding for: " << *load);
        auto sharedEncoding = getSharedEncIfAllUsersAreDotEnc(
                                  load, axisInfoAnalysis, targetFeatures,
                                  targetInfo, useAsyncCopy)
                                  .value_or(nullptr);
        loadToInfo[load] = {sharedEncoding, distance, use};
        LDBG("Populate loadInfo with shared encoding: " << sharedEncoding);
      }
    }
  }

  if (succeeded(mlir::ChainedDotSchedule::checkPreconditions(forOp, numStages,
                                                             loadToInfo))) {
    ChainedDotSchedule::updateSchedule(forOp, loadToInfo, schedule,
                                       axisInfoAnalysis, useAsyncCopy);
  } else {
    SingleDotSchedule::updateSchedule(forOp, loadToInfo, schedule,
                                      axisInfoAnalysis, useAsyncCopy,
                                      hasTDMLoad, waitAtTail);
  }

  dumpScheduleDebug(schedule, DEBUG_TYPE, "[lowerLoops]updated schedule:");

  schedule.serialize(forOp);
}

void lowerLoops(ModuleOp moduleOp, bool useAsyncCopy, bool usePingpong) {
  triton::AMD::ModuleAxisInfoAnalysis axisInfoAnalysis(moduleOp);
  SmallVector<scf::ForOp> loops;
  moduleOp->walk([&](scf::ForOp forOp) { loops.push_back(forOp); });
  if (loops.empty())
    return;
  for (auto forOp : loops) {
    lowerLoop(forOp, axisInfoAnalysis, useAsyncCopy, usePingpong);
  }
}

} // namespace mlir
