/**
 * SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
 * Copyright (c) 2015-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * See file LICENSE for terms.
 */

#include <bits/posix1_lim.h>
#include <coll_test.h>
#include <ctype.h>
#include <dlfcn.h>
#include <errno.h>
#include <sys/time.h>
#include "api/sharp.h"

#define MAX_IB_DEVICES (16)
#define DEBUG          (0)

#define SHARP_KBYTE (1ull << 10)
#define SHARP_MBYTE (1ull << 20)
#define SHARP_GBYTE (1ull << 30)
#define SHARP_TBYTE (1ull << 40)
#define SHARP_PBYTE (1ull << 50)

#define LOAD_SYM(handle, symbol, funcptr)                                  \
    do {                                                                   \
        cast = (void**)&funcptr;                                           \
        tmp = dlsym(handle, symbol);                                       \
        if (tmp == NULL) {                                                 \
            fprintf(stderr, "dlsym failed on %s - %s", symbol, dlerror()); \
            goto teardown;                                                 \
        }                                                                  \
        *cast = tmp;                                                       \
    } while (0)

int client_id;
int enable_sharp_coll = 1;
int sharp_world_rank;
int sharp_world_size;
int perf_test_iterations = 100;
int perf_test_skips = 10;
int perf_with_barrier = 1;
int nbc_count = 16;
int coll_root_rank = 0;
int test_flat_ppn = 1;
int use_cuda_dmabuf_flag = 0;
int use_data_direct_flag = 0;
int sharp_test_group_type = COMM_TYPE_WORLD;
int sharp_test_num_splits = 2; /* default odd even comm split */
coll_sharp_component_t coll_sharp_component;
char ib_dev_list[MAX_IB_DEVICES][16];
int num_ib_devices;
#ifdef MPI_SUPPORT
extern MPI_Datatype sharp_test_mpi_min_max_datatype[SHARP_DTYPE_NULL][SHARP_DTYPE_NULL];
#endif

const char* comm_type_names[] = {
    [COMM_TYPE_WORLD] = "Comm World",
    [COMM_TYPE_WORLD_DUP] = "Comm World Dup",
    [COMM_TYPE_WORLD_REVERSE] = "Comm World Reverse",
    [COMM_TYPE_N_SPLIT] = "Comm split (default: Odd/even)",
    [COMM_TYPE_N_SPLIT_REVERSE] = "Comm Split Reverse",
    [COMM_TYPE_HALF] = "Comm Split Half",
    [COMM_TYPE_PPN_COMM] = "Comm per PPN Rank",
    [COMM_TYPE_PPN_JOBS] = "Sharp job and Comm per PPN Rank",
};

static struct option long_options[] = {{"help", no_argument, NULL, 'h'},
                                       {"perf_with_barrier", required_argument, NULL, 'B'},
                                       {"collectives", required_argument, NULL, 'c'},
                                       {"rs_ag_mode", required_argument, NULL, 'R'},
                                       {"iov_count", required_argument, NULL, 'C'},
                                       {"ib-dev", required_argument, NULL, 'd'},
                                       {"data_layout", required_argument, NULL, 'D'},
                                       {"iters", required_argument, NULL, 'i'},
                                       {"jobid", required_argument, NULL, 'j'},
                                       {"host_alloc_type", required_argument, NULL, 'H'},
                                       {"mode", required_argument, NULL, 'm'},
                                       {"mem_type", required_argument, NULL, 'M'},
                                       {"nbc_count", required_argument, NULL, 'N'},
                                       {"flat_ppn", required_argument, NULL, 'p'},
                                       {"size", required_argument, NULL, 's'},
                                       {"root", required_argument, NULL, 'r'},
                                       {"datatype", required_argument, NULL, 'T'},
                                       {"group_type", required_argument, NULL, 't'},
                                       {"skips", required_argument, NULL, 'x'},
                                       {"stepfactor", required_argument, NULL, 'f'},
                                       {"use_cuda_dmabuf", required_argument, &use_cuda_dmabuf_flag, 1},
                                       {"use_data_direct", required_argument, &use_data_direct_flag, 1},
                                       {0, 0, 0, 0}};

static const char usage_string[] =
    "Usage:   sharp_coll_test[_mpi]  [OPTIONS]\n"
    "Options:\n"
    "-h, --help              Show this help message and exit\n"
    "-d, --ib-dev            Use IB device <dev:port> \n"
    "-j, --jobid             Explicit Job ID (Default: Random ID), valid range is between 1 and (2^63)-2\n"
    "-i, --iters             Number of iterations to run perf benchmark (Default: 100)\n"
    "-x, --skips             Number of warmup iterations to run perf benchmark (Default: 10)\n"
    "-M, --mem_type          Memory type(host,cuda,null) used in the communication buffers format: <src memtype>:<recv memtype> \n"
    "                        (Dafault: host cpu memory on non-GPU setup. cuda memory for GPU setup)\n"
    "--use_cuda_dmabuf       Use CUDA DMA-BUF for GPUDirect RDMA testing with CUDA memory\n"
    "--use_data_direct       Use Data-Direct CUDA DMA-BUF for GPUDirect RDMA testing\n"
#ifdef MPI_SUPPORT
    "-B, --perf_with_barrier Sync Allreduce with Barrier in Allreduce collective, Default:1 \n"
    "-c, --collectives       Comma separated list of collectives:[allreduce|iallreduce|reduce|ireduce|barrier|ibarrier|bcast|reduce_scatter|allgather|rs_ag|all|none] to run.\n"
    "                           Default: Allreduce\n"
    "-R, --rs_ag_mode        reduce-scatter and allgather chaining mode: seq|ovl (sequential or overlapped). Default: seq\n"
    "-H, --host_alloc_type   Host memory allocation method ( hugetlb, hugethp(MADVISE Transparent Hugepage support), malloc)\n"
    "-m, --mode              Test modes: <basic|complex|perf|all> . Default: perf\n"
    "-N, --nbc_count         Number of non-blocking operation posted before waiting for completion (max: 512)\n"
    "-p, --flat_ppn          Number of process per node default:1 \n"
    "-r, --root              Root for bcast, reduce operations \n"
    "-T, --datatype          Datatype of SHARP aggregation request(half_int, half_float, int, float, long, double) \n"
    "-t, --group_type        Set specific group type(world(default), world-dup, world-rev, half, split:<n>, split-rev:<n>, split:ppn-comm, split:ppn-jobs) to test\n"
