/**
 * Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See file LICENSE for terms.
 */

#include "config.h"
#include "allreduce.h"
#include "tl_ucp.h"
#include "tl_ucp_coll.h"
#include "core/ucc_progress_queue.h"
#include "utils/ucc_math.h"
#include "utils/ucc_coll_utils.h"
#include "../reduce_scatter/reduce_scatter.h"
#include "../allgather/allgather.h"

/*
 * Ring Allreduce - Implements allreduce as a composition of:
 *   1. Reduce-scatter ring: Each rank reduces and scatters data
 *   2. Allgather ring: Collects the scattered results to all ranks
 *
 * This approach uses UCC schedules to chain the two operations,
 * providing a clean, composable implementation of ring allreduce.
 */

static ucc_status_t
ucc_tl_ucp_allreduce_ring_start(ucc_coll_task_t *task)
{
    ucc_schedule_t  *schedule = ucc_derived_of(task, ucc_schedule_t);
    ucc_coll_args_t *args     = &schedule->super.bargs.args;
    size_t           count;
    size_t           data_size;
    ucc_status_t     status;

    UCC_TL_UCP_PROFILE_REQUEST_EVENT(task, "ucp_allreduce_ring_start", 0);

    /* Handle in-place vs out-of-place cases */
    if (UCC_IS_INPLACE(*args)) {
        /* Case 1: Allreduce is in-place (src == dst)
         * Data is already in dst buffer, no copy needed.
         * Reduce-scatter and allgather will both operate in-place on dst. */
    } else {
        /* Case 2: Allreduce is NOT in-place (src != dst)
         * Copy src to dst before starting reduce-scatter.
         *
         * Both sub-collectives run in-place on dst so each rank's reduced
         * segment lands at rank * (count/tsize).  The source data must
         * already reside in dst before the pipeline starts.
         *
         * Performance trade-off: this adds one full-buffer memory pass per
         * out-of-place call.  On large counts or high-latency memory types
         * (e.g. CUDA managed memory) the cost is non-trivial, but it avoids
         * threading src through both sub-collectives' ring logic while still
         * producing a correct result. */
        count     = args->dst.info.count;
        data_size = count * ucc_dt_size(args->dst.info.datatype);
        status    = ucc_mc_memcpy(args->dst.info.buffer,
                                  args->src.info.buffer,
                                  data_size,
                                  args->dst.info.mem_type,
                                  args->src.info.mem_type);
        if (ucc_unlikely(status != UCC_OK)) {
            return status;
        }
    }
    return ucc_schedule_start(task);
}

static ucc_status_t
ucc_tl_ucp_allreduce_ring_finalize(ucc_coll_task_t *task)
{
    ucc_schedule_t *schedule = ucc_derived_of(task, ucc_schedule_t);
    ucc_status_t    status;

    UCC_TL_UCP_PROFILE_REQUEST_EVENT(task, "ucp_allreduce_ring_done", 0);
    status = ucc_schedule_finalize(task);
    ucc_tl_ucp_put_schedule(schedule);
    return status;
}

