/*
 * Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
 *
 * See COPYRIGHT for license information
 */

#include "nvshmem.h"
#include "nvshmemx.h"
#include "gpu_coll.h"
#include "nvshmemi_coll.h"
#include <cstdio>
#include <cassert>

#ifdef __CUDA_ARCH__

#define GPU_HEAD_CHECK_OP_THREADGROUP(TYPENAME, TYPE, OP, dest, src, actual_src, nelems, myIdx, groupSize) \
    do {                                                                                         \
        int i, k;                                                                                \
        int subelems = sizeof(TYPE) / sizeof(uint32_t);                                          \
        volatile uint32_t *header = NULL;                                                        \
        TYPE tmp;                                                                                \
        uint32_t *tmp_ptr = (uint32_t *)&tmp;                                                    \
        uint32_t *payload = NULL;                                                                \
        for (i = myIdx; i < nelems; i += groupSize) {                                            \
            for (k = 0; k < subelems; k++) {                                                     \
                payload = (uint32_t *)((uint64_t *)src + (i * subelems) + k);                    \
                header = (uint32_t *)payload + 1;                                                \
                while (1 != *header)                                                             \
                    ;                                                                            \
                *header = 0;                                                                     \
                *(tmp_ptr + k) = *payload;                                                       \
            }                                                                                    \
            perform_gpu_rd_##OP(*((TYPE *)dest + i), *((TYPE *)actual_src + i), tmp);            \
        }                                                                                        \
    } while (0)

#define GPU_HEAD_CHECKALL_OP_THREADGROUP(TYPENAME, TYPE, OP, dest, src, actual_src, nelems, start, \
                                         stride, size, myIdx, groupSize)                           \
    do {                                                                                           \
        int i, j, k;                                                                               \
        int subelems = sizeof(TYPE) / sizeof(uint32_t);                                            \
        volatile uint32_t *header = NULL;                                                          \
        TYPE tmp;                                                                                  \
        uint32_t *tmp_ptr = (uint32_t *)&tmp;                                                      \
        uint32_t *payload = NULL;                                                                  \
        int my_active_set_pe = ((nvshmemi_device_state_d.mype - start) / stride);                  \
        TYPE *src_ptr = (TYPE *)actual_src;                                                        \
                                                                                                   \
        for (j = (my_active_set_pe - 1); j >= 0; j--) {                                            \
            for (i = myIdx; i < nelems; i += groupSize) {                                          \
                for (k = 0; k < subelems; k++) {                                                   \
                    payload = (uint32_t *)((uint64_t *)src + (i * subelems) + k +                  \
                                           (nelems * subelems * j));                               \
                    header = (uint32_t *)payload + 1;                                              \
                    while (1 != *header)                                                           \
                        ;                                                                          \
                    *header = 0;                                                                   \
                    *(tmp_ptr + k) = *payload;                                                     \
                }                                                                                  \
                perform_gpu_rd_##OP(*((TYPE *)dest + i), *((TYPE *)src_ptr + i), tmp);             \
            }                                                                                      \
            src_ptr = dest;                                                                        \
        }                                                                                          \
        for (j = size - 1; j > my_active_set_pe; j--) {                                            \
            for (i = myIdx; i < nelems; i += groupSize) {                                          \
                for (k = 0; k < subelems; k++) {                                                   \
                    payload = (uint32_t *)((uint64_t *)src + (i * subelems) + k +                  \
                                           (nelems * subelems * j));                               \
                    header = (uint32_t *)payload + 1;                                              \
                    while (1 != *header)                                                           \
                        ;                                                                          \
                    *header = 0;                                                                   \
                    *(tmp_ptr + k) = *payload;                                                     \
                }                                                                                  \
                perform_gpu_rd_##OP(*((TYPE *)dest + i), *((TYPE *)src_ptr + i), tmp);             \
            }                                                                                      \
            src_ptr = dest;                                                                        \
        }                                                                                          \
    } while (0)

#define GPU_LINEAR_REDUCE_THREADGROUP(SC, TYPENAME, TYPE, OP, x, y, z, nelems, myIdx, groupSize)     \
    do {                                                                               \
        int i;                                                                         \
        for (i = myIdx; i < nelems; i += groupSize) {                                  \
            perform_gpu_rd_##OP(*((TYPE *)z + i), *((TYPE *)x + i), *((TYPE *)y + i)); \
        }                                                                              \
    } while (0)

#define GPU_LINEAR_REDUCE_THREADGROUP_P2P_GET(SC, TYPENAME, TYPE, OP, x, y, next_rank, z, nelems, myIdx,   \
                                              groupSize, start, stride, size, pWrk, pSync) \
    do {                                                                                         \
        int i;                                                                                   \
        int group_nelems = ((nelems / groupSize) * groupSize);                                   \
        int excess = nelems - group_nelems;                                                      \
        long counter = 1;                                                                        \
        for (i = myIdx; i < group_nelems; i += groupSize) {                                      \
            nvshmem_##TYPENAME##_get((TYPE *)((TYPE *)pWrk + myIdx), (TYPE *)((TYPE *)y + i), \
                                 1, next_rank);                                                  \
            nvshmemxi_barrier_##SC(start, stride, size, (long *)pSync, &counter);                \
            perform_gpu_rd_##OP(*((TYPE *)z + i), *((TYPE *)x + i),                              \
                                *((TYPE *)pWrk + myIdx));                                 \
        }                                                                                        \
        if (excess) {                                                                            \
            if (i < nelems) {                                                                    \
                nvshmem_##TYPENAME##_get((TYPE *)((TYPE *)pWrk + myIdx),                  \
                                     (TYPE *)((TYPE *)y + i), 1, next_rank);                     \
            }                                                                                    \
            nvshmemxi_barrier_##SC(start, stride, size, (long *)pSync, &counter);                \
            if (i < nelems) {                                                                    \
                perform_gpu_rd_##OP(*((TYPE *)z + i), *((TYPE *)x + i),                          \
                                    *((TYPE *)pWrk + myIdx));                             \
            }                                                                                    \
        }                                                                                        \
        NVSHMEMI_SYNC_##SC();                                                                     \
        int end = start + stride * size;                                                         \
        for (i = myIdx; i < end; i += groupSize)                                                 \
            pSync[i] = NVSHMEMI_SYNC_VALUE;                                                       \
        NVSHMEMI_SYNC_##SC();                                                                     \
    } while (0)

