/*
 * Copyright 2020 NVIDIA Corporation.  All rights reserved.
 *
 * Please refer to the NVIDIA end user license agreement (EULA) associated
 * with this source code for terms and conditions that govern your use of
 * this software. Any use, reproduction, disclosure, or distribution of
 * this software and related documentation outside the terms of the EULA
 * is strictly prohibited.
 *
 */

/*
* This sample demonstrates peer-to-peer access of stream ordered memory 
* allocated with cudaMallocAsync and cudaMemPool family of APIs through simple
* kernel which does peer-to-peer to access & scales vector elements.
*/

// System includes
#include <stdio.h>
#include <assert.h>
#include <map>
#include <iostream>
#include <set>
#include <utility>

// CUDA runtime
#include <cuda_runtime.h>

// helper functions and utilities to work with CUDA
#include <helper_functions.h>
#include <helper_cuda.h>


// Simple kernel to demonstrate copying cudaMallocAsync memory via P2P to peer device
__global__ void copyP2PAndScale(const int *src, int *dst, int N)
{
    int idx = blockIdx.x*blockDim.x + threadIdx.x;

    if (idx < N)
    {
        // scale & store src vector.
        dst[idx] = 2 * src[idx];
    }
}

// Map of device version to device number
std::multimap<std::pair<int, int>, int> getIdenticalGPUs() {
    int numGpus = 0;
    checkCudaErrors(cudaGetDeviceCount(&numGpus));

    std::multimap<std::pair<int, int>, int> identicalGpus;

    for (int i = 0; i < numGpus; i++) {
        int isMemPoolSupported = 0;
        checkCudaErrors(cudaDeviceGetAttribute(&isMemPoolSupported, cudaDevAttrMemoryPoolsSupported, i));

        // Filter unsupported devices
        if (isMemPoolSupported) {
            int major = 0, minor = 0;
            checkCudaErrors(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, i));
            checkCudaErrors(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, i));
            identicalGpus.emplace(std::make_pair(major, minor), i);
        }
    }

    return identicalGpus;
}

std::pair<int, int> getP2PCapableGpuPair()
{
    constexpr size_t kNumGpusRequired = 2;

    auto gpusByArch = getIdenticalGPUs();

    auto it = gpusByArch.begin();
    auto end = gpusByArch.end();

    auto bestFit = std::make_pair(it, it);
    // use std::distance to find the largest number of GPUs amongst architectures
    auto distance = [](decltype(bestFit) p){return std::distance(p.first, p.second);};

    // Read each unique key/pair element in order
    for (; it != end; it = gpusByArch.upper_bound(it->first)) {
        // first and second are iterators bounded within the architecture group
        auto testFit = gpusByArch.equal_range(it->first);
        // Always use devices with highest architecture version or whichever has the most devices available
        if (distance(bestFit) <= distance(testFit))
            bestFit = testFit;
    }

    if (distance(bestFit) < kNumGpusRequired) {
        printf(
            "No Two or more GPUs with same architecture capable of cuda Memory Pools found."
            "\nWaiving the sample\n");
        exit(EXIT_WAIVED);
    }

    std::set<int> bestFitDeviceIds;

    // check & select peer-to-peer access capable GPU devices.
    int devIds[2];
    for (auto itr = bestFit.first; itr != bestFit.second; itr++) {
        int deviceId = itr->second;
        checkCudaErrors(cudaSetDevice(deviceId));

        std::for_each(itr, bestFit.second, [&deviceId, &bestFitDeviceIds, &kNumGpusRequired](decltype(*itr) mapPair) {
            if (deviceId != mapPair.second) {
                int access = 0;
                checkCudaErrors(cudaDeviceCanAccessPeer(&access, deviceId, mapPair.second));
                printf("Device=%d %s Access Peer Device=%d\n", deviceId, access ? "CAN" : "CANNOT", mapPair.second);
                if (access && bestFitDeviceIds.size() < kNumGpusRequired) {
                    bestFitDeviceIds.emplace(deviceId);
                    bestFitDeviceIds.emplace(mapPair.second);
                }
                else {
                    printf("Ignoring device %i (max devices exceeded)\n", mapPair.second);
                }
            }
        });

        if (bestFitDeviceIds.size() >= kNumGpusRequired) {
            printf("Selected p2p capable devices - ");
            int i=0;
            for (auto devicesItr = bestFitDeviceIds.begin(); devicesItr != bestFitDeviceIds.end(); devicesItr++) {
                devIds[i++] = *devicesItr;
                printf("deviceId = %d  ", *devicesItr);
            }
            printf("\n");
            break;
        }
    }

    // if bestFitDeviceIds.size() == 0 it means the GPUs in system are not p2p capable,
    // hence we add it without p2p capability check.
    if (!bestFitDeviceIds.size()) {
        printf("No Two or more Devices p2p capable found.. exiting..\n");
        exit(EXIT_WAIVED);
    }

    auto p2pGpuPair = std::make_pair(devIds[0], devIds[1]);

    return p2pGpuPair;
}