#endif
    "-s, --size              Set the minimum and/or the maximum message size. format:[MIN:]MAX Default:<4:32M>\n"
    "-f, --stepfactor        increment factor. multiplication factor between sizes. Set 1 to run test infinite loop. Default : 2";

sharp_conf_t sharp_conf = {
    .rank = 0,
    .size = 1,
    .jobid = 0,
    .test_mode = TEST_PERF,
    .sdata_layout = TEST_DATA_LAYOUT_CONTIG,
    .siov_count = 2,
    .rdata_layout = TEST_DATA_LAYOUT_CONTIG,
    .host_allocator_type = USE_MALLOC,
    .riov_count = 2,
    .run_colls = TEST_ALLREDUCE,
    .min_message_size = 4,
    .max_message_size = 33554432,
    .stepfactor = 2,
    .s_mem_type = SHARP_MEM_TYPE_HOST,
    .r_mem_type = SHARP_MEM_TYPE_HOST,
    .datatype = SHARP_DTYPE_INT,
    .run_ppn_comm_parallel = 0,
    .use_dma_buf = 1,
    .use_direct_nic = 0,
    .cuda_free = NULL,
    .cuda_malloc = NULL,
    .cuda_set_device = NULL,
    .cuda_get_device_count = NULL,
    .cuda_memcpy = NULL,
    .cuMemAlloc_ptr = NULL,
    .cuMemGetHandleForAddressRange_ptr = NULL,
    .cudart_lib = NULL,
    .cuda_lib = NULL,
};
int oob_bcast(void* comm_context, void* buf, int size, int root)
{
#ifdef MPI_SUPPORT
    MPI_Comm comm = (MPI_Comm)comm_context;
    MPI_Bcast(buf, size, MPI_BYTE, root, comm);
#endif
    return 0;
}

int oob_barrier(void* comm_context)
{
#ifdef MPI_SUPPORT
    MPI_Comm comm = (MPI_Comm)comm_context;
    MPI_Barrier(comm);
#endif
    return 0;
}
int oob_gather(void* comm_context, int root, void* sbuf, void* rbuf, int len)
{
    if (comm_context != NULL) {
#ifdef MPI_SUPPORT
        MPI_Comm comm = (MPI_Comm)comm_context;
        MPI_Gather(sbuf, len, MPI_BYTE, rbuf, len, MPI_BYTE, root, comm);
#endif
    } else {
        memcpy(rbuf, sbuf, len);
    }
    return 0;
}

int oob_progress(void)
{
    int flag = 0;
#ifdef MPI_SUPPORT
    MPI_Iprobe(MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &flag, MPI_STATUS_IGNORE);
#endif
    return flag;
}

#ifdef MPI_SUPPORT
static void setup_comm_manager_rank(MPI_Comm comm, coll_sharp_component_t* sharp)
{
    int name_length, size, i, rank;

    MPI_Comm_rank(comm, &rank);
    MPI_Comm_size(comm, &size);

    /* todo: do not use stack allocation but malloc */
    int name_len[size];
    int offsets[size];

    MPI_Get_processor_name(sharp->node_name, &name_length);

    /* collect hostname len from all ranks */
    MPI_Allgather(&name_length, 1, MPI_INT, &name_len[0], 1, MPI_INT, comm);

    /* calculate receive buffer byte count based on hostname len sum. */
    int bytes = 0;
    for (i = 0; i < size; ++i) {
        offsets[i] = bytes;
        bytes += name_len[i];
    }
    bytes++;

    char receive_buffer[bytes];
    receive_buffer[bytes - 1] = 0;

    // collect hostnames, form comma separated list
    MPI_Allgatherv(&sharp->node_name[0], name_length, MPI_CHAR, &receive_buffer[0], &name_len[0], &offsets[0], MPI_CHAR, comm);

    int node_rank_count = 0;
    int node_rank_min = rank;

    for (i = 0; i < size; i++) {
        int idx = offsets[i];
        int len = name_len[i];

        if (name_length != len) {
            continue;
        }

        if (0 == strncmp(sharp->node_name, &receive_buffer[idx], len)) {
            if (i == rank) {
                sharp->node_local_rank = node_rank_count;
#if DEBUG
                printf(stdout "rank:%d local rank:%d\n", i, sharp->node_local_rank);
#endif
            }
            node_rank_count++;
            if (node_rank_min > i) {
                node_rank_min = i;
            }
        }
    }

    sharp->is_leader = (rank == node_rank_min);
    sharp->ppn = node_rank_count;
}
#endif

static void sharp_component_init(struct coll_sharp_component_t* sharp)
{
    gethostname(sharp->node_name, sizeof(sharp->node_name));
    sharp->node_local_rank = 0;
    sharp->sharp_coll_context = NULL;
    sharp->is_leader = 1;
    sharp->ppn = 1;
}

