TheRock/rocm-systems/projects/rccl/test/AllocTests.cpp

/*************************************************************************
 * Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
 *
 * See LICENSE.txt for license information
 ************************************************************************/

#include <alloc.h>
#include <gtest/gtest.h>
#include <rccl/rccl.h>

#include "TestBed.hpp"
#include "common/ErrCode.hpp"
#include "common/ProcessIsolatedTestRunner.hpp"

template ncclResult_t ncclCudaMemcpy<float>(float*, float*, size_t);

namespace RcclUnitTesting
{
TEST(Alloc, ncclIbMallocDebugNonZero)
{
    void*  ptr  = nullptr;
    size_t size = 4096;

    ncclResult_t result = ncclIbMalloc(&ptr, size);

    EXPECT_EQ(result, ncclSuccess);
    ASSERT_NE(ptr, nullptr);

    char* char_ptr = static_cast<char*>(ptr);
    for(size_t i = 0; i < size; ++i)
    {
        ASSERT_EQ(char_ptr[i], 0);
    }

    free(ptr);
}

TEST(Alloc, ncclIbMallocDebugZeroSize)
{
    void*        ptr    = (void*)0xdeadbeef;
    ncclResult_t result = ncclIbMalloc(&ptr, 0);

    EXPECT_EQ(result, ncclSuccess);
    EXPECT_EQ(ptr, nullptr);
}

#if ROCM_VERSION < 71200
// These tests exercise the unsupported-fallback path of ncclCuMemHostAlloc/Free
// that returns ncclInternalError. On ROCm 7.12+ the real implementation is
// compiled in, so the fallback no longer exists and these tests are not applicable.
TEST(Alloc, ncclCuMemHostAlloc)
{
    RUN_ISOLATED_TEST(
        "ncclCuMemHostAlloc",
        []()
        {
            void*        ptr    = NULL;
            void*        handle = NULL;
            size_t       size   = 1024;
            ncclResult_t result = ncclCuMemHostAlloc(&ptr, handle, size);
            ASSERT_EQ(result, ncclInternalError);
        }
    );
}

TEST(Alloc, ncclCuMemHostFree)
{
    RUN_ISOLATED_TEST(
        "ncclCuMemHostFree",
        []()
        {
            void*        dummyPtr = reinterpret_cast<void*>(0x1234);
            ncclResult_t result   = ncclCuMemHostFree(dummyPtr);
            ASSERT_EQ(result, ncclInternalError);
        }
    );
}
#endif // ROCM_VERSION < 71200

#if ROCM_VERSION < 70000
// This test is only valid for ROCm versions < 7.0.0
// In ROCm 7.0.0+, the ncclCuMemAlloc signature changed
TEST(Alloc, ncclCuMemAlloc)
{
    RUN_ISOLATED_TEST(
        "ncclCuMemAlloc",
        []()
        {
            void*                      ptr    = reinterpret_cast<void*>(0x1234);
            void*                      handle = reinterpret_cast<void*>(0x5678);
            size_t                     size   = 1024;
            hipMemAllocationHandleType type   = hipMemHandleTypeNone;
            ncclResult_t               result = ncclCuMemAlloc(&ptr, &handle, type, size, /*manager=*/nullptr);
            EXPECT_EQ(result, ncclInternalError);
        }
    );
}

TEST(Alloc, ncclCuMemFree)
{
    RUN_ISOLATED_TEST(
        "ncclCuMemFree",
        []()
        {
            void*        dummyPtr = reinterpret_cast<void*>(0xdeadbeef);
            ncclResult_t result   = ncclCuMemFree(dummyPtr, /*manager=*/nullptr);
            EXPECT_EQ(result, ncclInternalError);
        }
    );
}

TEST(Alloc, ncclCuMemAllocAddr)
{
    RUN_ISOLATED_TEST(
        "ncclCuMemAllocAddr",
        []()
        {
            void* ptr = reinterpret_cast<void*>(0x1111);
            hipMemGenericAllocationHandle_t handle
                = reinterpret_cast<hipMemGenericAllocationHandle_t>(0x1234);
            size_t       size   = 4096;
            ncclResult_t result = ncclCuMemAllocAddr(&ptr, &handle, size);
            ASSERT_EQ(result, ncclInternalError);
        }
    );
}

TEST(Alloc, ncclCuMemFreeAddr)
{
    RUN_ISOLATED_TEST(
        "ncclCuMemFreeAddr",
        []()
        {
            void*        testPtr = reinterpret_cast<void*>(0xbeefcafe);
            ncclResult_t result  = ncclCuMemFreeAddr(testPtr, /*manager=*/nullptr);
            ASSERT_EQ(result, ncclInternalError);
        }
    );
}
#endif // ROCM_VERSION < 70000

TEST(Alloc, NcclCudaMemcpy)
{
    RUN_ISOLATED_TEST(
        "NcclCudaMemcpy",
        []()
        {
            // Initialize HIP device in forked process
            ASSERT_EQ(hipSetDevice(0), hipSuccess);

            constexpr size_t N     = 128;
            float *          d_src = nullptr, *d_dst = nullptr;
            float            h_src[N], h_dst[N];

            for(size_t i = 0; i < N; ++i)
                h_src[i] = static_cast<float>(i + 1);
            // Allocate device memory

            ASSERT_EQ(hipMalloc(&d_src, N * sizeof(float)), hipSuccess);
            ASSERT_EQ(hipMalloc(&d_dst, N * sizeof(float)), hipSuccess);

            // Copy from host to device (source buffer)
            ASSERT_EQ(
                hipMemcpy(d_src, h_src, N * sizeof(float), hipMemcpyHostToDevice),
                hipSuccess
            );

            // Perform the tested function
            ncclResult_t result = ncclCudaMemcpy<float>(d_dst, d_src, N);

            ASSERT_EQ(result, ncclSuccess);

            // Copy result back to host
            ASSERT_EQ(
                hipMemcpy(h_dst, d_dst, N * sizeof(float), hipMemcpyDeviceToHost),
                hipSuccess
            );

            // Check correctness
            for(size_t i = 0; i < N; ++i)
            {
                EXPECT_EQ(h_src[i], h_dst[i]) << "Mismatch at index " << i;
            }
            // Free memory
            ASSERT_EQ(hipFree(d_src), hipSuccess);
            ASSERT_EQ(hipFree(d_dst), hipSuccess);
        }
    );
}

TEST(Alloc, ZeroElementMemcpy)
{
    RUN_ISOLATED_TEST(
        "ZeroElementMemcpy",
        []()
        {
            // Initialize HIP device in forked process
            ASSERT_EQ(hipSetDevice(0), hipSuccess);

            float *d_src = nullptr, *d_dst = nullptr;
            ASSERT_EQ(hipMalloc(&d_src, sizeof(float)), hipSuccess);
            ASSERT_EQ(hipMalloc(&d_dst, sizeof(float)), hipSuccess);

            ncclResult_t result = ncclCudaMemcpy<float>(d_dst, d_src, 0);
            EXPECT_EQ(result, ncclSuccess) << "Zero-element copy should succeed (no-op)";

            ASSERT_EQ(hipFree(d_src), hipSuccess);
            ASSERT_EQ(hipFree(d_dst), hipSuccess);
        }
    );
}

TEST(Alloc, MemcpyNullSrcOrDstPointer)
{
    RUN_ISOLATED_TEST(
        "MemcpyNullSrcOrDstPointer",
        []()
        {
            // Initialize HIP device in forked process
            ASSERT_EQ(hipSetDevice(0), hipSuccess);

            constexpr size_t N       = 16;
            float*           d_valid = nullptr;
            ASSERT_EQ(hipMalloc(&d_valid, N * sizeof(float)), hipSuccess);

            // Case 1: src is nullptr
            ncclResult_t result = ncclCudaMemcpy<float>(d_valid, nullptr, N);
            EXPECT_EQ(result, ncclUnhandledCudaError)
                << "Expected ncclUnhandledCudaError when src is nullptr";

            // Case 2: dst is nullptr
            result = ncclCudaMemcpy<float>(nullptr, d_valid, N);
            EXPECT_EQ(result, ncclUnhandledCudaError)
                << "Expected ncclUnhandledCudaError when dst is nullptr";

            ASSERT_EQ(hipFree(d_valid), hipSuccess);
        }
    );
}
} // namespace RcclUnitTesting