/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/*
 * This sample demonstrates Inter Process Communication
 * using one process per GPU for computation.
 */

#include <cuda.h>
#include <stdio.h>
#include <stdlib.h>
#include <vector>
#define CUDA_DRIVER_API 1
#include "helper_cuda.h"
#include "helper_cuda_drvapi.h"
#include "helper_multiprocess.h"

static const char shmName[] = "streamOrderedAllocationIPCshm";
static const char ipcName[] = "streamOrderedAllocationIPC_pipe";
// For direct NVLINK and PCI-E peers, at max 8 simultaneous peers are allowed
// For NVSWITCH connected peers like DGX-2, simultaneous peers are not limited
// in the same way.
#define MAX_DEVICES (32)
#define DATA_SIZE   (64ULL << 20ULL) // 64MB

#if defined(__linux__)
#define cpu_atomic_add32(a, x) __sync_add_and_fetch(a, x)
#elif defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
#define cpu_atomic_add32(a, x) InterlockedAdd((volatile LONG *)a, x)
#else
#error Unsupported system
#endif

typedef struct shmStruct_st
{
    size_t                      nprocesses;
    int                         barrier;
    int                         sense;
    cudaMemAllocationHandleType handleType;
    int                         devices[MAX_DEVICES];
    cudaMemPoolPtrExportData    exportPtrData[MAX_DEVICES];
} shmStruct;

__global__ void simpleKernel(char *ptr, int sz, char val)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    for (; idx < sz; idx += (gridDim.x * blockDim.x)) {
        ptr[idx] = val;
    }
}

static void barrierWait(volatile int *barrier, volatile int *sense, unsigned int n)
{
    int count;

    // Check-in
    count = cpu_atomic_add32(barrier, 1);
    if (count == n) // Last one in
        *sense = 1;
    while (!*sense)
        ;

    // Check-out
    count = cpu_atomic_add32(barrier, -1);
    if (count == 0) // Last one out
        *sense = 0;
    while (*sense)
        ;
}

