Merge pull request #339 from InfiniTensor/issue/112_new

issue/112: infinirt测试框架

.github/workflows/build.yml

@@ -40,3 +40,13 @@ jobs:
 
     - name: Python Test
       run: python scripts/python_test.py --cpu
+
+    - name: run infinirt-test --cpu on Linux
+      if: matrix.os == 'ubuntu-latest'
+      run: |
+        ./build/linux/x86_64/release/infinirt-test --cpu
+
+    - name: run infinirt-test --cpu on Windows
+      if: matrix.os == 'windows-latest'
+      run: |
+        .\build\windows\x64\release\infinirt-test.exe --cpu

src/infinirt-test/main.cc

+#include "test.h"
+#include <infinirt.h>
+
+struct ParsedArgs {
+    infiniDevice_t device_type = INFINI_DEVICE_CPU;
+};
+
+void printUsage() {
+    std::cout << "Usage:" << std::endl
+              << "  infinirt-test [--<device>]" << std::endl
+              << std::endl
+              << "Options:" << std::endl
+              << "  --<device>   Specify the device type." << std::endl
+              << std::endl
+              << "Available devices:" << std::endl
+              << "  cpu         - Default" << std::endl
+              << "  nvidia" << std::endl
+              << "  cambricon" << std::endl
+              << "  ascend" << std::endl
+              << "  metax" << std::endl
+              << "  moore" << std::endl
+              << "  iluvatar" << std::endl
+              << "  kunlun" << std::endl
+              << "  sugon" << std::endl
+              << std::endl;
+    exit(EXIT_FAILURE);
+}
+
+ParsedArgs parseArgs(int argc, char *argv[]) {
+    ParsedArgs args;
+
+    if (argc < 2) {
+        return args; // 默认使用 CPU
+    }
+
+    std::string arg = argv[1];
+    if (arg == "--help" || arg == "-h") {
+        printUsage();
+    }
+
+    try {
+#define PARSE_DEVICE(FLAG, DEVICE) \
+    if (arg == FLAG) {             \
+        args.device_type = DEVICE; \
+    }
+        // clang-format off
+        PARSE_DEVICE("--cpu", INFINI_DEVICE_CPU)
+        else 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
+#undef PARSE_DEVICE
+    } catch (const std::exception &) {
+        printUsage();
+    }
+
+    return args;
+}
+
+int main(int argc, char *argv[]) {
+
+    ParsedArgs args = parseArgs(argc, argv);
+    std::cout << "Testing Device: " << args.device_type << std::endl;
+    infiniDevice_t device = args.device_type;
+
+    // 获取设备总数
+    std::vector<int> deviceCounts(INFINI_DEVICE_TYPE_COUNT, 0);
+    if (infinirtGetAllDeviceCount(deviceCounts.data()) != INFINI_STATUS_SUCCESS) {
+        std::cerr << "Failed to get total device count." << std::endl;
+        return 1;
+    }
+
+    int numDevices = deviceCounts[device];
+    std::cout << "Device Type: " << device << " | Available Devices: " << numDevices << std::endl;
+
+    if (numDevices == 0) {
+        std::cout << "Device type " << device << " has no available devices." << std::endl;
+        return 0;
+    }
+
+    for (int deviceId = 0; deviceId < numDevices; ++deviceId) {
+        if (!testSetDevice(device, deviceId)) {
+            return 1;
+        }
+
+        size_t dataSize[] = {1 << 10, 4 << 10, 2 << 20, 1L << 30};
+
+        for (size_t size : dataSize) {
+            if (!testMemcpy(device, deviceId, size)) {
+                return 1;
+            }
+        }
+    }
+
+    return 0;
+}

src/infinirt-test/test.cc

