Skip to content

Commit c846a91

Browse files
MC952-archMC952-arch
committed
Fix issues
1 parent 96148ef commit c846a91

20 files changed

Lines changed: 655 additions & 442 deletions

File tree

‎.github/workflows/unit-test.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ jobs:
7676
run: |
7777
cd /__w/FlagCX/FlagCX/test/unittest
7878
export MPI_HOME=/usr/local/mpi
79-
make -j$(nproc)
79+
make -j$(nproc) COMPILE_KERNEL=1 USE_NVIDIA=1
8080
8181
- name: Run single-cluster tests with mpirun
8282
run: |
@@ -104,4 +104,4 @@ jobs:
104104
cd /__w/FlagCX/FlagCX/test/unittest/kernel
105105
export MPI_HOME=/usr/local/mpi
106106
export PATH=$MPI_HOME/bin:$PATH
107-
mpirun -np 8 --allow-run-as-root -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_CLUSTER_SPLIT_LIST=2 -x FLAGCX_USE_HETERO_COMM=1 ./build/bin/kernel_mpi_tests
107+
mpirun -np 8 --allow-run-as-root -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 -x FLAGCX_P2P_DISABLE=1 -x FLAGCX_VMM_ENABLE=0 -x FLAGCX_DEBUG=TRACE -x FLAGCX_DEBUG_SUBSYS=ALL ./build/bin/kernel_mpi_tests

‎bindings/ir/flagcx_device_wrapper.h‎

Lines changed: 93 additions & 83 deletions
Large diffs are not rendered by default.

‎bindings/ir/flagcx_device_wrapper_impl.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -751,7 +751,7 @@ flagcxDevNetSignalCtrInc(const void *transOpaque, const void *teamOpaque,
751751
}
752752

753753
/* ================================================================
754-
* Category 13: Transport — One-Sided putValue<uint64_t> (16)
754+
* Category 13: Transport — One-Sided putValue<uint64_t> (4)
755755
* ================================================================ */
756756

757757
/* (None, None) */

‎bindings/ir/nvidia/Makefile‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ DEVICE_HOME ?= /usr/local/cuda
3333
CCL_HOME ?= /usr/local/nccl/build
3434

3535
# Detect CUDA version for arch selection
36-
CUDA_MAJOR ?= $(shell $(DEVICE_HOME)/bin/nvcc --version 2>/dev/null | grep -oP 'release \K[0-9]+' | head -1)
36+
CUDA_MAJOR ?= $(shell $(DEVICE_HOME)/bin/nvcc --version 2>/dev/null | sed -n 's/.*release \([0-9]*\)\..*/\1/p' | head -1)
3737
ifeq ($(CUDA_MAJOR),)
3838
CUDA_MAJOR := 12
3939
endif
@@ -132,7 +132,11 @@ $(LLVM_IR_FILE): $(OPTIMIZED_BC)
132132
$(LLVM_DIS) $< -o $@.tmp
133133
@echo "Cleaning LLVM IR (removing nvvm-reflect-ftz)..."
134134
@myVar="$$(cat $@.tmp | grep -E '!([0-9]+) = !\{[^"]*"nvvm-reflect-ftz"' | cut -d ' ' -f 1)"; \
135-
awk '!/nvvm-reflect-ftz/' $@.tmp | sed "/^!llvm\.module\.flags = /s/$$myVar, //" > $@
135+
if [ -n "$$myVar" ]; then \
136+
awk '!/nvvm-reflect-ftz/' $@.tmp | sed "/^!llvm\.module\.flags = /s/$$myVar, //" > $@; \
137+
else \
138+
cp $@.tmp $@; \
139+
fi
136140
@rm -f $@.tmp
137141

138142
$(FINAL_BC): $(LLVM_IR_FILE)

‎flagcx/adaptor/ccl/nccl_adaptor.cc‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ flagcxResult_t ncclAdaptorGetStagedBuffer(const flagcxInnerComm_t comm,
5454
return flagcxNotSupported;
5555
}
5656

