/**
 * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025-2026. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */

#include <cuda_runtime.h>
#include <ucp/ucp_test.h>

#include <ucp/api/device/ucp_device_types.h>

#include <common/cuda.h>
#include "cuda/test_kernels.h"

class test_ucp_device : public ucp_test {
public:
    static void get_test_variants(std::vector<ucp_test_variant> &variants);

    virtual void init() override;

private:
    static void get_base_variants(std::vector<ucp_test_variant> &variants);

protected:
    static constexpr size_t MAX_THREADS = 128;

    ucs_memory_type_t rx_mem_type() const {
        return static_cast<ucs_memory_type_t>(get_variant_value());
    }

    class mem_list {
    public:
        static constexpr uint64_t SEED_SRC = 0x1234;
        static constexpr uint64_t SEED_DST = 0x4321;

        enum mem_list_mode_t {
            MODE_DATA_ONLY,
            MODE_COUNTER_ONLY,
            MODE_WITH_GAP_ELEMENT
        };

        mem_list(test_ucp_device &test, size_t size, unsigned count,
                 ucs_memory_type_t mem_type = UCS_MEMORY_TYPE_CUDA,
                 mem_list_mode_t mode = MODE_DATA_ONLY);
        ~mem_list();

        void *src_ptr(unsigned index) const;

        uint64_t dst_ptr(unsigned index) const;

        std::vector<void*> src_ptrs() const;

        std::vector<uint64_t> dst_ptrs() const;

        void dst_counter_init(unsigned index, size_t remote_offset = 0);

        uint64_t
        dst_counter_read(unsigned index, size_t remote_offset = 0) const;

        ucp_device_local_mem_list_h local_handle() const;

        ucp_device_remote_mem_list_h remote_handle() const;

        void dst_pattern_check(unsigned index, uint64_t seed) const;

    private:
        std::vector<std::unique_ptr<mapped_buffer>> m_src, m_dst;
        std::vector<ucs::handle<ucp_rkey_h>> m_rkeys;
        ucp_device_local_mem_list_h m_local_mem_list_h;
        ucp_device_remote_mem_list_h m_remote_mem_list_h;
    };

    size_t counter_size();

    static void
    counter_init(const mapped_buffer &buffer, size_t remote_offset = 0);

    static uint64_t
    counter_read(const mapped_buffer &buffer, size_t remote_offset = 0);

    void launch_kernel(const test_ucp_device_kernel_params_t &params);
};


void test_ucp_device::get_base_variants(std::vector<ucp_test_variant> &variants)
{
    add_variant(variants,
                UCP_FEATURE_RMA | UCP_FEATURE_AMO64 | UCP_FEATURE_DEVICE);
}

void test_ucp_device::get_test_variants(std::vector<ucp_test_variant> &variants)
{
    add_variant_memtypes(variants, get_base_variants,
                         UCS_BIT(UCS_MEMORY_TYPE_CUDA) |
                         UCS_BIT(UCS_MEMORY_TYPE_HOST));
}

void test_ucp_device::init()
{
    m_env.push_back(new ucs::scoped_setenv("UCX_IB_GDA_MAX_SYS_LATENCY", "1us"));
    ucp_test::init();
    sender().connect(&receiver(), get_ep_params());
    if (!is_loopback()) {
        receiver().connect(&sender(), get_ep_params());
    }
}

