AI/rocm-systems
2
0 Garpu 0
Kode Masalah Tarik Permintaan Actions Packages Projects Rilis Wiki Kegiatan

Added Signalling Operations

Telusuri Sumber
This commit is contained in:
Yiltan Temucin committed 2024-11-22 15:36:17 -06:00
1 orang tua 308816bc5e
melakukan d8f44e4436
9 files changed
+425
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/ipc/context_ipc_device.cpp
+105
Melihat File
@@ -222,4 +222,109 @@ __device__ void IPCContext::internal_getmem_wave(void *dest,
ipcImpl_.ipcFence();
}
__device__ void IPCContext::putmem_signal(void *dest, const void *source, size_t nelems,
uint64_t *sig_addr, uint64_t signal, int sig_op,
int pe) {
putmem(dest, source, nelems, pe);
fence();
switch (sig_op) {
case ROC_SHMEM_SIGNAL_SET:
amo_set<uint64_t>(static_cast<void*>(sig_addr), signal, pe);
break;
case ROC_SHMEM_SIGNAL_ADD:
amo_add<uint64_t>(static_cast<void*>(dest), signal, pe);
break;
default:
DPRINTF("[%s] Invalid sig_op value (%d)\n", __func__, sig_op);
break;
}
}
__device__ void IPCContext::putmem_signal_wg(void *dest, const void *source, size_t nelems,
uint64_t *sig_addr, uint64_t signal, int sig_op,
int pe) {
putmem_wg(dest, source, nelems, pe);
fence();
if (is_thread_zero_in_block()) {
switch (sig_op) {
case ROC_SHMEM_SIGNAL_SET:
amo_set<uint64_t>(static_cast<void*>(sig_addr), signal, pe);
break;
case ROC_SHMEM_SIGNAL_ADD:
amo_add<uint64_t>(static_cast<void*>(dest), signal, pe);
break;
default:
DPRINTF("[%s] Invalid sig_op value (%d)\n", __func__, sig_op);
break;
}
}
}
__device__ void IPCContext::putmem_signal_wave(void *dest, const void *source, size_t nelems,
uint64_t *sig_addr, uint64_t signal, int sig_op,
int pe) {
putmem_wave(dest, source, nelems, pe);
fence();
if (is_thread_zero_in_wave()) {
switch (sig_op) {
case ROC_SHMEM_SIGNAL_SET:
amo_set<uint64_t>(static_cast<void*>(sig_addr), signal, pe);
break;
case ROC_SHMEM_SIGNAL_ADD:
amo_add<uint64_t>(static_cast<void*>(dest), signal, pe);
break;
default:
DPRINTF("[%s] Invalid sig_op value (%d)\n", __func__, sig_op);
break;
}
}
}
__device__ void IPCContext::putmem_signal_nbi(void *dest, const void *source, size_t nelems,
uint64_t *sig_addr, uint64_t signal, int sig_op,
int pe) {
putmem_signal(dest, source, nelems, sig_addr, signal, sig_op, pe);
}
__device__ void IPCContext::putmem_signal_nbi_wg(void *dest, const void *source, size_t nelems,
uint64_t *sig_addr, uint64_t signal, int sig_op,
int pe) {
putmem_signal_wg(dest, source, nelems, sig_addr, signal, sig_op, pe);
}
__device__ void IPCContext::putmem_signal_nbi_wave(void *dest, const void *source, size_t nelems,
uint64_t *sig_addr, uint64_t signal, int sig_op,
int pe) {
putmem_signal_wave(dest, source, nelems, sig_addr, signal, sig_op, pe);
}
__device__ uint64_t IPCContext::signal_fetch(const uint64_t *sig_addr) {
uint64_t *dst = const_cast<uint64_t*>(sig_addr);
return amo_fetch_add<uint64_t>(static_cast<void*>(dst), 0, my_pe);
}
__device__ uint64_t IPCContext::signal_fetch_wg(const uint64_t *sig_addr) {
__shared__ uint64_t value;
if (is_thread_zero_in_block()) {
uint64_t *dst = const_cast<uint64_t*>(sig_addr);
value = amo_fetch_add<uint64_t>(static_cast<void*>(dst), 0, my_pe);
}
__threadfence_block();
return value;
}
__device__ uint64_t IPCContext::signal_fetch_wave(const uint64_t *sig_addr) {
uint64_t value;
if (is_thread_zero_in_wave()) {
uint64_t *dst = const_cast<uint64_t*>(sig_addr);
value = amo_fetch_add<uint64_t>(static_cast<void*>(dst), 0, my_pe);
}
__threadfence_block();
value = __shfl(value, 0);
return value;
}
} // namespace rocshmem
src/ipc/context_ipc_device.hpp
+21
Melihat File
@@ -186,6 +186,27 @@ class IPCContext : public Context {
template <typename T>
__device__ void get_nbi_wave(T *dest, const T *source, size_t nelems, int pe);
#define IPC_CONTEXT_PUT_SIGNAL_DEC(SUFFIX) \
template <typename T> \
__device__ void put_signal##SUFFIX(T *dest, const T *source, size_t nelems, \
uint64_t *sig_addr, uint64_t signal, int sig_op, \
int pe); \
\
__device__ void putmem_signal##SUFFIX(void *dest, const void *source, size_t nelems, \
uint64_t *sig_addr, uint64_t signal, int sig_op, \
int pe);
IPC_CONTEXT_PUT_SIGNAL_DEC()
IPC_CONTEXT_PUT_SIGNAL_DEC(_wg)
IPC_CONTEXT_PUT_SIGNAL_DEC(_wave)
IPC_CONTEXT_PUT_SIGNAL_DEC(_nbi)
IPC_CONTEXT_PUT_SIGNAL_DEC(_nbi_wg)
IPC_CONTEXT_PUT_SIGNAL_DEC(_nbi_wave)
__device__ uint64_t signal_fetch(const uint64_t *sig_addr);
__device__ uint64_t signal_fetch_wg(const uint64_t *sig_addr);
__device__ uint64_t signal_fetch_wave(const uint64_t *sig_addr);
private:
//context class has IpcImpl object (ipcImpl_)
src/ipc/context_ipc_tmpl_device.hpp
+19
Melihat File
@@ -578,6 +578,25 @@ __device__ void IPCContext::get_nbi_wave(T *dest, const T *source,
getmem_nbi_wave(dest, source, nelems * sizeof(T), pe);
}
#define IPC_CONTEXT_PUT_SIGNAL_DEF(SUFFIX) \
template <typename T> \
__device__ void IPCContext::put_signal##SUFFIX(T *dest, const T *source, size_t nelems, \
uint64_t *sig_addr, uint64_t signal, int sig_op, \
int pe) { \
putmem_signal##SUFFIX(dest, source, nelems * sizeof(T), sig_addr, signal, sig_op, pe); \
} \
\
template <typename T> \
__device__ void IPCContext::put_signal_nbi##SUFFIX(T *dest, const T *source, size_t nelems, \
uint64_t *sig_addr, uint64_t signal, int sig_op, \
int pe) { \
putmem_signal##SUFFIX(dest, source, nelems * sizeof(T), sig_addr, signal, sig_op, pe); \
}
IPC_CONTEXT_PUT_SIGNAL_DEF()
IPC_CONTEXT_PUT_SIGNAL_DEF(_wg)
IPC_CONTEXT_PUT_SIGNAL_DEF(_wave)
} // namespace rocshmem
#endif // LIBRARY_SRC_IPC_CONTEXT_TMPL_DEVICE_HPP_