/* 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.
 */

#include "multidevicealloc_memmap.hpp"

static size_t round_up(size_t x, size_t y) { return ((x + y - 1) / y) * y; }

CUresult simpleMallocMultiDeviceMmap(CUdeviceptr                 *dptr,
                                     size_t                      *allocationSize,
                                     size_t                       size,
                                     const std::vector<CUdevice> &residentDevices,
                                     const std::vector<CUdevice> &mappingDevices,
                                     size_t                       align)
{
    CUresult status          = CUDA_SUCCESS;
    size_t   min_granularity = 0;
    size_t   stripeSize;

    // Setup the properties common for all the chunks
    // The allocations will be device pinned memory.
    // This property structure describes the physical location where the memory
    // will be allocated via cuMemCreate allong with additional properties In this
    // case, the allocation will be pinnded device memory local to a given device.
    CUmemAllocationProp prop = {};
    prop.type                = CU_MEM_ALLOCATION_TYPE_PINNED;
    prop.location.type       = CU_MEM_LOCATION_TYPE_DEVICE;

    // Get the minimum granularity needed for the resident devices
    // (the max of the minimum granularity of each participating device)
    for (int idx = 0; idx < residentDevices.size(); idx++) {
        size_t granularity = 0;

        // get the minnimum granularity for residentDevices[idx]
        prop.location.id = residentDevices[idx];
        status           = cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM);
        if (status != CUDA_SUCCESS) {
            goto done;
        }
        if (min_granularity < granularity) {
            min_granularity = granularity;
        }
    }

    // Get the minimum granularity needed for the accessing devices
    // (the max of the minimum granularity of each participating device)
    for (size_t idx = 0; idx < mappingDevices.size(); idx++) {
        size_t granularity = 0;

        // get the minnimum granularity for mappingDevices[idx]
        prop.location.id = mappingDevices[idx];
        status           = cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM);
        if (status != CUDA_SUCCESS) {
            goto done;
        }
        if (min_granularity < granularity) {
            min_granularity = granularity;
        }
    }

    // Round up the size such that we can evenly split it into a stripe size tha
    // meets the granularity requirements Essentially size = N *
    // residentDevices.size() * min_granularity is the requirement, since each
    // piece of the allocation will be stripeSize = N * min_granularity and the
    // min_granularity requirement applies to each stripeSize piece of the
    // allocation.
    size       = round_up(size, residentDevices.size() * min_granularity);
    stripeSize = size / residentDevices.size();

    // Return the rounded up size to the caller for use in the free
    if (allocationSize) {
        *allocationSize = size;
    }

    // Reserve the required contiguous VA space for the allocations
    status = cuMemAddressReserve(dptr, size, align, 0, 0);
    if (status != CUDA_SUCCESS) {
        goto done;
    }

    // Create and map the backings on each gpu
    // note: reusing CUmemAllocationProp prop from earlier with prop.type &
    // prop.location.type already specified.
    for (size_t idx = 0; idx < residentDevices.size(); idx++) {
        CUresult status2 = CUDA_SUCCESS;

        // Set the location for this chunk to this device
        prop.location.id = residentDevices[idx];

        // Create the allocation as a pinned allocation on this device
        CUmemGenericAllocationHandle allocationHandle;
        status = cuMemCreate(&allocationHandle, stripeSize, &prop, 0);
        if (status != CUDA_SUCCESS) {
            goto done;
        }

        // Assign the chunk to the appropriate VA range and release the handle.
        // After mapping the memory, it can be referenced by virtual address.
        // Since we do not need to make any other mappings of this memory or export
        // it, we no longer need and can release the allocationHandle. The
        // allocation will be kept live until it is unmapped.
        status = cuMemMap(*dptr + (stripeSize * idx), stripeSize, 0, allocationHandle, 0);

        // the handle needs to be released even if the mapping failed.
        status2 = cuMemRelease(allocationHandle);
        if (status == CUDA_SUCCESS) {
            // cuMemRelease should not have failed here
            // as the handle was just allocated successfully
            // however return an error if it does.
            status = status2;
        }

        // Cleanup in case of any mapping failures.
        if (status != CUDA_SUCCESS) {
            goto done;
        }
    }

    {
        // Each accessDescriptor will describe the mapping requirement for a single
        // device
        std::vector<CUmemAccessDesc> accessDescriptors;
        accessDescriptors.resize(mappingDevices.size());

        // Prepare the access descriptor array indicating where and how the backings
        // should be visible.
        for (size_t idx = 0; idx < mappingDevices.size(); idx++) {
            // Specify which device we are adding mappings for.
            accessDescriptors[idx].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
            accessDescriptors[idx].location.id   = mappingDevices[idx];

            // Specify both read and write access.
            accessDescriptors[idx].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
        }

        // Apply the access descriptors to the whole VA range.
        status = cuMemSetAccess(*dptr, size, &accessDescriptors[0], accessDescriptors.size());
        if (status != CUDA_SUCCESS) {
            goto done;
        }
    }

done:
    if (status != CUDA_SUCCESS) {
        if (*dptr) {
            simpleFreeMultiDeviceMmap(*dptr, size);
        }
    }

    return status;
}

CUresult simpleFreeMultiDeviceMmap(CUdeviceptr dptr, size_t size)
{
    CUresult status = CUDA_SUCCESS;

    // Unmap the mapped virtual memory region
    // Since the handles to the mapped backing stores have already been released
    // by cuMemRelease, and these are the only/last mappings referencing them,
    // The backing stores will be freed.
    // Since the memory has been unmapped after this call, accessing the specified
    // va range will result in a fault (unitll it is remapped).
    status = cuMemUnmap(dptr, size);
    if (status != CUDA_SUCCESS) {
        return status;
    }
    // Free the virtual address region.  This allows the virtual address region
    // to be reused by future cuMemAddressReserve calls.  This also allows the
    // virtual address region to be used by other allocation made through
    // opperating system calls like malloc & mmap.
    status = cuMemAddressFree(dptr, size);
    if (status != CUDA_SUCCESS) {
        return status;
    }

    return status;
}