int coll_component_open()
{
    int ret, d;
    struct sharp_coll_init_spec init_spec = {0};
    struct timeval tval;
    char device_list[256] = "";

    coll_sharp_component.sharp_coll_context = NULL;

    if (!enable_sharp_coll)
        return 0;

    gettimeofday(&tval, NULL);
    srand((int)tval.tv_usec);

    init_spec.progress_func = oob_progress;
    /* coverity[dont_call] */
    if (sharp_conf.jobid && sharp_test_group_type != COMM_TYPE_PPN_JOBS) {
        init_spec.job_id = sharp_conf.jobid;
    } else {
        init_spec.job_id = (gethostid() << 32) | (rand() ^ getpid());
    }

#ifdef MPI_SUPPORT
    MPI_Bcast(&init_spec.job_id, sizeof(init_spec.job_id), MPI_BYTE, 0, (MPI_Comm)coll_sharp_component.mpi_comm);
#endif
    init_spec.world_rank = coll_sharp_component.rank;
    init_spec.world_size = coll_sharp_component.size;
    init_spec.enable_thread_support = 0;
    if (test_flat_ppn || sharp_test_group_type == COMM_TYPE_PPN_JOBS) {
        init_spec.world_local_rank = 0;
        init_spec.group_channel_idx = 0;
    } else {
        init_spec.world_local_rank = coll_sharp_component.node_local_rank;
        init_spec.group_channel_idx = coll_sharp_component.node_local_rank;
        if (sharp_test_group_type == COMM_TYPE_PPN_COMM) {
            init_spec.group_channel_idx = 0;
        }
    }

    init_spec.oob_colls.barrier = oob_barrier;
    init_spec.oob_colls.bcast = oob_bcast;
    init_spec.oob_colls.gather = oob_gather;
    init_spec.oob_ctx = (void*)coll_sharp_component.mpi_comm;
    init_spec.config = sharp_coll_default_config;

    if (test_flat_ppn && coll_sharp_component.ppn > num_ib_devices &&
        (sharp_test_group_type != COMM_TYPE_PPN_COMM && sharp_test_group_type != COMM_TYPE_PPN_JOBS))
    {
        fprintf(stderr, "PPN > 1 not supported with this group type \n");
        return -1;
    }

    // split device list between all local ranks
    if (num_ib_devices % coll_sharp_component.ppn != 0) {
        fprintf(stderr, "Can't split device list equally between local ranks");
        return -1;
    }

    int devices_per_ppn = num_ib_devices / coll_sharp_component.ppn;
    d = 0;
    while (1) {
        strcat(device_list, ib_dev_list[coll_sharp_component.node_local_rank * devices_per_ppn + d]);
        d++;
        if (d == devices_per_ppn) {
            break;
        }
        strcat(device_list, ",");
    }

    init_spec.config.ib_dev_list = device_list;

    if (init_spec.world_size == 1) {
        init_spec.config.flags |= SHARP_COLL_JOB_POLLING_KEEPALIVE;
    }

#if DEBUG
    fprintf(stdout,
            "[rank:%d, localrank:%d] Using IB device :%s world_local_rank:%d group_channel_idx:%d \n",
            sharp_world_rank,
            coll_sharp_component.node_local_rank,
            init_spec.config.ib_dev_list,
            init_spec.world_local_rank,
            init_spec.group_channel_idx);
#endif

    ret = sharp_coll_init(&init_spec, &coll_sharp_component.sharp_coll_context);
    if (ret != SHARP_COLL_SUCCESS) {
        return -1;
    }

    ret = sharp_coll_caps_query(coll_sharp_component.sharp_coll_context, &coll_sharp_component.sharp_caps);

    coll_sharp_component.src_shmid = 0;
    coll_sharp_component.dst_shmid = 0;

    return ret;
}

int coll_module_enable(coll_sharp_module_t* sharp_module)
{
    int ret;
    struct sharp_coll_comm_init_spec comm_spec;

    sharp_module->sharp_coll_comm = NULL;

    if (!enable_sharp_coll)
        return 0;

    comm_spec.rank = sharp_module->rank;
    comm_spec.size = sharp_module->size;
    comm_spec.oob_ctx = sharp_module->mpi_comm;
    comm_spec.group_world_ranks = NULL;
    ret = sharp_coll_comm_init(coll_sharp_component.sharp_coll_context, &comm_spec, &sharp_module->sharp_coll_comm);
    if (ret < 0) {
        fprintf(stdout, "sharp communicator creation failed: %s\n", sharp_coll_strerror(ret));
        return -1;
    }
    return 0;
}

void coll_module_destroy(coll_sharp_module_t* sharp_module)
{
    if (!enable_sharp_coll || !sharp_module->sharp_coll_comm)
        return;
    sharp_coll_comm_destroy(sharp_module->sharp_coll_comm);
    sharp_module->sharp_coll_comm = NULL;
    sharp_module->mpi_comm = NULL;
}
void coll_component_close(void)
{
    if (!enable_sharp_coll || !coll_sharp_component.sharp_coll_context)
        return;
    sharp_coll_finalize(coll_sharp_component.sharp_coll_context);
}

static void sharp_run_basic(coll_sharp_module_t* sharp_comm)
{
    if (sharp_conf.run_colls & TEST_ALLREDUCE)
        coll_test_sharp_allreduce_basic(sharp_comm, 0);
    if (sharp_conf.run_colls & TEST_IALLREDUCE)
        coll_test_sharp_iallreduce_basic(sharp_comm);
    if (sharp_conf.run_colls & TEST_REDUCE)
        coll_test_sharp_allreduce_basic(sharp_comm, 1);
    if (sharp_conf.run_colls & TEST_BARRIER)
        coll_test_sharp_barrier_basic(sharp_comm);
    if (sharp_conf.run_colls & TEST_IBARRIER)
        coll_test_sharp_ibarrier_basic(sharp_comm);
    if (sharp_conf.run_colls & TEST_BCAST)
        coll_test_sharp_bcast_basic(sharp_comm);
    if (sharp_conf.run_colls & TEST_REDUCE_SCATTER)
        coll_test_sharp_reduce_scatter_basic(sharp_comm);
    if (sharp_conf.run_colls & TEST_ALLGATHER)
        coll_test_sharp_allgather_basic(sharp_comm);
}

static void sharp_run_complex(coll_sharp_module_t* sharp_comm)
{
    if (sharp_conf.run_colls & TEST_ALLREDUCE)
        coll_test_sharp_allreduce_complex(sharp_comm);
    if (sharp_conf.run_colls & TEST_IALLREDUCE)
        coll_test_sharp_iallreduce_complex(sharp_comm);
    if (sharp_conf.run_colls & TEST_BARRIER)
        coll_test_sharp_barrier_complex(sharp_comm);
}
static void sharp_run_perf(coll_sharp_module_t* sharp_comm)
{
    if (sharp_conf.run_colls & TEST_ALLREDUCE)
        coll_test_sharp_allreduce_perf(sharp_comm, 0);
    if (sharp_conf.run_colls & TEST_IALLREDUCE)
        coll_test_sharp_iallreduce_perf(sharp_comm);
    if (sharp_conf.run_colls & TEST_REDUCE)
        coll_test_sharp_allreduce_perf(sharp_comm, 1);
    if (sharp_conf.run_colls & TEST_BARRIER)
        coll_test_sharp_barrier_perf(sharp_comm);
    if (sharp_conf.run_colls & TEST_BCAST)
        coll_test_sharp_bcast_perf(sharp_comm);
    if (sharp_conf.run_colls & TEST_REDUCE_SCATTER)
        coll_test_sharp_reduce_scatter_perf(sharp_comm);
    if (sharp_conf.run_colls & TEST_ALLGATHER)
        coll_test_sharp_allgather_perf(sharp_comm);
    if (sharp_conf.run_colls & TEST_RS_AG)
        coll_test_sharp_rs_ag_perf(sharp_comm, sharp_conf.rs_ag_overlapped);
}

