#ifndef TRITON_TRITONGPU_TRANSFORM_PIPELINE_WARPSPECIALIZATION_H_
#define TRITON_TRITONGPU_TRANSFORM_PIPELINE_WARPSPECIALIZATION_H_

#include "mlir/Support/LogicalResult.h"

namespace mlir {
class LoopLikeOpInterface;
namespace scf {
class ForOp;
} // namespace scf
namespace triton::gpu {
// This is the final step to prepare a loop for warp specialization. This takes
// a loop with a partition schedule and rewrites the loop such that all SSA
// dependencies between partitions are passed through shared memory and
// multibuffers them according to partition stages.
LogicalResult rewritePartitionDependencies(scf::ForOp &loop);
// Given a loop where the partitions' inputs and outputs have been fully
// rewritten to be reference semantic, partitiong the loop into a
// `ttg.warp_specialize` by duplicating the loop for each partition and
// rematerializing, as necessary, operations in the root partition.
LogicalResult partitionLoop(LoopLikeOpInterface loop);
} // namespace triton::gpu
} // namespace mlir

#endif // TRITON_TRITONGPU_TRANSFORM_PIPELINE_WARPSPECIALIZATION_H_