test_ucp_device::mem_list::mem_list(test_ucp_device &test, size_t size,
                                    unsigned count, ucs_memory_type_t mem_type,
                                    mem_list_mode_t mode) :
    m_local_mem_list_h(nullptr),
    m_remote_mem_list_h(nullptr)
{
    bool has_counter  = (mode == MODE_COUNTER_ONLY);
    size_t gap_count  = (mode == MODE_WITH_GAP_ELEMENT) ? 1 : 0;
    size_t data_count = (has_counter) ? count - 1 : count;
    ucs_status_t status;

    // Prepare src and dst buffers
    for (auto i = 0; i < data_count; ++i) {
        m_src.emplace_back(new mapped_buffer(size, test.sender(), 0, mem_type));
        m_dst.emplace_back(new mapped_buffer(size, test.receiver(), 0,
                                             test.rx_mem_type()));
        m_rkeys.push_back(m_dst.back()->rkey(test.sender()));
        m_src.back()->pattern_fill(SEED_SRC, size);
        m_dst.back()->pattern_fill(SEED_DST, size);
    }

    if (has_counter) {
        m_dst.emplace_back(new mapped_buffer(size, test.receiver(), 0,
                                             test.rx_mem_type()));
        m_rkeys.push_back(m_dst.back()->rkey(test.sender()));
        m_dst.back()->pattern_fill(SEED_DST, size);
    }

    ucp_device_mem_list_params_t params;

    if (mode != MODE_COUNTER_ONLY) {
        // Initialize local elements
        std::vector<ucp_device_mem_list_elem_t> local_elems(data_count);
        for (auto i = 0; i < data_count; ++i) {
            local_elems[i].field_mask =
                    UCP_DEVICE_MEM_LIST_ELEM_FIELD_MEMH |
                    UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR;
            local_elems[i].memh       = m_src[i]->memh();
            local_elems[i].local_addr = m_src[i]->ptr();
        }

        // Initialize parameters
        params.field_mask   = UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENTS |
                              UCP_DEVICE_MEM_LIST_PARAMS_FIELD_NUM_ELEMENTS |
                              UCP_DEVICE_MEM_LIST_PARAMS_FIELD_WORKER |
                              UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENT_SIZE;
        params.element_size = sizeof(local_elems[0]);
        params.num_elements = data_count;
        params.elements     = local_elems.data();
        params.worker       = test.sender().worker();
        // Create memory list (with retry on connection)
        status = ucp_device_local_mem_list_create(&params, &m_local_mem_list_h);

        ASSERT_UCS_OK(status);
    }

    // Initialize remote elements
    size_t remote_count = count + gap_count;
    std::vector<ucp_device_mem_list_elem_t> remote_elems(remote_count);
    for (auto i = 0; i < gap_count; ++i) {
        remote_elems[i].field_mask = 0;
    }

    for (auto i = gap_count; i < data_count + gap_count; ++i) {
        auto data_index             = i - gap_count;
        remote_elems[i].field_mask  = UCP_DEVICE_MEM_LIST_ELEM_FIELD_EP |
                                      UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR |
                                      UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY;
        remote_elems[i].ep          = test.sender().ep();
        remote_elems[i].remote_addr = reinterpret_cast<uint64_t>(
                m_dst[data_index]->ptr());
        remote_elems[i].rkey        = m_rkeys[data_index];
    }

    if (has_counter) {
        remote_elems[data_count + gap_count].field_mask =
                UCP_DEVICE_MEM_LIST_ELEM_FIELD_EP |
                UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR |
                UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY;
        remote_elems[data_count + gap_count].ep = test.sender().ep();
        remote_elems[data_count + gap_count].remote_addr =
                reinterpret_cast<uint64_t>(m_dst[data_count]->ptr());
        remote_elems[data_count + gap_count].rkey = m_rkeys[data_count];
    }

    // Initialize parameters
    params.field_mask   = UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENTS |
                          UCP_DEVICE_MEM_LIST_PARAMS_FIELD_NUM_ELEMENTS |
                          UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENT_SIZE;
    params.element_size = sizeof(remote_elems[0]);
    params.num_elements = remote_count;
    params.elements     = remote_elems.data();

    // Create memory list (with retry on connection)
    do {
        test.progress();
        status = ucp_device_remote_mem_list_create(&params,
                                                   &m_remote_mem_list_h);
    } while (status == UCS_ERR_NOT_CONNECTED);

    if (status == UCS_ERR_NO_DEVICE) {
        UCS_TEST_SKIP_R("Skipping test if no device lanes exists.");
    } else {
        ASSERT_UCS_OK(status);
    }
}