#define GPU_LINEAR_REDUCE_THREADGROUP_P2P_PUT(SC, TYPENAME, TYPE, OP, x, y, next_rank, z, offset, nelems,  \
                                              myIdx, groupSize, start, stride, size, pWrk, pSync) \
    do {                                                                                         \
        int i;                                                                                   \
        int group_nelems = ((nelems / groupSize) * groupSize);                                   \
        int excess = nelems - group_nelems;                                                      \
        long counter = 1;                                                                        \
        for (i = myIdx; i < group_nelems; i += groupSize) {                                      \
            nvshmem_##TYPENAME##_put_nbi(                                                        \
                (TYPE *)((TYPE *)pWrk + (myIdx + ((offset & 1) * groupSize))),                   \
                (const TYPE *)((TYPE *)y + i), 1, next_rank);                                    \
            nvshmemxi_barrier_##SC(start, stride, size, (long *)pSync, &counter);                \
            perform_gpu_rd_##OP(*((TYPE *)z + i), *((TYPE *)x + i),                              \
                                *((TYPE *)pWrk + (myIdx + ((offset & 1) * groupSize))));         \
        }                                                                                        \
        offset++;                                                                                \
        if (excess) {                                                                            \
            if (i < nelems) {                                                                    \
                nvshmem_##TYPENAME##_put_nbi(                                                    \
                    (TYPE *)((TYPE *)pWrk + (myIdx + ((offset & 1) * groupSize))),               \
                    (const TYPE *)((TYPE *)y + i), 1, next_rank);                                \
            }                                                                                    \
            nvshmemxi_barrier_##SC(start, stride, size, (long *)pSync, &counter);                \
            if (i < nelems) {                                                                    \
                perform_gpu_rd_##OP(                                                             \
                    *((TYPE *)z + i), *((TYPE *)x + i),                                          \
                    *((TYPE *)pWrk + (myIdx + ((offset & 1) * groupSize))));                     \
            }                                                                                    \
        }                                                                                        \
        offset++;                                                                                \
        int end = start + stride * size;                                                         \
        NVSHMEMI_SYNC_##SC();                                                                     \
        for (i = myIdx; i < end; i += groupSize)                                                 \
            pSync[i] = NVSHMEMI_SYNC_VALUE;                                                      \
        NVSHMEMI_SYNC_##SC();                                                                     \
    } while (0)

#define GPU_LINEAR_REDUCE_THREADGROUP_P2P_PUT_DIRECT(SC, TYPENAME, TYPE, OP, x, y, next_rank, z, offset,  \
                                                     nelems, myIdx, groupSize, start,           \
                                                     stride, size, pWrk, pSync)                 \
    do {                                                                                        \
        int i;                                                                                  \
        int group_nelems = ((nelems / groupSize) * groupSize);                                  \
        int excess = nelems - group_nelems;                                                     \
        long counter = 1;                                                                       \
        for (i = myIdx; i < group_nelems; i += groupSize) {                                     \
            *((TYPE *)((TYPE *)pWrk + (myIdx + ((offset & 1) * groupSize)))) =                  \
                *((TYPE *)((TYPE *)y + i));                                                     \
            nvshmemxi_barrier_##SC(start, stride, size, pSync, &counter);                       \
            perform_gpu_rd_##OP(*((TYPE *)z + i), *((TYPE *)x + i),                             \
                                *((TYPE *)pWrk + (myIdx + ((offset & 1) * groupSize))));        \
        }                                                                                       \
        offset++;                                                                               \
        if (excess) {                                                                           \
            if (i < nelems) {                                                                   \
                *((TYPE *)((TYPE *)pWrk + (myIdx + ((offset & 1) * groupSize)))) =              \
                    *((TYPE *)((TYPE *)y + i));                                                 \
            }                                                                                   \
            nvshmemxi_barrier_##SC(start, stride, size, pSync, &counter);                       \
            if (i < nelems) {                                                                   \
                perform_gpu_rd_##OP(                                                            \
                    *((TYPE *)z + i), *((TYPE *)x + i),                                         \
                    *((TYPE *)pWrk + (myIdx + ((offset & 1) * groupSize))));                    \
            }                                                                                   \
        }                                                                                       \
        offset++;                                                                               \
        int end = start + stride * size;                                                        \
        NVSHMEMI_SYNC_##SC();                                                                     \
        for (i = myIdx; i < end; i += groupSize)                                                \
            pSync[i] = NVSHMEMI_SYNC_VALUE;                                                     \
        NVSHMEMI_SYNC_##SC();                                                                     \
    } while (0)

