issue/193 接入通信库，支持英伟达平台，增加通信库测试，修复运行时bug

README.md

@@ -11,8 +11,6 @@
 ![GitHub contributors](https://img.shields.io/github/contributors/InfiniTensor/InfiniCore)
 ![GitHub commit activity](https://img.shields.io/github/commit-activity/m/InfiniTensor/InfiniCore)
 
-文档
-
 InfiniCore 是一个跨平台统一编程工具集，为不同芯片平台的功能（包括计算、运行时、通信等）提供统一 C 语言接口。目前支持的硬件和后端包括：
 
 - CPU；
@@ -51,6 +49,7 @@ python scripts/install.py [XMAKE_CONFIG_FLAGS]
 | `--moore-gpu=[y\|n]`     | 是否编译摩尔线程 GPU 接口实现 | n
 | `--sugon-dcu=[y\|n]`     | 是否编译曙光 DCU 接口实现     | n
 | `--kunlun-xpu=[y\|n]`    | 是否编译昆仑 XPU 接口实现     | n
+| `--ccl=[y\|n]`           | 是否编译 InfiniCCL 通信库接口实现     | n
 
 ### 手动安装
 
@@ -112,6 +111,18 @@ python scripts/python_test.py [--cpu | --nvidia | --cambricon | --ascend]
 
 详见 `test/infiniop-test` 目录
 
+#### 通信库（InfiniCCL）测试
+
+编译（需要先安装InfiniCCL）：
+```shell
+xmake build infiniccl-test
+```
+
+在英伟达平台运行测试（会自动使用所有可见的卡）：
+```shell
+infiniccl-test --nvidia
+```
+
 ## 开发指南
 
 ### 代码格式化

include/infiniccl.h

+#ifndef __INFINICCL_API_H__
+#define __INFINICCL_API_H__
+
+#include "infinirt.h"
+
+typedef enum {
+    INFINICCL_SUM = 0,
+    INFINICCL_PROD = 1,
+    INFINICCL_MAX = 2,
+    INFINICCL_MIN = 3,
+    INFINICCL_AVG = 4,
+} infinicclReduceOp_t;
+
+struct InfinicclComm;
+
+typedef struct InfinicclComm *infinicclComm_t;
+
+__C __export infiniStatus_t infinicclCommInitAll(
+    infiniDevice_t device_type,
+    infinicclComm_t *comms,
+    int ndevice,
+    const int *device_ids);
+
+__C __export infiniStatus_t infinicclCommDestroy(infinicclComm_t comm);
+
+__C __export infiniStatus_t infinicclAllReduce(
+    void *sendbuf,
+    void *recvbuf,
+    size_t count,
+    infiniDtype_t dataype,
+    infinicclReduceOp_t op,
+    infinicclComm_t comm,
+    infinirtStream_t stream);
+
+#endif

src/infiniccl-test/infiniccl_test.cpp

+#include "infiniccl_test.hpp"
+
+#include <chrono>
+#include <cstring>
+#include <iostream>
+#include <numeric>
+#include <pthread.h>
+#include <vector>
+
+#define TEST_INFINI(API__) CHECK_API_OR(API__, INFINI_STATUS_SUCCESS, return 1)
+#define TEST_INFINI_THREAD(API__) CHECK_API_OR(API__, INFINI_STATUS_SUCCESS, return nullptr)
+
+const size_t MAX_COUNT = 100ULL * 1024 * 1024;
+
+const size_t TEST_COUNTS[] = {
+    128,
+    1024,
+    4 * 1024,
+    MAX_COUNT,
+};
+
+const infiniDtype_t TEST_DTYPES[] = {INFINI_DTYPE_F32, INFINI_DTYPE_F16};
+
+const size_t WARM_UPS = 10;
+
+const size_t ITERATIONS = 100;
+
+struct ThreadArgs {
+    int rank;
+    int device_id;
+    infinicclComm_t comm;
+    infiniDevice_t device_type;
+    infiniDtype_t dtype;
+    size_t count;
+    const void *data;
+    const void *ans;
+    int *result;
+    double *time;
+};
+
+void setData(infiniDtype_t dtype, void *data, size_t count, float val) {
+    switch (dtype) {
+    case INFINI_DTYPE_F32:
+        for (size_t i = 0; i < count; i++) {
+            ((float *)data)[i] = val;
+        }
+        break;
+
+    case INFINI_DTYPE_F16:
+        for (size_t i = 0; i < count; i++) {
+            ((fp16_t *)data)[i] = utils::cast<fp16_t>(val);
+        }
+        break;
+    default:
+        std::abort();
+        break;
+    }
+}
+
+template <typename T>
+int checkData(const T *actual_, const T *expected_, size_t count) {
+    int failed = 0;
+    for (size_t i = 0; i < count; i++) {
+        if constexpr (std::is_same<T, fp16_t>::value) {
+            float actual = utils::cast<float>(actual_[i]);
+            float expected = utils::cast<float>(expected_[i]);
+            if (std::abs(actual - expected) > 1e-4) {
+                failed += 1;
+            }
+        } else {
+            if (std::abs(actual_[i] - expected_[i]) > 1e-4) {
+                failed += 1;
+            }
+        }
+    }
+    return failed;
+}
+
+int checkData(const void *actual, const void *expected, infiniDtype_t dtype, size_t count) {
+    switch (dtype) {
+    case INFINI_DTYPE_F32:
+        return checkData((const float *)actual, (const float *)expected, count);
+    case INFINI_DTYPE_F16:
+        return checkData((const fp16_t *)actual, (const fp16_t *)expected, count);
+    default:
+        std::abort();
+        return 1;
+    }
+}
+
+void *testAllReduceThread(void *arg) {
+    ThreadArgs *args = (ThreadArgs *)arg;
+    *(args->result) = 1;
+    TEST_INFINI_THREAD(infinirtSetDevice(args->device_type, args->device_id));
+    void *output = std::malloc(args->count * infiniSizeOf(args->dtype));
+    std::memset(output, 0, args->count * infiniSizeOf(args->dtype));
+    void *buf;
+    TEST_INFINI_THREAD(infinirtMalloc(&buf, args->count * infiniSizeOf(args->dtype)));
+    TEST_INFINI_THREAD(infinirtMemcpy(buf, args->data, args->count * infiniSizeOf(args->dtype), INFINIRT_MEMCPY_H2D));
+    TEST_INFINI_THREAD(infinicclAllReduce(buf, buf, args->count, args->dtype, INFINICCL_SUM, args->comm, NULL));
+    TEST_INFINI_THREAD(infinirtDeviceSynchronize());
+    TEST_INFINI_THREAD(infinirtMemcpy(output, buf, args->count * infiniSizeOf(args->dtype), INFINIRT_MEMCPY_D2H));
+
+    if (checkData(output, args->ans, args->dtype, args->count) != 0) {
+        std::free(output);
+        infinirtFree(buf);
+        return nullptr;
+    }
+    for (size_t i = 0; i < WARM_UPS; i++) {
+        TEST_INFINI_THREAD(infinicclAllReduce(buf, buf, args->count, args->dtype, INFINICCL_SUM, args->comm, NULL));
+    }
+    TEST_INFINI_THREAD(infinirtDeviceSynchronize());
+
+    // measure time
+    auto start = std::chrono::high_resolution_clock::now();
+    for (size_t i = 0; i < ITERATIONS; i++) {
+        TEST_INFINI_THREAD(infinicclAllReduce(buf, buf, args->count, args->dtype, INFINICCL_SUM, args->comm, NULL));
+    }
+    TEST_INFINI_THREAD(infinirtDeviceSynchronize());
+    auto end = std::chrono::high_resolution_clock::now();
+    double elapsed_ms = std::chrono::duration<double, std::milli>(end - start).count();
+    *args->time = elapsed_ms / ITERATIONS;
+
+    *args->result = 0;
+
+    std::free(output);
+    infinirtFree(buf);
+    return nullptr;
+}
+
+int testAllReduce(infiniDevice_t device_type, int ndevice) {
+    std::vector<ThreadArgs> thread_args(ndevice);
+    std::vector<infinicclComm_t> comms(ndevice);
+    std::vector<pthread_t> threads(ndevice);
+    std::vector<int> device_ids(ndevice);
+    std::vector<int> results(ndevice);
+    std::vector<double> times(ndevice);
+    void *data = std::malloc(MAX_COUNT * sizeof(float)); // Use float as max dtype size
+    void *ans = std::malloc(MAX_COUNT * sizeof(float));
+
+    for (int i = 0; i < ndevice; i++) {
+        device_ids[i] = i;
+    }
+    TEST_INFINI(infinicclCommInitAll(device_type, comms.data(), ndevice, device_ids.data()));
+
+    for (infiniDtype_t dtype : TEST_DTYPES) {
+        setData(dtype, data, MAX_COUNT, 1.0f);
+        setData(dtype, ans, MAX_COUNT, 1.0f * ndevice);
+        for (size_t count : TEST_COUNTS) {
+            std::cout << "Testing AllReduce with " << count << " elements of " << infiniDtypeToString(dtype) << std::endl;
+            for (int rank = 0; rank < ndevice; rank++) {
+                thread_args[rank] = {rank, device_ids[rank], comms[rank], device_type, dtype, count, data, ans, &results[rank], &times[rank]};
+                pthread_create(&threads[rank], NULL, testAllReduceThread, &thread_args[rank]);
+            }
+            for (int rank = 0; rank < ndevice; rank++) {
+                pthread_join(threads[rank], NULL);
+            }
+            int failed = std::accumulate(results.begin(), results.end(), 0);
+            for (int rank = 0; rank < ndevice; rank++) {
+                if (results[rank] != 0) {
+                    std::cout << "Rank " << rank << ": incorrect results." << std::endl;
+                } else {
+                    std::cout << "Rank " << rank << ": " << times[rank] << " ms." << std::endl;
+                }
+            }
+
+            if (failed > 0) {
+                std::cout << "Failed with " << failed << " errors." << std::endl
+                          << std::endl;
+                std::free(data);
+                std::free(ans);
+                return 1;
+            }
+            std::cout << std::endl;
+        }
+    }
+
+    std::free(data);
+    std::free(ans);
+    return 0;
+}

src/infiniccl-test/infiniccl_test.hpp

+#ifndef INFINICCL_TEST_HPP
+#define INFINICCL_TEST_HPP
+
+#include <infiniccl.h>
+
+#include "../utils.h"
+
+int testAllReduce(infiniDevice_t device_type, int ndevice);
+
+#endif // INFINICCL_TEST_HPP

src/infiniccl-test/main.cpp

+#include "infiniccl_test.hpp"
+
+#include <iostream>
+
+struct ParsedArgs {
+    infiniDevice_t device_type;
+};
+
+void printUsage() {
+    std::cout << "Usage:" << std::endl
+              << std::endl;
+    std::cout << "infiniccl-test --<device>" << std::endl
+              << std::endl;
+    std::cout << "  --<device>" << std::endl;
+    std::cout << "    Specify the device type --(nvidia|cambricon|ascend|metax|moore|iluvatar|kunlun|sugon)." << std::endl
+              << std::endl;
+    std::cout << "The program will run tests on all visible devices of the specified device type."
+              << " Use Environmental Variables such as CUDA_VSIBLE_DEVICES to limit visible device IDs.";
+    exit(-1);
+}
+
+#define PARSE_DEVICE(FLAG, DEVICE) \
+    if (arg == FLAG) {             \
+        args.device_type = DEVICE; \
+    }
+
+ParsedArgs parseArgs(int argc, char *argv[]) {
+    if (argc != 2) {
+        printUsage();
+    }
+
+    if (std::string(argv[1]) == "--help" || std::string(argv[1]) == "-h") {
+        printUsage();
+    }
+
+    ParsedArgs args;
+    try {
+        std::string arg = argv[1];
+        // clang-format off
+        PARSE_DEVICE("--nvidia", INFINI_DEVICE_NVIDIA)
+        else PARSE_DEVICE("--cambricon", INFINI_DEVICE_CAMBRICON)
+        else PARSE_DEVICE("--ascend", INFINI_DEVICE_ASCEND)
+        else PARSE_DEVICE("--metax", INFINI_DEVICE_METAX)
+        else PARSE_DEVICE("--moore", INFINI_DEVICE_MOORE)
+        else PARSE_DEVICE("--iluvatar", INFINI_DEVICE_ILUVATAR)
+        else PARSE_DEVICE("--kunlun", INFINI_DEVICE_KUNLUN)
+        else PARSE_DEVICE("--sugon", INFINI_DEVICE_SUGON)
+        else {
+            printUsage();
+        }
+        // clang-format on
+
+    } catch (const std::exception &) {
+        printUsage();
+    }
+
+    return args;
+}
+
+int main(int argc, char *argv[]) {
+    ParsedArgs args = parseArgs(argc, argv);
+    int ndevice = 0;
+    if (infinirtGetDeviceCount(args.device_type, &ndevice) != INFINI_STATUS_SUCCESS) {
+        std::cout << "Failed to get device count" << std::endl;
+        return -1;
+    }
+    if (ndevice == 0) {
+        std::cout << "No devices found. Tests skipped." << std::endl;
+        return 0;
+    } else {
+        std::cout << "Found " << ndevice << " devices. Running tests..." << std::endl;
+    }
+
+    int failed = 0;
+    failed += testAllReduce(args.device_type, ndevice);
+    return failed;
+}

src/infiniccl/cuda/infiniccl_cuda.cu

+#include "infiniccl_cuda.h"
+
+#include <cuda_runtime.h>
+#include <iostream>
+#include <nccl.h>
+#include <vector>
+
+#include "../../utils.h"
+
+#define CHECK_NCCL(API__) CHECK_INTERNAL(API__, ncclSuccess)
+
+inline cudaStream_t getCudaStream(infinirtStream_t stream) {
+    if (stream == nullptr) {
+        return 0;
+    }
+    return static_cast<cudaStream_t>(stream);
+}
+
+inline ncclDataType_t getNcclDtype(infiniDtype_t datatype) {
+    switch (datatype) {
+    case INFINI_DTYPE_F32:
+        return ncclFloat;
+    case INFINI_DTYPE_F16:
+        return ncclHalf;
+    default:
+        std::abort();
+        return ncclHalf;
+    }
+}
+
+inline ncclRedOp_t getNcclRedOp(infinicclReduceOp_t op) {
+    switch (op) {
+    case INFINICCL_SUM:
+        return ncclSum;
+    case INFINICCL_PROD:
+        return ncclProd;
+    case INFINICCL_MAX:
+        return ncclMax;
+    case INFINICCL_MIN:
+        return ncclMin;
+    case INFINICCL_AVG:
+        return ncclAvg;
+    default:
+        std::abort();
+        return ncclSum;
+    }
+}
+
+inline ncclComm_t getNcclComm(infinicclComm_t comm) {
+    return static_cast<ncclComm_t>(comm->comm);
+}
+
+namespace infiniccl::cuda {
+
+infiniStatus_t commInitAll(
+    infinicclComm_t *comms,
+    int ndevice,
+    const int *device_ids) {
+
+    std::vector<ncclComm_t> nccl_comms(ndevice);
+    CHECK_NCCL(ncclCommInitAll(nccl_comms.data(), ndevice, (int const *)device_ids));
+
+    for (int i = 0; i < ndevice; i++) {
+        comms[i] = new InfinicclComm{INFINI_DEVICE_NVIDIA, device_ids[i], (void *)(nccl_comms[i])};
+    }
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t commDestroy(infinicclComm_t comm) {
+    CHECK_NCCL(ncclCommDestroy(getNcclComm(comm)));
+    delete comm;
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t allReduce(
+    void *sendbuf,
+    void *recvbuf,
+    size_t count,
+    infiniDtype_t datatype,
+    infinicclReduceOp_t op,
+    infinicclComm_t comm,
+    infinirtStream_t stream) {
+
+    if (datatype != INFINI_DTYPE_F32 && datatype != INFINI_DTYPE_F16) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+
+    CHECK_NCCL(ncclAllReduce(sendbuf, recvbuf, count, getNcclDtype(datatype),
+                             getNcclRedOp(op), getNcclComm(comm), getCudaStream(stream)));
+
+    return INFINI_STATUS_SUCCESS;
+}
+} // namespace infiniccl::cuda

src/infiniccl/cuda/infiniccl_cuda.h

+#ifndef INFINICCL_CUDA_H_
+#define INFINICCL_CUDA_H_
+
+#include "../infiniccl_impl.h"
+
+// Windows does not support CUDA
+#if defined(ENABLE_CUDA_API) && defined(ENABLE_CCL) && !defined(_WIN32)
+INFINICCL_DEVICE_API_IMPL(cuda)
+#else
+INFINICCL_DEVICE_API_NOOP(cuda)
+#endif
+
+#endif /* INFINICCL_CUDA_H_ */

src/infiniccl/infiniccl.cc

+#include "infiniccl.h"
+
+#include "./cuda/infiniccl_cuda.h"
+
+__C infiniStatus_t infinicclCommInitAll(
+    infiniDevice_t device_type,
+    infinicclComm_t *comms,
+    int ndevice,
+    const int *device_ids) {
+
+#define COMM_INIT_ALL(CASE_, NAMESPACE_) \
+    case CASE_:                          \
+        return infiniccl::NAMESPACE_::commInitAll(comms, ndevice, device_ids);
+
+    switch (device_type) {
+        COMM_INIT_ALL(INFINI_DEVICE_NVIDIA, cuda)
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+
+#undef COMM_INIT_ALL
+}
+
+__C infiniStatus_t infinicclCommDestroy(infinicclComm_t comm) {
+    if (comm == nullptr) {
+        return INFINI_STATUS_SUCCESS;
+    }
+
+#define COMM_DESTROY(CASE_, NAMESPACE_) \
+    case CASE_:                         \
+        return infiniccl::NAMESPACE_::commDestroy(comm);
+
+    switch (comm->device_type) {
+        COMM_DESTROY(INFINI_DEVICE_NVIDIA, cuda)
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+#undef COMM_DESTROY
+}
+
+__C infiniStatus_t infinicclAllReduce(
+    void *sendbuf,
+    void *recvbuf,
+    size_t count,
+    infiniDtype_t dataype,
+    infinicclReduceOp_t op,
+    infinicclComm_t comm,
+    infinirtStream_t stream) {
+
+    if (comm == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+
+#define ALL_REDUCE(CASE_, NAMESPACE_) \
+    case CASE_:                       \
+        return infiniccl::NAMESPACE_::allReduce(sendbuf, recvbuf, count, dataype, op, comm, stream);
+
+    switch (comm->device_type) {
+        ALL_REDUCE(INFINI_DEVICE_NVIDIA, cuda)
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+
+#undef ALL_REDUCE
+}

src/infiniccl/infiniccl_impl.h

+#ifndef INFINICCL_IMPL_H
+#define INFINICCL_IMPL_H
+
+#include "infiniccl.h"
+
+struct InfinicclComm {
+    infiniDevice_t device_type;
+    int device_id; // the actual device ID, not rank number
+    void *comm;    // the actual communicator
+};
+
+#define INFINICCL_DEVICE_API(NAMSPACE, IMPL)               \
+    namespace infiniccl::NAMSPACE {                        \
+    infiniStatus_t commInitAll(                            \
+        infinicclComm_t *comms,                            \
+        int ndevice,                                       \
+        const int *device_ids) IMPL;                       \
+                                                           \
+    infiniStatus_t commDestroy(infinicclComm_t comm) IMPL; \
+                                                           \
+    infiniStatus_t allReduce(                              \
+        void *sendbuf,                                     \
+        void *recvbuf,                                     \
+        size_t count,                                      \
+        infiniDtype_t datatype,                            \
+        infinicclReduceOp_t op,                            \
+        infinicclComm_t comm,                              \
+        infinirtStream_t stream) IMPL;                     \
+    };
+
+#define INFINICCL_DEVICE_API_IMPL(NAMSPACE) \
+    INFINICCL_DEVICE_API(NAMSPACE, )
+
+#define INFINICCL_DEVICE_API_NOOP(NAMSPACE) \
+    INFINICCL_DEVICE_API(NAMSPACE, { return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; })
+
+#endif // INFINICCL_IMPL_H

src/infinirt/infinirt.cc

@@ -44,10 +44,10 @@ __C infiniStatus_t infinirtGetDevice(infiniDevice_t *device_ptr, int *device_id_
     return INFINI_STATUS_SUCCESS;
 }
 
-#define INFINIRT_CALL_DEVICE_API_AND(API, PARAMS, ACTION)   \
+#define INFINIRT_CALL_DEVICE_API_AND(DEVICE_TYPE, API, PARAMS, ACTION) \
     {                                                                  \
         infiniStatus_t _status;                                        \
-        switch (CURRENT_DEVICE_TYPE) {                      \
+        switch (DEVICE_TYPE) {                                         \
         case INFINI_DEVICE_CPU:                                        \
             _status = infinirt::cpu::API PARAMS;                       \
             break;                                                     \
@@ -73,17 +73,17 @@ __C infiniStatus_t infinirtGetDevice(infiniDevice_t *device_ptr, int *device_id_
         return _status;                                                \
     }
 
-#define INFINIRT_CALL_DEVICE_API(API, PARAMS) INFINIRT_CALL_DEVICE_API_AND(API, PARAMS, )
+#define INFINIRT_CALL_DEVICE_API(API, PARAMS) INFINIRT_CALL_DEVICE_API_AND(CURRENT_DEVICE_TYPE, API, PARAMS, )
 
 __C infiniStatus_t infinirtGetDeviceCount(infiniDevice_t device, int *count) {
     if (count == nullptr) {
         return INFINI_STATUS_NULL_POINTER;
     }
-    INFINIRT_CALL_DEVICE_API(getDeviceCount, (count));
+    INFINIRT_CALL_DEVICE_API_AND(device, getDeviceCount, (count), {});
 }
 
 __C infiniStatus_t infinirtSetDevice(infiniDevice_t device, int device_id) {
-    INFINIRT_CALL_DEVICE_API_AND(setDevice, (device_id),
+    INFINIRT_CALL_DEVICE_API_AND(device, setDevice, (device_id),
                                  { CURRENT_DEVICE_TYPE = device;
                                    CURRENT_DEVICE_ID = device_id; });
 }

xmake.lua

@@ -118,6 +118,18 @@ if has_config("kunlun-xpu") then
     includes("xmake/kunlun.lua")
 end
 
+-- InfiniCCL
+option("ccl")
+set_default(false)
+    set_default(false)
+    set_showmenu(true)
+    set_description("Wether to complie implementations for InfiniCCL")
+option_end()
+
+if has_config("ccl") then
+    add_defines("ENABLE_CCL")
+end
+
 target("infini-utils")
     set_kind("static")
     on_install(function (target) end)
@@ -220,10 +232,25 @@ target("infiniop")
     add_installfiles("include/infinicore.h", {prefixdir = "include"})
 target_end()
 
+target("infiniccl")
+    set_kind("shared")
+    add_deps("infinirt")
+
+    if has_config("nv-gpu") then
+        add_deps("infiniccl-cuda")
+    end
+    
+    set_languages("cxx17")
+
+    add_files("src/infiniccl/*.cc")
+    add_installfiles("include/infiniccl.h", {prefixdir = "include"})
+
+    set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
+target_end()
 
 target("all")
     set_kind("phony")
-    add_deps("infiniop", "infinirt")
+    add_deps("infiniop", "infinirt", "infiniccl")
     after_build(function (target) print(YELLOW .. "[Congratulations!] Now you can install the libraries with \"xmake install\"" .. NC) end)
 target_end()
 

xmake/cuda.lua

@@ -58,3 +58,34 @@ target("infinirt-cuda")
     set_languages("cxx17")
     add_files("../src/infinirt/cuda/*.cu")
 target_end()
+
+target("infiniccl-cuda")
+    set_kind("static")
+    add_deps("infinirt")
+    on_install(function (target) end)
+    if has_config("ccl") then
+        set_policy("build.cuda.devlink", true)
+        set_toolchains("cuda")
+        add_links("cudart")
+
+        if not is_plat("windows") then
+            add_cuflags("-Xcompiler=-fPIC")
+            add_culdflags("-Xcompiler=-fPIC")
+            add_cxflags("-fPIC")
+
+            local nccl_root = os.getenv("NCCL_ROOT")
+            if nccl_root then
+                add_includedirs(nccl_root .. "/include")
+                add_links(nccl_root .. "/lib/libnccl.so")
+            else
+                add_links("nccl") -- Fall back to default nccl linking
+            end
+
+            add_files("../src/infiniccl/cuda/*.cu")
+        else
+            print("[Warning] NCCL is not supported on Windows")
+        end
+    end
+    set_languages("cxx17")
+    
+target_end()

xmake/test.lua

@@ -34,3 +34,20 @@ target("infiniop-test")
 
     set_installdir(INFINI_ROOT)
 target_end()
+
+target("infiniccl-test")
+    set_kind("binary")
+    add_deps("infini-utils")
+    set_default(false)
+
+    set_warnings("all", "error")
+    set_languages("cxx17")
+
+    local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini")
+    add_includedirs(INFINI_ROOT.."/include")
+    add_linkdirs(INFINI_ROOT.."/lib")
+    add_links("infinirt", "infiniccl")
+    add_files(os.projectdir().."/src/infiniccl-test/*.cpp")
+
+    set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
+target_end()