Skip to content

Commit 02fbd40

Browse files
yangchaossyangchaoss
committed
test(kunlunxin): add IPC memory handle validation
1 parent 0beba7a commit 02fbd40

5 files changed

Lines changed: 249 additions & 3 deletions

File tree

‎flagcx/adaptor/device/kunlunxin_adaptor.cc‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -267,7 +267,16 @@ flagcxResult_t kunlunAdaptorEventQuery(flagcxEvent_t event) {
267267

268268
flagcxResult_t kunlunAdaptorIpcMemHandleCreate(flagcxIpcMemHandle_t *handle,
269269
size_t *size) {
270-
flagcxCalloc(handle, 1);
270+
if (handle == NULL) {
271+
return flagcxInvalidArgument;
272+
}
273+
274+
*handle = NULL;
275+
flagcxResult_t result = flagcxCalloc(handle, 1);
276+
if (result != flagcxSuccess) {
277+
return result;
278+
}
279+
271280
if (size != NULL) {
272281
*size = sizeof(cudaIpcMemHandle_t);
273282
}

‎test/make.inc‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,14 @@ ifeq ($(USE_NVIDIA), 1)
7676
DEVICE_COMPILER := $(DEVICE_HOME)/bin/nvcc
7777
ADAPTOR_DEFINE := -DUSE_NVIDIA_ADAPTOR
7878
PLATFORM_IR_DIR := $(PROJECT_ROOT)/bindings/ir/nvidia
79+
else ifeq ($(USE_KUNLUNXIN), 1)
80+
DEVICE_HOME ?= /usr/local/xpu
81+
DEVICE_LIB := $(DEVICE_HOME)/so
82+
DEVICE_INCLUDE := -I$(DEVICE_HOME)/include
83+
DEVICE_LINK := -L$(DEVICE_LIB) -lxpurt -lcudart
84+
ADAPTOR_DEFINE := -DUSE_KUNLUNXIN_ADAPTOR
85+
CCL_HOME ?= /usr/local/xccl
86+
CCL_INCLUDE := -I$(CCL_HOME)/include
7987
else ifeq ($(USE_DU), 1)
8088
DEVICE_HOME ?= $(CUDA_PATH)
8189
DEVICE_LIB := $(DEVICE_HOME)/lib64

‎test/unittest/adaptor/Makefile‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,12 +35,12 @@ $(MPI_TARGET): $(MPI_OBJS)
3535
$(OBJDIR)/test_%.o: test_%.cpp
3636
@mkdir -p $(OBJDIR)
3737
@echo "Compiling $<"
38-
@$(CXX) $< -o $@ -c $(CXXFLAGS) $(FLAGCX_INCLUDES) $(GTEST_INCLUDE) $(FIXTURE_INCLUDE)
38+
@$(CXX) $< -o $@ -c $(CXXFLAGS) $(ADAPTOR_DEFINE) $(FLAGCX_INCLUDES) $(DEVICE_INCLUDE) $(CCL_INCLUDE) $(GTEST_INCLUDE) $(FIXTURE_INCLUDE)
3939

4040
$(OBJDIR)/coll_%.o: coll_%.cpp
4141
@mkdir -p $(OBJDIR)
4242
@echo "Compiling $<"
43-
@$(CXX) $< -o $@ -c $(CXXFLAGS) $(FLAGCX_INCLUDES) $(GTEST_INCLUDE) $(FIXTURE_INCLUDE) $(MPI_INCLUDE)
43+
@$(CXX) $< -o $@ -c $(CXXFLAGS) $(ADAPTOR_DEFINE) $(FLAGCX_INCLUDES) $(GTEST_INCLUDE) $(FIXTURE_INCLUDE) $(MPI_INCLUDE) $(DEVICE_INCLUDE) $(CCL_INCLUDE)
4444

4545
clean:
4646
@rm -rf $(BUILDDIR)

‎test/unittest/adaptor/coll_ipc_mem_handle.cpp‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
/*************************************************************************
2+
* Copyright (c) 2026. All Rights Reserved.
3+
* Cross-process IPC memory handle test (requires MPI and KunlunXin XPU).
4+
************************************************************************/
5+
6+
#include <gtest/gtest.h>
7+
#include <mpi.h>
8+
9+
#include "adaptor.h"
10+
#include "flagcx.h"
11+
12+
13+
namespace {
14+
15+
#define ASSERT_FLAGCX_SUCCESS(expr) \
16+
do { \
17+
flagcxResult_t result = (expr); \
18+
if (result != flagcxSuccess) { \
19+
ADD_FAILURE() << #expr << " returned " << static_cast<int>(result); \
20+
MPI_Abort(MPI_COMM_WORLD, static_cast<int>(result)); \
21+
return; \
22+
} \
23+
} while (0)
24+
25+
#define ASSERT_MPI_SUCCESS(expr) \
26+
do { \
27+
int result = (expr); \
28+
if (result != MPI_SUCCESS) { \
29+
ADD_FAILURE() << #expr << " returned " << result; \
30+
MPI_Abort(MPI_COMM_WORLD, result); \
31+
return; \
32+
} \
33+
} while (0)
34+
35+
class IpcMemHandleMpiTest : public ::testing::Test {
36+
protected:
37+
void SetUp() override {
38+
flagcxDeviceHandleInit(&devHandle);
39+
ASSERT_NE(devHandle, nullptr);
40+
41+
int deviceCount = 0;
42+
ASSERT_FLAGCX_SUCCESS(devHandle->getDeviceCount(&deviceCount));
43+
if (deviceCount <= 0) {
44+
ADD_FAILURE() << "No visible XPU device";
45+
MPI_Abort(MPI_COMM_WORLD, 1);
46+
return;
47+
}
48+
ASSERT_FLAGCX_SUCCESS(devHandle->setDevice(0));
49+
}
50+
51+
void TearDown() override {
52+
if (devHandle != nullptr) {
53+
flagcxDeviceHandleFree(devHandle);
54+
}
55+
}
56+
57+
flagcxDeviceHandle_t devHandle = nullptr;
58+
};
59+
60+
TEST_F(IpcMemHandleMpiTest, CrossProcessLifecycle) {
61+
int rank = -1;
62+
int worldSize = 0;
63+
ASSERT_MPI_SUCCESS(MPI_Comm_rank(MPI_COMM_WORLD, &rank));
64+
ASSERT_MPI_SUCCESS(MPI_Comm_size(MPI_COMM_WORLD, &worldSize));
65+
66+
if (worldSize != 2) {
67+
GTEST_SKIP() << "CrossProcessLifecycle requires exactly 2 MPI ranks";
68+
}
69+
70+
constexpr size_t bufferSize = 4096;
71+
constexpr int expectedValue = 0x12345678;
72+
73+
if (rank == 0) {
74+
void *devPtr = nullptr;
75+
ASSERT_FLAGCX_SUCCESS(devHandle->deviceMalloc(
76+
&devPtr, bufferSize, flagcxMemDevice, nullptr));
77+
ASSERT_NE(devPtr, nullptr);
78+
79+
int hostValue = expectedValue;
80+
ASSERT_FLAGCX_SUCCESS(devHandle->deviceMemcpy(
81+
devPtr, &hostValue, sizeof(hostValue), flagcxMemcpyHostToDevice,
82+
nullptr));
83+
ASSERT_FLAGCX_SUCCESS(devHandle->streamSynchronize(nullptr));
84+
85+
flagcxIpcMemHandle_t handle = nullptr;
86+
size_t handleSize = 0;
87+
ASSERT_FLAGCX_SUCCESS(
88+
devHandle->ipcMemHandleCreate(&handle, &handleSize));
89+
ASSERT_NE(handle, nullptr);
90+
ASSERT_GT(handleSize, static_cast<size_t>(0));
91+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleGet(handle, devPtr));
92+
93+
ASSERT_MPI_SUCCESS(MPI_Send(&handleSize, sizeof(handleSize), MPI_BYTE, 1,
94+
0, MPI_COMM_WORLD));
95+
ASSERT_MPI_SUCCESS(MPI_Send(handle, static_cast<int>(handleSize), MPI_BYTE,
96+
1, 1, MPI_COMM_WORLD));
97+
98+
int acknowledgement = 0;
99+
ASSERT_MPI_SUCCESS(MPI_Recv(&acknowledgement, 1, MPI_INT, 1, 2,
100+
MPI_COMM_WORLD, MPI_STATUS_IGNORE));
101+
ASSERT_EQ(acknowledgement, 1);
102+
103+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleFree(handle));
104+
ASSERT_FLAGCX_SUCCESS(
105+
devHandle->deviceFree(devPtr, flagcxMemDevice, nullptr));
106+
} else {
107+
size_t receivedHandleSize = 0;
108+
ASSERT_MPI_SUCCESS(MPI_Recv(&receivedHandleSize, sizeof(receivedHandleSize),
109+
MPI_BYTE, 0, 0, MPI_COMM_WORLD,
110+
MPI_STATUS_IGNORE));
111+
112+
flagcxIpcMemHandle_t handle = nullptr;
113+
size_t localHandleSize = 0;
114+
ASSERT_FLAGCX_SUCCESS(
115+
devHandle->ipcMemHandleCreate(&handle, &localHandleSize));
116+
ASSERT_NE(handle, nullptr);
117+
if (localHandleSize != receivedHandleSize) {
118+
ADD_FAILURE() << "IPC handle size mismatch: local=" << localHandleSize
119+
<< ", remote=" << receivedHandleSize;
120+
MPI_Abort(MPI_COMM_WORLD, 1);
121+
return;
122+
}
123+
124+
ASSERT_MPI_SUCCESS(MPI_Recv(handle, static_cast<int>(receivedHandleSize),
125+
MPI_BYTE, 0, 1, MPI_COMM_WORLD,
126+
MPI_STATUS_IGNORE));
127+
128+
void *mappedPtr = nullptr;
129+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleOpen(handle, &mappedPtr));
130+
ASSERT_NE(mappedPtr, nullptr);
131+
132+
int receivedValue = 0;
133+
ASSERT_FLAGCX_SUCCESS(devHandle->deviceMemcpy(
134+
&receivedValue, mappedPtr, sizeof(receivedValue),
135+
flagcxMemcpyDeviceToHost, nullptr));
136+
ASSERT_FLAGCX_SUCCESS(devHandle->streamSynchronize(nullptr));
137+
EXPECT_EQ(receivedValue, expectedValue);
138+
139+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleClose(mappedPtr));
140+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleFree(handle));
141+
142+
int acknowledgement = 1;
143+
ASSERT_MPI_SUCCESS(
144+
MPI_Send(&acknowledgement, 1, MPI_INT, 0, 2, MPI_COMM_WORLD));
145+
}
146+
147+
ASSERT_MPI_SUCCESS(MPI_Barrier(MPI_COMM_WORLD));
148+
}
149+
150+
} // namespace
151+
152+
153+
int main(int argc, char **argv) {
154+
int mpiResult = MPI_Init(&argc, &argv);
155+
if (mpiResult != MPI_SUCCESS) {
156+
return mpiResult;
157+
}
158+
159+
::testing::InitGoogleTest(&argc, argv);
160+
int testResult = RUN_ALL_TESTS();
161+
MPI_Finalize();
162+
return testResult;
163+
}