void coll_module_op(coll_sharp_module_t* sharp_comm)
{
    if (sharp_conf.test_mode == TEST_BASIC) {
        sharp_run_basic(sharp_comm);
    } else if (sharp_conf.test_mode == TEST_COMPLEX) {
        sharp_run_complex(sharp_comm);
    } else if (sharp_conf.test_mode == TEST_PERF) {
        sharp_run_perf(sharp_comm);
    } else {
        sharp_run_basic(sharp_comm);
        sharp_run_complex(sharp_comm);
        sharp_run_perf(sharp_comm);
    }
}

void sharp_setenv(char* sharp_var, char* env_var, char* default_val)
{
    char* tmp = getenv(env_var);

    if (NULL == tmp) {
        if (NULL != default_val) {
            tmp = default_val;
        } else {
            return;
        }
    }
    setenv(sharp_var, tmp, 0);
}

void sharp_env2int(char* env_var, int* val, int default_val)
{
    long tmp;
    char *endptr, *str;

    *val = default_val;

    str = getenv(env_var);
    if (NULL == str)
        return;

    tmp = strtol(str, &endptr, 10);
    if (*endptr != '\0') {
        fprintf(stdout, "Invalid  %s environment value\n", env_var);
    } else {
        *val = tmp;
    }
}

#ifdef MPI_SUPPORT
void mpi_get_communicator(MPI_Comm* comm)
{
    int size, rank;

    MPI_Comm_size(MPI_COMM_WORLD, &size);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    switch (sharp_test_group_type) {
        case COMM_TYPE_WORLD:
            *comm = MPI_COMM_WORLD;
            break;
        case COMM_TYPE_WORLD_DUP:
            /* dup of comm_world */
            MPI_Comm_dup(MPI_COMM_WORLD, comm);
            break;
        case COMM_TYPE_WORLD_REVERSE:
            /*reverse comm */
            MPI_Comm_split(MPI_COMM_WORLD, 0, size - rank, comm);
            break;
        case COMM_TYPE_N_SPLIT:
            /* subset of world. odd, even groups */
            MPI_Comm_split(MPI_COMM_WORLD, (rank % sharp_test_num_splits), rank, comm);
            break;
        case COMM_TYPE_N_SPLIT_REVERSE:
            /* subset of world. odd, even groups in reserse*/
            MPI_Comm_split(MPI_COMM_WORLD, (rank % sharp_test_num_splits), (size - rank), comm);
            break;
        case COMM_TYPE_HALF:
            MPI_Comm_split(MPI_COMM_WORLD, (rank < (size / 2)), rank, comm);
            break;
        case COMM_TYPE_PPN_COMM:
        case COMM_TYPE_PPN_JOBS:
            MPI_Comm_split(MPI_COMM_WORLD, coll_sharp_component.node_local_rank, rank, comm);
            break;
        case COMM_TYPE_RANDOM:
            /* TODO: random comm with random size */
        default:
            *comm = MPI_COMM_WORLD;
    }
}

void mpi_free_comm(MPI_Comm* comm)
{
    if (*comm != MPI_COMM_WORLD && *comm != MPI_COMM_NULL) {
        MPI_Comm_free(comm);
    }
}
#endif

void setup_sharp_env(sharp_conf_t* sharp_conf)
{
    sharp_env2int("ENABLE_SHARP_COLL", &enable_sharp_coll, enable_sharp_coll);
    sharp_env2int("SHARP_RUN_PPN_COMM_PARALLEL", &sharp_conf->run_ppn_comm_parallel, 0);
    setenv("SHARP_COLL_ENABLE_SAT", "1", 0);
}

static size_t sharp_coll_test_power_floor(size_t x)
{
    size_t power = 1;
    while (x >>= 1)
        power <<= 1;
    return power;
}

static size_t sharp_coll_test_power_ceil(size_t x)
{
    if (x <= 1)
        return 1;
    size_t power = 2;
    x--;
    while (x >>= 1)
        power <<= 1;
    return power;
}

size_t sharp_test_string_quantity_prefix_value(char prefix)
{
    switch (prefix) {
        case 'B':
            return 1;
        case 'K':
            return SHARP_KBYTE;
        case 'M':
            return SHARP_MBYTE;
        case 'G':
            return SHARP_GBYTE;
        case 'T':
            return SHARP_TBYTE;
        default:
            return 0;
    }
}
/*
 * Parssing memory units. valid syntax : <number>[b|kb|mb|gb].
 * The memory units ( 'b' , 'k', 'm', "g', 't') can be specificed in both upper and lower case.
 */
int sharp_test_read_memunits(const void* str, size_t* dest, ssize_t low, ssize_t high)
{
    char units[3];
    int num_fields;
    ssize_t value;
    size_t bytes;

    /* Special value: infinity */
    if (!strcasecmp(str, "inf")) {
        *(size_t*)dest = SSIZE_MAX;
        return 0;
    }

    memset(units, 0, sizeof(units));
    num_fields = sscanf(str, "%zd%c%c", &value, &units[0], &units[1]);
    if (num_fields == 1) {
        bytes = 1;
    } else if ((num_fields == 2) || (num_fields == 3)) {
        bytes = sharp_test_string_quantity_prefix_value(toupper(units[0]));
        if (!bytes || ((num_fields == 3) && toupper(units[1]) != 'B')) {
            goto err_memunits;
        }
    } else {
        goto err_memunits;
    }

    value *= bytes;

    if (value < low || value > high) {
        fprintf(stderr, "%s %lu - %lu\n", "Invalid range of memory size. expected range ", low, high);
        return 1;
    }

    *dest = value;
    return 0;
err_memunits:
    fprintf(stderr, "%s: %s\n", (char*)str, "Invalid memory units. syntax:<number>[b|kb|mb|gb] or \"inf\"");
    return 1;
}