test_ucp_device::mem_list::~mem_list()
{
    if (m_local_mem_list_h) {
        ucp_device_mem_list_release(m_local_mem_list_h);
    }

    if (m_remote_mem_list_h) {
        ucp_device_mem_list_release(m_remote_mem_list_h);
    }
}

void *test_ucp_device::mem_list::src_ptr(unsigned index) const
{
    return m_src[index]->ptr();
}

uint64_t test_ucp_device::mem_list::dst_ptr(unsigned index) const
{
    return reinterpret_cast<uint64_t>(m_dst[index]->ptr());
}

std::vector<void*> test_ucp_device::mem_list::src_ptrs() const
{
    std::vector<void*> result;
    for (const auto &src : m_src) {
        result.push_back(src->ptr());
    }
    return result;
}

std::vector<uint64_t> test_ucp_device::mem_list::dst_ptrs() const
{
    std::vector<uint64_t> result;
    for (const auto &dst : m_dst) {
        result.push_back(reinterpret_cast<uint64_t>(dst->ptr()));
    }
    return result;
}

void test_ucp_device::mem_list::dst_counter_init(unsigned index,
                                                 size_t remote_offset)
{
    test_ucp_device::counter_init(*m_dst[index], remote_offset);
}

uint64_t test_ucp_device::mem_list::dst_counter_read(unsigned index,
                                                     size_t remote_offset) const
{
    return test_ucp_device::counter_read(*m_dst[index], remote_offset);
}

void test_ucp_device::mem_list::dst_pattern_check(unsigned index,
                                                  uint64_t seed) const
{
    m_dst[index]->pattern_check(seed, m_dst[index]->size());
}

ucp_device_local_mem_list_h test_ucp_device::mem_list::local_handle() const
{
    return m_local_mem_list_h;
}

ucp_device_remote_mem_list_h test_ucp_device::mem_list::remote_handle() const
{
    return m_remote_mem_list_h;
}

size_t test_ucp_device::counter_size()
{
    ucp_context_attr_t attr;
    attr.field_mask = UCP_ATTR_FIELD_DEVICE_COUNTER_SIZE;
    ASSERT_UCS_OK(ucp_context_query(receiver().ucph(), &attr));
    return attr.device_counter_size;
}

void test_ucp_device::counter_init(const mapped_buffer &buffer,
                                   size_t remote_offset)
{
    ucp_device_counter_params_t params;
    params.field_mask = UCP_DEVICE_COUNTER_PARAMS_FIELD_MEMH;
    params.memh       = buffer.memh();
    ASSERT_UCS_OK(ucp_device_counter_init(buffer.worker(), &params,
                                          UCS_PTR_BYTE_OFFSET(buffer.ptr(),
                                                              remote_offset)));
}

uint64_t
test_ucp_device::counter_read(const mapped_buffer &buffer, size_t remote_offset)
{
    ucp_device_counter_params_t params;
    params.field_mask = UCP_DEVICE_COUNTER_PARAMS_FIELD_MEMH;
    params.memh       = buffer.memh();
    return ucp_device_counter_read(buffer.worker(), &params,
                                   UCS_PTR_BYTE_OFFSET(buffer.ptr(),
                                                       remote_offset));
}

UCS_TEST_P(test_ucp_device, create_success)
{
    mem_list list(*this, 4 * UCS_MBYTE, 4);
    EXPECT_NE(nullptr, list.local_handle());
    EXPECT_NE(nullptr, list.remote_handle());
}