#define NVSHMEMXI_GPU_RDXN_THREADGROUP_ZCOPY_GET_BAR_DIRECT(                                      \
    SC, TYPENAME, TYPE, OP, dest, source, nreduce, start, stride, size, pWrk, pSync)              \
    do {                                                                                          \
        int next_rank = -1;                                                                       \
        int src_offset = -1;                                                                      \
        int next_offset = -1;                                                                     \
        char *base = NULL;                                                                        \
        char *peer_base = NULL;                                                                   \
        char *peer_source = NULL;                                                                 \
        NVSHMEMI_DECL_THREAD_IDX_##SC();                                                          \
        NVSHMEMI_DECL_THREADGROUP_SIZE_##SC();                                                    \
        int i;                                                                                    \
        int my_active_set_pe = ((nvshmemi_device_state_d.mype - start) / stride);                 \
                                                                                                  \
        base = (char *)((void *)__ldg(                                                            \
            (const long long unsigned *)nvshmemi_device_state_d.peer_heap_base +                  \
            nvshmemi_device_state_d.mype));                                                       \
        src_offset = ((char *)source - base);                                                     \
                                                                                                  \
        next_rank = start + ((my_active_set_pe + 1) % size) * stride;                             \
        next_offset = src_offset;                                                                 \
        peer_base = (char *)((void *)__ldg(                                                       \
            (const long long unsigned *)nvshmemi_device_state_d.peer_heap_base + next_rank));     \
        peer_source = peer_base + next_offset;                                                    \
        GPU_LINEAR_REDUCE_THREADGROUP(SC, TYPENAME, TYPE, OP, (void *)source, peer_source, dest,  \
                                      nreduce, myIdx, groupSize);                                 \
                                                                                                  \
        for (i = 2; i < size; i++) {                                                              \
            next_rank = start + ((my_active_set_pe + i) % size) * stride;                         \
            next_offset = src_offset;                                                             \
            peer_base = (char *)((void *)__ldg(                                                   \
                (const long long unsigned *)nvshmemi_device_state_d.peer_heap_base + next_rank)); \
            peer_source = peer_base + next_offset;                                                \
            GPU_LINEAR_REDUCE_THREADGROUP(SC, TYPENAME, TYPE, OP, dest, peer_source, dest,        \
                                          nreduce, myIdx, groupSize);                             \
        }                                                                                         \
        nvshmemxi_barrier_##SC(start, stride, size, pSync, NULL);                                 \
    } while (0)

// barrier limited
#define NVSHMEMXI_GPU_RDXN_THREADGROUP_GET_BAR(SC, TYPENAME, TYPE, OP, dest, source, nreduce,      \
                                               start, stride, size, pWrk, pSync)                   \
    do {                                                                                           \
        int next_rank = -1;                                                                        \
        size_t src_offset;                                                                         \
        size_t next_offset;                                                                        \
        char *base = NULL;                                                                         \
        char *peer_base = NULL;                                                                    \
        char *peer_source = NULL;                                                                  \
        NVSHMEMI_DECL_THREAD_IDX_##SC();                                                           \
        NVSHMEMI_DECL_THREADGROUP_SIZE_##SC();                                                     \
        int i;                                                                                     \
        int my_active_set_pe = ((nvshmemi_device_state_d.mype - start) / stride);                  \
                                                                                                   \
        next_rank = start + ((my_active_set_pe + 1) % size) * stride;                              \
        GPU_LINEAR_REDUCE_THREADGROUP_P2P_GET(SC, TYPENAME, TYPE, OP, source, source, next_rank,   \
                                              dest, nreduce, myIdx, groupSize, start, stride,      \
                                              size, pWrk, pSync);                                  \
                                                                                                   \
        for (i = 2; i < size; i++) {                                                               \
            next_rank = start + ((my_active_set_pe + i) % size) * stride;                          \
            GPU_LINEAR_REDUCE_THREADGROUP_P2P_GET(SC, TYPENAME, TYPE, OP, dest, source, next_rank, \
                                                  dest, nreduce, myIdx, groupSize, start, stride,  \
                                                  size, pWrk, pSync);                              \
        }                                                                                          \
        nvshmemxi_barrier_##SC(start, stride, size, (long *)pSync, NULL);                          \
    } while (0)


// barrier limited
#define NVSHMEMXI_GPU_RDXN_THREADGROUP_PUT_BAR(SC, TYPENAME, TYPE, OP, dest, source, nreduce, start,   \
                                               stride, size, pWrk, pSync)              \
    do {                                                                                        \
        int next_rank = -1;                                                                     \
        int counter = 0;                                                                        \
        NVSHMEMI_DECL_THREAD_IDX_##SC();                                                        \
        NVSHMEMI_DECL_THREADGROUP_SIZE_##SC();                                                  \
        int i;                                                                                  \
        int my_active_set_pe = ((nvshmemi_device_state_d.mype - start) / stride);               \
                                                                                                \
        next_rank = start + ((my_active_set_pe + 1) % size) * stride;                           \
        GPU_LINEAR_REDUCE_THREADGROUP_P2P_PUT(SC, TYPE, OP, source, source, next_rank, dest,    \
                                              counter, nreduce, myIdx, groupSize, start,        \
                                              stride, size, pSync);                     \
                                                                                                \
        for (i = 2; i < size; i++) {                                                            \
            next_rank = start + ((my_active_set_pe + i) % size) * stride;                       \
            GPU_LINEAR_REDUCE_THREADGROUP_P2P_PUT(SC, TYPENAME, TYPE, OP, dest, source, next_rank, dest,  \
                                                  counter, nreduce, myIdx, groupSize, start,    \
                                                  stride, size, pSync);                 \
        }                                                                                       \
        NVSHMEMI_SYNC_##SC();                                                                   \
    } while (0)

