Skip to content

Commit 50eaee6

Browse files
PanZezhong1725PanZezhong1725
authored
Merge pull request #339 from InfiniTensor/issue/112_new
issue/112: infinirt测试框架
2 parents 639d921 + 3c74a2e commit 50eaee6

File tree

6 files changed

+231
-0
lines changed

6 files changed

+231
-0
lines changed

‎.github/workflows/build.yml‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,3 +40,13 @@ jobs:
4040
4141
- name: Python Test
4242
run: python scripts/python_test.py --cpu
43+
44+
- name: run infinirt-test --cpu on Linux
45+
if: matrix.os == 'ubuntu-latest'
46+
run: |
47+
./build/linux/x86_64/release/infinirt-test --cpu
48+
49+
- name: run infinirt-test --cpu on Windows
50+
if: matrix.os == 'windows-latest'
51+
run: |
52+
.\build\windows\x64\release\infinirt-test.exe --cpu

‎src/infinirt-test/main.cc‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
#include "test.h"
2+
#include <infinirt.h>
3+
4+
struct ParsedArgs {
5+
infiniDevice_t device_type = INFINI_DEVICE_CPU;
6+
};
7+
8+
void printUsage() {
9+
std::cout << "Usage:" << std::endl
10+
<< " infinirt-test [--<device>]" << std::endl
11+
<< std::endl
12+
<< "Options:" << std::endl
13+
<< " --<device> Specify the device type." << std::endl
14+
<< std::endl
15+
<< "Available devices:" << std::endl
16+
<< " cpu - Default" << std::endl
17+
<< " nvidia" << std::endl
18+
<< " cambricon" << std::endl
19+
<< " ascend" << std::endl
20+
<< " metax" << std::endl
21+
<< " moore" << std::endl
22+
<< " iluvatar" << std::endl
23+
<< " kunlun" << std::endl
24+
<< " sugon" << std::endl
25+
<< std::endl;
26+
exit(EXIT_FAILURE);
27+
}
28+
29+
ParsedArgs parseArgs(int argc, char *argv[]) {
30+
ParsedArgs args;
31+
32+
if (argc < 2) {
33+
return args; // 默认使用 CPU
34+
}
35+
36+
std::string arg = argv[1];
37+
if (arg == "--help" || arg == "-h") {
38+
printUsage();
39+
}
40+
41+
try {
42+
#define PARSE_DEVICE(FLAG, DEVICE) \
43+
if (arg == FLAG) { \
44+
args.device_type = DEVICE; \
45+
}
46+
// clang-format off
47+
PARSE_DEVICE("--cpu", INFINI_DEVICE_CPU)
48+
else PARSE_DEVICE("--nvidia", INFINI_DEVICE_NVIDIA)
49+
else PARSE_DEVICE("--cambricon", INFINI_DEVICE_CAMBRICON)
50+
else PARSE_DEVICE("--ascend", INFINI_DEVICE_ASCEND)
51+
else PARSE_DEVICE("--metax", INFINI_DEVICE_METAX)
52+
else PARSE_DEVICE("--moore", INFINI_DEVICE_MOORE)
53+
else PARSE_DEVICE("--iluvatar", INFINI_DEVICE_ILUVATAR)
54+
else PARSE_DEVICE("--kunlun", INFINI_DEVICE_KUNLUN)
55+
else PARSE_DEVICE("--sugon", INFINI_DEVICE_SUGON)
56+
else {
57+
printUsage();
58+
}
59+
// clang-format on
60+
#undef PARSE_DEVICE
61+
} catch (const std::exception &) {
62+
printUsage();
63+
}
64+
65+
return args;
66+
}
67+
68+
int main(int argc, char *argv[]) {
69+
70+
ParsedArgs args = parseArgs(argc, argv);
71+
std::cout << "Testing Device: " << args.device_type << std::endl;
72+
infiniDevice_t device = args.device_type;
73+
74+
// 获取设备总数
75+
std::vector<int> deviceCounts(INFINI_DEVICE_TYPE_COUNT, 0);
76+
if (infinirtGetAllDeviceCount(deviceCounts.data()) != INFINI_STATUS_SUCCESS) {
77+
std::cerr << "Failed to get total device count." << std::endl;
78+
return 1;
79+
}
80+
81+
int numDevices = deviceCounts[device];
82+
std::cout << "Device Type: " << device << " | Available Devices: " << numDevices << std::endl;
83+
84+
if (numDevices == 0) {
85+
std::cout << "Device type " << device << " has no available devices." << std::endl;
86+
return 0;
87+
}
88+
89+
for (int deviceId = 0; deviceId < numDevices; ++deviceId) {
90+
if (!testSetDevice(device, deviceId)) {
91+
return 1;
92+
}
93+
94+
size_t dataSize[] = {1 << 10, 4 << 10, 2 << 20, 1L << 30};
95+
96+
for (size_t size : dataSize) {
97+
if (!testMemcpy(device, deviceId, size)) {
98+
return 1;
99+
}
100+
}
101+
}
102+
103+
return 0;
104+
}

