#include "PartitionAttrs.h"
#include "mlir/Analysis/TopologicalSortUtils.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/RegionUtils.h"
#include "nvidia/include/Dialect/NVWS/IR/Dialect.h"
#include "nvidia/include/Dialect/NVWS/Transforms/Passes.h"
#include "triton/Dialect/TritonGPU/IR/Dialect.h"
#include "triton/Dialect/TritonGPU/Transforms/Partition.h"
#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h"
#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
#include "triton/Dialect/TritonGPU/Transforms/WarpSpecialization.h"
#include "llvm/ADT/SCCIterator.h"

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

namespace {

struct WarpGroupBuilder : public OpBuilder {
  WarpGroupBuilder(Block *block, Block::iterator insertPoint,
                   size_t partitionId)
      : OpBuilder(block, insertPoint), partitionId(partitionId) {}

  IRMapping mapping;
  size_t partitionId;
};

// This is computed per loop and partition
enum class LoopVarCategory {
  // The given loop variable is not used by the given partition. For example,
  // the use-D flag for MMA is only used by the MMA partition, and thus
  // is `Unused` for any other partition.
  Unused,
  // The given loop variable is used by the given partition. For example, a loop
  // index might be used to compute a relevant stage or phase value for the
  // given partition.
  Used,
  // The results of warp_group op are defined to be those of the first
  // partition. If the original loop results include a tensor which is computed
  // only by a non-default partition, such tensor cannot be returned from the
  // first partition and and must be passed through shared memory. The
  // corresponding loop variable falls into this category.
  // Recognizing this category is necessary for the first partition. For other
  // partitions, some loop variables might be assigned this category, but that
  // information is not used.
  TensorResultFromOtherPartition,
};

SetVector<int> getResultPartitionIds(Operation *op, int index) {
  return getPartitionOutputs(op)[index];
}

SetVector<int> getIfOpResultPartitionIds(scf::IfOp ifOp, Value value) {
  for (auto result : ifOp.getResults()) {
    if (result == value) {
      auto pos = result.getResultNumber();
      return getResultPartitionIds(ifOp, pos);
    }
  }
  llvm_unreachable("value is not a result of if-stmt");
}

bool isTensorResultComputedBy(LoopLikeOpInterface loop, size_t resultIdx,
                              const Partition *partition,
                              const PartitionSet &partitions) {
  auto value = loop.getYieldedValues()[resultIdx];
  if (!isa<RankedTensorType>(value.getType()))
    return false;
  auto defOp = value.getDefiningOp();
  auto partitionIds = getPartitionIds(defOp);
  if (auto ifOp = dyn_cast<scf::IfOp>(defOp)) {
    partitionIds = getIfOpResultPartitionIds(ifOp, value);
  }
  return llvm::is_contained(partitionIds, partition->getIndex());
}

SmallVector<LoopVarCategory> classifyLoopVars(LoopLikeOpInterface loop,
                                              const Partition *partition,
                                              const PartitionSet &partitions) {
  auto isTensorResultFromOtherPartition = [&](int i) {
    for (auto otherPartition : partitions.getPartitions()) {
      if (&otherPartition == partition) {
        continue;
      }
      if (isTensorResultComputedBy(loop, i, &otherPartition, partitions)) {
        return true;
      }
    }
    return false;
  };

  SmallVector<LoopVarCategory> categories(loop.getRegionIterArgs().size());
  for (auto [i, arg] : llvm::enumerate(loop.getRegionIterArgs())) {
    auto partitionIds = getResultPartitionIds(loop, i);
    if (llvm::is_contained(partitionIds, partition->getIndex())) {
      categories[i] = LoopVarCategory::Used;
    } else if (isTensorResultFromOtherPartition(i) &&
               !loop->getResult(i).use_empty()) {
      categories[i] = LoopVarCategory::TensorResultFromOtherPartition;
    } else {
      categories[i] = LoopVarCategory::Unused;
    }
  }

  return categories;
}

std::pair<SmallVector<size_t>, SmallVector<std::optional<size_t>>>
getLoopVarIndicesToKeep(LoopLikeOpInterface loop, const Partition *partition,
                        ArrayRef<LoopVarCategory> loopVarCategories) {
  SmallVector<size_t> indices;
  // The null index means an invalid index, the corresponding loop variable in
  // the original loop is removed in the cloned loop
  SmallVector<std::optional<size_t>> reverseIndices(
      loop.getRegionIterArgs().size(), std::nullopt);
  for (auto [i, arg] : llvm::enumerate(loop.getRegionIterArgs())) {
    if (loopVarCategories[i] == LoopVarCategory::Used) {
      reverseIndices[i] = indices.size();
      indices.push_back(i);
    }
  }
  return std::make_pair(indices, reverseIndices);
}

std::pair<SmallVector<size_t>, SmallVector<std::optional<size_t>>>
getLoopVarIndicesToKeep(LoopLikeOpInterface loop, const Partition *partition,
                        const PartitionSet &partitions) {
  auto loopVarCategories = classifyLoopVars(loop, partition, partitions);
  return getLoopVarIndicesToKeep(loop, partition, loopVarCategories);
}

void mapRange(ValueRange fromRange, ValueRange toRange, IRMapping &mapping) {
  for (auto [from, to] : llvm::zip(fromRange, toRange)) {
    mapping.map(from, to);
  }
}

void cloneOpsInBlock(Block *block, SmallVector<WarpGroupBuilder> &builders,
                     const PartitionSet &partitions);

void cloneForOp(scf::ForOp forOp, SmallVector<WarpGroupBuilder> &builders,
                const PartitionSet &partitions) {
  auto forOpPartitions = getPartitionIds(forOp);

  SmallVector<scf::ForOp> newForOps;
  for (int i : forOpPartitions) {
    auto &b = builders[i];
    auto partition = partitions.getPartition(i);
    auto [newLoopIndices, _] =
        getLoopVarIndicesToKeep(forOp, partition, partitions);
    auto lb = b.mapping.lookupOrDefault(forOp.getLowerBound());
    auto ub = b.mapping.lookupOrDefault(forOp.getUpperBound());
    auto step = b.mapping.lookupOrDefault(forOp.getStep());
    SmallVector<Value> initArgs;
    for (auto idx : newLoopIndices) {
      initArgs.push_back(b.mapping.lookupOrDefault(forOp.getInitArgs()[idx]));
    }
    auto newForOp =
        scf::ForOp::create(b, forOp.getLoc(), lb, ub, step, initArgs);
    newForOp->setAttrs(forOp->getAttrs());
    if (forOp->hasAttr(kPartitionOutputsAttrName)) {
      newForOp->removeAttr(kPartitionOutputsAttrName);
    }
    newForOps.push_back(newForOp);

    b.mapping.map(forOp.getInductionVar(), newForOp.getInductionVar());

    auto oldIterArgs = forOp.getRegionIterArgs();
    auto newIterArgs = newForOp.getRegionIterArgs();
    for (auto [newIdx, oldIdx] : llvm::enumerate(newLoopIndices)) {
      b.mapping.map(oldIterArgs[oldIdx], newIterArgs[newIdx]);
      b.mapping.map(forOp.getResult(oldIdx), newForOp.getResult(newIdx));
    }

    b.setInsertionPointToStart(newForOp.getBody());
  }

  cloneOpsInBlock(forOp.getBody(), builders, partitions);

  for (auto [i, newForOp] : llvm::zip(forOpPartitions, newForOps)) {
    builders[i].setInsertionPointAfter(newForOp);
    newForOp.walk([&](Operation *op) { op->removeAttr(kPartitionAttrName); });
    newForOp->removeAttr(kPartitionStagesAttrName);
  }
}

void cloneWhileOp(scf::WhileOp whileOp, SmallVector<WarpGroupBuilder> &builders,
                  const PartitionSet &partitions) {
  auto whilePartitions = getPartitionIds(whileOp);
  SmallVector<scf::WhileOp> newWhileOps;

  for (int i : whilePartitions) {
    auto &b = builders[i];
    auto partition = partitions.getPartition(i);
    auto [stateIndices, _] =
        getLoopVarIndicesToKeep(whileOp, partition, partitions);

    SmallVector<Value> inits;
    SmallVector<Type> resultTypes;
    for (size_t oldIdx : stateIndices) {
      inits.push_back(b.mapping.lookupOrDefault(whileOp.getInits()[oldIdx]));
      resultTypes.push_back(whileOp.getResult(oldIdx).getType());
    }

    auto newWhile =
        scf::WhileOp::create(b, whileOp.getLoc(), resultTypes, inits);
    newWhile->setAttrs(whileOp->getAttrs());
    newWhile->removeAttr(kPartitionOutputsAttrName);

    auto before = &newWhile.getBefore().emplaceBlock();
    auto after = &newWhile.getAfter().emplaceBlock();
    for (auto [newIdx, oldIdx] : llvm::enumerate(stateIndices)) {
      auto oldBefore = whileOp.getBeforeArguments()[oldIdx];
      auto newBefore =
          before->addArgument(oldBefore.getType(), oldBefore.getLoc());
      b.mapping.map(oldBefore, newBefore);

      auto oldAfter = whileOp.getAfterArguments()[oldIdx];
      auto newAfter = after->addArgument(oldAfter.getType(), oldAfter.getLoc());
      b.mapping.map(oldAfter, newAfter);

      b.mapping.map(whileOp.getResult(oldIdx), newWhile.getResult(newIdx));
    }

    newWhileOps.push_back(newWhile);
    b.setInsertionPointToStart(before);
  }

  cloneOpsInBlock(&whileOp.getBefore().front(), builders, partitions);
  for (auto [i, newWhile] : llvm::zip(whilePartitions, newWhileOps)) {
    builders[i].setInsertionPointToStart(&newWhile.getAfter().front());
  }
  cloneOpsInBlock(&whileOp.getAfter().front(), builders, partitions);

  for (auto [i, newWhile] : llvm::zip(whilePartitions, newWhileOps)) {
    builders[i].setInsertionPointAfter(newWhile);
    newWhile.walk([&](Operation *op) { op->removeAttr(kPartitionAttrName); });
    newWhile->removeAttr(kPartitionStagesAttrName);
  }
}

void cloneIfOp(scf::IfOp ifOp, SmallVector<WarpGroupBuilder> &builders,
               const PartitionSet &partitions) {
  auto partitionIndices = getPartitionIds(ifOp);

  SmallVector<scf::IfOp> newIfOps;
  for (size_t idx : partitionIndices) {
    auto &b = builders[idx];
    auto cond = b.mapping.lookupOrDefault(ifOp.getCondition());
    SmallVector<Type> newIfResultTypes;
    SmallVector<int> newIfResultIndices;
    for (auto pos = 0; pos < ifOp.getResultTypes().size(); ++pos) {
      auto partitionIds = getResultPartitionIds(ifOp, pos);
      if (llvm::is_contained(partitionIds, b.partitionId)) {
        newIfResultTypes.push_back(ifOp.getResult(pos).getType());
        newIfResultIndices.push_back(pos);
      }
    }
    auto newIfOp = scf::IfOp::create(b, ifOp.getLoc(), newIfResultTypes, cond,
                                     ifOp.elseBlock() ? true : false);
    newIfOp->setAttrs(ifOp->getAttrs());
    if (ifOp->hasAttr(kPartitionOutputsAttrName)) {
      newIfOp->removeAttr(kPartitionOutputsAttrName);
    }
    newIfOps.push_back(newIfOp);

    for (auto [newIdx, oldIdx] : llvm::enumerate(newIfResultIndices)) {
      b.mapping.map(ifOp.getResult(oldIdx), newIfOp.getResult(newIdx));
    }
    assert(ifOp.thenBlock()->getNumArguments() == 0);

    b.setInsertionPointToStart(newIfOp.thenBlock());
  }

  cloneOpsInBlock(ifOp.thenBlock(), builders, partitions);

  if (auto elseBlock = ifOp.elseBlock()) {
    for (auto [idx, newIfOp] : llvm::zip(partitionIndices, newIfOps)) {
      builders[idx].setInsertionPointToStart(newIfOp.elseBlock());
    }
    cloneOpsInBlock(elseBlock, builders, partitions);
  }

  for (auto [idx, newIfOp] : llvm::zip(partitionIndices, newIfOps)) {
    builders[idx].setInsertionPointAfter(newIfOp);
  }
}

void cloneReduceOp(triton::ReduceOp reduceOp,
                   SmallVector<WarpGroupBuilder> &builders,
                   const PartitionSet &partitions) {
  auto partitionIndices = getPartitionIds(reduceOp);

  SmallVector<ReduceOp> newReduceOps;
  for (size_t idx : partitionIndices) {
    auto &b = builders[idx];

    SmallVector<Value> srcs;
    for (auto src : reduceOp.getSrcs()) {
      srcs.push_back(b.mapping.lookupOrDefault(src));
    }
    auto axis = reduceOp.getAxis();
    auto newReduceOp =
        triton::ReduceOp::create(b, reduceOp.getLoc(), srcs, axis);
    newReduceOp->setAttrs(reduceOp->getAttrs());
    if (reduceOp->hasAttr(kPartitionOutputsAttrName)) {
      newReduceOp->removeAttr(kPartitionOutputsAttrName);
    }
    newReduceOps.push_back(newReduceOp);

    mapRange(reduceOp.getResults(), newReduceOp.getResults(), b.mapping);

    auto &region = newReduceOp.getRegion();
    Block *block = &region.emplaceBlock();
    for (auto arg : reduceOp.getRegion().getBlocks().front().getArguments()) {
      auto newArg = block->addArgument(arg.getType(), arg.getLoc());
      b.mapping.map(arg, newArg);
    }

    b.setInsertionPointToStart(block);
  }

  cloneOpsInBlock(reduceOp.getBody(), builders, partitions);

  for (auto [idx, newReduceOp] : llvm::zip(partitionIndices, newReduceOps)) {
    builders[idx].setInsertionPointAfter(newReduceOp);
  }
}

void cloneOp(Operation *op, SmallVector<WarpGroupBuilder> &builders,
             const SetVector<int> &partitionIndices) {
  if (op->getNumRegions() != 0) {
    llvm::report_fatal_error(
        "Ops are expected to be regionless at this point.");
  }

  if (auto arefPhi = dyn_cast<nvws::ArefPhiOp>(op)) {
    auto getValuePartitions = [](Value value) {
      if (auto result = dyn_cast<OpResult>(value)) {
        auto def = result.getDefiningOp();
        if (!def->hasAttr(kPartitionOutputsAttrName))
          return getPartitionIds(def);
        return getPartitionOutputs(def)[result.getResultNumber()];
      }
      auto arg = cast<BlockArgument>(value);
      auto forOp = cast<scf::ForOp>(arg.getOwner()->getParentOp());
      return getPartitionOutputs(forOp)[arg.getArgNumber() - 1];
    };
    auto localPartitions = getValuePartitions(arefPhi.getLocal());
    auto remotePartitions = getValuePartitions(arefPhi.getRemote());
    for (size_t idx : partitionIndices) {
      auto selected = localPartitions.contains(idx) ? arefPhi.getLocal()
                                                    : arefPhi.getRemote();
      assert(
          (localPartitions.contains(idx) != remotePartitions.contains(idx)) &&
          "expected exactly one aref.phi operand per partition");
      builders[idx].mapping.map(
          arefPhi.getResult(), builders[idx].mapping.lookupOrDefault(selected));
    }
  } else {
    for (size_t idx : partitionIndices) {
      auto &builder = builders[idx];
      auto newOp = builder.clone(*op, builder.mapping);
      mapRange(op->getResults(), newOp->getResults(), builder.mapping);
    }
  }
}

void cloneOpsInBlock(Block *block, SmallVector<WarpGroupBuilder> &builders,
                     const PartitionSet &partitions) {
  for (auto &op_ : *block) {
    auto op = &op_;

    if (auto forOp = dyn_cast<scf::ForOp>(op)) {
      cloneForOp(forOp, builders, partitions);
    } else if (auto whileOp = dyn_cast<scf::WhileOp>(op)) {
      cloneWhileOp(whileOp, builders, partitions);
    } else if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
      cloneIfOp(ifOp, builders, partitions);
    } else if (auto reduceOp = dyn_cast<triton::ReduceOp>(op)) {
      cloneReduceOp(reduceOp, builders, partitions);
    } else if (auto conditionOp = dyn_cast<scf::ConditionOp>(op)) {
      for (int idx : getPartitionIds(conditionOp)) {
        auto &builder = builders[idx];
        auto condition =
            builder.mapping.lookupOrDefault(conditionOp.getCondition());
        scf::ConditionOp::create(builder, conditionOp.getLoc(), condition,
                                 builder.getInsertionBlock()->getArguments());
      }
    } else if (auto yieldOp = dyn_cast<scf::YieldOp>(op)) {
      if (yieldOp.getOperands().empty()) {
        continue;
      }
      // empty yield has no partition annotations
      assert(hasPartition(op));
      auto partitionIndices = getPartitionIds(op);

      for (size_t idx : partitionIndices) {
        auto &builder = builders[idx];
        SmallVector<size_t> newOperandIndices;
        if (auto loopOp =
                dyn_cast<LoopLikeOpInterface>(yieldOp->getParentOp())) {
          newOperandIndices =
              getLoopVarIndicesToKeep(
                  loopOp, partitions.getPartition(builder.partitionId),
                  partitions)
                  .first;
        } else {
          auto ifOp = cast<scf::IfOp>(yieldOp->getParentOp());
          for (size_t i = 0; i < yieldOp.getOperands().size(); ++i) {
            auto ids = getResultPartitionIds(ifOp, i);
            if (llvm::is_contained(ids, builder.partitionId)) {
              newOperandIndices.push_back(i);
            }
          }
        }

        if (newOperandIndices.empty())
          continue;

        SmallVector<Value> newYieldOperands;
        for (size_t i : newOperandIndices) {
          newYieldOperands.push_back(
              builder.mapping.lookupOrDefault(yieldOp.getOperand(i)));
        }

        scf::YieldOp::create(builder, op->getLoc(), newYieldOperands);
      }
    } else {
      assert(hasPartition(op));
      auto partitionIndices = getPartitionIds(op);
      cloneOp(op, builders, partitionIndices);
    }
  }
}

} // namespace