#define GPU_RDXN_SEGMENT_THREADGROUP(SC, TYPENAME, TYPE, OP, dest, source, nelems, start, stride, \
                                     size, pWrk, pSync)                                           \
    do {                                                                                          \
        int type_size = sizeof(TYPE);                                                             \
        size_t msg_len = nelems * type_size;                                                         \
        int next_rank = -1;                                                                       \
        NVSHMEMI_DECL_THREAD_IDX_##SC();                                                          \
        NVSHMEMI_DECL_THREADGROUP_SIZE_##SC();                                                    \
        TYPE *op1 = NULL, *op2 = NULL;                                                            \
        int i;                                                                                    \
        size_t j;                                                                                 \
        volatile TYPE *tmp_operand;                                                               \
        size_t remainder = 0;                                                                     \
        size_t rnds_floor = 0;                                                                    \
        size_t offset = 0;                                                                        \
        int pe_offset = 0;                                                                        \
        int pes_per_round = 0;                                                                    \
        int round = 0;                                                                            \
        size_t exchange_size = 0;                                                                 \
        int my_active_set_pe = ((nvshmemi_device_state_d.mype - start) / stride);                 \
        size_t nvshm_gpu_rdxn_seg_size = NVSHMEMI_REDUCE_MIN_WRKDATA_SIZE;                        \
                                                                                                  \
        tmp_operand = (TYPE *)pWrk;                                                               \
        nvshmemx_##TYPENAME##_put_nbi_##SC((TYPE *)dest, (const TYPE *)source, nelems,            \
                                           nvshmemi_device_state_d.mype);                         \
                                                                                                  \
        rnds_floor = msg_len / nvshm_gpu_rdxn_seg_size;                                           \
        remainder = msg_len % nvshm_gpu_rdxn_seg_size;                                            \
                                                                                                  \
        long sync_counter = NVSHMEMI_SYNC_VALUE + 1;                                              \
        for (j = 0; j < rnds_floor; j++) {                                                        \
            exchange_size = nvshm_gpu_rdxn_seg_size;                                              \
            for (i = 1; i < size; i++) {                                                          \
                next_rank = start + ((my_active_set_pe + i) % size) * stride;                     \
                nvshmemx_##TYPENAME##_put_nbi_##SC((TYPE *)tmp_operand, (const TYPE *)source + offset,      \
                                               (exchange_size / sizeof(TYPE)), next_rank);        \
                nvshmemxi_barrier_##SC(start, stride, size, (long *)pSync, &sync_counter);        \
                op1 = (TYPE *)dest + offset;                                                      \
                op2 = (TYPE *)tmp_operand;                                                        \
                GPU_LINEAR_REDUCE_THREADGROUP(SC, TYPENAME, TYPE, OP, op1, op2, op1,              \
                                              (exchange_size / sizeof(TYPE)), myIdx, groupSize);  \
                nvshmemxi_sync_##SC(start, stride, size, (long *)pSync, &sync_counter);           \
            }                                                                                     \
            offset += (exchange_size / sizeof(TYPE));                                             \
        }                                                                                         \
        if (remainder != 0) {                                                                     \
            exchange_size = remainder;                                                            \
            pes_per_round = nvshm_gpu_rdxn_seg_size / remainder;                                  \
            pe_offset = 1;                                                                        \
            do {                                                                                  \
                round = 0;                                                                        \
                for (i = pe_offset; ((round < pes_per_round) && (i < size)); i++) {               \
                    next_rank = start + ((my_active_set_pe + i) % size) * stride;                 \
                    nvshmemx_##TYPENAME##_put_nbi_##SC(                                           \
                        (TYPE *)((TYPE *)tmp_operand + (round * (exchange_size / sizeof(TYPE)))), \
                        (TYPE *)source + offset, (exchange_size / sizeof(TYPE)), next_rank);      \
                    round++;                                                                      \
                    pe_offset++;                                                                  \
                }                                                                                 \
                nvshmemxi_barrier_##SC(start, stride, size, (long *)pSync, &sync_counter);        \
                for (i = 0; i < round; i++) {                                                     \
                    op1 = (TYPE *)dest + offset;                                                  \
                    op2 = (TYPE *)((TYPE *)tmp_operand + (i * (exchange_size / sizeof(TYPE))));   \
                    GPU_LINEAR_REDUCE_THREADGROUP(SC, TYPENAME, TYPE, OP, op1, op2, op1,          \
                                                  (exchange_size / sizeof(TYPE)), myIdx,          \
                                                  groupSize);                                     \
                }                                                                                 \
                nvshmemxi_sync_##SC(start, stride, size, (long *)pSync, &sync_counter);           \
            } while (pe_offset < size);                                                           \
        }                                                                                         \
        NVSHMEMI_SYNC_##SC();                                                                     \
        for (i = myIdx; i < size; i += groupSize) {                                               \
            pSync[start + i * stride] = NVSHMEMI_SYNC_VALUE;                                      \
        }                                                                                         \
        NVSHMEMI_SYNC_##SC();                                                                     \
    } while (0)

