Add ipc unit_tests

[ROCm/rocshmem commit: 58c5a98b5d]
This commit is contained in:
Brandon Potter
2024-08-06 12:11:08 -07:00
parent a542e0a332
commit bda68c964e
8 changed files with 1367 additions and 3 deletions
+9 -2
View File
@@ -50,7 +50,6 @@ __host__ void IpcOnImpl::ipcHostInit(int my_pe, const HEAP_BASES_T &heap_bases,
/*
* Figure out how this process' rank among local processes.
*/
int shm_rank;
MPI_Comm_rank(shmcomm, &shm_rank);
/*
@@ -92,7 +91,6 @@ __host__ void IpcOnImpl::ipcHostInit(int my_pe, const HEAP_BASES_T &heap_bases,
void **ipc_base_uncast = reinterpret_cast<void **>(&ipc_base[i]);
CHECK_HIP(hipIpcOpenMemHandle(ipc_base_uncast, vec_ipc_handle[i],
hipIpcMemLazyEnablePeerAccess));
// TODO(bpotter): add some error checking here if happens to fail
} else {
ipc_base[i] = base_heap;
}
@@ -110,6 +108,15 @@ __host__ void IpcOnImpl::ipcHostInit(int my_pe, const HEAP_BASES_T &heap_bases,
free(vec_ipc_handle);
}
__host__ void IpcOnImpl::ipcHostStop() {
for (size_t i = 0; i < shm_size; i++) {
if (i != shm_rank) {
CHECK_HIP(hipIpcCloseMemHandle(ipc_bases[i]));
}
}
CHECK_HIP(hipFree(ipc_bases));
}
__device__ void IpcOnImpl::ipcCopy(void *dst, void *src, size_t size) {
memcpy(dst, src, size);
}
+6
View File
@@ -42,6 +42,8 @@ class IpcOnImpl {
using HEAP_BASES_T = std::vector<char *, StdAllocatorHIP<char *>>;
public:
int shm_rank{0};
uint32_t shm_size{0};
char **ipc_bases{nullptr};
@@ -49,6 +51,8 @@ class IpcOnImpl {
__host__ void ipcHostInit(int my_pe, const HEAP_BASES_T &heap_bases,
MPI_Comm thread_comm);
__host__ void ipcHostStop();
__device__ bool isIpcAvailable(int my_pe, int target_pe) {
return my_pe / shm_size == target_pe / shm_size;
}
@@ -115,6 +119,8 @@ class IpcOffImpl {
__host__ void ipcHostInit(int my_pe, const HEAP_BASES_T &heap_bases,
MPI_Comm thread_comm) {}
__host__ void ipcHostStop() {}
__device__ bool isIpcAvailable(int my_pe, int target_pe) { return false; }
__device__ void ipcGpuInit(Backend *roc_shmem_handle, Context *ctx,