UCS_TEST_P(test_ucp_device, create_local_fail)
{
    ucp_device_local_mem_list_h handle = nullptr;

    scoped_log_handler wrap_err(wrap_errors_logger);

    // Empty params
    ucp_device_mem_list_params_t empty_params;
    empty_params.field_mask = UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENTS;
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_local_mem_list_create(&empty_params, &handle));
    EXPECT_EQ(nullptr, handle);

    // Empty mem list
    ucp_device_mem_list_params_t invalid_params{};
    invalid_params.field_mask   = UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENTS |
                                  UCP_DEVICE_MEM_LIST_PARAMS_FIELD_WORKER |
                                  UCP_DEVICE_MEM_LIST_PARAMS_FIELD_NUM_ELEMENTS |
                                  UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENT_SIZE;
    invalid_params.elements     = nullptr;
    invalid_params.worker       = sender().worker();
    invalid_params.num_elements = 0;
    invalid_params.element_size = sizeof(ucp_device_mem_list_elem_t);
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_local_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);

    // Zero element size
    ucp_device_mem_list_elem_t dummy_elem{};
    invalid_params.elements     = &dummy_elem;
    invalid_params.num_elements = 1;
    invalid_params.element_size = 0;
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_local_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);

    invalid_params.element_size = sizeof(ucp_device_mem_list_elem_t);
    mapped_buffer src_cuda(4096, sender(), 0, UCS_MEMORY_TYPE_CUDA);
    mapped_buffer src_host(4096, sender(), 0, UCS_MEMORY_TYPE_HOST);

    ucp_device_mem_list_elem_t local_elems[2]{};
    for (int i = 0; i < 2; i++) {
        local_elems[i].field_mask = UCP_DEVICE_MEM_LIST_ELEM_FIELD_MEMH |
                                    UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR |
                                    UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH;
        local_elems[i].memh       = src_cuda.memh();
        local_elems[i].local_addr = src_cuda.ptr();
    }

    // Missing worker
    invalid_params.worker   = nullptr;
    invalid_params.elements = local_elems;
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_local_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);

    // Mismatched local sys_dev
    invalid_params.worker       = sender().worker();
    local_elems[1].memh         = src_host.memh(); // Different sys_dev
    local_elems[1].local_addr   = src_host.ptr();
    invalid_params.num_elements = 2;
    invalid_params.elements     = local_elems;
    EXPECT_EQ(UCS_ERR_UNSUPPORTED,
              ucp_device_local_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);
}

UCS_TEST_P(test_ucp_device, create_remote_fail)
{
    ucp_device_remote_mem_list_h handle = nullptr;

    scoped_log_handler wrap_err(wrap_errors_logger);

    // Empty params
    ucp_device_mem_list_params_t empty_params{};
    empty_params.field_mask = UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENTS;
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_remote_mem_list_create(&empty_params, &handle));
    EXPECT_EQ(nullptr, handle);

    // Empty mem list
    ucp_device_mem_list_params_t invalid_params{};
    invalid_params.field_mask   = UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENTS |
                                  UCP_DEVICE_MEM_LIST_PARAMS_FIELD_NUM_ELEMENTS |
                                  UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENT_SIZE;
    invalid_params.elements     = nullptr;
    invalid_params.num_elements = 0;
    invalid_params.element_size = sizeof(ucp_device_mem_list_elem_t);
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_remote_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);

    // Zero element size
    ucp_device_mem_list_elem_t dummy_elem{};
    invalid_params.elements     = &dummy_elem;
    invalid_params.num_elements = 1;
    invalid_params.element_size = 0;
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_remote_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);

    invalid_params.element_size = sizeof(ucp_device_mem_list_elem_t);
    mapped_buffer dst1(4096, receiver(), 0, UCS_MEMORY_TYPE_CUDA);
    mapped_buffer dst2(4096, receiver(), 0, UCS_MEMORY_TYPE_HOST);
    auto rkey1          = dst1.rkey(sender());
    auto rkey2          = dst2.rkey(sender());
    ucp_rkey_h rkeys[2] = {rkey1, rkey2};

    ucp_device_mem_list_elem_t remote_elems[2]{};
    for (int i = 0; i < 2; i++) {
        remote_elems[i].field_mask  = UCP_DEVICE_MEM_LIST_ELEM_FIELD_EP |
                                      UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR |
                                      UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY;
        remote_elems[i].ep          = sender().ep();
        remote_elems[i].remote_addr = reinterpret_cast<uint64_t>(dst1.ptr());
        remote_elems[i].rkey        = rkeys[i];

        if (i == 1) {
            remote_elems[i].field_mask = 0;
        }
    }

    // Missing rkey (always required)
    remote_elems[0].field_mask &= ~UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY;
    invalid_params.num_elements = 1;
    invalid_params.elements     = remote_elems;
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_remote_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);
    remote_elems[0].field_mask |= UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY; // Restore
    remote_elems[0].rkey        = nullptr;
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_remote_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);
    remote_elems[0].rkey = rkey1;

    // Missing ep (always required)
    remote_elems[0].field_mask &= ~UCP_DEVICE_MEM_LIST_ELEM_FIELD_EP;
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_remote_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);
    remote_elems[0].field_mask |= UCP_DEVICE_MEM_LIST_ELEM_FIELD_EP; // Restore
    remote_elems[0].ep          = nullptr;
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_remote_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);

    // Only gap element
    invalid_params.num_elements = 1;
    invalid_params.elements     = remote_elems + 1;
    EXPECT_EQ(UCS_ERR_INVALID_PARAM,
              ucp_device_remote_mem_list_create(&invalid_params, &handle));
    EXPECT_EQ(nullptr, handle);
}