static void childProcess(int id)
{
    volatile shmStruct *shm = NULL;
    cudaStream_t        stream;
    sharedMemoryInfo    info;
    size_t              procCount, i;
    int                 blocks  = 0;
    int                 threads = 128;
    cudaDeviceProp      prop;
    std::vector<void *> ptrs;
    pid_t               pid;
    char                pidString[20] = {0};
    char                lshmName[40]  = {0};

    std::vector<char> verification_buffer(DATA_SIZE);

    pid = getppid();
    snprintf(pidString, sizeof(pidString), "%d", pid);
    strcat(lshmName, shmName);
    strcat(lshmName, pidString);

    printf("CP: lshmName = %s\n", lshmName);

    ipcHandle *ipcChildHandle = NULL;
    checkIpcErrors(ipcOpenSocket(ipcChildHandle));

    if (sharedMemoryOpen(lshmName, sizeof(shmStruct), &info) != 0) {
        printf("Failed to create shared memory slab\n");
        exit(EXIT_FAILURE);
    }
    shm       = (volatile shmStruct *)info.addr;
    procCount = shm->nprocesses;

    barrierWait(&shm->barrier, &shm->sense, (unsigned int)(procCount + 1));

    // Receive all allocation handles shared by Parent.
    std::vector<ShareableHandle> shHandle(shm->nprocesses);
    checkIpcErrors(ipcRecvShareableHandles(ipcChildHandle, shHandle));

    checkCudaErrors(cudaSetDevice(shm->devices[id]));
    checkCudaErrors(cudaGetDeviceProperties(&prop, shm->devices[id]));
    checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
    checkCudaErrors(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, simpleKernel, threads, 0));
    blocks *= prop.multiProcessorCount;

    std::vector<cudaMemPool_t> pools(shm->nprocesses);

    cudaMemAllocationHandleType handleType = shm->handleType;

    // Import mem pools from all the devices created in the master
    // process using shareable handles received via socket
    // and import the pointer to the allocated buffer using
    // exportData filled in shared memory by the master process.
    for (i = 0; i < procCount; i++) {
        checkCudaErrors(cudaMemPoolImportFromShareableHandle(&pools[i], (void *)shHandle[i], handleType, 0));

        cudaMemAccessFlags accessFlags;
        cudaMemLocation    location;
        location.type = cudaMemLocationTypeDevice;
        location.id   = shm->devices[id];
        checkCudaErrors(cudaMemPoolGetAccess(&accessFlags, pools[i], &location));
        if (accessFlags != cudaMemAccessFlagsProtReadWrite) {
            cudaMemAccessDesc desc;
            memset(&desc, 0, sizeof(cudaMemAccessDesc));
            desc.location.type = cudaMemLocationTypeDevice;
            desc.location.id   = shm->devices[id];
            desc.flags         = cudaMemAccessFlagsProtReadWrite;
            checkCudaErrors(cudaMemPoolSetAccess(pools[i], &desc, 1));
        }

        // Import the allocation from each memory pool by iterating over exportData
        // until import is success.
        for (int j = 0; j < procCount; j++) {
            void *ptr = NULL;
            // Import the allocation using the opaque export data retrieved through
            // the shared memory".
            cudaError_t ret =
                cudaMemPoolImportPointer(&ptr, pools[i], (cudaMemPoolPtrExportData *)&shm->exportPtrData[j]);

            if (ret == cudaSuccess) {
                // Pointer import is successful hence add it to the ptrs bag.
                ptrs.push_back(ptr);
                break;
            }
            else {
                // Reset failure error received from cudaMemPoolImportPointer
                // for further try.
                cudaGetLastError();
            }
        }
        // Since we have imported allocations shared by the parent with us, we can
        // close this ShareableHandle.
        checkIpcErrors(ipcCloseShareableHandle(shHandle[i]));
    }

    // Since we have imported allocations shared by the parent with us, we can
    // close the socket.
    checkIpcErrors(ipcCloseSocket(ipcChildHandle));

    // At each iteration of the loop, each sibling process will push work on
    // their respective devices accessing the next peer mapped buffer allocated
    // by the master process (these can come from other sibling processes as
    // well). To coordinate each process' access, we force the stream to wait for
    // the work already accessing this buffer.
    for (i = 0; i < procCount; i++) {
        size_t bufferId = (i + id) % procCount;

        // Push a simple kernel on it
        simpleKernel<<<blocks, threads, 0, stream>>>((char *)ptrs[bufferId], DATA_SIZE, id);
        checkCudaErrors(cudaGetLastError());
        checkCudaErrors(cudaStreamSynchronize(stream));

        // Wait for all my sibling processes to push this stage of their work
        // before proceeding to the next. This prevents siblings from racing
        // ahead and clobbering the recorded event or waiting on the wrong
        // recorded event.
        barrierWait(&shm->barrier, &shm->sense, (unsigned int)procCount);
        if (id == 0) {
            printf("Step %lld done\n", (unsigned long long)i);
        }
    }

    // Now wait for my buffer to be ready so I can copy it locally and verify it
    checkCudaErrors(cudaMemcpyAsync(&verification_buffer[0], ptrs[id], DATA_SIZE, cudaMemcpyDeviceToHost, stream));

    // And wait for all the queued up work to complete
    checkCudaErrors(cudaStreamSynchronize(stream));

    printf("Process %d: verifying...\n", id);

    // The contents should have the id of the sibling just after me
    char compareId = (char)((id + 1) % procCount);
    for (unsigned long long j = 0; j < DATA_SIZE; j++) {
        if (verification_buffer[j] != compareId) {
            printf("Process %d: Verification mismatch at %lld: %d != %d\n",
                   id,
                   j,
                   (int)verification_buffer[j],
                   (int)compareId);
        }
    }

    // Clean up!
    for (i = 0; i < procCount; i++) {
        // Free the memory before the exporter process frees it
        checkCudaErrors(cudaFreeAsync(ptrs[i], stream));
    }

    // And wait for all the queued up work to complete
    checkCudaErrors(cudaStreamSynchronize(stream));
    checkCudaErrors(cudaStreamDestroy(stream));

    printf("Process %d complete!\n", id);
}