int parse_opts(int argc, char** argv, sharp_conf_t* sharp_conf)
{
    int opt, option_index;
    int res = PARSE_ARGS_OK;
    long tmp_id = 0;
    size_t val1, val2;
    char *endptr, *saveptr, *tmp_ptr = NULL;

    while (res == PARSE_ARGS_OK) {
        option_index = 0;
        endptr = "\0";
#ifdef MPI_SUPPORT
        opt = getopt_long(argc, argv, "ht:d:m:c:R:i:j:H:x:B:N:M:p:s:r:D:T:C:f:", long_options, &option_index);
#else
        opt = getopt_long(argc, argv, "hd:i:j:x:M:s:f:", long_options, &option_index);
#endif
        if (opt == -1)   // no more args to parse
            break;
        switch (opt) {
            case 'h':
                return PARSE_ARGS_HELP;
            case 'c':
                sharp_conf->run_colls = 0;
                saveptr = optarg;
                for (endptr = strtok_r(optarg, ",", &saveptr); endptr; endptr = strtok_r(NULL, ",", &saveptr)) {
                    if (strcmp(endptr, "allreduce") == 0) {
                        sharp_conf->run_colls |= TEST_ALLREDUCE;
                    } else if (strcmp(endptr, "iallreduce") == 0) {
                        sharp_conf->run_colls |= TEST_IALLREDUCE;
                    } else if (strcmp(endptr, "reduce") == 0) {
                        sharp_conf->run_colls |= TEST_REDUCE;
                    } else if (strcmp(endptr, "ireduce") == 0) {
                        sharp_conf->run_colls |= TEST_IREDUCE;
                    } else if (strcmp(endptr, "barrier") == 0) {
                        sharp_conf->run_colls |= TEST_BARRIER;
                    } else if (strcmp(endptr, "ibarrier") == 0) {
                        sharp_conf->run_colls |= TEST_IBARRIER;
                    } else if (strcmp(endptr, "bcast") == 0) {
                        sharp_conf->run_colls |= TEST_BCAST;
                    } else if (strcmp(endptr, "reduce_scatter") == 0) {
                        sharp_conf->run_colls |= TEST_REDUCE_SCATTER;
                    } else if (strcmp(endptr, "allgather") == 0) {
                        sharp_conf->run_colls |= TEST_ALLGATHER;
                    } else if (strcmp(endptr, "rs_ag") == 0) {
                        sharp_conf->run_colls |= TEST_RS_AG;
                    } else if (strcmp(endptr, "none") == 0) {
                        sharp_conf->run_colls = 0;
                    } else if (strcmp(endptr, "all") == 0) {
                        sharp_conf->run_colls |=
                            TEST_ALLREDUCE | TEST_IALLREDUCE | TEST_BARRIER | TEST_ALLGATHER | TEST_REDUCE_SCATTER | TEST_IBARRIER;
                    } else {
                        fprintf(stderr, "Invalid --collectives option\n");
                        res = INVALID_ARG;
                    }
                }
                break;
            case 'R':
                if (strcmp(optarg, "seq") == 0) {
                    sharp_conf->rs_ag_overlapped = 0;
                } else if (strcmp(optarg, "ovl") == 0) {
                    sharp_conf->rs_ag_overlapped = 1;
                } else {
                    fprintf(stderr, "Invalid --rs_ag_mode option. Use seq|ovl\n");
                    res = INVALID_ARG;
                }
                break;
            case 'C':
                tmp_id = strtol(optarg, &endptr, 10);
                if (*endptr != '\0' || tmp_id < 0 || tmp_id > SHARP_COLL_DATA_MAX_IOV) {
                    fprintf(stderr, "Invalid --iov_count value specified\n");
                    res = INVALID_ARG;
                } else {
                    sharp_conf->siov_count = tmp_id;
                    sharp_conf->riov_count = tmp_id;
                }
                break;
            case 'D':
                if (strcmp(optarg, "contig") == 0) {
                    sharp_conf->sdata_layout = TEST_DATA_LAYOUT_CONTIG;
                    sharp_conf->rdata_layout = TEST_DATA_LAYOUT_CONTIG;
                } else if (strcmp(optarg, "iov") == 0) {
                    sharp_conf->sdata_layout = TEST_DATA_LAYOUT_IOV;
                    sharp_conf->rdata_layout = TEST_DATA_LAYOUT_IOV;
                    res = INVALID_ARG;
                } else {
                    fprintf(stderr, "Invalid --data_layout option\n");
                }
                break;
            case 't':
                saveptr = optarg;
                endptr = strtok_r(optarg, ":", &saveptr);
                if (endptr) {
                    if (strcmp(endptr, "world") == 0) {
                        sharp_test_group_type = COMM_TYPE_WORLD;
                    } else if (strcmp(endptr, "world-dup") == 0) {
                        sharp_test_group_type = COMM_TYPE_WORLD_DUP;
                    } else if (strcmp(endptr, "world-rev") == 0) {
                        sharp_test_group_type = COMM_TYPE_WORLD_REVERSE;
                    } else if (strcmp(endptr, "split") == 0) {
                        endptr = strtok_r(NULL, ",", &saveptr);
                        if (endptr) {
                            if (!strcmp(endptr, "ppn-jobs")) {
                                sharp_test_group_type = COMM_TYPE_PPN_JOBS;
                                test_flat_ppn = 0;
                            } else if (!strcmp(endptr, "ppn-comm")) {
                                sharp_test_group_type = COMM_TYPE_PPN_COMM;
                                test_flat_ppn = 0;
                            } else {
                                sharp_test_group_type = COMM_TYPE_N_SPLIT;
                                sharp_test_num_splits = strtol(endptr, &tmp_ptr, 10);
                                if (*tmp_ptr != '\0') {
                                    fprintf(stderr, "Invalid --group_type option\n");
                                    res = INVALID_ARG;
                                }
                            }
                        }
                    } else if (strcmp(endptr, "split-rev") == 0) {
                        sharp_test_group_type = COMM_TYPE_N_SPLIT_REVERSE;
                        endptr = strtok_r(NULL, ",", &saveptr);
                        if (endptr) {
                            sharp_test_num_splits = strtol(endptr, &tmp_ptr, 10);
                            if (*tmp_ptr != '\0') {
                                fprintf(stderr, "Invalid --group_type option\n");
                                res = INVALID_ARG;
                            }
                        }
                    } else if (strcmp(endptr, "half") == 0) {
                        sharp_test_group_type = COMM_TYPE_HALF;
                    } else {
                        fprintf(stderr, "Invalid --group_type option\n");
                        res = INVALID_ARG;
                    }
                } else {
                    fprintf(stderr, "Invalid --group_type option\n");
                    res = INVALID_ARG;
                }
                break;
            case 'T':
                if (0 == strcmp(optarg, "half_int"))
                    sharp_conf->datatype = SHARP_DTYPE_SHORT;
                else if (0 == strcmp(optarg, "half_float"))
                    sharp_conf->datatype = SHARP_DTYPE_FLOAT_SHORT;
                else if (0 == strcmp(optarg, "int"))
                    sharp_conf->datatype = SHARP_DTYPE_INT;
                else if (0 == strcmp(optarg, "float"))
                    sharp_conf->datatype = SHARP_DTYPE_FLOAT;
                else if (0 == strcmp(optarg, "long"))
                    sharp_conf->datatype = SHARP_DTYPE_LONG;
                else if (0 == strcmp(optarg, "double"))
                    sharp_conf->datatype = SHARP_DTYPE_DOUBLE;
                else if (0 == strcmp(optarg, "byte"))
                    sharp_conf->datatype = SHARP_DTYPE_INT8;
                else {
                    fprintf(stderr, "Invalid --datatype option\n");
                    res = INVALID_ARG;
                }
                break;
            case 'd':
                saveptr = optarg;
                for (num_ib_devices = 0, endptr = strtok_r(optarg, ",", &saveptr); endptr && (num_ib_devices < MAX_IB_DEVICES);
                     endptr = strtok_r(NULL, ",", &saveptr))
                {
                    strcpy(ib_dev_list[num_ib_devices++], endptr);
                }
                break;
            case 'm':
                if (0 == strcmp(optarg, "basic"))
                    sharp_conf->test_mode = TEST_BASIC;
                else if (0 == strcmp(optarg, "complex"))
                    sharp_conf->test_mode = TEST_COMPLEX;
                else if (0 == strcmp(optarg, "perf"))
                    sharp_conf->test_mode = TEST_PERF;
                else if (0 == strcmp(optarg, "all"))
                    sharp_conf->test_mode = TEST_ALL;
                else {
                    fprintf(stderr, "Invalid --mode option\n");
                    res = INVALID_ARG;
                }
                break;
            case 'i':
                tmp_id = strtol(optarg, &endptr, 10);
                if (*endptr != '\0' || tmp_id < 0 || tmp_id > INT_MAX) {
                    fprintf(stderr, "Invalid  number of interations specified\n");
                    res = INVALID_ARG;
                } else
                    perf_test_iterations = tmp_id;
                break;
            case 'f':
                tmp_id = strtol(optarg, &endptr, 10);
                if (*endptr != '\0' || tmp_id <= 0 || tmp_id > 2) {
                    fprintf(stderr, "Invalid step factor value. acceptd values are 1(infinte) and 2\n");
                    res = INVALID_ARG;
                } else
                    sharp_conf->stepfactor = tmp_id;
                break;
            case 'j':
                tmp_id = strtol(optarg, &endptr, 10);
                // Check for overflow and invalid values
                // errno is set to ERANGE if the value is too large
                if (*endptr != '\0' || tmp_id <= 0 || tmp_id == LONG_MAX || errno == ERANGE) {
                    fprintf(stderr, "Invalid job id value is provided, valid values are between 1 and (2^63)-2\n");
                    res = INVALID_ARG;
                } else {
                    sharp_conf->jobid = tmp_id;
                }
                break;
            case 'H':
                if (0 == strcmp(optarg, "malloc"))
                    sharp_conf->host_allocator_type = USE_MALLOC;
                else if (0 == strcmp(optarg, "hugetlb"))
                    sharp_conf->host_allocator_type = USE_HUGETLB;
                else if (0 == strcmp(optarg, "hugethp"))
                    sharp_conf->host_allocator_type = USE_HUGETHP;
                else {
                    fprintf(stderr, "Invalid --host_alloc_type option\n");
                    res = INVALID_ARG;
                }
                break;
            case 'r':
                tmp_id = strtol(optarg, &endptr, 10);
                if (*endptr != '\0' || tmp_id < 0)
                    res = INVALID_ARG;
                else
                    coll_root_rank = tmp_id;
                break;
            case 'x':
                tmp_id = strtol(optarg, &endptr, 10);
                if (*endptr != '\0' || tmp_id < 0 || tmp_id > INT_MAX) {
                    fprintf(stderr, "Invalid  number of warmup interations specified\n");
                    res = INVALID_ARG;
                } else
                    perf_test_skips = tmp_id;
                break;
            case 'N':
                tmp_id = strtol(optarg, &endptr, 10);
                if (*endptr != '\0' || tmp_id < 0 || tmp_id > 512)
                    res = INVALID_ARG;
                else
                    nbc_count = tmp_id;
                break;
            case 'p':
                tmp_id = strtol(optarg, &endptr, 10);
                if (*endptr != '\0' || tmp_id < 0)
                    res = INVALID_ARG;
                else
                    test_flat_ppn = tmp_id;
                break;
            case 'B':
                tmp_id = strtol(optarg, &endptr, 10);
                if (*endptr != '\0' || tmp_id < 0)
                    res = INVALID_ARG;
                else
                    perf_with_barrier = tmp_id;
                break;
            case 'M':
                saveptr = optarg;
                endptr = strtok_r(optarg, ":", &saveptr);
                if (endptr) {
                    const char* p_invalid_value = "Invalid memory type value. Accepted values are: host, cuda, and null.\n";
                    if (0 == strcasecmp(endptr, "CUDA")) {
                        sharp_conf->s_mem_type = SHARP_MEM_TYPE_CUDA;
                    } else if (0 == strcasecmp(endptr, "HOST")) {
                        sharp_conf->s_mem_type = SHARP_MEM_TYPE_HOST;
                    } else if (0 == strcasecmp(endptr, "NULL")) {
                        sharp_conf->s_mem_type = SHARP_MEM_TYPE_LAST;
                    } else {
                        fprintf(stderr, "%s", p_invalid_value);
                        res = INVALID_ARG;
                        break;
                    }

                    endptr = strtok_r(NULL, ":", &saveptr);
                    if (endptr) {
                        if (0 == strcasecmp(endptr, "CUDA")) {
                            sharp_conf->r_mem_type = SHARP_MEM_TYPE_CUDA;
                        } else if (0 == strcasecmp(endptr, "HOST")) {
                            sharp_conf->r_mem_type = SHARP_MEM_TYPE_HOST;
                        } else if (0 == strcasecmp(endptr, "NULL")) {
                            sharp_conf->r_mem_type = SHARP_MEM_TYPE_LAST;
                        } else {
                            fprintf(stderr, "%s", p_invalid_value);
                            res = INVALID_ARG;
                            break;
                        }
                    } else {
                        sharp_conf->r_mem_type = sharp_conf->s_mem_type;
                    }
                }
                if (!sharp_conf->cudart_lib &&
                    (sharp_conf->s_mem_type == SHARP_MEM_TYPE_CUDA || sharp_conf->r_mem_type == SHARP_MEM_TYPE_CUDA))
                {
                    fprintf(stderr, "CUDA runtime libs not found. CUDA memory not supported\n");
                    res = INVALID_ARG;
                }
                break;
            case 's':
                val1 = val2 = 0;
                saveptr = optarg;
                endptr = strtok_r(optarg, ":", &saveptr);
                if (endptr) {
                    if (sharp_test_read_memunits(endptr, &val1, 4, SSIZE_MAX)) {
                        res = INVALID_ARG;
                        break;
                    }
                    sharp_conf->min_message_size = sharp_coll_test_power_ceil(val1);
                    sharp_conf->max_message_size = sharp_coll_test_power_floor(val1);
                    endptr = strtok_r(NULL, ":", &saveptr);
                    if (endptr) {
                        if (sharp_test_read_memunits(endptr, &val2, sharp_conf->min_message_size, SSIZE_MAX)) {
                            res = INVALID_ARG;
                            break;
                        }
                        sharp_conf->max_message_size = sharp_coll_test_power_floor(val2);
                    }
                } else {
                    res = INVALID_ARG;
                }

                break;
            case 0:
                if (use_cuda_dmabuf_flag) {
                    tmp_id = strtol(optarg, &endptr, 10);
                    if (*endptr != '\0' || tmp_id < 0)
                        res = INVALID_ARG;
                    else
                        sharp_conf->use_dma_buf = tmp_id;
                    use_cuda_dmabuf_flag = 0;
                }
                if (use_data_direct_flag) {
                    tmp_id = strtol(optarg, &endptr, 10);
                    if (*endptr != '\0' || tmp_id < 0)
                        res = INVALID_ARG;
                    else
                        sharp_conf->use_direct_nic = tmp_id;
                    use_data_direct_flag = 0;
                }
                break;
            default:
                res = INVALID_ARG;
        }
    }

    if (opt == -1 && argv[optind] != NULL)
        res = INVALID_ARG;
    return res;
}

