/* Copyright 2020 National Research Foundation (SARAO)
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU Lesser General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option) any
 * later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
 * details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/* This example sends data directly from a GPU to the network, without
 * passing through the CPU, using GPUDirect RDMA. It requires
 * - An NVIDIA GPU
 * - An NVIDIA (Mellanox) NIC
 * - The nv_peer_mem kernel module (https://github.com/Mellanox/nv_peer_memory)
 * - It doesn't play nicely with IOMMU remapping. Adding "iommu=pt" to the
 *   kernel command line seems to help.
 *
 * Run it by passing the IP address of the NIC on the command line. It will
 * send data to 239.255.88.88:8888 (multicast). You can verify that the
 * correct data is being sent by running
 * spead2_recv --ibv --bind $ip 239.255.88.88:8888 --verify
 * It should print a message as it receives each heap.
 *
 * If the sender reports "ibv_reg_mr failed: Bad address" it probably means
 * that nv_peer_mem is not set up correctly.
 *
 * It should be noted that this example is intended for exposition rather
 * than high performance. For example, a real application would most likely
 * overlap computation on one buffer with transmission of another.
 */

#include <cuda.h>
#include <spead2/send_stream.h>
#include <spead2/send_udp_ibv.h>
#include <cstdint>
#include <iostream>

/// Crash out if a CUDA call fails
#define CUDA_CHECK(cmd)             \
    ({                              \
        cudaError_t result = (cmd); \
        if (result != cudaSuccess)  \
        {                           \
            std::cerr << "CUDA error: " << #cmd << "\n"; \
            std::exit(1);           \
        }                           \
        result;                     \
    })

// Fill the output with random numbers, as generated by std::minstd_rand
__global__ void fill_random(std::uint32_t *out, std::uint64_t offset)
{
    constexpr std::uint32_t m = 2147483647;
    constexpr std::uint32_t a = 48271;
    std::uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    std::uint64_t z = idx + offset;
    std::uint64_t x = a;
    std::uint64_t apow = a;
    while (z > 0)
    {
        if (z & 1)
            x = x * apow % m;
        apow = apow * apow % m;
        z >>= 1;
    }
    // Convert to big endian
#if __BYTE_ORDER == __LITTLE_ENDIAN
    x = __byte_perm(x, x, 0x0123);
#elif __BYTE_ORDER != __BIG_ENDIAN
# error "Unhandled __BYTE_ORDER"
#endif
    out[idx] = x;
}

int main(int argc, const char * const *argv)
{
    if (argc != 2)
    {
        std::cerr << "Usage: gpudirect_example <NIC IP address>\n";
        return 2;
    }

    const std::size_t elements = 1048576;   // Elements per heap
    const std::size_t size = elements * sizeof(std::uint32_t);
    const std::size_t block_size = 256;     // Elements per CUDA block
    std::uint32_t *dout;                    // Device pointer
    CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&dout), size));

    // Set up the stream
    spead2::thread_pool tp;
    spead2::send::stream_config config;
    config.set_rate(1e8);                   // Low rate to allow spead2_recv --verify to keep up
    spead2::send::udp_ibv_config ibv_config;
    ibv_config.add_endpoint(
        boost::asio::ip::udp::endpoint(
            boost::asio::ip::address::from_string("239.255.88.88"),
            8888));
    ibv_config.set_interface_address(boost::asio::ip::address::from_string(argv[1]));
    // The nv_peer_mem kernel module recognises that dout is a device pointer
    ibv_config.add_memory_region(dout, size);

    auto empty_callback = [](const boost::system::error_code &ec, spead2::item_pointer_t bytes) {
        if (ec)
            std::cerr << ec << '\n';
    };
    spead2::send::udp_ibv_stream stream(tp, config, ibv_config);
    spead2::flavour flavour(4, 64, 48, 0);
    spead2::descriptor desc;
    desc.id = 0x1000;
    desc.name = "random";
    desc.description = "Random numbers";
    desc.shape.push_back(elements);
    desc.format.emplace_back('u', 32);

    std::uint64_t offset = 0;   // Number of elements already sent
    bool first = true;          // Whether this is the first heap
    while (true)
    {
        // Fill the buffer on the GPU
        fill_random<<<elements / block_size, block_size>>>(dout, offset);
        offset += elements;
        CUDA_CHECK(cudaDeviceSynchronize());   // Ensure data is ready for transmission
        // Transmit it
        spead2::send::heap heap(flavour);
        if (first)
        {
            heap.add_descriptor(desc);
            first = false;
        }
        heap.add_item(0x1000, dout, size, false);
        stream.async_send_heap(heap, empty_callback);
        stream.flush();            // Ensure data is free to be written by the GPU
        std::cout << "Sent a heap\n";
    }
}