static void parentProcess(char *app)
{
    sharedMemoryInfo            info;
    int                         devCount, i;
    volatile shmStruct         *shm = NULL;
    std::vector<void *>         ptrs;
    std::vector<Process>        processes;
    cudaMemAllocationHandleType handleType = cudaMemHandleTypeNone;
    pid_t                       pid;
    char                        pidString[20] = {0};
    char                        lshmName[40]  = {0};

    pid = getpid();
    snprintf(pidString, sizeof(pidString), "%d", pid);
    strcat(lshmName, shmName);
    strcat(lshmName, pidString);

    printf("PP: lshmName = %s\n", lshmName);

    checkCudaErrors(cudaGetDeviceCount(&devCount));
    std::vector<CUdevice> devices(devCount);
    for (i = 0; i < devCount; i++) {
        cuDeviceGet(&devices[i], i);
    }

    if (sharedMemoryCreate(lshmName, sizeof(*shm), &info) != 0) {
        printf("Failed to create shared memory slab\n");
        exit(EXIT_FAILURE);
    }
    shm = (volatile shmStruct *)info.addr;
    memset((void *)shm, 0, sizeof(*shm));

    // Pick all the devices that can access each other's memory for this test
    // Keep in mind that CUDA has minimal support for fork() without a
    // corresponding exec() in the child process, but in this case our
    // spawnProcess will always exec, so no need to worry.
    for (i = 0; i < devCount; i++) {
        bool           allPeers = true;
        cudaDeviceProp prop;
        checkCudaErrors(cudaGetDeviceProperties(&prop, i));

        int isMemPoolSupported = 0;
        checkCudaErrors(cudaDeviceGetAttribute(&isMemPoolSupported, cudaDevAttrMemoryPoolsSupported, i));
        // CUDA IPC is only supported on devices with unified addressing
        if (!isMemPoolSupported) {
            printf("Device %d does not support cuda memory pools, skipping...\n", i);
            continue;
        }
        int supportedHandleTypes = 0;
        checkCudaErrors(cudaDeviceGetAttribute(&supportedHandleTypes, cudaDevAttrMemoryPoolSupportedHandleTypes, i));
        if (supportedHandleTypes == 0) {
            printf("Device %d does not support Memory pool based IPC, skipping...\n", i);
            continue;
        }

        if (handleType == cudaMemHandleTypeNone) {
            if (supportedHandleTypes & cudaMemHandleTypePosixFileDescriptor) {
                handleType = cudaMemHandleTypePosixFileDescriptor;
            }
            else if (supportedHandleTypes & cudaMemHandleTypeWin32) {
                handleType = cudaMemHandleTypeWin32;
            }
            else {
                printf("Device %d does not support any supported handle types, skipping...\n", i);
                continue;
            }
        }
        else {
            if ((supportedHandleTypes & handleType) != handleType) {
                printf("Mixed handle types are not supported, waiving test\n");
                exit(EXIT_WAIVED);
            }
        }
        // This sample requires two processes accessing each device, so we need
        // to ensure exclusive or prohibited mode is not set
        int computeMode;
        checkCudaErrors(cudaDeviceGetAttribute(&computeMode, cudaDevAttrComputeMode, i));
        if (computeMode != cudaComputeModeDefault) {
            printf("Device %d is in an unsupported compute mode for this sample\n", i);
            continue;
        }
#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
        // CUDA IPC on Windows is only supported on TCC
        if (!prop.tccDriver) {
            printf("Device %d is not in TCC mode\n", i);
            continue;
        }
#endif

        for (int j = 0; j < shm->nprocesses; j++) {
            int canAccessPeerIJ, canAccessPeerJI;
            checkCudaErrors(cudaDeviceCanAccessPeer(&canAccessPeerJI, shm->devices[j], i));
            checkCudaErrors(cudaDeviceCanAccessPeer(&canAccessPeerIJ, i, shm->devices[j]));
            if (!canAccessPeerIJ || !canAccessPeerJI) {
                allPeers = false;
                break;
            }
        }
        if (allPeers) {
            // Enable peers here.  This isn't necessary for IPC, but it will
            // setup the peers for the device.  For systems that only allow 8
            // peers per GPU at a time, this acts to remove devices from CanAccessPeer
            for (int j = 0; j < shm->nprocesses; j++) {
                checkCudaErrors(cudaSetDevice(i));
                checkCudaErrors(cudaDeviceEnablePeerAccess(shm->devices[j], 0));
                checkCudaErrors(cudaSetDevice(shm->devices[j]));
                checkCudaErrors(cudaDeviceEnablePeerAccess(i, 0));
            }
            shm->devices[shm->nprocesses++] = i;
            if (shm->nprocesses >= MAX_DEVICES)
                break;
        }
        else {
            printf("Device %d is not peer capable with some other selected peers, "
                   "skipping\n",
                   i);
        }
    }

    if (shm->nprocesses == 0) {
        printf("No CUDA devices support IPC\n");
        exit(EXIT_WAIVED);
    }

    if (handleType == cudaMemHandleTypeNone) {
        printf("No supported handle types found, waiving test\n");
        exit(EXIT_WAIVED);
    }

    std::vector<ShareableHandle> shareableHandles(shm->nprocesses);
    std::vector<cudaStream_t>    streams(shm->nprocesses);
    std::vector<cudaMemPool_t>   pools(shm->nprocesses);

    // Now allocate memory for each process and fill the shared
    // memory buffer with the export data and get memPool handles to communicate
    for (i = 0; i < shm->nprocesses; i++) {
        void *ptr = NULL;
        checkCudaErrors(cudaSetDevice(shm->devices[i]));
        checkCudaErrors(cudaStreamCreateWithFlags(&streams[i], cudaStreamNonBlocking));
        // Allocate an explicit pool with IPC capabilities
        cudaMemPoolProps poolProps;
        memset(&poolProps, 0, sizeof(cudaMemPoolProps));
        poolProps.allocType   = cudaMemAllocationTypePinned;
        poolProps.handleTypes = handleType;

        poolProps.location.type = cudaMemLocationTypeDevice;
        poolProps.location.id   = shm->devices[i];

        checkCudaErrors(cudaMemPoolCreate(&pools[i], &poolProps));

        // Query the shareable handle for the pool
        // Allocate memory in a stream from the pool just created
        checkCudaErrors(cudaMallocAsync(&ptr, DATA_SIZE, pools[i], streams[i]));

        checkCudaErrors(cudaMemPoolExportToShareableHandle(&shareableHandles[i], pools[i], handleType, 0));

        // Memset handle to 0 to make sure call to `cudaMemPoolImportPointer` in
        // childProcess will fail if the following call to
        // `cudaMemPoolExportPointer` fails.
        memset((void *)&shm->exportPtrData[i], 0, sizeof(cudaMemPoolPtrExportData));
        // Get the opaque ‘bag-of-bits’ representing the allocation
        checkCudaErrors(cudaMemPoolExportPointer((cudaMemPoolPtrExportData *)&shm->exportPtrData[i], ptr));
        ptrs.push_back(ptr);
    }

    shm->handleType = handleType;

    // Launch the child processes!
    for (i = 0; i < shm->nprocesses; i++) {
        char        devIdx[10];
        char *const args[] = {app, devIdx, NULL};
        Process     process;

        SPRINTF(devIdx, "%d", i);

        if (spawnProcess(&process, app, args)) {
            printf("Failed to create process\n");
            exit(EXIT_FAILURE);
        }

        processes.push_back(process);
    }

    barrierWait(&shm->barrier, &shm->sense, (unsigned int)(shm->nprocesses + 1));

    ipcHandle *ipcParentHandle = NULL;
    checkIpcErrors(ipcCreateSocket(ipcParentHandle, ipcName, processes));
    checkIpcErrors(ipcSendShareableHandles(ipcParentHandle, shareableHandles, processes));

    // Close the shareable handles as they are not needed anymore.
    for (int i = 0; i < shm->nprocesses; i++) {
        checkIpcErrors(ipcCloseShareableHandle(shareableHandles[i]));
    }
    checkIpcErrors(ipcCloseSocket(ipcParentHandle));

    // And wait for them to finish
    for (i = 0; i < processes.size(); i++) {
        if (waitProcess(&processes[i]) != EXIT_SUCCESS) {
            printf("Process %d failed!\n", i);
            exit(EXIT_FAILURE);
        }
    }

    // Clean up!
    for (i = 0; i < shm->nprocesses; i++) {
        checkCudaErrors(cudaSetDevice(shm->devices[i]));
        checkCudaErrors(cudaFreeAsync(ptrs[i], streams[i]));
        checkCudaErrors(cudaStreamSynchronize(streams[i]));
        checkCudaErrors(cudaMemPoolDestroy(pools[i]));
    }

    sharedMemoryClose(&info);
}

// Host code
int main(int argc, char **argv)
{
#if defined(__arm__) || defined(__aarch64__) || defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
    printf("Not supported on ARM or Windows\n");
    return EXIT_WAIVED;
#else
    if (argc == 1) {
        parentProcess(argv[0]);
    }
    else {
        childProcess(atoi(argv[1]));
    }
    return EXIT_SUCCESS;
#endif
}