int memPoolP2PCopy()
{
    int *dev0_srcVec, *dev1_dstVec; // Device buffers
    cudaStream_t stream1, stream2;
    cudaMemPool_t memPool;
    cudaEvent_t waitOnStream1;

    // Allocate CPU memory.
    size_t nelem = 1048576;
    size_t bytes = nelem*sizeof(int);

    int *a        = (int*) malloc(bytes);
    int *output   = (int*) malloc(bytes);

    /* Initialize the vectors. */
    for (int n = 0; n < nelem; n++) {
        a[n] = rand() / (int)RAND_MAX;
    }

    auto p2pDevices = getP2PCapableGpuPair();
    printf("selected devices = %d & %d\n", p2pDevices.first, p2pDevices.second);
    checkCudaErrors(cudaSetDevice(p2pDevices.first));
    checkCudaErrors(cudaEventCreate(&waitOnStream1));

    checkCudaErrors(cudaStreamCreateWithFlags(&stream1, cudaStreamNonBlocking));

    // Get the default mempool for device p2pDevices.first from the pair
    checkCudaErrors(cudaDeviceGetDefaultMemPool(&memPool, p2pDevices.first));

    // Allocate memory in a stream from the pool set above.
    checkCudaErrors(cudaMallocAsync(&dev0_srcVec, bytes, stream1));

    checkCudaErrors(cudaMemcpyAsync(dev0_srcVec, a, bytes, cudaMemcpyHostToDevice, stream1));
    checkCudaErrors(cudaEventRecord(waitOnStream1, stream1));

    checkCudaErrors(cudaSetDevice(p2pDevices.second));
    checkCudaErrors(cudaStreamCreateWithFlags(&stream2, cudaStreamNonBlocking));

    // Allocate memory in p2pDevices.second device
    checkCudaErrors(cudaMallocAsync(&dev1_dstVec, bytes, stream2));

    // Setup peer mappings for p2pDevices.second device
    cudaMemAccessDesc desc;
    memset(&desc, 0, sizeof(cudaMemAccessDesc));
    desc.location.type = cudaMemLocationTypeDevice;
    desc.location.id = p2pDevices.second;
    desc.flags = cudaMemAccessFlagsProtReadWrite;
    checkCudaErrors(cudaMemPoolSetAccess(memPool, &desc, 1));


    printf("> copyP2PAndScale kernel running ...\n");
    dim3 block(256);
    dim3 grid((unsigned int)ceil(nelem/(int)block.x));
    checkCudaErrors(cudaStreamWaitEvent(stream2, waitOnStream1));
    copyP2PAndScale<<<grid, block, 0, stream2>>>(dev0_srcVec, dev1_dstVec, nelem);

    checkCudaErrors(cudaMemcpyAsync(output, dev1_dstVec, bytes, cudaMemcpyDeviceToHost, stream2));
    checkCudaErrors(cudaFreeAsync(dev0_srcVec, stream2));
    checkCudaErrors(cudaFreeAsync(dev1_dstVec, stream2));
    checkCudaErrors(cudaStreamSynchronize(stream2));

    /* Compare the results */
    printf("> Checking the results from copyP2PAndScale() ...\n");

    for (int n = 0; n < nelem; n++) {
        if ((2*a[n]) != output[n]) {
            printf("mismatch i = %d expected = %d val = %d\n",n, 2*a[n], output[n]);
            return EXIT_FAILURE;
        }
    }

    free(a);
    free(output);
    checkCudaErrors(cudaStreamDestroy(stream1));
    checkCudaErrors(cudaStreamDestroy(stream2));
    printf("PASSED\n");

    return EXIT_SUCCESS;
}

int main(int argc, char **argv)
{
    int ret = memPoolP2PCopy();
    return ret;
}