UCS_TEST_P(test_ucp_device, get_local_mem_list_length)
{
    constexpr unsigned num_elements = 8;
    mem_list list(*this, 1 * UCS_KBYTE, num_elements);
    EXPECT_EQ(num_elements,
              ucp_device_get_mem_list_length(list.local_handle()));
}

UCS_TEST_P(test_ucp_device, get_remote_mem_list_length)
{
    constexpr unsigned num_elements = 8;
    mem_list list(*this, 1 * UCS_KBYTE, num_elements);
    EXPECT_EQ(num_elements,
              ucp_device_get_mem_list_length(list.remote_handle()));
}

UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_device, rc_gda, "rc,rc_gda")


class test_ucp_device_kernel : public test_ucp_device {
public:
    static void get_test_variants(std::vector<ucp_test_variant> &variants)
    {
        // TODO move to UCS
        static const char *ucs_device_level_names[] = {"thread", "warp",
                                                       "block", "grid"};
        add_variant_values(variants, test_ucp_device::get_test_variants,
                           UCS_BIT(UCS_DEVICE_LEVEL_THREAD) |
                                   UCS_BIT(UCS_DEVICE_LEVEL_WARP),
                           ucs_device_level_names);
    }

protected:
    ucs_device_level_t get_device_level() const
    {
        return static_cast<ucs_device_level_t>(get_variant_value(1));
    }

    test_ucp_device_kernel_params_t init_params(unsigned num_iters = 1)
    {
        test_ucp_device_kernel_params_t params = {};
        params.num_threads                     = get_num_threads();
        params.num_blocks                      = 1;
        params.level                           = get_device_level();
        params.num_iters                       = num_iters;
        return params;
    }

    virtual unsigned get_num_threads() const
    {
        switch (get_device_level()) {
        case UCS_DEVICE_LEVEL_THREAD:
            return 1;
        case UCS_DEVICE_LEVEL_WARP:
            return UCS_DEVICE_NUM_THREADS_IN_WARP;
        default:
            return 1;
        }
    }

    test_ucp_device_kernel_result_t
    launch_kernel(const test_ucp_device_kernel_params_t &params)
    {
        auto result = launch_test_ucp_device_kernel(params);
        ASSERT_UCS_OK(result.status);
        return result;
    }
};

UCS_TEST_P(test_ucp_device_kernel, local_counter)
{
    mapped_buffer counter_buffer(counter_size(), receiver(), 0,
                                 UCS_MEMORY_TYPE_CUDA);
    const uint64_t value = 1764;

    // Perform the write
    auto params                  = init_params(1);
    params.operation             = TEST_UCP_DEVICE_KERNEL_COUNTER_WRITE;
    params.local_counter.address = counter_buffer.ptr();
    params.local_counter.value   = value;
    launch_kernel(params);

    EXPECT_TRUE(mem_buffer::compare(&value, counter_buffer.ptr(), sizeof(value),
                                    counter_buffer.mem_type()));

    // Check counter value using device API
    params.operation = TEST_UCP_DEVICE_KERNEL_COUNTER_READ;
    launch_kernel(params);
}

UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_device_kernel, rc_gda,
                                        "rc,rc_gda")


class test_ucp_device_xfer : public test_ucp_device_kernel {
public:
    static void get_test_variants(std::vector<ucp_test_variant> &variants)
    {
        add_variant_values(variants, test_ucp_device_kernel::get_test_variants,
                           NODELAY_WITH_REQ, "nodelay_with_req");
        add_variant_values(variants, test_ucp_device_kernel::get_test_variants,
                           NODELAY_WITHOUT_REQ, "nodelay_without_req");
        add_variant_values(variants, test_ucp_device_kernel::get_test_variants,
                           LAZY_WITHOUT_REQ, "lazy_without_req");
        add_variant_values(variants, test_ucp_device_kernel::get_test_variants,
                           MULTI_CHANNEL, "multi_channel");
    }

    virtual void init() override
    {
        if (get_send_mode() == MULTI_CHANNEL) {
            m_env.push_back(
                    new ucs::scoped_setenv("UCX_RC_GDA_NUM_CHANNELS", "32"));
        }
        test_ucp_device::init();
    }

protected:
    typedef enum {
        NODELAY_WITH_REQ,
        NODELAY_WITHOUT_REQ,
        LAZY_WITHOUT_REQ,
        MULTI_CHANNEL,
    } send_mode_t;

    test_ucp_device_kernel_params_t init_params()
    {
        // TODO: Test different sizes and alignment
        test_ucp_device_kernel_params_t params;
        params.num_threads = get_num_threads();
        params.num_blocks  = 1;
        params.level       = get_device_level();
        params.num_iters   = get_num_iters();
        params.num_channels = 1;
        switch (get_send_mode()) {
        case MULTI_CHANNEL:
            params.num_channels = 32;
            // fall through, rest args from nodelay_with_req
        case NODELAY_WITH_REQ:
            params.with_no_delay = true;
            params.with_request  = true;
            break;
        case NODELAY_WITHOUT_REQ:
            params.with_no_delay = true;
            params.with_request  = false;
            break;
        case LAZY_WITHOUT_REQ:
            params.with_no_delay = false;
            params.with_request  = false;
            break;
        default:
            UCS_TEST_ABORT("Invalid send mode");
        }
        return params;
    }

    test_ucp_device_kernel_params_t
    init_params_put(mem_list &list, unsigned mem_list_index, size_t size,
                    test_ucp_device_operation_t operation =
                            TEST_UCP_DEVICE_KERNEL_PUT)
    {
        auto params               = init_params();
        params.operation          = operation;
        params.local_mem_list     = list.local_handle();
        params.remote_mem_list    = list.remote_handle();
        params.put.mem_list_index = mem_list_index;
        params.put.address        = list.src_ptr(mem_list_index);
        params.put.remote_address = list.dst_ptr(mem_list_index);
        params.put.length         = size;

        return params;
    }

    test_ucp_device_kernel_params_t
    init_params_counter_inc(mem_list &list, unsigned mem_list_index,
                            size_t remote_offset = 0,
                            test_ucp_device_operation_t operation =
                                    TEST_UCP_DEVICE_KERNEL_COUNTER_INC)
    {
        auto params                       = init_params();
        params.operation                  = operation;
        params.remote_mem_list            = list.remote_handle();
        params.counter_inc.mem_list_index = mem_list_index;
        params.counter_inc.inc_value      = 1;
        params.counter_inc.remote_address = list.dst_ptr(mem_list_index);
        params.counter_inc.remote_offset  = remote_offset;
        return params;
    }

    send_mode_t get_send_mode() const
    {
        return static_cast<send_mode_t>(get_variant_value(2));
    }

    virtual unsigned get_num_threads() const override
    {
        switch (get_device_level()) {
        case UCS_DEVICE_LEVEL_THREAD:
            /* When using thread-level use less threads to shorten the test */
            return 8;
        case UCS_DEVICE_LEVEL_WARP:
            return 128;
        default:
            return 1;
        }
    }