// barrier limited
#define NVSHMEMXI_GPU_RDXN_THREADGROUP_PUT_BAR_DIRECT(SC, TYPENAME, TYPE, OP, dest, source,        \
                                                      nreduce, start, stride, size, pWrk, pSync)   \
    do {                                                                                           \
        int next_rank = -1;                                                                        \
        size_t src_offset = -1;                                                                    \
        size_t next_offset = -1;                                                                   \
        char *base = NULL;                                                                         \
        char *peer_base = NULL;                                                                    \
        char *peer_source = NULL;                                                                  \
        int counter = 0;                                                                           \
        int my_active_set_pe = ((nvshmemi_device_state_d.mype - start) / stride);                  \
        NVSHMEMI_DECL_THREAD_IDX_##SC();                                                           \
        NVSHMEMI_DECL_THREADGROUP_SIZE_##SC();                                                     \
        int i;                                                                                     \
                                                                                                   \
        base = (char *)((void *)__ldg(                                                             \
            (const long long unsigned *)nvshmemi_device_state_d.peer_heap_base +                   \
            nvshmemi_device_state_d.mype));                                                        \
        src_offset = ((char *)source - base);                                                      \
                                                                                                   \
        next_rank = start + ((my_active_set_pe + 1) % size) * stride;                              \
        next_offset = src_offset;                                                                  \
        peer_base = (char *)((void *)__ldg(                                                        \
            (const long long unsigned *)nvshmemi_device_state_d.peer_heap_base + next_rank));      \
        peer_source = peer_base + next_offset;                                                     \
        GPU_LINEAR_REDUCE_THREADGROUP_P2P_PUT_DIRECT(SC, TYPENAME, TYPE, OP, source, peer_source,  \
                                                     next_rank, dest, counter, nreduce, myIdx,     \
                                                     groupSize, start, stride, size, pWrk, pSync); \
                                                                                                   \
        for (i = 2; i < size; i++) {                                                               \
            next_rank = start + ((my_active_set_pe + i) % size) * stride;                          \
            next_offset = src_offset;                                                              \
            peer_base = (char *)((void *)__ldg(                                                    \
                (const long long unsigned *)nvshmemi_device_state_d.peer_heap_base + next_rank));  \
            peer_source = peer_base + next_offset;                                                 \
            GPU_LINEAR_REDUCE_THREADGROUP_P2P_PUT_DIRECT(                                          \
                SC, TYPENAME, TYPE, OP, dest, peer_source, next_rank, dest, counter, nreduce,      \
                myIdx, groupSize, start, stride, size, pWrk, pSync);                               \
        }                                                                                          \
        NVSHMEMI_SYNC_##SC();                                                                      \
    } while (0)

#define NVSHMEMXI_GPU_RDXN_THREADGROUP_PUTRING(SC, TYPENAME, TYPE, OP, dest, source, nelems, start,      \
                                               stride, size, pWrk, pSync)                                \
    do {                                                                                                 \
        int next_rank = -1;                                                                              \
        int prev_rank = -1;                                                                              \
        int i;                                                                                           \
        size_t j;                                                                                        \
        int PE_end = start + (stride * size);                                                            \
        uint32_t tmp[2];                                                                                 \
        uint64_t *tmp_rdxn = pWrk;                                                                       \
        uint64_t *tmp_rdxn_2 = &tmp_rdxn[1];                                                             \
        uint32_t payld;                                                                                  \
        NVSHMEMI_DECL_THREAD_IDX_##SC();                                                                 \
        NVSHMEMI_DECL_THREADGROUP_SIZE_##SC();                                                           \
        volatile uint32_t *my_notif_ptr = NULL;                                                          \
        size_t subelems = sizeof(TYPE) / sizeof(uint32_t);                                               \
        tmp[1] = 1;                                                                                      \
        next_rank = (nvshmemi_mype_d != (PE_end - stride)) ? (nvshmemi_device_state_d.mype + stride) : start;         \
        prev_rank = (nvshmemi_mype_d != start) ? (nvshmemi_device_state_d.mype - stride) : (PE_end - stride);         \
        my_notif_ptr = (uint32_t *)((uint32_t *)((uint64_t *)pWrk + (nelems * subelems)) + myIdx);       \
                                                                                                         \
        for (j = myIdx; j < nelems * subelems; j += groupSize) {                                         \
            payld = *((uint32_t *)source + j);                                                           \
            tmp[0] = payld;                                                                              \
            *tmp_rdxn = *((uint64_t *)tmp);                                                              \
            nvshmemx_signal_op((pWrk + j), *tmp_rdxn, NVSHMEM_SIGNAL_SET, next_rank);                    \
        }                                                                                                \
        GPU_HEAD_CHECK_OP_THREADGROUP(TYPENAME, TYPE, OP, dest, pWrk, source, nelems, myIdx, groupSize); \
        /* sync needed on volta (intermittent hangs seen otherwise) */                                   \
        NVSHMEMI_SYNC_##SC();                                                                            \
                                                                                                         \
        for (i = 1; i < (size - 1); i++) {                                                               \
            __threadfence_system();                                                                      \
            /* Don't want notification to overtake data transfer */                                      \
            *tmp_rdxn_2 = 1;                                                                             \
            nvshmemx_signal_op(pWrk + (nelems * subelems) + myIdx, *tmp_rdxn_2, NVSHMEM_SIGNAL_SET,      \
                               prev_rank);                                                               \
            while (1 != *my_notif_ptr)                                                                   \
                ;                                                                                        \
            *my_notif_ptr = 0;                                                                           \
            for (j = myIdx; j < nelems * subelems; j += groupSize) {                                     \
                payld = *((uint32_t *)dest + j);                                                         \
                tmp[0] = payld;                                                                          \
                *tmp_rdxn = *((uint64_t *)tmp);                                                          \
                nvshmemx_signal_op((pWrk + j), *tmp_rdxn, NVSHMEM_SIGNAL_SET, next_rank);                \
            }                                                                                            \
            GPU_HEAD_CHECK_OP_THREADGROUP(TYPENAME, TYPE, OP, dest, pWrk, source, nelems, myIdx,         \
                                          groupSize);                                                    \
            NVSHMEMI_SYNC_##SC();                                                                        \
        }                                                                                                \
    } while (0)