LogicalResult triton::gpu::partitionLoop(LoopLikeOpInterface loop) {
  FailureOr<PartitionSet> partitionsOr = PartitionSet::fromLoop(loop);
  if (failed(partitionsOr))
    return failure();
  PartitionSet partitions = std::move(*partitionsOr);

  // Only the root node should have consumers at this point.
  for (const Partition &partition : partitions.getPartitions()) {
    bool failed = false;
    auto callback = [&](OpResult output, OpOperand &use, unsigned distance) {
      if (isa<nvws::ArefPhiOp>(use.getOwner()))
        return;
      auto partitionIds = getPartitionIds(use.getOwner());
      if (llvm::is_contained(partitionIds, partition.getIndex()))
        return;

      // check if consumer partition set is a subset of the producer partitions
      auto defOpPartitionIds = getPartitionIds(output.getDefiningOp());
      bool isValidSubset = std::all_of(
          partitionIds.begin(), partitionIds.end(), [&](int consumerId) {
            return llvm::is_contained(defOpPartitionIds, consumerId);
          });

      if (isValidSubset)
        return; // Valid: consumer ⊆ producer

      failed = true;
      InFlightDiagnostic diag =
          mlir::emitWarning(output.getLoc(), "non-root partition #")
          << partition.getIndex() << " has direct SSA consumer";

      for (auto partitionId : partitionIds) {
        diag.attachNote(use.getOwner()->getLoc())
            << "use at distance " << distance << " in partition #"
            << partitionId << " here";
      }
    };
    partition.iterateUses(loop, callback);
    if (failed)
      return failure();
  }

  // There is nothing to do if the loop has 1 or fewer partitions.
  if (llvm::size(partitions.getPartitions()) <= 1)
    return success();

  auto numPartitions = partitions.getNumPartitions();
  auto defaultPartition = partitions.getPartition((int)0);
  auto loopVarCategories = classifyLoopVars(loop, defaultPartition, partitions);
  auto [loopVarIndices, newResultIndices] =
      getLoopVarIndicesToKeep(loop, defaultPartition, loopVarCategories);

  ImplicitLocOpBuilder topBuilder(loop.getLoc(), loop);
  SmallVector<Value> tensorResultAllocs(loop.getRegionIterArgs().size());
  for (auto [i, res] : llvm::enumerate(loop->getResults())) {
    if (loopVarCategories[i] ==
        LoopVarCategory::TensorResultFromOtherPartition) {
      auto ty = cast<RankedTensorType>(res.getType());
      auto memdesc = MemDescType::get(
          ty.getShape(), ty.getElementType(), getSharedEncoding(ty),
          SharedMemorySpaceAttr::get(ty.getContext()), /*mutable=*/true);
      tensorResultAllocs[i] = LocalAllocOp::create(topBuilder, memdesc);
    }
  }

  SmallVector<Type> resultTypes;
  for (auto i : loopVarIndices) {
    resultTypes.push_back(loop->getResultTypes()[i]);
  }

  SmallVector<int32_t> numWarps(numPartitions, lookupNumWarps(loop));
  auto wgOp = nvws::WarpGroupOp::create(topBuilder, resultTypes, numWarps,
                                        numPartitions);

  SmallVector<WarpGroupBuilder> builders;
  for (Region &region : wgOp.getPartitionRegions()) {
    auto partitionId = builders.size();
    auto &block = region.emplaceBlock();
    builders.push_back(WarpGroupBuilder(&block, block.end(), partitionId));
  }

  SmallVector<Operation *> opsToErase;
  for (auto &op_ : *loop->getBlock()) {
    auto op = &op_;
    if (!hasPartition(op))
      continue;
    assert(hasWarpSpecializeTag(op));
    if (*getWarpSpecializeTag(op) != partitions.getTag())
      continue;
    if (op == loop) {
      if (auto forOp = dyn_cast<scf::ForOp>(loop.getOperation())) {
        cloneForOp(forOp, builders, partitions);
      } else {
        cloneWhileOp(cast<scf::WhileOp>(loop.getOperation()), builders,
                     partitions);
      }
      opsToErase.push_back(loop.getOperation());
    } else {
      cloneOp(op, builders, getPartitionIds(op));
      opsToErase.push_back(op);
    }
  }

  for (auto [b, region, partition] : llvm::zip(
           builders, wgOp.getPartitionRegions(), partitions.getPartitions())) {
    if (!llvm::is_contained(getPartitionIds(loop.getOperation()),
                            b.partitionId)) {
      nvws::WarpGroupYieldOp::create(b, wgOp.getLoc(), SmallVector<Value>{});
      continue;
    }
    Operation *newLoop = nullptr;
    if (isa<scf::ForOp>(loop.getOperation())) {
      newLoop = *region.front().getOps<scf::ForOp>().begin();
    } else {
      newLoop = *region.front().getOps<scf::WhileOp>().begin();
    }
    auto outputs = newLoop->getResults();

    if (b.partitionId == 0) {
      nvws::WarpGroupYieldOp::create(b, wgOp.getLoc(), outputs);
    } else {
      // Tensor results computed by non-default partitions are communicated back
      // via SMEM.
      // The calls to getLoopVarIndicesToKeep and isTensorResultComputedBy
      // below are unnecessary if we can encode the partition index and the
      // corresponding result tensor index of the cloned loop in
      // LoopVarCategory::TensorResultFromOtherPartition. In the absence of such
      // language support, we end up computing the same information multiple
      // times.
      auto [_, reverseIndices] =
          getLoopVarIndicesToKeep(loop, &partition, partitions);
      for (size_t i = 0; i < loop.getRegionIterArgs().size(); ++i) {
        if (loopVarCategories[i] ==
                LoopVarCategory::TensorResultFromOtherPartition &&
            isTensorResultComputedBy(loop, i, &partition, partitions)) {
          assert(reverseIndices[i] && "A valid index is expected.");
          auto result = newLoop->getResult(*reverseIndices[i]);
          LocalStoreOp::create(b, wgOp.getLoc(), result, tensorResultAllocs[i]);
        }
      }
      nvws::WarpGroupReturnOp::create(b, wgOp.getLoc());
    }
  }

  topBuilder.setInsertionPointAfter(wgOp);

  for (auto [i, res] : llvm::enumerate(loop->getResults())) {
    if (res.use_empty())
      continue;

    if (loopVarCategories[i] ==
        LoopVarCategory::TensorResultFromOtherPartition) {
      auto ty = cast<RankedTensorType>(loop->getResult(i).getType());
      auto output = LocalLoadOp::create(topBuilder, ty, tensorResultAllocs[i]);
      LocalDeallocOp::create(topBuilder, tensorResultAllocs[i]);
      res.replaceAllUsesWith(output);
    } else if (llvm::any_of(res.getUsers(), [&](Operation *user) {
                 return !hasPartition(user) ||
                        (isa<scf::ForOp, scf::WhileOp>(user) &&
                         hasWarpSpecializeTag(user));
               })) {
      // If some users are in the root partition (no partition attribute) or
      // used by another warp-specialized loop, we need to replace their uses
      // with the corresponding result from the warp group operation
      assert(newResultIndices[i] && "A valid index is expected.");
      res.replaceAllUsesWith(wgOp.getResult(*newResultIndices[i]));
    }
  }

  for (auto op : llvm::reverse(opsToErase))
    op->erase();

  return success();
}

//===----------------------------------------------------------------------===//
// Pass Definition
//===----------------------------------------------------------------------===//

namespace mlir::triton::gpu {
#define GEN_PASS_DEF_TRITONGPUPARTITIONLOOPS
#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc"
} // namespace mlir::triton::gpu

namespace {
struct PartitionLoops
    : triton::gpu::impl::TritonGPUPartitionLoopsBase<PartitionLoops> {
  using TritonGPUPartitionLoopsBase::TritonGPUPartitionLoopsBase;

  void runOnOperation() override;
};
} // namespace

void PartitionLoops::runOnOperation() {
  // Collect loops to warp specialize. This pass expects each loop to already
  // be annotated with partitions.
  SmallVector<LoopLikeOpInterface> loops;
  getOperation().walk([&](LoopLikeOpInterface loop) {
    if (loop->hasAttrOfType<ArrayAttr>(kPartitionStagesAttrName))
      loops.push_back(loop);
  });

  for (LoopLikeOpInterface loop : loops) {
    if (failed(partitionLoop(loop)))
      return signalPassFailure();
  }
}