+#include "test.h"
+#include <cstring>
+#include <infinirt.h>
+#include <iostream>
+
+bool testMemcpy(infiniDevice_t device, int deviceId, size_t dataSize) {
+
+    std::cout << "==============================================\n"
+              << "Testing memcpy on Device ID: " << deviceId << "\n"
+              << "==============================================" << std::endl;
+
+    // 分配主机内存
+    std::cout << "[Device " << deviceId << "] Allocating host memory: " << dataSize * sizeof(float) << " bytes" << std::endl;
+    std::vector<float> hostData(dataSize, 1.23f);
+    std::vector<float> hostCopy(dataSize, 0.0f);
+
+    // 分配设备内存
+    void *deviceSrc = nullptr, *deviceDst = nullptr;
+    size_t dataSizeInBytes = dataSize * sizeof(float);
+
+    std::cout << "[Device " << deviceId << "] Allocating device memory: " << dataSizeInBytes << " bytes" << std::endl;
+    if (infinirtMalloc(&deviceSrc, dataSizeInBytes) != INFINI_STATUS_SUCCESS) {
+        std::cerr << "[Device " << deviceId << "] Failed to allocate device memory for deviceSrc." << std::endl;
+        return false;
+    }
+
+    if (infinirtMalloc(&deviceDst, dataSizeInBytes) != INFINI_STATUS_SUCCESS) {
+        std::cerr << "[Device " << deviceId << "] Failed to allocate device memory for deviceDst." << std::endl;
+        infinirtFree(deviceSrc);
+        return false;
+    }
+
+    // 复制数据到设备
+    std::cout << "[Device " << deviceId << "] Copying data from host to device..." << std::endl;
+    if (infinirtMemcpy(deviceSrc, hostData.data(), dataSizeInBytes, INFINIRT_MEMCPY_H2D) != INFINI_STATUS_SUCCESS) {
+        std::cerr << "[Device " << deviceId << "] Failed to copy data from host to device." << std::endl;
+        infinirtFree(deviceSrc);
+        infinirtFree(deviceDst);
+        return false;
+    }
+
+    // 设备内存间复制
+    std::cout << "[Device " << deviceId << "] Copying data between device memory (D2D)..." << std::endl;
+    if (infinirtMemcpy(deviceDst, deviceSrc, dataSizeInBytes, INFINIRT_MEMCPY_D2D) != INFINI_STATUS_SUCCESS) {
+        std::cerr << "[Device " << deviceId << "] Failed to copy data from device to device." << std::endl;
+        infinirtFree(deviceSrc);
+        infinirtFree(deviceDst);
+        return false;
+    }
+
+    // 设备数据复制回主机
+    std::cout << "[Device " << deviceId << "] Copying data from device back to host..." << std::endl;
+    if (infinirtMemcpy(hostCopy.data(), deviceDst, dataSizeInBytes, INFINIRT_MEMCPY_D2H) != INFINI_STATUS_SUCCESS) {
+        std::cerr << "[Device " << deviceId << "] Failed to copy data from device to host." << std::endl;
+        infinirtFree(deviceSrc);
+        infinirtFree(deviceDst);
+        return false;
+    }
+
+    // 数据验证
+    std::cout << "[Device " << deviceId << "] Validating copied data..." << std::endl;
+    if (std::memcmp(hostData.data(), hostCopy.data(), dataSizeInBytes) != 0) {
+        std::cerr << "[Device " << deviceId << "] Data mismatch between hostData and hostCopy." << std::endl;
+        infinirtFree(deviceSrc);
+        infinirtFree(deviceDst);
+        return false;
+    }
+
+    std::cout << "[Device " << deviceId << "] Data copied correctly!" << std::endl;
+
+    // 释放设备内存
+    std::cout << "[Device " << deviceId << "] Freeing device memory..." << std::endl;
+    infinirtFree(deviceSrc);
+    infinirtFree(deviceDst);
+
+    std::cout << "[Device " << deviceId << "] Memory copy test PASSED!" << std::endl;
+
+    return true;
+}
+
+bool testSetDevice(infiniDevice_t device, int deviceId) {
+
+    std::cout << "Setting device " << device << " with ID: " << deviceId << std::endl;
+
+    infiniStatus_t status = infinirtSetDevice(device, deviceId);
+
+    if (status != INFINI_STATUS_SUCCESS) {
+        std::cerr << "Failed to set device " << device << " with ID " << deviceId << std::endl;
+        return false;
+    }
+
+    return true;
+}

src/infinirt-test/test.h

+#ifndef __INFINIRT_TEST_H__
+#define __INFINIRT_TEST_H__
+#include "../utils.h"
+
+bool testSetDevice(infiniDevice_t device, int deviceId);
+bool testMemcpy(infiniDevice_t device, int deviceId, size_t dataSize);
+
+#endif

src/infinirt/infinirt.cc

@@ -23,6 +23,10 @@ __C infiniStatus_t infinirtGetAllDeviceCount(int *count_array) {
         return INFINI_STATUS_NULL_POINTER;
     }
     for (size_t i = 0; i < INFINI_DEVICE_TYPE_COUNT; i++) {
+        if (i == INFINI_DEVICE_ILUVATAR || i == INFINI_DEVICE_KUNLUN || i == INFINI_DEVICE_SUGON) {
+            count_array[i] = 0;
+            continue;
+        }
         auto status = infinirtGetDeviceCount(static_cast<infiniDevice_t>(i), &count_array[i]);
         if (status != INFINI_STATUS_SUCCESS) {
             return status;

xmake/test.lua

@@ -51,3 +51,15 @@ target("infiniccl-test")
 
     set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
 target_end()
+
+target("infinirt-test")
+    set_kind("binary")
+    add_deps("infinirt")
+    on_install(function (target) end)
+
+    set_languages("cxx17")
+    set_warnings("all", "error")
+
+    add_files(os.projectdir().."/src/infinirt-test/*.cc")
+    set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
+target_end()