#define NVSHMEMXI_GPU_RDXN_THREADGROUP_PUTRING_DIRECT(                                             \
    SC, TYPENAME, TYPE, OP, dest, source, nelems, start, stride, size, pWrk, pSync)                \
    do {                                                                                           \
        size_t offset;                                                                             \
        char *notif_pwrk_dest;                                                                     \
        char *round_pwrk_dest;                                                                     \
        int next_rank = -1;                                                                        \
        int prev_rank = -1;                                                                        \
        int i;                                                                                     \
        size_t j;                                                                                  \
        int PE_end = start + (stride * size);                                                      \
        uint32_t tmp[2];                                                                           \
        uint32_t payld;                                                                            \
        uint32_t *notif_ptr = NULL;                                                                \
        volatile uint32_t *my_notif_ptr = NULL;                                                    \
        NVSHMEMI_DECL_THREAD_IDX_##SC();                                                           \
        NVSHMEMI_DECL_THREADGROUP_SIZE_##SC();                                                     \
        size_t subelems = sizeof(TYPE) / sizeof(uint32_t);                                         \
        tmp[1] = 1;                                                                                \
        offset =                                                                                   \
            (char *)pWrk - (char *)(__ldg((const long long unsigned *)nvshmemi_device_state_d.peer_heap_base +  \
                                          nvshmemi_mype_d));                                       \
        my_notif_ptr = (uint32_t *)((uint32_t *)((uint64_t *)pWrk + (nelems * subelems)) + myIdx); \
        next_rank = (nvshmemi_device_state_d.mype != (PE_end - stride)) ? (nvshmemi_device_state_d.mype + stride) : start;   \
        round_pwrk_dest =                                                                          \
            (char *)(__ldg((const long long unsigned *)nvshmemi_device_state_d.peer_heap_base + next_rank)) +   \
            offset;                                                                                \
        prev_rank = (nvshmemi_device_state_d.mype != start) ? (nvshmemi_device_state_d.mype - stride) : (PE_end - stride);   \
        notif_pwrk_dest =                                                                          \
            (char *)(__ldg((const long long unsigned *)nvshmemi_device_state_d.peer_heap_base + prev_rank)) +   \
            offset;                                                                                \
        notif_ptr =                                                                                \
            (uint32_t *)((uint32_t *)((uint64_t *)notif_pwrk_dest + (nelems * subelems)) + myIdx); \
                                                                                                   \
        for (j = myIdx; j < nelems * subelems; j += groupSize) {                                   \
            payld = *((uint32_t *)source + j);                                                     \
            tmp[0] = payld;                                                                        \
            *((uint64_t *)round_pwrk_dest + j) = *((uint64_t *)tmp);                               \
        }                                                                                          \
        GPU_HEAD_CHECK_OP_THREADGROUP(TYPENAME, TYPE, OP, dest, pWrk, source, nelems, myIdx,       \
                                      groupSize);                                                  \
        /* sync needed on volta (intermittent hangs seen otherwise) */                             \
        NVSHMEMI_SYNC_##SC();                                                                      \
                                                                                                   \
        for (i = 1; i < (size - 1); i++) {                                                         \
            __threadfence_system();                                                                \
            /* Don't want notification to overtake data transfer */                                \
            *notif_ptr = 1;                                                                        \
            while (1 != *my_notif_ptr)                                                             \
                ;                                                                                  \
            *my_notif_ptr = 0;                                                                     \
            for (j = myIdx; j < nelems * subelems; j += groupSize) {                               \
                payld = *((uint32_t *)dest + j);                                                   \
                tmp[0] = payld;                                                                    \
                *((uint64_t *)round_pwrk_dest + j) = *((uint64_t *)tmp);                           \
            }                                                                                      \
            GPU_HEAD_CHECK_OP_THREADGROUP(TYPENAME, TYPE, OP, dest, pWrk, source, nelems, myIdx,   \
                                          groupSize);                                              \
            NVSHMEMI_SYNC_##SC();                                                                  \
        }                                                                                          \
    } while (0)

#define NVSHMEMXI_GPU_RDXN_THREADGROUP_PUTALL(SC, TYPENAME, TYPE, OP, dest, source, nelems, start, \
                                              stride, size, pWrk, pSync)                           \
    do {                                                                                           \
        int i;                                                                                     \
        size_t j;                                                                                  \
        int PE_end = start + (stride * size);                                                      \
        NVSHMEMI_DECL_THREAD_IDX_##SC();                                                           \
        NVSHMEMI_DECL_THREADGROUP_SIZE_##SC();                                                     \
        uint64_t *tmp_rdxn = pWrk;                                                                 \
        uint32_t tmp[2];                                                                           \
        uint32_t payld;                                                                            \
        size_t subelems = sizeof(TYPE) / sizeof(uint32_t);                                         \
        int my_active_set_pe = ((nvshmemi_device_state_d.mype - start) / stride);                  \
        tmp[1] = 1;                                                                                \
                                                                                                   \
        for (j = myIdx; j < nelems * subelems; j += groupSize) {                                   \
            payld = *((uint32_t *)source + j);                                                     \
            tmp[0] = payld;                                                                        \
            *tmp_rdxn = *((uint64_t *)tmp);                                                        \
            for (i = start; i < nvshmemi_device_state_d.mype; i += stride) {                       \
                nvshmemx_signal_op((pWrk + j + (nelems * subelems * my_active_set_pe)),            \
                                     *tmp_rdxn, NVSHMEM_SIGNAL_SET, i);                            \
            }                                                                                      \
            for (i = nvshmemi_device_state_d.mype + stride; i < PE_end; i += stride) {             \
                nvshmemx_signal_op((pWrk + j + (nelems * subelems * my_active_set_pe)),            \
                                     *tmp_rdxn, NVSHMEM_SIGNAL_SET, i);                            \
            }                                                                                      \
        }                                                                                          \
        GPU_HEAD_CHECKALL_OP_THREADGROUP(TYPENAME, TYPE, OP, dest, pWrk, source, nelems, start,    \
                                         stride, size, myIdx, groupSize);                          \
        __threadfence();                                                                           \
        NVSHMEMI_SYNC_##SC();                                                                      \
    } while (0)