void print_parser_msg(int parse_res)
{
    if (parse_res == PARSE_ARGS_HELP)
        printf("%s\n", usage_string);
    else if (parse_res == INVALID_ARG)
        fprintf(stderr,
                "ERROR. Either a non-existing command was specified, "
                "or an option with a missing/invalid argument was given\n\n");
}

const char* sharp_get_host_name(void)
{
    static char hostname[256] = {0};

    if (*hostname == 0) {
        gethostname(hostname, sizeof(hostname));
        strtok(hostname, ".");
    }
    return hostname;
}

static void coll_test_load_cuda(void)
{
    void* tmp;
    void** cast;

    sharp_conf.cuda_lib = dlopen("libcuda.so.1", RTLD_NOW);
    if (!sharp_conf.cuda_lib) {
        // fprintf(stderr, "Failed to open libcuda.so\n");
        goto teardown;
    }

    LOAD_SYM(sharp_conf.cuda_lib, "cuMemAlloc_v2", sharp_conf.cuMemAlloc_ptr);
    LOAD_SYM(sharp_conf.cuda_lib, "cuMemGetHandleForAddressRange", sharp_conf.cuMemGetHandleForAddressRange_ptr);

    sharp_conf.cudart_lib = dlopen("libcudart.so", RTLD_NOW);
    if (!sharp_conf.cudart_lib) {
        // fprintf(stderr, "Failed to open libcudart.so\n");
        goto teardown;
    }

    LOAD_SYM(sharp_conf.cudart_lib, "cudaGetDeviceCount", sharp_conf.cuda_get_device_count);
    LOAD_SYM(sharp_conf.cudart_lib, "cudaSetDevice", sharp_conf.cuda_set_device);
    LOAD_SYM(sharp_conf.cudart_lib, "cudaMalloc", sharp_conf.cuda_malloc);
    LOAD_SYM(sharp_conf.cudart_lib, "cudaMemcpy", sharp_conf.cuda_memcpy);
    LOAD_SYM(sharp_conf.cudart_lib, "cudaFree", sharp_conf.cuda_free);

    sharp_conf.s_mem_type = SHARP_MEM_TYPE_CUDA;
    sharp_conf.r_mem_type = SHARP_MEM_TYPE_CUDA;

    return;
teardown:
    if (sharp_conf.cudart_lib != NULL)
        dlclose(sharp_conf.cudart_lib);
    if (sharp_conf.cuda_lib != NULL)
        dlclose(sharp_conf.cuda_lib);
    sharp_conf.cudart_lib = NULL;
    sharp_conf.cuda_lib = NULL;
    return;
}

