/*
 * Copyright (c) 2025 NVIDIA Corporation & Affiliates. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining
 * a copy of this software and associated documentation files
 * (the "Software"), to deal in the Software without restriction,
 * including without limitation the rights to use, copy, modify, merge,
 * publish, distribute, sublicense, and/or sell copies of the Software,
 * and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#include "Utilities.h"
#include "lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionAttrs.h"
#include "mlir/Analysis/TopologicalSortUtils.h"

#include "Utilities.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/AttrTypeSubElements.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/Passes.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/PartitionBuilder.h"
#include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h"
#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
#include "llvm/Support/ErrorHandling.h"

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

#define DEBUG_TYPE "nvws-lower-aref"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")

namespace mlir {
namespace triton {

#define GEN_PASS_DEF_NVWSASSIGNSTAGEPHASE
#include "nvidia/include/Dialect/NVWS/Transforms/Passes.h.inc"
namespace {
template <class T> struct AssignStagePhase {
  struct StagePhase {
    Value stage;
    Value phase;
    Value token;
  };
  Value aref;
  SetVector<int> partitionIds;
  DenseMap<std::pair<Operation *, Value>, int> tokToStagePosMap;

  AssignStagePhase(Value aref, SetVector<int> partitionIds)
      : aref(aref), partitionIds(std::move(partitionIds)) {}

  bool owns(Operation *op) const {
    if (!hasPartition(op))
      return true;
    auto ids = getPartitionIds(op);
    return llvm::all_of(partitionIds, [&](int id) { return ids.contains(id); });
  }

  SetVector<int> getOwnerPartitions() const { return partitionIds; }

  T getTypedOp(Operation *op) {
    if (auto opT = dyn_cast<T>(op)) {
      if (opT.getAref() == aref) {
        if (owns(op))
          return opT;
      }
    }
    return {};
  }
  bool isBufferUsed(ArefBufferOp bufOp, Value token) {
    if (!bufOp)
      return false;
    if (bufOp.getAref() != this->aref)
      return false;
    return token == bufOp.getToken();
  }

  bool analyzeArefUseInBlock(Block *block, Value token) {
    for (auto &op : *block) {
      if (getTypedOp(&op) || isBufferUsed(dyn_cast<ArefBufferOp>(op), token)) {
        return true;
      } else if (auto forOp = dyn_cast<scf::ForOp>(op)) {
        Value newTok;
        if (auto pos = findValuePosInRange(forOp.getInitArgs(), token)) {
          newTok = forOp.getRegionIterArgs()[*pos];
        }
        if (analyzeArefUseInBlock(forOp.getBody(), newTok))
          return true;
      } else if (auto whileOp = dyn_cast<scf::WhileOp>(op)) {
        Value afterTok = token;
        if (auto pos = findValuePosInRange(whileOp.getInits(), token))
          afterTok = whileOp.getAfterArguments()[*pos];
        // The before region computes the condition in the same partitions as
        // the While and does not use arefs, so only inspect the after region.
        if (analyzeArefUseInBlock(&whileOp.getAfter().front(), afterTok))
          return true;
      } else if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
        if (analyzeArefUseInBlock(ifOp.thenBlock(), token))
          return true;
        if (ifOp.elseBlock() && analyzeArefUseInBlock(ifOp.elseBlock(), token))
          return true;
      }
    }
    return false;
  }

  void assignArefIndexInForOp(scf::ForOp forOp, StagePhase &index) {

    Value newTok;
    if (auto pos = findValuePosInRange(forOp.getInitArgs(), index.token)) {
      newTok = forOp.getRegionIterArgs()[*pos];
    }
    // find uses of arefs in forOp body
    if (!analyzeArefUseInBlock(forOp.getBody(), newTok))
      return;

    // add extra iterArgs to the forOp
    SmallVector<Value> extraIterArgs{index.stage, index.phase};
    SmallVector<Value *> arefIndexRefs{&index.stage, &index.phase};
    llvm::MapVector<int, Value *> arefTokenRefs;

    if (auto pos = findValuePosInRange(forOp.getInitArgs(), index.token)) {
      // keep reference of the token position to latest token value
      // we will need it update with the value returned from forOp
      arefTokenRefs[*pos] = &index.token;
      // update token value with iter argument
      index.token = forOp.getRegionIterArgs()[*pos];
    }
    // create new forOp with extra iterArgs
    OpBuilder builder(forOp);
    size_t nArgs = forOp.getRegionIterArgs().size();

    assert(hasPartition(forOp));
    auto forOpIds = getPartitionIds(forOp);
    auto forOpOutputsIds = getPartitionOutputs(forOp);
    forOp = addIterArgsToLoop(builder, forOp, extraIterArgs);

    // update arefIndex with iterArgs in the forOp body
    for (size_t idx = nArgs; idx < forOp.getRegionIterArgs().size(); ++idx)
      *arefIndexRefs[idx - nArgs] = forOp.getRegionIterArgs()[idx];

    // assign arefIndex in the forOp body
    auto indexInBlock = assignArefIndexInBlock(forOp.getBody(), index);

    // update yieldOp to return new indexes
    SmallVector<Value> extraYieldArgs;
    // associate token with stage positional argument in the iterArgs & yieldOp
    // we will need this in propagateStage function that will assign stage
    // to arefBuffer and arefExit ops
    extraYieldArgs.push_back(indexInBlock.stage);
    if (index.phase)
      extraYieldArgs.push_back(indexInBlock.phase);
    appendToForOpYield(forOp, extraYieldArgs);
    tokToStagePosMap[{forOp, index.token}] = nArgs;
    tokToStagePosMap[{forOp.getBody()->getTerminator(), indexInBlock.token}] =
        nArgs;

    // update partitions of the forOp
    for (auto arg : extraYieldArgs) {
      SetVector<int> argIds;
      if (auto defOp = arg.getDefiningOp()) {
        if (defOp->getNumRegions() == 0) {
          // if there is defOp, use partitions of defOp
          assert(hasPartition(defOp));
          argIds = getPartitionIds(defOp);
        } else {
          // if op has region, it returns result, get partition from result
          auto pos = findValuePosInRange(defOp->getResults(), arg);
          argIds = getPartitionOutputs(defOp)[*pos];
        }
      } else {
        // otherwise it is a block-arg, use partitions of users
        for (auto user : arg.getUsers()) {
          if (isa<scf::YieldOp>(user))
            continue;
          assert(hasPartition(user));
          auto ids = getPartitionIds(user);
          argIds.insert(ids.begin(), ids.end());
        }
      }
      // Unchanged stage/phase inherits this aref put/get group's owners.
      if (argIds.empty())
        argIds = getOwnerPartitions();
      forOpIds.insert(argIds.begin(), argIds.end());
      forOpOutputsIds.push_back(argIds);
    }
    setPartition(forOp, forOpIds);
    setPartitionOutputs(forOp, forOpOutputsIds);

    // update arefIndex with results from newForOp
    for (size_t idx = nArgs; idx < forOp.getRegionIterArgs().size(); ++idx)
      *arefIndexRefs[idx - nArgs] = forOp.getResult(idx);
    for (auto [idx, arefTokenRef] : arefTokenRefs)
      *arefTokenRef = forOp.getResult(idx);
  }

  void assignArefIndexInWhileOp(scf::WhileOp whileOp, StagePhase &index) {

    Value newTok;
    if (auto pos = findValuePosInRange(whileOp.getInits(), index.token))
      newTok = whileOp.getAfterArguments()[*pos];
    // find uses of arefs in the While body, only inspect after region
    if (!analyzeArefUseInBlock(&whileOp.getAfter().front(), newTok))
      return;

    // add extra iterArgs to the While
    SmallVector<Value> extraIterArgs{index.stage, index.phase};
    SmallVector<Value *> arefIndexRefs{&index.stage, &index.phase};
    llvm::MapVector<int, Value *> arefTokenRefs;
    if (auto pos = findValuePosInRange(whileOp.getInits(), index.token))
      arefTokenRefs[*pos] = &index.token;

    // create new While with extra iterArgs
    OpBuilder builder(whileOp);
    auto nArgs = whileOp.getNumResults();
    auto whileOpIds = getPartitionIds(whileOp);
    auto whileOpOutputsIds = getPartitionOutputs(whileOp);
    auto oldWhileOp = whileOp;
    whileOp = replaceWhileOpWithNewSignature(
        builder, whileOp, extraIterArgs,
        TypeRange{index.stage.getType(), index.phase.getType()});
    oldWhileOp.erase();

    // forward the new state from the before region to the body
    whileOp.getConditionOp().getArgsMutable().append(
        whileOp.getBeforeArguments().drop_front(nArgs));

    // update arefIndex with iterArgs in the While body
    for (size_t idx = nArgs; idx < whileOp.getAfterArguments().size(); ++idx)
      *arefIndexRefs[idx - nArgs] = whileOp.getAfterArguments()[idx];
    for (auto [idx, arefTokenRef] : arefTokenRefs)
      *arefTokenRef = whileOp.getAfterArguments()[idx];

    // assign arefIndex in the While body
    auto indexInBlock =
        assignArefIndexInBlock(&whileOp.getAfter().front(), index);

    // update yieldOp to return new indexes
    auto yield = whileOp.getYieldOp();
    SmallVector<Value> extraYieldArgs{indexInBlock.stage, indexInBlock.phase};
    yield.getResultsMutable().append(extraYieldArgs);
    for (auto [idx, _] : arefTokenRefs) {
      if (yield.getOperand(idx) == indexInBlock.token)
        tokToStagePosMap[{yield, indexInBlock.token}] = nArgs;
    }

    // update partitions of the While
    auto ownerIds = getOwnerPartitions();
    whileOpIds.insert(ownerIds.begin(), ownerIds.end());
    whileOpOutputsIds.resize(whileOpOutputsIds.size() + extraYieldArgs.size(),
                             ownerIds);
    setPartition(whileOp, whileOpIds);
    setPartitionOutputs(whileOp, whileOpOutputsIds);

    // update arefIndex with results from new While
    for (size_t idx = nArgs; idx < whileOp.getNumResults(); ++idx)
      *arefIndexRefs[idx - nArgs] = whileOp.getResult(idx);
    for (auto [idx, arefTokenRef] : arefTokenRefs)
      *arefTokenRef = whileOp.getResult(idx);
  }

  void assignArefIndexInIfOp(scf::IfOp ifOp, StagePhase &index) {

    auto useInThenBlock = analyzeArefUseInBlock(ifOp.thenBlock(), index.token);
    auto useInElseBlock =
        ifOp.elseBlock() ? analyzeArefUseInBlock(ifOp.elseBlock(), index.token)
                         : false;
    if (!useInThenBlock && !useInElseBlock)
      return;

    // add extra results to the ifOp
    SmallVector<Type> extraIfResults{index.stage.getType(),
                                     index.phase.getType()};
    SmallVector<Value *> arefIndexRefs{&index.stage, &index.phase};

    // create new ifOp with extra results
    OpBuilder builder(ifOp);
    size_t nArgs = ifOp.getResults().size();
    auto newIfOp = replaceIfOpWithNewSignature(builder, ifOp, extraIfResults);

    // assign arefIndex in then-body
    auto thenIndex = assignArefIndexInBlock(newIfOp.thenBlock(), index);

    // assign arefIndex in else-body
    auto elseIndex = newIfOp.elseBlock()
                         ? assignArefIndexInBlock(newIfOp.elseBlock(), index)
                         : index;

    // update yieldOp to return new indexes
    auto thenYieldOp = newIfOp.thenYield();
    auto elseYieldOp = newIfOp.elseYield();
    // insert new indexes to the yieldOp
    llvm::MapVector<int, Value *> arefTokenRefs;

    // find token pos in yieldOp and make a reference to  arefIndexMap value
    if (auto pos =
            findValuePosInRange(thenYieldOp->getOperands(), index.token)) {
      arefTokenRefs[*pos] = &index.token;
    }
    if (auto pos =
            findValuePosInRange(elseYieldOp->getOperands(), index.token)) {
      arefTokenRefs[*pos] = &index.token;
    }

    tokToStagePosMap[{newIfOp.thenYield(), thenIndex.token}] =
        thenYieldOp.getNumOperands();
    tokToStagePosMap[{newIfOp.elseYield(), elseIndex.token}] =
        elseYieldOp.getNumOperands();
    thenYieldOp->insertOperands(thenYieldOp.getNumOperands(), thenIndex.stage);
    elseYieldOp->insertOperands(elseYieldOp.getNumOperands(), elseIndex.stage);
    thenYieldOp->insertOperands(thenYieldOp.getNumOperands(), thenIndex.phase);
    elseYieldOp->insertOperands(elseYieldOp.getNumOperands(), elseIndex.phase);

    assert(hasPartition(ifOp));
    auto ifOpIds = getPartitionIds(ifOp);
    auto ifOpOutputsIds = getPartitionOutputs(ifOp);
    ifOp.erase();

    SetVector<int> stageIds;
    // at least one of the then/else block must have producing op
    for (auto arg : {thenIndex.stage, elseIndex.stage}) {
      if (auto defOp = arg.getDefiningOp()) {
        auto argIds = getPartitionIds(defOp);
        stageIds.insert(argIds.begin(), argIds.end());
      }
    }
    SetVector<int> phaseIds;
    for (auto arg : {thenIndex.phase, elseIndex.phase}) {
      if (auto defOp = arg.getDefiningOp()) {
        auto argIds = getPartitionIds(defOp);
        phaseIds.insert(argIds.begin(), argIds.end());
      }
    }
    ifOpOutputsIds.push_back(stageIds);
    ifOpOutputsIds.push_back(phaseIds);
    setPartition(newIfOp, ifOpIds);
    setPartitionOutputs(newIfOp, ifOpOutputsIds);

    // update arefIndex with results from newIfOp
    for (size_t idx = nArgs; idx < newIfOp.getResults().size(); ++idx)
      *arefIndexRefs[idx - nArgs] = newIfOp.getResult(idx);
    for (auto [idx, arefTokenRef] : arefTokenRefs)
      *arefTokenRef = newIfOp.getResult(idx);
  }

  StagePhase assignArefIndexInBlock(Block *block, StagePhase index) {
    for (auto &op : llvm::make_early_inc_range(*block)) {
      if (auto opT = getTypedOp(&op)) {
        ImplicitLocOpBuilder builder(opT.getLoc(), opT);
        std::optional<SetVector<int>> partitionIds;
        if (hasPartition(&op))
          partitionIds = getPartitionIds(&op);
        auto wsTag = getWarpSpecializeTag(&op);
        auto stageCluster = getStageCluster(&op);

        auto createInto = [&](auto opTy, auto... args) {
          using ty = decltype(opTy);
          auto op = triton::gpu::createInto<ty>(
              builder, builder.getLoc(), partitionIds, stageCluster,
              std::forward<decltype(args)>(args)...);
          if (wsTag)
            setWarpSpecializeTag(op, *wsTag);
          return op;
        };

        auto nextStage = createInto(arith::AddIOp{}, index.stage,
                                    createInto(arith::ConstantIntOp{}, 1, 32));
        auto arefBuf = opT.getAref()
                           .template getDefiningOp<nvws::ArefCreateOp>()
                           .getOperand(0);
        auto depth = getArefDepth(cast<MemDescType>(arefBuf.getType()));

        auto cnd =
            createInto(arith::CmpIOp{}, arith::CmpIPredicate::eq, nextStage,
                       createInto(arith::ConstantIntOp{}, depth, 32));
        auto zero = createInto(arith::ConstantIntOp{}, 0, 32);
        index.stage = createInto(arith::SelectOp{}, cnd, zero, nextStage);

        auto nextPhase = createInto(arith::XOrIOp{}, index.phase,
                                    createInto(arith::ConstantIntOp{}, 1, 32));
        index.phase =
            createInto(arith::SelectOp{}, cnd, nextPhase, index.phase);

        index.token = opT.getToken();
        opT.getStageMutable().assign(index.stage);
        opT.getPhaseMutable().assign(index.phase);
      } else if (auto forOp = dyn_cast<scf::ForOp>(op)) {
        assignArefIndexInForOp(forOp, index);
      } else if (auto whileOp = dyn_cast<scf::WhileOp>(op)) {
        assignArefIndexInWhileOp(whileOp, index);
      } else if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
        assignArefIndexInIfOp(ifOp, index);
      }
    }

    return index;
  }

  void propagateStage(Value token, Value stage,
                      DenseSet<Operation *> &visited) {
    for (auto &tokUse : token.getUses()) {
      auto owner = tokUse.getOwner();
      if (visited.contains(owner))
        continue;
      visited.insert(owner);
      if (auto stageOp = dyn_cast<ArefStageInterface>(owner)) {
        if (auto blk = dyn_cast<BlockArgument>(stage)) {
          assert(hasPartition(stageOp));
          auto stageOpIds = getPartitionIds(stageOp);
          auto forOp = cast<scf::ForOp>(blk.getOwner()->getParentOp());
          auto pos = findValuePosInRange(forOp.getRegionIterArgs(), stage);
          assert(pos);

          // update op partitions
          assert(hasPartition(forOp));
          auto forOpIds = getPartitionIds(forOp);
          forOpIds.insert(stageOpIds.begin(), stageOpIds.end());
          setPartition(forOp, forOpIds);

          auto forOpOutputsIds = getPartitionOutputs(forOp);
          forOpOutputsIds[*pos + 0].insert(stageOpIds.begin(),
                                           stageOpIds.end());
          forOpOutputsIds[*pos + 1].insert(stageOpIds.begin(),
                                           stageOpIds.end());
          setPartitionOutputs(forOp, forOpOutputsIds);
        }
        stageOp.setStage(stage);
      } else if (auto forOp = dyn_cast<scf::ForOp>(owner)) {
        auto tokPos = tokUse.getOperandNumber() - forOp.getNumControlOperands();
        auto iterTok = forOp.getRegionIterArg(tokPos);
        auto stagePos = tokToStagePosMap.at({forOp, iterTok});
        propagateStage(iterTok, forOp.getRegionIterArgs()[stagePos], visited);
      } else if (auto whileOp = dyn_cast<scf::WhileOp>(owner)) {
        auto tokPos = tokUse.getOperandNumber();
        auto stagePos = *findValuePosInRange(whileOp.getInits(), stage);
        propagateStageAcrossWhile(whileOp, tokPos, stagePos, visited);
      } else if (auto yieldOp = dyn_cast<scf::YieldOp>(owner)) {
        auto tokPos = tokUse.getOperandNumber();
        auto stagePos = tokToStagePosMap.at({yieldOp, token});
        auto parentOp = yieldOp->getParentOp();
        if (auto whileOp = dyn_cast<scf::WhileOp>(parentOp)) {
          propagateStageAcrossWhile(whileOp, tokPos, stagePos, visited);
        } else {
          propagateStage(parentOp->getResult(tokPos),
                         parentOp->getResult(stagePos), visited);
        }
      }
    }
  }

  void propagateStageAcrossWhile(scf::WhileOp whileOp, unsigned tokenPos,
                                 unsigned stagePos,
                                 DenseSet<Operation *> &visited) {
    propagateStage(whileOp.getAfterArguments()[tokenPos],
                   whileOp.getAfterArguments()[stagePos], visited);
    propagateStage(whileOp.getResult(tokenPos), whileOp.getResult(stagePos),
                   visited);
  }

  static LogicalResult run(ArefCreateOp arefOp) {
    SmallVector<SetVector<int>> ownerGroups;
    for (auto user : arefOp->getUsers()) {
      // Track stage/phase once per distinct aref owner group.
      if (!isa<T>(user) || !hasPartition(user))
        continue;
      auto ids = getPartitionIds(user);
      bool exists = llvm::any_of(ownerGroups, [&](const SetVector<int> &group) {
        return ids.size() == group.size() &&
               llvm::all_of(ids, [&](int id) { return group.contains(id); });
      });
      if (!exists)
        ownerGroups.push_back(std::move(ids));
    }
    if (ownerGroups.empty()) {
      // No partition metadata means the aref is outside ttg.ws. Use one dummy
      // owner so it still receives a valid stage and phase.
      SetVector<int> singleton;
      singleton.insert(0);
      ownerGroups.push_back(std::move(singleton));
    }
    llvm::sort(ownerGroups, [](const auto &lhs, const auto &rhs) {
      return lhs.getArrayRef() < rhs.getArrayRef();
    });

    // initialize indexes
    StagePhase index;
    ImplicitLocOpBuilder b(arefOp.getLoc(), arefOp);
    b.setInsertionPointAfter(arefOp);
    auto depth =
        getArefDepth(cast<MemDescType>(arefOp.getOperand(0).getType()));
    index.stage = arith::ConstantIntOp::create(b, depth - 1, 32);

    static_assert(std::is_same_v<T, ArefPutEnterOp> ||
                      std::is_same_v<T, ArefGetEnterOp>,
                  "ArefPutEnterOp or ArefGetEnterOp expected");
    auto initPhase = std::is_same_v<T, ArefPutEnterOp> ? 0 : 1;
    index.phase = arith::ConstantIntOp::create(b, initPhase, 32);

    for (const SetVector<int> &ownerGroup : ownerGroups) {
      // Assign stage/phase to enter/exit Ops in partitions aref is used
      AssignStagePhase arefIndex(arefOp.getResult(), ownerGroup);

      // Assign stage/phase to enter operations.
      arefIndex.assignArefIndexInBlock(arefOp->getBlock(), index);

      // Propagate stage to exit operations following each enter token.
      for (auto user : arefOp->getUsers()) {
        auto enterOp = dyn_cast<T>(user);
        if (!enterOp || !arefIndex.owns(enterOp))
          continue;
        DenseSet<Operation *> visited;
        arefIndex.propagateStage(enterOp.getToken(), enterOp.getStage(),
                                 visited);
      }
    }

    return success();
  }
};

void updateOutputWithDefaultPartition(Operation *op, int pos) {
  auto opIds = getPartitionIds(op);
  opIds.insert(0);
  setPartition(op, opIds);

  auto opOutputsIds = getPartitionOutputs(op);
  opOutputsIds[pos].insert(0);
  setPartitionOutputs(op, opOutputsIds);
}

void visitBackwardSlice(scf::ForOp wsLoop, Value value,
                        std::function<void(Operation *)> callback,
                        DenseSet<Value> &visited) {
  if (!visited.insert(value).second)
    return;

  if (auto blockArg = dyn_cast<BlockArgument>(value)) {
    if (auto forOp = dyn_cast<scf::ForOp>(blockArg.getOwner()->getParentOp())) {
      if (forOp->hasAttr(kWarpSpecializeAttrName))
        return;
      auto pos = findValuePosInRange(forOp.getRegionIterArgs(), value);
      assert(pos);
      visitBackwardSlice(wsLoop, forOp.getInitArgs()[*pos], callback, visited);
    }
  } else if (auto defOp = value.getDefiningOp();
             isa<scf::IfOp, scf::ForOp>(defOp)) {
    auto pos = findValuePosInRange(defOp->getResults(), value);
    assert(pos);
    updateOutputWithDefaultPartition(defOp, *pos);
    if (auto ifOp = dyn_cast<scf::IfOp>(defOp)) {
      visitBackwardSlice(wsLoop, ifOp.thenYield()->getOperand(*pos), callback,
                         visited);
      if (ifOp.elseBlock())
        visitBackwardSlice(wsLoop, ifOp.elseYield()->getOperand(*pos), callback,
                           visited);
      visitBackwardSlice(wsLoop, ifOp.getCondition(), callback, visited);
    } else {
      auto forOp = cast<scf::ForOp>(defOp);
      visitBackwardSlice(wsLoop,
                         forOp.getBody()->getTerminator()->getOperand(*pos),
                         callback, visited);
      // visit control operands of for-op
      for (int idx = 0; idx < forOp.getNumControlOperands(); ++idx) {
        auto control = forOp.getOperand(idx);
        visitBackwardSlice(wsLoop, control, callback, visited);
      }
    }
  } else if (wsLoop.getBody()->findAncestorOpInBlock(*defOp)) {
    callback(defOp);
    for (auto operand : defOp->getOperands()) {
      visitBackwardSlice(wsLoop, operand, callback, visited);
    }
  }
}

LogicalResult assignStagePhase(triton::FuncOp funcOp) {
  SmallVector<ArefCreateOp> arefOps;
  funcOp.walk([&](ArefCreateOp arefOp) { arefOps.push_back(arefOp); });
  for (auto arefOp : arefOps) {
    if (failed(AssignStagePhase<ArefPutEnterOp>::run(arefOp)))
      return failure();
    if (failed(AssignStagePhase<ArefGetEnterOp>::run(arefOp)))
      return failure();
  }

  auto callback = [&](Operation *op) {
    if (!isa<scf::YieldOp, scf::IfOp, scf::ForOp, triton::ReduceOp>(op)) {
      assert(hasPartition(op));
      auto partitionIds = getPartitionIds(op);
      partitionIds.insert(0);
      setPartition(op, partitionIds);
    }
  };

  funcOp.walk([&](scf::ForOp forOp) {
    DenseSet<Value> visited;
    if (forOp->hasAttr(kWarpSpecializeAttrName)) {
      for (auto result : forOp.getResults()) {
        // if result is of scalar type and is used outside of for-op, visit
        // all dependencies and assign default partition to them
        if (isa<IntegerType, FloatType>(result.getType()) &&
            !result.use_empty()) {
          auto arg = forOp.getBody()->getTerminator()->getOperand(
              result.getResultNumber());
          // Check if any users of this scalar result lack ttg.partition, or if
          // it is used in another warp-specialized loop. If so, the scalar is
          // consumed by the root partition outside the warp-specialized loop,
          // requiring us to assign the default partition to all operations that
          // compute this result.
          bool assignDefaultPartition =
              llvm::any_of(result.getUsers(), [&](Operation *user) {
                return !hasPartition(user) ||
                       (isa<scf::ForOp>(user) && hasWarpSpecializeTag(user));
              });
          if (assignDefaultPartition) {
            updateOutputWithDefaultPartition(forOp, result.getResultNumber());
            visitBackwardSlice(forOp, arg, callback, visited);
          }
        }
      }
    }
  });
  return success();
}

// ----------------------------------------------------------------------------

} // anonymous namespace

class NVWSAssignStagePhase
    : public impl::NVWSAssignStagePhaseBase<NVWSAssignStagePhase> {
public:
  void runOnOperation() override {
    mlir::ModuleOp m = getOperation();

    m.walk([&](triton::FuncOp funcOp) {
      if (failed(assignStagePhase(funcOp)))
        signalPassFailure();
    });
  }
}; // namespace triton

} // namespace triton
} // namespace mlir