‎src/infinirt-test/test.cc‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
#include "test.h"
2+
#include <cstring>
3+
#include <infinirt.h>
4+
#include <iostream>
5+
6+
bool testMemcpy(infiniDevice_t device, int deviceId, size_t dataSize) {
7+
8+
std::cout << "==============================================\n"
9+
<< "Testing memcpy on Device ID: " << deviceId << "\n"
10+
<< "==============================================" << std::endl;
11+
12+
// 分配主机内存
13+
std::cout << "[Device " << deviceId << "] Allocating host memory: " << dataSize * sizeof(float) << " bytes" << std::endl;
14+
std::vector<float> hostData(dataSize, 1.23f);
15+
std::vector<float> hostCopy(dataSize, 0.0f);
16+
17+
// 分配设备内存
18+
void *deviceSrc = nullptr, *deviceDst = nullptr;
19+
size_t dataSizeInBytes = dataSize * sizeof(float);
20+
21+
std::cout << "[Device " << deviceId << "] Allocating device memory: " << dataSizeInBytes << " bytes" << std::endl;
22+
if (infinirtMalloc(&deviceSrc, dataSizeInBytes) != INFINI_STATUS_SUCCESS) {
23+
std::cerr << "[Device " << deviceId << "] Failed to allocate device memory for deviceSrc." << std::endl;
24+
return false;
25+
}
26+
27+
if (infinirtMalloc(&deviceDst, dataSizeInBytes) != INFINI_STATUS_SUCCESS) {
28+
std::cerr << "[Device " << deviceId << "] Failed to allocate device memory for deviceDst." << std::endl;
29+
infinirtFree(deviceSrc);
30+
return false;
31+
}
32+
33+
// 复制数据到设备
34+
std::cout << "[Device " << deviceId << "] Copying data from host to device..." << std::endl;
35+
if (infinirtMemcpy(deviceSrc, hostData.data(), dataSizeInBytes, INFINIRT_MEMCPY_H2D) != INFINI_STATUS_SUCCESS) {
36+
std::cerr << "[Device " << deviceId << "] Failed to copy data from host to device." << std::endl;
37+
infinirtFree(deviceSrc);
38+
infinirtFree(deviceDst);
39+
return false;
40+
}
41+
42+
// 设备内存间复制
43+
std::cout << "[Device " << deviceId << "] Copying data between device memory (D2D)..." << std::endl;
44+
if (infinirtMemcpy(deviceDst, deviceSrc, dataSizeInBytes, INFINIRT_MEMCPY_D2D) != INFINI_STATUS_SUCCESS) {
45+
std::cerr << "[Device " << deviceId << "] Failed to copy data from device to device." << std::endl;
46+
infinirtFree(deviceSrc);
47+
infinirtFree(deviceDst);
48+
return false;
49+
}
50+
51+
// 设备数据复制回主机
52+
std::cout << "[Device " << deviceId << "] Copying data from device back to host..." << std::endl;
53+
if (infinirtMemcpy(hostCopy.data(), deviceDst, dataSizeInBytes, INFINIRT_MEMCPY_D2H) != INFINI_STATUS_SUCCESS) {
54+
std::cerr << "[Device " << deviceId << "] Failed to copy data from device to host." << std::endl;
55+
infinirtFree(deviceSrc);
56+
infinirtFree(deviceDst);
57+
return false;
58+
}
59+
60+
// 数据验证
61+
std::cout << "[Device " << deviceId << "] Validating copied data..." << std::endl;
62+
if (std::memcmp(hostData.data(), hostCopy.data(), dataSizeInBytes) != 0) {
63+
std::cerr << "[Device " << deviceId << "] Data mismatch between hostData and hostCopy." << std::endl;
64+
infinirtFree(deviceSrc);
65+
infinirtFree(deviceDst);
66+
return false;
67+
}
68+
69+
std::cout << "[Device " << deviceId << "] Data copied correctly!" << std::endl;
70+
71+
// 释放设备内存
72+
std::cout << "[Device " << deviceId << "] Freeing device memory..." << std::endl;
73+
infinirtFree(deviceSrc);
74+
infinirtFree(deviceDst);
75+
76+
std::cout << "[Device " << deviceId << "] Memory copy test PASSED!" << std::endl;
77+
78+
return true;
79+
}
80+
81+
bool testSetDevice(infiniDevice_t device, int deviceId) {
82+
83+
std::cout << "Setting device " << device << " with ID: " << deviceId << std::endl;
84+
85+
infiniStatus_t status = infinirtSetDevice(device, deviceId);
86+
87+
if (status != INFINI_STATUS_SUCCESS) {
88+
std::cerr << "Failed to set device " << device << " with ID " << deviceId << std::endl;
89+
return false;
90+
}
91+
92+
return true;
93+
}

‎src/infinirt-test/test.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#ifndef __INFINIRT_TEST_H__
2+
#define __INFINIRT_TEST_H__
3+
#include "../utils.h"
4+
5+
bool testSetDevice(infiniDevice_t device, int deviceId);
6+
bool testMemcpy(infiniDevice_t device, int deviceId, size_t dataSize);
7+
8+
#endif

‎src/infinirt/infinirt.cc‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,10 @@ __C infiniStatus_t infinirtGetAllDeviceCount(int *count_array) {
2323
return INFINI_STATUS_NULL_POINTER;
2424
}
2525
for (size_t i = 0; i < INFINI_DEVICE_TYPE_COUNT; i++) {
26+
if (i == INFINI_DEVICE_ILUVATAR || i == INFINI_DEVICE_KUNLUN || i == INFINI_DEVICE_SUGON) {
27+
count_array[i] = 0;
28+
continue;
29+
}
2630
auto status = infinirtGetDeviceCount(static_cast<infiniDevice_t>(i), &count_array[i]);
2731
if (status != INFINI_STATUS_SUCCESS) {
2832
return status;

‎xmake/test.lua‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,3 +51,15 @@ target("infiniccl-test")
5151

5252
set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
5353
target_end()
54+
55+
target("infinirt-test")
56+
set_kind("binary")
57+
add_deps("infinirt")
58+
on_install(function (target) end)
59+
60+
set_languages("cxx17")
61+
set_warnings("all", "error")
62+
63+
add_files(os.projectdir().."/src/infinirt-test/*.cc")
64+
set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
65+
target_end()

0 commit comments

Comments
 (0)