#include <cuda_runtime.h>
#include <nccl.h>

#include <cstdio>
#include <cstdlib>
#include <vector>

#define CUDA_CHECK(call) do { if ((call) != cudaSuccess) std::exit(2); } while (0)
#define NCCL_CHECK(call) do { if ((call) != ncclSuccess) std::exit(3); } while (0)

int main() {
    int device_count = 0;
    CUDA_CHECK(cudaGetDeviceCount(&device_count));
    if (device_count < 2) {
        std::fprintf(stderr, "two GPUs required\n");
        return 77;
    }

    constexpr int ranks = 2;
    constexpr int elements = 1024;
    const int devices[ranks] = {0, 1};
    std::vector<ncclComm_t> communicators(ranks);
    std::vector<cudaStream_t> streams(ranks);
    std::vector<float*> buffers(ranks);
    NCCL_CHECK(ncclCommInitAll(communicators.data(), ranks, devices));

    for (int rank = 0; rank < ranks; ++rank) {
        CUDA_CHECK(cudaSetDevice(devices[rank]));
        CUDA_CHECK(cudaStreamCreate(&streams[rank]));
        CUDA_CHECK(cudaMalloc(&buffers[rank], elements * sizeof(float)));
        std::vector<float> host(elements, static_cast<float>(rank + 1));
        CUDA_CHECK(cudaMemcpyAsync(buffers[rank], host.data(), elements * sizeof(float),
                                   cudaMemcpyHostToDevice, streams[rank]));
    }

    NCCL_CHECK(ncclGroupStart());
    for (int rank = 0; rank < ranks; ++rank) {
        NCCL_CHECK(ncclAllReduce(buffers[rank], buffers[rank], elements, ncclFloat,
                                 ncclSum, communicators[rank], streams[rank]));
    }
    NCCL_CHECK(ncclGroupEnd());

    for (int rank = 0; rank < ranks; ++rank) {
        CUDA_CHECK(cudaSetDevice(devices[rank]));
        float first = 0.0F;
        CUDA_CHECK(cudaMemcpyAsync(&first, buffers[rank], sizeof(first),
                                   cudaMemcpyDeviceToHost, streams[rank]));
        CUDA_CHECK(cudaStreamSynchronize(streams[rank]));
        std::printf("rank=%d first=%.1f expected=3.0\n", rank, first);
        CUDA_CHECK(cudaFree(buffers[rank]));
        CUDA_CHECK(cudaStreamDestroy(streams[rank]));
        NCCL_CHECK(ncclCommDestroy(communicators[rank]));
    }
}