57-
#if NCCL_VERSION_CODE > NCCL_VERSION(2, 28, 0)
57+
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
5858
static flagcxResult_t ncclDevCommCreateHelper(ncclComm_t comm,
5959
ncclDevCommRequirements *reqs,
6060
ncclDevComm *devComm) {
@@ -94,7 +94,7 @@ static flagcxResult_t ncclDevCommDestroyHelper(ncclComm_t comm,
9494
dlclose(handle);
9595
return (flagcxResult_t)ret;
9696
}
97-
#endif // NCCL_VERSION_CODE > NCCL_VERSION(2, 28, 0)
97+
#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
9898

9999
const char *ncclAdaptorGetErrorString(flagcxResult_t result) {
100100
return ncclGetErrorString((ncclResult_t)result);
@@ -401,7 +401,7 @@ flagcxResult_t ncclAdaptorGroupEnd() { return (flagcxResult_t)ncclGroupEnd(); }
401401
flagcxResult_t ncclAdaptorDevCommCreate(flagcxInnerComm_t comm,
402402
const flagcxDevCommRequirements *reqs,
403403
flagcxInnerDevComm_t *devComm) {
404-
#if NCCL_VERSION_CODE > NCCL_VERSION(2, 28, 0)
404+
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
405405
flagcxInnerDevComm_t inner =
406406
(flagcxInnerDevComm_t)malloc(sizeof(struct flagcxInnerDevComm));
407407
if (!inner)
@@ -436,7 +436,7 @@ flagcxResult_t ncclAdaptorDevCommCreate(flagcxInnerComm_t comm,
436436

437437
flagcxResult_t ncclAdaptorDevCommDestroy(flagcxInnerComm_t comm,
438438
flagcxInnerDevComm_t devComm) {
439-
#if NCCL_VERSION_CODE > NCCL_VERSION(2, 28, 0)
439+
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
440440
if (!devComm)
441441
return flagcxSuccess;
442442
flagcxResult_t ret = ncclDevCommDestroyHelper(comm->base, &devComm->base);
@@ -449,7 +449,7 @@ flagcxResult_t ncclAdaptorDevCommDestroy(flagcxInnerComm_t comm,
449449

450450
flagcxResult_t ncclAdaptorDevCommReqsInit(flagcxInnerComm_t comm,
451451
flagcxDevCommRequirements *reqs) {
452-
#if NCCL_VERSION_CODE > NCCL_VERSION(2, 28, 0)
452+
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
453453
const char *winEnv = flagcxGetEnv("NCCL_WIN_ENABLE");
454454
const char *cuMemEnv = flagcxGetEnv("NCCL_CUMEM_ENABLE");
455455
const char *crossNicEnv = flagcxGetEnv("NCCL_CROSS_NIC");

‎flagcx/adaptor/flagcx_device.cc‎

Lines changed: 62 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -659,6 +659,7 @@ flagcxResult_t flagcxDevCommCreate(flagcxComm_t comm,
659659
return flagcxSystemError;
660660
}
661661
memset(handle, 0, sizeof(struct flagcxDevCommInternal));
662+
pthread_mutex_init(&handle->cachedPtrMutex, NULL);
662663
handle->barrierIpcIndex = -1;
663664

664665
// ---- Baseline: always ----
@@ -691,6 +692,7 @@ flagcxResult_t flagcxDevCommCreate(flagcxComm_t comm,
691692
if (ret != flagcxSuccess && ret != flagcxNotSupported &&
692693
ret != flagcxInvalidArgument) {
693694
WARN("flagcxDevCommCreate: vendor devCommCreate failed (%d)", ret);
695+
pthread_mutex_destroy(&handle->cachedPtrMutex);
694696
free(handle);
695697
return ret;
696698
}
@@ -719,6 +721,7 @@ flagcxResult_t flagcxDevCommCreate(flagcxComm_t comm,
719721
WARN("flagcxDevCommCreate: IPC barrier setup failed (%d), "
720722
"barriers unavailable",
721723
res);
724+
pthread_mutex_destroy(&handle->cachedPtrMutex);
722725
free(handle);
723726
return res;
724727
}
@@ -988,9 +991,10 @@ flagcxResult_t flagcxDevCommDestroy(flagcxComm_t comm,
988991

989992
// Device pointer cache cleanup
990993
if (devComm->cachedDevicePtr) {
991-
deviceAdaptor->deviceFree(devComm->cachedDevicePtr, flagcxMemDevice, NULL);
994+
flagcxCommDeferFree(comm, devComm->cachedDevicePtr, flagcxMemDevice);
992995
}
993996

997+
pthread_mutex_destroy(&devComm->cachedPtrMutex);
994998
free(devComm->localRankToRank);
995999
free(devComm);
9961000
return flagcxSuccess;
@@ -1015,6 +1019,7 @@ flagcxResult_t flagcxDevMemCreate(flagcxComm_t comm, void *buff, size_t size,
10151019
return flagcxSystemError;
10161020
}
10171021
memset(handle, 0, sizeof(struct flagcxDevMemInternal));
1022+
pthread_mutex_init(&handle->cachedPtrMutex, NULL);
10181023

10191024
// ---- Baseline: always ----
10201025
handle->rawPtr = buff;
@@ -1093,6 +1098,7 @@ flagcxResult_t flagcxDevMemCreate(flagcxComm_t comm, void *buff, size_t size,
10931098
auto *kWin = new (std::nothrow) typename DeviceAPI::Window{};
10941099
if (kWin == nullptr) {
10951100
WARN("flagcxDevMemCreate: failed to allocate DeviceAPI::Window");
1101+
pthread_mutex_destroy(&handle->cachedPtrMutex);
10961102
free(handle);
10971103
return flagcxSystemError;
10981104
}
@@ -1140,9 +1146,10 @@ flagcxResult_t flagcxDevMemDestroy(flagcxComm_t comm, flagcxDevMem_t devMem) {
11401146

11411147
// Free cached device pointer if present
11421148
if (devMem->cachedDevicePtr) {
1143-
deviceAdaptor->deviceFree(devMem->cachedDevicePtr, flagcxMemDevice, NULL);
1149+
flagcxCommDeferFree(comm, devMem->cachedDevicePtr, flagcxMemDevice);
11441150
}
11451151

1152+
pthread_mutex_destroy(&devMem->cachedPtrMutex);
11461153
free(devMem);
11471154
return flagcxSuccess;
11481155
}
@@ -1156,8 +1163,11 @@ flagcxResult_t flagcxDevCommGetDevicePtr(flagcxDevComm_t devComm,
11561163
if (!devComm || !devPtr)
11571164
return flagcxInvalidArgument;
11581165

1166+
pthread_mutex_lock(&devComm->cachedPtrMutex);
1167+
11591168
if (devComm->cachedDevicePtr) {
11601169
*devPtr = devComm->cachedDevicePtr;
1170+
pthread_mutex_unlock(&devComm->cachedPtrMutex);
11611171
return flagcxSuccess;
11621172
}
11631173

@@ -1166,24 +1176,39 @@ flagcxResult_t flagcxDevCommGetDevicePtr(flagcxDevComm_t devComm,
11661176

11671177
// Allocate device memory and copy
11681178
void *dPtr = nullptr;
1169-
FLAGCXCHECK(deviceAdaptor->deviceMalloc(&dPtr, sizeof(flagcxDevComm),
1170-
flagcxMemDevice, NULL));
1171-
FLAGCXCHECK(
1179+
flagcxResult_t res = flagcxSuccess;
1180+
FLAGCXCHECKGOTO(deviceAdaptor->deviceMalloc(&dPtr, sizeof(flagcxDevComm),
1181+
flagcxMemDevice, NULL),
1182+
res, fail);
1183+
FLAGCXCHECKGOTO(
11721184
deviceAdaptor->deviceMemcpy(dPtr, &hostCopy, sizeof(flagcxDevComm),
1173-
flagcxMemcpyHostToDevice, NULL, NULL));
1185+
flagcxMemcpyHostToDevice, NULL, NULL),
1186+
res, fail);
11741187

11751188
devComm->cachedDevicePtr = dPtr;
11761189
*devPtr = dPtr;
1190+
pthread_mutex_unlock(&devComm->cachedPtrMutex);
11771191
return flagcxSuccess;
1192+
1193+
fail:
1194+
pthread_mutex_unlock(&devComm->cachedPtrMutex);
1195+
if (dPtr) {
1196+
deviceAdaptor->deviceFree(dPtr, flagcxMemDevice, NULL);
1197+
}
1198+
return res;
11781199
}
11791200

11801201
flagcxResult_t flagcxDevCommFreeDevicePtr(flagcxDevComm_t devComm) {
11811202
if (!devComm)
11821203
return flagcxInvalidArgument;
1183-
if (devComm->cachedDevicePtr) {
1184-
FLAGCXCHECK(deviceAdaptor->deviceFree(devComm->cachedDevicePtr,
1185-
flagcxMemDevice, NULL));
1186-
devComm->cachedDevicePtr = nullptr;
1204+
1205+
pthread_mutex_lock(&devComm->cachedPtrMutex);
1206+
void *ptr = devComm->cachedDevicePtr;
1207+
devComm->cachedDevicePtr = nullptr;
1208+
pthread_mutex_unlock(&devComm->cachedPtrMutex);
1209+
1210+
if (ptr) {
1211+
FLAGCXCHECK(deviceAdaptor->deviceFree(ptr, flagcxMemDevice, NULL));
11871212
}
11881213
return flagcxSuccess;
11891214
}
@@ -1192,8 +1217,11 @@ flagcxResult_t flagcxDevMemGetDevicePtr(flagcxDevMem_t devMem, void **devPtr) {
11921217
if (!devMem || !devPtr)
11931218
return flagcxInvalidArgument;
11941219

1220+
pthread_mutex_lock(&devMem->cachedPtrMutex);
1221+
11951222
if (devMem->cachedDevicePtr) {
11961223
*devPtr = devMem->cachedDevicePtr;
1224+
pthread_mutex_unlock(&devMem->cachedPtrMutex);
11971225
return flagcxSuccess;
11981226
}
11991227

@@ -1202,24 +1230,39 @@ flagcxResult_t flagcxDevMemGetDevicePtr(flagcxDevMem_t devMem, void **devPtr) {
12021230

12031231
// Allocate device memory and copy
12041232
void *dPtr = nullptr;
1205-
FLAGCXCHECK(deviceAdaptor->deviceMalloc(&dPtr, sizeof(flagcxDevMem),
1206-
flagcxMemDevice, NULL));
1207-
FLAGCXCHECK(deviceAdaptor->deviceMemcpy(dPtr, &hostCopy, sizeof(flagcxDevMem),
1208-
flagcxMemcpyHostToDevice, NULL,
1209-
NULL));
1233+
flagcxResult_t res = flagcxSuccess;
1234+
FLAGCXCHECKGOTO(deviceAdaptor->deviceMalloc(&dPtr, sizeof(flagcxDevMem),
1235+
flagcxMemDevice, NULL),
1236+
res, fail);
1237+
FLAGCXCHECKGOTO(
1238+
deviceAdaptor->deviceMemcpy(dPtr, &hostCopy, sizeof(flagcxDevMem),
1239+
flagcxMemcpyHostToDevice, NULL, NULL),
1240+
res, fail);
12101241

12111242
devMem->cachedDevicePtr = dPtr;
12121243
*devPtr = dPtr;
1244+
pthread_mutex_unlock(&devMem->cachedPtrMutex);
12131245
return flagcxSuccess;
1246+
1247+
fail:
1248+
pthread_mutex_unlock(&devMem->cachedPtrMutex);
1249+
if (dPtr) {
1250+
deviceAdaptor->deviceFree(dPtr, flagcxMemDevice, NULL);
1251+
}
1252+
return res;
12141253
}
12151254

12161255
flagcxResult_t flagcxDevMemFreeDevicePtr(flagcxDevMem_t devMem) {
12171256
if (!devMem)
12181257
return flagcxInvalidArgument;
1219-
if (devMem->cachedDevicePtr) {
1220-
FLAGCXCHECK(deviceAdaptor->deviceFree(devMem->cachedDevicePtr,
1221-
flagcxMemDevice, NULL));
1222-
devMem->cachedDevicePtr = nullptr;
1258+
1259+
pthread_mutex_lock(&devMem->cachedPtrMutex);
1260+
void *ptr = devMem->cachedDevicePtr;
1261+
devMem->cachedDevicePtr = nullptr;
1262+
pthread_mutex_unlock(&devMem->cachedPtrMutex);
1263+
1264+
if (ptr) {
1265+
FLAGCXCHECK(deviceAdaptor->deviceFree(ptr, flagcxMemDevice, NULL));
12231266
}
12241267
return flagcxSuccess;
12251268
}

‎flagcx/adaptor/include/device_api/flagcx_device_internal.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
#include "flagcx_kernel.h"
1818
#include "shmutils.h"
19+
#include <pthread.h>
1920

2021
// Forward declaration for typed vendor device comm handle
2122
struct flagcxInnerDevComm;
@@ -99,6 +100,7 @@ struct flagcxDevCommInternal {
99100

100101
// ---- Device pointer cache (for Triton integration) ----
101102
void *cachedDevicePtr; // Lazily allocated by flagcxDevCommGetDevicePtr
103+
pthread_mutex_t cachedPtrMutex; // Protects lazy init of cachedDevicePtr
102104
};
103105

104106
// ============================================================
@@ -135,6 +137,7 @@ struct flagcxDevMemInternal {
135137

136138
// ---- Device pointer cache (for Triton integration) ----
137139
void *cachedDevicePtr; // Lazily allocated by flagcxDevMemGetDevicePtr
140+
pthread_mutex_t cachedPtrMutex; // Protects lazy init of cachedDevicePtr
138141
};
139142
#ifndef FLAGCX_DEV_MEM_T_DEFINED
140143
#define FLAGCX_DEV_MEM_T_DEFINED

‎flagcx/adaptor/include/device_api/nvidia_comm_traits.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
// ============================================================
2828
// NVIDIA Vendor Backend (NCCL device API)
2929
// ============================================================
30-
#if NCCL_VERSION_CODE > NCCL_VERSION(2, 28, 0) && \
30+
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) && \
3131
!defined(FLAGCX_FORCE_FALLBACK)
3232

3333
#include "nccl_device.h"

‎flagcx/adaptor/include/nvidia_adaptor.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
#include <cuda.h>
66
#include <cuda_runtime.h>
77
#include <map>
8-
#if NCCL_VERSION_CODE > NCCL_VERSION(2, 28, 0)
8+
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
99
#include "nccl_device.h"
1010

1111
#define NCCL_ADAPTOR_DEVICE_CTA_COUNT 36
@@ -20,7 +20,7 @@ struct flagcxInnerDevComm {
2020
typedef void ncclDevComm;
2121
struct flagcxInnerDevComm {};
2222

23-
#endif // NCCL_VERSION_CODE > NCCL_VERSION(2, 28, 0)
23+
#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
2424

2525
#if NCCL_VERSION_CODE > NCCL_VERSION(2, 27, 0)
2626

0 commit comments

Comments
 (0)