‎test/unittest/adaptor/test_device_adaptor.cpp‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
* Single device adaptor test - no multi-GPU or MPI required
44
************************************************************************/
55

6+
#include <cstdlib>
67
#include <cstring>
78
#include <gtest/gtest.h>
89
#include <iostream>
@@ -11,6 +12,10 @@
1112
#include "flagcx.h"
1213
#include "topo.h"
1314

15+
#ifdef USE_KUNLUNXIN_ADAPTOR
16+
#include "kunlunxin_adaptor.h"
17+
#endif
18+
1419
class DeviceAdaptorTest : public ::testing::Test {
1520
protected:
1621
void SetUp() override {
@@ -441,6 +446,67 @@ TEST_F(DeviceAdaptorTest, StreamCopyAndFree) {
441446
// Clean up the original
442447
devHandle->streamDestroy(tempStream);
443448
}
449+
#ifdef USE_KUNLUNXIN_ADAPTOR
450+
// Test: ipcMemHandleCreate allocates a wrapper and reports its serialized size.
451+
// Get/Open/Close/Free are intentionally not called in this test.
452+
TEST_F(DeviceAdaptorTest, IpcMemHandleCreate) {
453+
ASSERT_NE(devHandle->ipcMemHandleCreate, nullptr);
454+
455+
flagcxIpcMemHandle_t handle = nullptr;
456+
size_t handleSize = 0;
457+
EXPECT_EQ(devHandle->ipcMemHandleCreate(&handle, &handleSize),
458+
flagcxSuccess);
459+
EXPECT_NE(handle, nullptr);
460+
EXPECT_GT(handleSize, static_cast<size_t>(0));
461+
462+
flagcxIpcMemHandle_t handleWithoutSize = nullptr;
463+
EXPECT_EQ(devHandle->ipcMemHandleCreate(&handleWithoutSize, nullptr),
464+
flagcxSuccess);
465+
EXPECT_NE(handleWithoutSize, nullptr);
466+
467+
EXPECT_EQ(devHandle->ipcMemHandleCreate(nullptr, &handleSize),
468+
flagcxInvalidArgument);
469+
470+
// Direct cleanup keeps this test independent of ipcMemHandleFree.
471+
std::free(handle);
472+
std::free(handleWithoutSize);
473+
}
474+
475+
// Test: ipcMemHandleGet exports a handle from a live device allocation.
476+
// The wrapper is allocated directly so Create/Open/Close/Free are not involved.
477+
TEST_F(DeviceAdaptorTest, IpcMemHandleGet) {
478+
ASSERT_NE(devHandle->ipcMemHandleGet, nullptr);
479+
480+
constexpr size_t bufferSize = 4096;
481+
void *devPtr = nullptr;
482+
ASSERT_EQ(devHandle->deviceMalloc(&devPtr, bufferSize, flagcxMemDevice,
483+
stream),
484+
flagcxSuccess);
485+
ASSERT_NE(devPtr, nullptr);
486+
487+
flagcxIpcMemHandle_t handle = nullptr;
488+
flagcxCalloc(&handle, 1);
489+
ASSERT_NE(handle, nullptr);
490+
491+
EXPECT_EQ(devHandle->ipcMemHandleGet(handle, devPtr), flagcxSuccess);
492+
EXPECT_EQ(devHandle->ipcMemHandleGet(nullptr, devPtr),
493+
flagcxInvalidArgument);
494+
EXPECT_EQ(devHandle->ipcMemHandleGet(handle, nullptr),
495+
flagcxInvalidArgument);
496+
497+
std::free(handle);
498+
EXPECT_EQ(devHandle->deviceFree(devPtr, flagcxMemDevice, stream),
499+
flagcxSuccess);
500+
}
501+
502+
// Test: ipcMemHandleClose rejects a null mapped pointer.
503+
// Successful Close is covered by the MPI lifecycle test.
504+
TEST_F(DeviceAdaptorTest, IpcMemHandleClose) {
505+
ASSERT_NE(devHandle->ipcMemHandleClose, nullptr);
506+
EXPECT_EQ(devHandle->ipcMemHandleClose(nullptr), flagcxInvalidArgument);
507+
}
508+
509+
#endif
444510

445511
int main(int argc, char **argv) {
446512
::testing::InitGoogleTest(&argc, argv);

0 commit comments

Comments
 (0)