Skip to content

Commit 655031c

Browse files
committed
resolve adaptor test and merge pr516
1 parent 828f57b commit 655031c

5 files changed

Lines changed: 345 additions & 33 deletions

File tree

flagcx/adaptor/device/kunlunxin_adaptor.cc

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -270,7 +270,16 @@ flagcxResult_t kunlunAdaptorEventQuery(flagcxEvent_t event) {
270270

271271
flagcxResult_t kunlunAdaptorIpcMemHandleCreate(flagcxIpcMemHandle_t *handle,
272272
size_t *size) {
273-
flagcxCalloc(handle, 1);
273+
if (handle == NULL) {
274+
return flagcxInvalidArgument;
275+
}
276+
277+
*handle = NULL;
278+
flagcxResult_t result = flagcxCalloc(handle, 1);
279+
if (result != flagcxSuccess) {
280+
return result;
281+
}
282+
274283
if (size != NULL) {
275284
*size = sizeof(cudaIpcMemHandle_t);
276285
}
@@ -355,7 +364,7 @@ flagcxResult_t kunlunAdaptorGetDeviceProperties(struct flagcxDevProps *props,
355364
}
356365

357366
props->pciDomainId = static_cast<int>(domain);
358-
props->pciBusId = static_cast<int>(bus);
367+
props->pciBusId = static_cast<int>(bus);
359368
props->pciDeviceId = static_cast<int>(slot);
360369

361370
return flagcxSuccess;

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: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ UNIT_TARGET := $(BINDIR)/adaptor_unit_tests
99
MPI_SRCS := $(wildcard coll_*.cpp)
1010
MPI_OBJS := $(MPI_SRCS:%.cpp=$(OBJDIR)/%.o)
1111
MPI_TARGET := $(BINDIR)/adaptor_mpi_tests
12+
MPIRUN ?= mpirun
13+
MPI_NP ?= 8
1214

1315
.PHONY: all clean run-unit run-mpi run
1416

@@ -49,6 +51,6 @@ run-unit: $(UNIT_TARGET)
4951
@$(UNIT_TARGET)
5052

5153
run-mpi: $(MPI_TARGET)
52-
@mpirun --allow-run-as-root -np 8 $(MPI_TARGET)
54+
@$(MPIRUN) --allow-run-as-root -np $(MPI_NP) $(MPI_TARGET)
5355

5456
run: run-unit run-mpi
Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
/*************************************************************************
2+
* Copyright (c) 2026. All Rights Reserved.
3+
* Cross-process IPC memory handle test.
4+
*
5+
* Run exactly two MPI ranks on the same host. Device IPC handles are
6+
* host-local and cannot be transferred between different nodes.
7+
************************************************************************/
8+
9+
#include <gtest/gtest.h>
10+
#include <mpi.h>
11+
12+
#include "adaptor.h"
13+
#include "flagcx.h"
14+
15+
16+
namespace {
17+
18+
#define ASSERT_FLAGCX_SUCCESS(expr) \
19+
do { \
20+
flagcxResult_t result = (expr); \
21+
if (result != flagcxSuccess) { \
22+
ADD_FAILURE() << #expr << " returned " << static_cast<int>(result); \
23+
MPI_Abort(MPI_COMM_WORLD, static_cast<int>(result)); \
24+
return; \
25+
} \
26+
} while (0)
27+
28+
#define ASSERT_MPI_SUCCESS(expr) \
29+
do { \
30+
int result = (expr); \
31+
if (result != MPI_SUCCESS) { \
32+
ADD_FAILURE() << #expr << " returned " << result; \
33+
MPI_Abort(MPI_COMM_WORLD, result); \
34+
return; \
35+
} \
36+
} while (0)
37+
38+
#define ASSERT_MPI_TRUE(condition) \
39+
do { \
40+
if (!(condition)) { \
41+
ADD_FAILURE() << "MPI assertion failed: " << #condition; \
42+
MPI_Abort(MPI_COMM_WORLD, 1); \
43+
return; \
44+
} \
45+
} while (0)
46+
47+
class IpcMemHandleMpiTest : public ::testing::Test {
48+
protected:
49+
void SetUp() override {
50+
flagcxDeviceHandleInit(&devHandle);
51+
ASSERT_MPI_TRUE(devHandle != nullptr);
52+
53+
int deviceCount = 0;
54+
ASSERT_FLAGCX_SUCCESS(devHandle->getDeviceCount(&deviceCount));
55+
if (deviceCount <= 0) {
56+
ADD_FAILURE() << "No visible device";
57+
MPI_Abort(MPI_COMM_WORLD, 1);
58+
return;
59+
}
60+
ASSERT_FLAGCX_SUCCESS(devHandle->setDevice(0));
61+
}
62+
63+
void TearDown() override {
64+
if (devHandle != nullptr) {
65+
flagcxDeviceHandleFree(devHandle);
66+
}
67+
}
68+
69+
flagcxDeviceHandle_t devHandle = nullptr;
70+
};
71+
72+
TEST_F(IpcMemHandleMpiTest, CrossProcessLifecycle) {
73+
int rank = -1;
74+
int worldSize = 0;
75+
ASSERT_MPI_SUCCESS(MPI_Comm_rank(MPI_COMM_WORLD, &rank));
76+
ASSERT_MPI_SUCCESS(MPI_Comm_size(MPI_COMM_WORLD, &worldSize));
77+
78+
if (worldSize != 2) {
79+
GTEST_SKIP() << "CrossProcessLifecycle requires exactly 2 MPI ranks; "
80+
<< "run with: make MPI_NP=2 run-mpi";
81+
}
82+
83+
constexpr size_t bufferSize = 4096;
84+
constexpr int expectedValue = 0x12345678;
85+
86+
int localApisAvailable =
87+
devHandle->ipcMemHandleCreate != nullptr &&
88+
devHandle->ipcMemHandleGet != nullptr &&
89+
devHandle->ipcMemHandleOpen != nullptr &&
90+
devHandle->ipcMemHandleClose != nullptr &&
91+
devHandle->ipcMemHandleFree != nullptr;
92+
int allApisAvailable = 0;
93+
ASSERT_MPI_SUCCESS(MPI_Allreduce(&localApisAvailable, &allApisAvailable, 1,
94+
MPI_INT, MPI_MIN, MPI_COMM_WORLD));
95+
if (!allApisAvailable) {
96+
GTEST_SKIP() << "IPC memory handle APIs are not available";
97+
}
98+
99+
// Every rank creates its receive/export storage before communication. This
100+
// also provides a coordinated runtime capability check for stub backends.
101+
flagcxIpcMemHandle_t handle = nullptr;
102+
size_t localHandleSize = 0;
103+
flagcxResult_t createResult =
104+
devHandle->ipcMemHandleCreate(&handle, &localHandleSize);
105+
if (createResult != flagcxSuccess &&
106+
createResult != flagcxNotSupported) {
107+
ADD_FAILURE() << "ipcMemHandleCreate returned "
108+
<< static_cast<int>(createResult);
109+
MPI_Abort(MPI_COMM_WORLD, static_cast<int>(createResult));
110+
return;
111+
}
112+
int localSupported = createResult != flagcxNotSupported;
113+
int allSupported = 0;
114+
ASSERT_MPI_SUCCESS(MPI_Allreduce(&localSupported, &allSupported, 1, MPI_INT,
115+
MPI_MIN, MPI_COMM_WORLD));
116+
if (!allSupported) {
117+
if (createResult == flagcxSuccess && handle != nullptr) {
118+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleFree(handle));
119+
}
120+
GTEST_SKIP() << "IPC memory handles are not supported";
121+
}
122+
ASSERT_FLAGCX_SUCCESS(createResult);
123+
ASSERT_MPI_TRUE(handle != nullptr);
124+
ASSERT_MPI_TRUE(localHandleSize > 0);
125+
126+
if (rank == 0) {
127+
void *devPtr = nullptr;
128+
ASSERT_FLAGCX_SUCCESS(devHandle->deviceMalloc(
129+
&devPtr, bufferSize, flagcxMemDevice, nullptr));
130+
ASSERT_MPI_TRUE(devPtr != nullptr);
131+
132+
int hostValue = expectedValue;
133+
ASSERT_FLAGCX_SUCCESS(devHandle->deviceMemcpy(
134+
devPtr, &hostValue, sizeof(hostValue), flagcxMemcpyHostToDevice,
135+
nullptr));
136+
ASSERT_FLAGCX_SUCCESS(devHandle->streamSynchronize(nullptr));
137+
138+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleGet(handle, devPtr));
139+
140+
ASSERT_MPI_SUCCESS(MPI_Send(&localHandleSize, sizeof(localHandleSize),
141+
MPI_BYTE, 1, 0, MPI_COMM_WORLD));
142+
ASSERT_MPI_SUCCESS(
143+
MPI_Send(handle, static_cast<int>(localHandleSize), MPI_BYTE, 1, 1,
144+
MPI_COMM_WORLD));
145+
146+
int acknowledgement = 0;
147+
ASSERT_MPI_SUCCESS(MPI_Recv(&acknowledgement, 1, MPI_INT, 1, 2,
148+
MPI_COMM_WORLD, MPI_STATUS_IGNORE));
149+
ASSERT_MPI_TRUE(acknowledgement == 1);
150+
151+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleFree(handle));
152+
ASSERT_FLAGCX_SUCCESS(
153+
devHandle->deviceFree(devPtr, flagcxMemDevice, nullptr));
154+
} else {
155+
size_t receivedHandleSize = 0;
156+
ASSERT_MPI_SUCCESS(MPI_Recv(&receivedHandleSize, sizeof(receivedHandleSize),
157+
MPI_BYTE, 0, 0, MPI_COMM_WORLD,
158+
MPI_STATUS_IGNORE));
159+
160+
if (localHandleSize != receivedHandleSize) {
161+
ADD_FAILURE() << "IPC handle size mismatch: local=" << localHandleSize
162+
<< ", remote=" << receivedHandleSize;
163+
MPI_Abort(MPI_COMM_WORLD, 1);
164+
return;
165+
}
166+
167+
ASSERT_MPI_SUCCESS(MPI_Recv(handle, static_cast<int>(receivedHandleSize),
168+
MPI_BYTE, 0, 1, MPI_COMM_WORLD,
169+
MPI_STATUS_IGNORE));
170+
171+
void *mappedPtr = nullptr;
172+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleOpen(handle, &mappedPtr));
173+
ASSERT_MPI_TRUE(mappedPtr != nullptr);
174+
175+
int receivedValue = 0;
176+
ASSERT_FLAGCX_SUCCESS(devHandle->deviceMemcpy(
177+
&receivedValue, mappedPtr, sizeof(receivedValue),
178+
flagcxMemcpyDeviceToHost, nullptr));
179+
ASSERT_FLAGCX_SUCCESS(devHandle->streamSynchronize(nullptr));
180+
ASSERT_MPI_TRUE(receivedValue == expectedValue);
181+
182+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleClose(mappedPtr));
183+
ASSERT_FLAGCX_SUCCESS(devHandle->ipcMemHandleFree(handle));
184+
185+
int acknowledgement = 1;
186+
ASSERT_MPI_SUCCESS(
187+
MPI_Send(&acknowledgement, 1, MPI_INT, 0, 2, MPI_COMM_WORLD));
188+
}
189+
190+
ASSERT_MPI_SUCCESS(MPI_Barrier(MPI_COMM_WORLD));
191+
}
192+
193+
} // namespace
194+
195+
196+
int main(int argc, char **argv) {
197+
int mpiResult = MPI_Init(&argc, &argv);
198+
if (mpiResult != MPI_SUCCESS) {
199+
return mpiResult;
200+
}
201+
202+
::testing::InitGoogleTest(&argc, argv);
203+
int testResult = RUN_ALL_TESTS();
204+
MPI_Finalize();
205+
return testResult;
206+
}

0 commit comments

Comments
 (0)