#define NVSHMEMXI_GPU_RDXN_THREADGROUP_PUTALL_DIRECT(SC, TYPENAME, TYPE, OP, dest, source, nelems, \
                                                     start, stride, size, pWrk, pSync)             \
    do {                                                                                           \
        size_t offset;                                                                             \
        char *round_pwrk_dest;                                                                     \
        int i, j;                                                                                  \
        int PE_end = start + (stride * size);                                                      \
        NVSHMEMI_DECL_THREAD_IDX_##SC();                                                           \
        NVSHMEMI_DECL_THREADGROUP_SIZE_##SC();                                                     \
        uint32_t tmp[2];                                                                           \
        uint32_t payld;                                                                            \
        size_t subelems = sizeof(TYPE) / sizeof(uint32_t);                                         \
        int my_active_set_pe = ((nvshmemi_device_state_d.mype - start) / stride);                  \
        tmp[1] = 1;                                                                                \
        offset =                                                                                   \
            (char *)pWrk -                                                                         \
            (char *)(__ldg((const long long unsigned *)nvshmemi_device_state_d.peer_heap_base +    \
                           nvshmemi_device_state_d.mype));                                         \
                                                                                                   \
        for (j = myIdx; j < nelems * subelems; j += groupSize) {                                   \
            payld = *((uint32_t *)source + j);                                                     \
            tmp[0] = payld;                                                                        \
            for (i = nvshmemi_device_state_d.mype + stride; i < PE_end; i += stride) {             \
                round_pwrk_dest =                                                                  \
                    (char *)(__ldg(                                                                \
                        (const long long unsigned *)nvshmemi_device_state_d.peer_heap_base + i)) + \
                    offset;                                                                        \
                *((uint64_t *)round_pwrk_dest + j + (nelems * subelems * my_active_set_pe)) =      \
                    *((uint64_t *)tmp);                                                            \
            }                                                                                      \
            for (i = start; i < nvshmemi_device_state_d.mype; i += stride) {                       \
                round_pwrk_dest =                                                                  \
                    (char *)(__ldg(                                                                \
                        (const long long unsigned *)nvshmemi_device_state_d.peer_heap_base + i)) + \
                    offset;                                                                        \
                *((uint64_t *)round_pwrk_dest + j + (nelems * subelems * my_active_set_pe)) =      \
                    *((uint64_t *)tmp);                                                            \
            }                                                                                      \
        }                                                                                          \
        GPU_HEAD_CHECKALL_OP_THREADGROUP(TYPENAME, TYPE, OP, dest, pWrk, source, nelems, start,    \
                                         stride, size, myIdx, groupSize);                          \
        __threadfence();                                                                           \
        NVSHMEMI_SYNC_##SC();                                                                      \
    } while (0)

#ifdef NVSHMEM_GPU_COLL_USE_LDST
#ifdef NVSHMEM_DISABLE_COLL_POLL
#define NVSHMEMXI_GPU_RDXN_THREADGROUP(SC, TYPENAME, TYPE, OP, dest, source, nelems, start, stride, \
                                       size, pWrk, pSync)                                           \
    do {                                                                                            \
        NVSHMEMXI_GPU_RDXN_THREADGROUP_ZCOPY_GET_BAR_DIRECT(                                        \
            SC, TYPENAME, TYPE, OP, dest, source, nreduce, start, stride, size, pWrk, pSync);       \
    } while (0)
#else
#define NVSHMEMXI_GPU_RDXN_THREADGROUP(SC, TYPENAME, TYPE, OP, dest, source, nelems, start, stride, \
                                       size, pWrk, pSync)                                           \
    do {                                                                                            \
        size_t subelems = sizeof(TYPE) / sizeof(uint32_t);                                          \
        size_t pwrk_req_sz_allgather = ((subelems * nelems) * sizeof(uint64_t)) * size;             \
        size_t wrk_size = NVSHMEMI_REDUCE_MIN_WRKDATA_SIZE * sizeof(TYPE);                             \
        if (subelems && (pwrk_req_sz_allgather <= wrk_size)) {                                      \
            NVSHMEMXI_GPU_RDXN_THREADGROUP_PUTALL_DIRECT(SC, TYPENAME, TYPE, OP, dest, source, nreduce,      \
                                                         start, stride, size, pWrk,                 \
                                                         pSync);                                    \
        } else {                                                                                    \
            NVSHMEMXI_GPU_RDXN_THREADGROUP_ZCOPY_GET_BAR_DIRECT(SC, TYPENAME, TYPE, OP, dest, source,        \
                                                                nreduce, start, stride,             \
                                                                size, pWrk, pSync);                 \
        }                                                                                           \
    } while (0)
#endif
#else
#ifdef NVSHMEM_DISABLE_COLL_POLL
#define NVSHMEMXI_GPU_RDXN_THREADGROUP(SC, TYPENAME, TYPE, OP, dest, source, nelems, start, stride, \
                                       size, pWrk, pSync)                                           \
    do {                                                                                            \
        GPU_RDXN_SEGMENT_THREADGROUP(SC, TYPENAME, TYPE, OP, dest, source, nelems, start, stride,   \
                                     size, pWrk, pSync);                                            \
    } while (0)