    unsigned get_multi_elem_count() const
    {
        /* When using thread-level use less threads to shorten the test */
        switch (get_device_level()) {
        case UCS_DEVICE_LEVEL_THREAD:
            /* Thread level uses less elements to not overflow the QP */
            return 16;
        case UCS_DEVICE_LEVEL_WARP:
            return 32;
        default:
            return 1;
        }
    }

    size_t get_num_iters() const
    {
        return 10;
    }

    size_t get_num_ops_multiplier() const
    {
        switch (get_device_level()) {
        case UCS_DEVICE_LEVEL_THREAD:
            return get_num_threads();
        case UCS_DEVICE_LEVEL_WARP:
            return get_num_threads() / 32;
        default:
            return 1;
        }
    }

    void wait_for_counter(const mem_list &list, unsigned counter_index,
                          size_t remote_offset = 0)
    {
        const size_t multiplier = get_num_ops_multiplier();
        uint64_t target_value   = get_num_iters() * multiplier;

        wait_for_cond(
                [&list, counter_index, target_value, remote_offset]() {
                    return list.dst_counter_read(counter_index,
                                                 remote_offset) == target_value;
                },
                [] {});
        EXPECT_EQ(get_num_iters() * multiplier,
                  list.dst_counter_read(counter_index, remote_offset))
                << "multiplier: " << multiplier;
    }
};

UCS_TEST_P(test_ucp_device_xfer, put)
{
    static constexpr size_t size = 32 * UCS_KBYTE;
    mem_list list(*this, size, 6, UCS_MEMORY_TYPE_CUDA,
                  mem_list::MODE_WITH_GAP_ELEMENT);

    // Perform the transfer
    static constexpr unsigned mem_list_index = 3;
    auto params = init_params_put(list, mem_list_index, size,
                                  TEST_UCP_DEVICE_KERNEL_PUT);
    params.put.remote_mem_list_index = 4;
    launch_kernel(params);

    // Check proper index received data
    list.dst_pattern_check(mem_list_index - 1, mem_list::SEED_DST);
    list.dst_pattern_check(mem_list_index, mem_list::SEED_SRC);
    list.dst_pattern_check(mem_list_index + 1, mem_list::SEED_DST);
}

/* TODO: Enable these tests in CI */
UCS_TEST_SKIP_COND_P(test_ucp_device_xfer, put_stress_test,
                     RUNNING_ON_VALGRIND)
{
#ifdef __SANITIZE_ADDRESS__
    UCS_TEST_SKIP_R("Skipping stress test under ASAN");
#endif

    constexpr size_t size             = 8;
    constexpr unsigned mem_list_index = 0;
    mem_list list(*this, size, 1, UCS_MEMORY_TYPE_CUDA,
                  mem_list::MODE_WITH_GAP_ELEMENT);

    // Perform the transfer
    auto params = init_params_put(list, mem_list_index, size,
                                  TEST_UCP_DEVICE_KERNEL_PUT);
    params.put.remote_mem_list_index = 1;
    params.num_iters                    = 1000;
    params.num_blocks                   = 1;
    params.num_threads                  = MAX_THREADS;
    launch_kernel(params);

    // Check proper index received data
    list.dst_pattern_check(mem_list_index, mem_list::SEED_SRC);
}

UCS_TEST_P(test_ucp_device_xfer, counter)
{
    const size_t remote_offset = 8;
    const size_t size          = counter_size() + remote_offset;
    mem_list list(*this, size, 1, UCS_MEMORY_TYPE_CUDA,
                  mem_list::MODE_COUNTER_ONLY);

    constexpr unsigned mem_list_index = 0;
    list.dst_counter_init(mem_list_index, remote_offset);

    auto params = init_params_counter_inc(list, mem_list_index, remote_offset,
                                          TEST_UCP_DEVICE_KERNEL_COUNTER_INC);
    launch_kernel(params);

    // Check destination
    wait_for_counter(list, mem_list_index, remote_offset);
}

UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_device_xfer, rc_gda,
                                        "rc,rc_gda")