ucc_status_t ucc_tl_ucp_allreduce_ring_init(ucc_base_coll_args_t *coll_args,
                                            ucc_base_team_t *team,
                                            ucc_coll_task_t **task_h)
{
    ucc_tl_ucp_team_t     *tl_team = ucc_derived_of(team, ucc_tl_ucp_team_t);
    ucc_rank_t             tsize   = UCC_TL_TEAM_SIZE(tl_team);
    size_t                 count   = coll_args->args.dst.info.count;
    ucc_schedule_t        *schedule;
    ucc_coll_task_t       *rs_task;
    ucc_coll_task_t       *ag_task;
    ucc_base_coll_args_t   rs_args;
    ucc_base_coll_args_t   ag_args;
    ucc_status_t           status;

    /* Check for predefined datatypes */
    if (!ucc_coll_args_is_predefined_dt(&coll_args->args, UCC_RANK_INVALID)) {
        tl_error(team->context->lib,
                 "user defined datatype is not supported");
        return UCC_ERR_NOT_SUPPORTED;
    }

    if (UCC_TL_UCP_TEAM_LIB(tl_team)->cfg.reduce_avg_pre_op &&
        coll_args->args.op == UCC_OP_AVG) {
        return UCC_ERR_NOT_SUPPORTED;
    }

    /* Check that count is divisible by team size for ring algorithm */
    if (count % tsize != 0) {
        tl_debug(team->context->lib,
                 "ring requires count (%zu) divisible by team size (%u)",
                 count, tsize);
        return UCC_ERR_NOT_SUPPORTED;
    }

    /* Check for same memory types (required by allreduce) */
    ALLREDUCE_TASK_CHECK(coll_args->args, tl_team);

    /* Allocate schedule */
    status = ucc_tl_ucp_get_schedule(tl_team, coll_args,
                                     (ucc_tl_ucp_schedule_t **)&schedule);
    if (ucc_unlikely(UCC_OK != status)) {
        goto out;
    }

    /*
     * Step 1: Reduce-Scatter Ring
     */
    rs_args                     = *coll_args;
    rs_args.args.coll_type      = UCC_COLL_TYPE_REDUCE_SCATTER;
    rs_args.args.dst.info.count = count;
    rs_args.args.mask          |= UCC_COLL_ARGS_FIELD_FLAGS;
    rs_args.args.flags         |= UCC_COLL_ARGS_FLAG_IN_PLACE;


    status = ucc_tl_ucp_reduce_scatter_ring_init(&rs_args, team, &rs_task);
    if (ucc_unlikely(UCC_OK != status)) {
        tl_error(team->context->lib, "failed to init reduce_scatter_ring task");
        goto err_schedule;
    }

    /*
     * Step 2: Allgather Ring
     */
    ag_args                      = *coll_args;
    ag_args.args.coll_type       = UCC_COLL_TYPE_ALLGATHER;
    ag_args.args.src.info.buffer = NULL;
    ag_args.args.src.info.count  = 0;
    ag_args.args.dst.info.buffer = coll_args->args.dst.info.buffer;
    ag_args.args.dst.info.count  = count;
    ag_args.args.mask           |= UCC_COLL_ARGS_FIELD_FLAGS;
    ag_args.args.flags          |= UCC_COLL_ARGS_FLAG_IN_PLACE;

    status = ucc_tl_ucp_allgather_ring_init(&ag_args, team, &ag_task);
    if (ucc_unlikely(UCC_OK != status)) {
        tl_error(team->context->lib, "failed to init allgather_ring task");
        goto err_rs;
    }

    /* ucc_schedule_add_task always links the task into schedule->tasks[]
     * before returning, even on error.  Jump targets must therefore reflect
     * the state *after* the call, not only on success. */
    UCC_CHECK_GOTO(ucc_schedule_add_task(schedule, rs_task), err_ag_rs_owned,
                   status);
    UCC_CHECK_GOTO(
        ucc_event_manager_subscribe(
            &schedule->super,
            UCC_EVENT_SCHEDULE_STARTED,
            rs_task,
            ucc_task_start_handler),
        err_ag_rs_owned,
        status);

    /* Same rule for ag_task: once add_task is called it is in the schedule. */
    UCC_CHECK_GOTO(ucc_schedule_add_task(schedule, ag_task), err_schedule,
                   status);
    UCC_CHECK_GOTO(
        ucc_event_manager_subscribe(
            rs_task, UCC_EVENT_COMPLETED, ag_task, ucc_task_start_handler),
        err_schedule,
        status);

    schedule->super.post     = ucc_tl_ucp_allreduce_ring_start;
    schedule->super.finalize = ucc_tl_ucp_allreduce_ring_finalize;

    *task_h = &schedule->super;
    return UCC_OK;

err_ag_rs_owned:
    /* rs_task is in the schedule; ucc_schedule_finalize will clean it up */
    ag_task->finalize(ag_task);
    goto err_schedule;
err_rs:
    rs_task->finalize(rs_task);
err_schedule:
    ucc_schedule_finalize(&schedule->super);
    ucc_tl_ucp_put_schedule(schedule);
out:
    return status;
}