#else
#define NVSHMEMXI_GPU_RDXN_THREADGROUP(SC, TYPENAME, TYPE, OP, dest, source, nelems, start, stride, \
                                       size, pWrk, pSync)                                           \
    do {                                                                                           \
        size_t subelems = sizeof(TYPE) / sizeof(uint32_t);                                         \
        size_t pwrk_req_sz_allgather = ((subelems * nelems) * sizeof(uint64_t)) * size;            \
        NVSHMEMI_DECL_THREADGROUP_SIZE_##SC();                                                     \
        size_t pwrk_req_sz_ring =                                                                     \
            ((subelems * nelems) * sizeof(uint64_t)) + (groupSize * sizeof(uint32_t));             \
        size_t wrk_size = NVSHMEMI_REDUCE_MIN_WRKDATA_SIZE * sizeof(TYPE);                            \
        if (subelems && pwrk_req_sz_allgather <= wrk_size) {                                       \
            NVSHMEMXI_GPU_RDXN_THREADGROUP_PUTALL(SC, TYPENAME, TYPE, OP, dest, source, nreduce, start,   \
                                                  stride, size, pWrk, pSync);             \
        } else if (subelems && pwrk_req_sz_ring <= wrk_size) {                                     \
            NVSHMEMXI_GPU_RDXN_THREADGROUP_PUTRING(SC, TYPENAME, TYPE, OP, dest, source, nreduce, start,  \
                                                   stride, size, pWrk, pSync);            \
        } else {                                                                                   \
            NVSHMEMXI_GPU_RDXN_THREADGROUP_PUT_BAR(SC, TYPENAME, TYPE, OP, dest, source, nreduce, start,  \
                                                   stride, size, pWrk, pSync);            \
        }                                                                                          \
    } while (0)
#endif
#endif


#define DEFN_NVSHMEMXI_TYPENAME_OP_REDUCE_THREADGROUP(SC, TYPENAME, TYPE, OP)                     \
    __device__ void nvshmemxi_##TYPENAME##_##OP##_reduce_##SC(TYPE * dest, const TYPE *source,    \
          size_t nreduce, int start, int stride, int size, uint64_t *pWrk, uint64_t *pSync) {     \
        NVSHMEMXI_GPU_RDXN_THREADGROUP(SC, TYPENAME, TYPE, OP, dest, source, nreduce, start,      \
                                       stride, size, pWrk, pSync);                                \
    }

#define DEFN_NVSHMEMI_REDUCE_THREADGROUP(SC)                                                                 \
NVSHMEMI_REPT_FOR_BITWISE_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMXI_TYPENAME_OP_REDUCE_THREADGROUP, SC, and)    \
NVSHMEMI_REPT_FOR_BITWISE_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMXI_TYPENAME_OP_REDUCE_THREADGROUP, SC, or)     \
NVSHMEMI_REPT_FOR_BITWISE_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMXI_TYPENAME_OP_REDUCE_THREADGROUP, SC, xor)    \
                                                                                                             \
NVSHMEMI_REPT_FOR_STANDARD_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMXI_TYPENAME_OP_REDUCE_THREADGROUP, SC, max)   \
NVSHMEMI_REPT_FOR_STANDARD_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMXI_TYPENAME_OP_REDUCE_THREADGROUP, SC, min)   \
                                                                                                             \
NVSHMEMI_REPT_FOR_ARITH_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMXI_TYPENAME_OP_REDUCE_THREADGROUP, SC, sum)      \
NVSHMEMI_REPT_FOR_ARITH_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMXI_TYPENAME_OP_REDUCE_THREADGROUP, SC, prod)

DEFN_NVSHMEMI_REDUCE_THREADGROUP(warp);
DEFN_NVSHMEMI_REDUCE_THREADGROUP(block);
#undef DEFN_NVSHMEMXI_TYPENAME_OP_REDUCE_THREADGROUP
#undef DEFN_NVSHMEMI_REDUCE_THREADGROUP

#define DEFN_NVSHMEMX_TYPENAME_OP_REDUCE_THREADGROUP(SC, TYPENAME, TYPE, OP)                      \
    __device__ int nvshmemx_##TYPENAME##_##OP##_reduce_##SC(nvshmem_team_t team, TYPE *dest,      \
                                                           const TYPE *source, size_t nreduce) {  \
        nvshmemi_team_t *teami = nvshmemi_device_state_d.team_pool[team];                         \
        uint64_t *pWrk = (uint64_t *) nvshmemi_team_get_psync(teami, REDUCE);                     \
        uint64_t *pSync = (uint64_t *) ((uint64_t *)pWrk + NVSHMEMI_REDUCE_MIN_WRKDATA_SIZE);     \
        nvshmemx_barrier_##SC(team);                                                              \
        nvshmemxi_##TYPENAME##_##OP##_reduce_##SC(dest, source, nreduce, teami->start,            \
                                       teami->stride, teami->size, pWrk, pSync);                  \
        return 0;                                                                                 \
    }

#define DEFN_NVSHMEM_REDUCE_THREADGROUP(SC)                                                                 \
NVSHMEMI_REPT_FOR_BITWISE_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMX_TYPENAME_OP_REDUCE_THREADGROUP, SC, and)    \
NVSHMEMI_REPT_FOR_BITWISE_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMX_TYPENAME_OP_REDUCE_THREADGROUP, SC, or)     \
NVSHMEMI_REPT_FOR_BITWISE_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMX_TYPENAME_OP_REDUCE_THREADGROUP, SC, xor)    \
                                                                                                            \
NVSHMEMI_REPT_FOR_STANDARD_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMX_TYPENAME_OP_REDUCE_THREADGROUP, SC, max)   \
NVSHMEMI_REPT_FOR_STANDARD_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMX_TYPENAME_OP_REDUCE_THREADGROUP, SC, min)   \
                                                                                                            \
NVSHMEMI_REPT_FOR_ARITH_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMX_TYPENAME_OP_REDUCE_THREADGROUP, SC, sum)      \
NVSHMEMI_REPT_FOR_ARITH_REDUCE_TYPES_WITH_SCOPE(DEFN_NVSHMEMX_TYPENAME_OP_REDUCE_THREADGROUP, SC, prod)

DEFN_NVSHMEM_REDUCE_THREADGROUP(warp);
DEFN_NVSHMEM_REDUCE_THREADGROUP(block);
#undef DEFN_NVSHMEMX_TYPENAME_OP_REDUCE_THREADGROUP
#undef DEFN_NVSHMEM_REDUCE_THREADGROUP

#endif