int main(int argc, char** argv)
{
    int rank, size, parse_res, mpi_rank, local_rank;
    int ret = 0;
    coll_sharp_module_t sharp_comm;
#ifdef MPI_SUPPORT
    int dtype, tag_dtype, comm_size;
    MPI_Comm mpi_comm;
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
#else
    rank = 0;
    size = 1;
#endif
    sharp_conf.size = sharp_world_rank = rank;
    sharp_conf.rank = sharp_world_size = size;

    coll_test_load_cuda();

    parse_res = parse_opts(argc, argv, &sharp_conf);
    if (parse_res != PARSE_ARGS_OK) {
        if (rank == 0) {
            print_parser_msg(parse_res);
            if (parse_res != PARSE_ARGS_HELP) {
                print_parser_msg(PARSE_ARGS_HELP);
                ret = -1;
            }
        }
        goto out;
    }

    if (sharp_conf.s_mem_type == SHARP_MEM_TYPE_LAST || sharp_conf.r_mem_type == SHARP_MEM_TYPE_LAST) {
        if (sharp_conf.test_mode != TEST_PERF) {
            fprintf(stderr, "mem type \"null\" is supported only with with perf test mode. please specify \"-m perf\" options\n");
        }

        sharp_conf.min_message_size = MAX(16384, sharp_conf.min_message_size);
    }

    if (!num_ib_devices) {
        fprintf(stderr, "Please specify IB device with -d test option\n");
        goto out;
    }

    if (getenv("SLEEP")) {
        sleep(atoi(getenv("SLEEP")));
    }

    sharp_component_init(&coll_sharp_component);
#ifdef MPI_SUPPORT
    setup_comm_manager_rank(MPI_COMM_WORLD, &coll_sharp_component);
    for (dtype = 0; dtype < SHARP_DTYPE_NULL; dtype++) {
        for (tag_dtype = 0; tag_dtype < SHARP_DTYPE_NULL; tag_dtype++) {
            sharp_test_mpi_min_max_datatype[dtype][tag_dtype] = MPI_DATATYPE_NULL;
        }
    }
    sharp_test_mpi_min_max_datatype[SHARP_DTYPE_FLOAT][SHARP_DTYPE_INT] = MPI_FLOAT_INT;
    sharp_test_mpi_min_max_datatype[SHARP_DTYPE_INT][SHARP_DTYPE_INT] = MPI_2INT;
#endif

    if ((sharp_conf.s_mem_type == SHARP_MEM_TYPE_CUDA || sharp_conf.r_mem_type == SHARP_MEM_TYPE_CUDA)) {
        int num_gpus, gpu_index, gpu_index_stride;
        int cerr;
        cerr = sharp_conf.cuda_get_device_count(&num_gpus);
        if (cerr != 0) {
            fprintf(stdout, "cudaGetDeviceCount failed %d\n", cerr);
            goto out;
        }

        gpu_index_stride = num_gpus / coll_sharp_component.ppn;
        gpu_index = (coll_sharp_component.node_local_rank * gpu_index_stride) % num_gpus;

        cerr = sharp_conf.cuda_set_device(gpu_index);
        if (cerr != 0) {
            fprintf(stdout, "cudaSetDevice failed\n");
            goto out;
        }
#if DEBUG
        fprintf(stdout, "[Host:%s Rank:%d, GPU:%d] Testing with cuda buffers\n", coll_sharp_component.node_name, rank, gpu_index);
#endif
    }

    setup_sharp_env(&sharp_conf);

#ifdef MPI_SUPPORT
    mpi_get_communicator(&mpi_comm);
    MPI_Comm_rank(mpi_comm, &mpi_rank);
    MPI_Comm_size(mpi_comm, &comm_size);
    coll_sharp_component.rank = sharp_comm.rank = mpi_rank;
    coll_sharp_component.size = sharp_comm.size = comm_size;
    coll_sharp_component.mpi_comm = sharp_comm.mpi_comm =
        (void*)((sharp_test_group_type == COMM_TYPE_PPN_JOBS) ? mpi_comm : MPI_COMM_WORLD);
#else
    if (enable_sharp_coll == 0) {
        fprintf(stderr, "MPI is not supported");
        return -1;
    }
    mpi_rank = 0;
    coll_sharp_component.rank = sharp_comm.rank = 0;
    coll_sharp_component.size = sharp_comm.size = 1;
    sharp_comm.mpi_comm = NULL;
#endif

    if (coll_component_open()) {
        if (mpi_rank == 0)
            fprintf(stdout, "SHArP coll failed to initialize..\n");
        goto out;
    }

    if (coll_module_enable(&sharp_comm)) {
        fprintf(stdout, "sharp comm create failed for :  %s\n", comm_type_names[sharp_test_group_type]);
        ret = -1;
        goto close_component;
    }

    // After getting the sharp comm, we can tell the Sharp configuration and compare them with the coll_test configuration.
    // If only "reduce" is selected, we need to set the min message size to the minimum size of the reduce operation.
    if (sharp_conf.run_colls == TEST_REDUCE) {
        unsigned long sat_threshold = coll_sharp_component.sharp_caps.sat_threshold;
        if (sharp_conf.min_message_size < sat_threshold) {
            sharp_conf.min_message_size = sat_threshold;
        }
        if (sharp_conf.max_message_size < sat_threshold) {
            sharp_conf.max_message_size = sat_threshold;
        }
    }

    if (!sharp_conf.run_ppn_comm_parallel && (sharp_test_group_type == COMM_TYPE_PPN_COMM || sharp_test_group_type == COMM_TYPE_PPN_JOBS)) {
        for (local_rank = 0; local_rank < coll_sharp_component.ppn; local_rank++) {
            if (coll_sharp_component.node_local_rank == local_rank) {
                if (sharp_test_group_type != COMM_TYPE_WORLD && mpi_rank == 0) {
                    fprintf(stdout, "MPI COMM Type : %s\n", comm_type_names[sharp_test_group_type]);
                }
                coll_module_op(&sharp_comm);
            }
#ifdef MPI_SUPPORT
            MPI_Barrier(MPI_COMM_WORLD);
#endif
        }

    } else {
        if (sharp_test_group_type != COMM_TYPE_WORLD && mpi_rank == 0) {
            fprintf(stdout, "MPI COMM Type : %s\n", comm_type_names[sharp_test_group_type]);
        }
        coll_module_op(&sharp_comm);
    }
    coll_module_destroy(&sharp_comm);

#ifdef MPI_SUPPORT
    mpi_free_comm(&mpi_comm);
    MPI_Barrier(MPI_COMM_WORLD);
#endif
close_component:
    coll_component_close();
out:
#ifdef MPI_SUPPORT
    MPI_Finalize();
#endif
    return ret;
}
