diff --git a/projects/rocshmem/CMakeLists.txt b/projects/rocshmem/CMakeLists.txt index 05ef4acedfe6..cfd12d03c68c 100644 --- a/projects/rocshmem/CMakeLists.txt +++ b/projects/rocshmem/CMakeLists.txt @@ -47,7 +47,8 @@ option(PROFILE "Enable statistics and timing support" OFF) option(USE_RO "Enable RO conduit" OFF) option(USE_IPC "Enable IPC support (using HIP)" ON) option(USE_GDA "Enable GDA conduit" OFF) -option(USE_SDMA "Enable SDMA transport (MI300X+)." OFF) +option(USE_SDMA "Enable SDMA transport (MI300X+)" OFF) +option(USE_TDM "Enable TDM (Tensor Data Mover) LDS-staged IPC transfers (MI4xx / gfx1250+)" OFF) option(USE_THREADS "Enable workgroup threads to share network queues" OFF) option(USE_WF_COAL "Enable wavefront message coalescing" OFF) option(USE_FUNC_CALL "Force compiler to use function calls on library API" OFF) diff --git a/projects/rocshmem/cmake/rocshmem_config.h.in b/projects/rocshmem/cmake/rocshmem_config.h.in index 10ec550c7458..5180bb44001e 100644 --- a/projects/rocshmem/cmake/rocshmem_config.h.in +++ b/projects/rocshmem/cmake/rocshmem_config.h.in @@ -30,6 +30,7 @@ #cmakedefine USE_IPC #cmakedefine USE_GDA #cmakedefine USE_SDMA +#cmakedefine USE_TDM #cmakedefine USE_THREADS #cmakedefine USE_SHARED_CTX #cmakedefine USE_WF_COAL diff --git a/projects/rocshmem/docs/env_variables.rst b/projects/rocshmem/docs/env_variables.rst index b19c82641dd3..87c338c03c43 100644 --- a/projects/rocshmem/docs/env_variables.rst +++ b/projects/rocshmem/docs/env_variables.rst @@ -257,3 +257,28 @@ control the behavior of rocSHMEM. - ``0`` - | ``0``: Apply wf_id spreading only for the default (shared) context. | ``1``: Apply wf_id spreading for all contexts. + + * - | ``ROCSHMEM_TDM_TILE_BYTES`` + | TDM (Tensor Data Mover) tile size in bytes for LDS-staged work-group + | or wave put/get on gfx1250 (requires ``USE_TDM`` build option). + | Doubled internally for double buffering; see + | ``rocshmem_query_tdm_lds_bytes()``. Wave put/get splits the + | registered LDS buffer evenly across the block's waves, so each + | wave's effective tile is smaller than this value. + - ``0`` + - | ``0`` (default): auto-size the tile to use all LDS available per + | block on this device. + | Nonzero: use this tile size explicitly. Must be a multiple of 8 + | bytes (the TDM element size); non-multiples are rounded down with + | a warning. If the requested size (doubled) doesn't fit in this + | device's shared memory per block, falls back to the auto-sized + | value instead. + + * - | ``ROCSHMEM_TDM_THRESHOLD`` + | Minimum transfer size in bytes to route through TDM for work-group + | or wave put/get on gfx1250 (requires ``USE_TDM`` build option). + | Transfers smaller than this use the normal (non-TDM) copy path + | instead, since TDM's issue latency isn't amortized for small + | transfers. + - ``16384`` + - Size in bytes. diff --git a/projects/rocshmem/include/rocshmem/rocshmem.hpp b/projects/rocshmem/include/rocshmem/rocshmem.hpp index 93107bcb5227..1f7334ee0bb3 100644 --- a/projects/rocshmem/include/rocshmem/rocshmem.hpp +++ b/projects/rocshmem/include/rocshmem/rocshmem.hpp @@ -196,6 +196,19 @@ __host__ int rocshmem_set_attr_uniqueid_args(int rank, int nranks, */ __host__ void rocshmem_query_thread(int *provided); +/** + * @brief Bytes of dynamic LDS the TDM (Tensor Data Mover) fast path wants + * for work-group or wave put/get on gfx1250 -- size the kernel's dynamic + * LDS and the buffer passed to rocshmem_set_tdm_lds() with this. Reflects + * ROCSHMEM_TDM_TILE_BYTES, doubled for double buffering (0/unset + * auto-sizes from this device's available LDS). Wave put/get splits this + * buffer across the block's waves, so its effective tile is smaller. + * Returns 0 if rocSHMEM wasn't built with USE_TDM. + * + * @return Number of bytes of dynamic LDS required. + */ +__host__ size_t rocshmem_query_tdm_lds_bytes(); + /** * @brief Function that dumps internal stats to stdout. */ @@ -818,6 +831,45 @@ __host__ void rocshmem_global_exit(int status); */ __device__ void rocshmem_query_thread(int *provided); +/** + * @brief Bytes of dynamic LDS the TDM (Tensor Data Mover) fast path wants + * for work-group or wave put/get on gfx1250 -- size the kernel's dynamic + * LDS and the buffer passed to rocshmem_set_tdm_lds() with this. Reflects + * ROCSHMEM_TDM_TILE_BYTES, doubled for double buffering (0/unset + * auto-sizes from this device's available LDS). Wave put/get splits this + * buffer across the block's waves, so its effective tile is smaller. + * Returns 0 if rocSHMEM wasn't built with USE_TDM. + * + * @return Number of bytes of dynamic LDS required. + */ +__device__ size_t rocshmem_query_tdm_lds_bytes(); + +/** + * @brief Registers the LDS buffer this work-group lets rocSHMEM use for its + * TDM (Tensor Data Mover) fast path, for both work-group and wave put/get. + * Publishes the registration to the rest of the block via an internal + * barrier before returning. rocshmem_wg_ctx_create() resets it to + * "unregistered" when creating the default context, so a work-group that + * never calls this simply never uses TDM. Creating an additional context + * via rocshmem_wg_team_create_ctx() does not reset it, so one call per + * kernel is enough. + * + * Work-group put/get uses the whole buffer; wave put/get splits it into + * one fixed, non-overlapping slice per wave the block could launch, so + * concurrent waves never tile through the same LDS bytes. + * + * @param[in] lds Workgroup-shared LDS buffer at least lds_bytes long + * (e.g. the caller's own extern __shared__ array). + * At least 8-byte aligned; 256-byte alignment + * recommended for peak TDM throughput. + * @param[in] lds_bytes Size of \p lds in bytes. See + * rocshmem_query_tdm_lds_bytes() for rocSHMEM's + * preferred size; a smaller value still works (uses + * a smaller tile), and 0 disables TDM for this + * workgroup. + */ +__device__ void rocshmem_set_tdm_lds(void *lds, size_t lds_bytes); + /** * @brief Creates an OpenSHMEM context. By design, the context is private * to the calling work-group. diff --git a/projects/rocshmem/include/rocshmem/rocshmem_RMA_X.hpp b/projects/rocshmem/include/rocshmem/rocshmem_RMA_X.hpp index 7be101adf37d..87c888f30434 100644 --- a/projects/rocshmem/include/rocshmem/rocshmem_RMA_X.hpp +++ b/projects/rocshmem/include/rocshmem/rocshmem_RMA_X.hpp @@ -233,6 +233,11 @@ __device__ ATTR_NO_INLINE void rocshmem_ulonglong_put_wg( * granularity. However, all threads in a wave must participate in the * call using the same parameters. * + * On gfx1250 built with USE_TDM, this call uses the faster TDM (Tensor + * Data Mover) copy path if the workgroup registered an LDS buffer via + * rocshmem_set_tdm_lds(), split evenly across the block's waves; + * otherwise it falls back to the plain copy path. + * * @param[in] ctx Context with which to perform this operation. * @param[in] dest Destination address. Must be an address on the symmetric * heap. @@ -259,6 +264,11 @@ __device__ ATTR_NO_INLINE void rocshmem_putmem_wave(void *dest, * (WG) granularity. However, all threads in the workgroup must participate in * the call using the same parameters. * + * On gfx1250 built with USE_TDM, this call uses the faster TDM (Tensor Data + * Mover) copy path if this workgroup registered an LDS buffer via + * rocshmem_set_tdm_lds() (see also rocshmem_query_tdm_lds_bytes()). + * Without registration, it falls back to the plain copy path. + * * @param[in] ctx Context with which to perform this operation. * @param[in] dest Destination address. Must be an address on the symmetric * heap. @@ -481,6 +491,11 @@ __device__ ATTR_NO_INLINE void rocshmem_ulonglong_get_wg( * granularity. However, all threads in a the wave must participate in the * call using the same parameters * + * On gfx1250 built with USE_TDM, this call uses the faster TDM (Tensor + * Data Mover) copy path if the workgroup registered an LDS buffer via + * rocshmem_set_tdm_lds(), split evenly across the block's waves; + * otherwise it falls back to the plain copy path. + * * @param[in] ctx Context with which to perform this operation. * @param[in] dest Destination address. Must be an address on the symmetric * heap. @@ -506,6 +521,11 @@ __device__ ATTR_NO_INLINE void rocshmem_getmem_wave(void *dest, * (WG) granularity. However, all threads in the workgroup must participate * in the call using the same parameters * + * On gfx1250 built with USE_TDM, this call uses the faster TDM (Tensor Data + * Mover) copy path if this workgroup registered an LDS buffer via + * rocshmem_set_tdm_lds() (see also rocshmem_query_tdm_lds_bytes()). + * Without registration, it falls back to the plain copy path. + * * @param[in] ctx Context with which to perform this operation. * @param[in] dest Destination address. Must be an address on the symmetric * heap. diff --git a/projects/rocshmem/src/backend_bc.hpp b/projects/rocshmem/src/backend_bc.hpp index 1dffafb0cdb4..bb6c161a72e1 100644 --- a/projects/rocshmem/src/backend_bc.hpp +++ b/projects/rocshmem/src/backend_bc.hpp @@ -325,6 +325,14 @@ class Backend { */ int hip_dev_id{0}; + /** + * @brief Host-side copy of the auto-resolved TDM tile size (bytes). + * + * Mirrors constmem_t::tdm_tile_bytes (GPU constant memory) so host code + * can query the resolved tile size without a device->host memcpy. + */ + uint32_t tdm_resolved_tile_bytes{0}; + /** * @brief Add ctx from the list of user-created ctxs */ diff --git a/projects/rocshmem/src/constmem.cpp b/projects/rocshmem/src/constmem.cpp index 802f78f51723..efb0167cbd84 100644 --- a/projects/rocshmem/src/constmem.cpp +++ b/projects/rocshmem/src/constmem.cpp @@ -1,6 +1,7 @@ #include "constmem.hpp" #include "backend_bc.hpp" #include "envvar.hpp" +#include "tdm.hpp" #if defined(USE_GDA) #include "gda/backend_gda.hpp" #endif @@ -50,6 +51,74 @@ void init_constant_memory(void) { reinterpret_cast(backend->heap.get_local_heap_base()); constmem_values.heap_size = backend->heap.get_size(); + constmem_values.tdm_tile_bytes = envvar::tdm::tile_bytes; + constmem_values.tdm_threshold = envvar::tdm::threshold; + +#if defined(USE_TDM) + { + int device_id = 0; + CHECK_HIP(hipGetDevice(&device_id)); + int max_shared_mem_per_block = 0; + CHECK_HIP(hipDeviceGetAttribute(&max_shared_mem_per_block, + hipDeviceAttributeMaxSharedMemoryPerBlock, + device_id)); + + // tile_dim0 is a 16-bit field (see tdm.hpp GROUP1), so the tile can + // never encode more than 65535 elements regardless of available LDS. + constexpr uint32_t element_bytes = 1u << tdm::FlatCopyElementLog2; + constexpr uint32_t max_encodable_tile_bytes = 65535u * element_bytes; + // Reserve 16KB for rocSHMEM's own (small) static __shared__ usage plus + // headroom for the caller's, then halve the rest for double buffering. + constexpr size_t reserve_bytes = 16 * 1024; + const size_t avail_bytes = + (static_cast(max_shared_mem_per_block) > reserve_bytes) + ? static_cast(max_shared_mem_per_block) - reserve_bytes + : 0; + const uint32_t avail_tile_bytes = + static_cast((avail_bytes / 2) / element_bytes) * element_bytes; + const uint32_t fittable_tile_bytes = + (avail_tile_bytes < max_encodable_tile_bytes) ? avail_tile_bytes + : max_encodable_tile_bytes; + + if (constmem_values.tdm_tile_bytes == 0) { + constmem_values.tdm_tile_bytes = fittable_tile_bytes; + if (constmem_values.tdm_tile_bytes == 0) { + LOG_WARN( + "Not enough LDS on this device (%d bytes/block) to auto-size a TDM " + "tile after reserving %zu bytes; TDM disabled. Set " + "ROCSHMEM_TDM_TILE_BYTES explicitly to override.", + max_shared_mem_per_block, reserve_bytes); + } else { + LOG_INFO("TDM tile size auto-sized to %u bytes (device max shared mem/block: %d)", + constmem_values.tdm_tile_bytes, max_shared_mem_per_block); + } + } else { + const uint32_t aligned_tile_bytes = + constmem_values.tdm_tile_bytes - (constmem_values.tdm_tile_bytes % element_bytes); + if (aligned_tile_bytes != constmem_values.tdm_tile_bytes) { + LOG_WARN( + "ROCSHMEM_TDM_TILE_BYTES=%u is not a multiple of the %u-byte TDM " + "element size; rounding down to %u bytes.", + constmem_values.tdm_tile_bytes, element_bytes, aligned_tile_bytes); + constmem_values.tdm_tile_bytes = aligned_tile_bytes; + } + const size_t tdm_lds_bytes = tdm::lds_bytes_for_tile(constmem_values.tdm_tile_bytes); + if (tdm_lds_bytes > static_cast(max_shared_mem_per_block)) { + // Don't abort: fall back to whatever tile size does fit (possibly 0, + // which disables TDM for this run) rather than taking the job down. + LOG_WARN( + "ROCSHMEM_TDM_TILE_BYTES=%u needs %zu bytes of LDS (double-buffered) " + "but this device only has %d bytes of shared memory per block; " + "falling back to %u bytes.", + constmem_values.tdm_tile_bytes, tdm_lds_bytes, max_shared_mem_per_block, + fittable_tile_bytes); + constmem_values.tdm_tile_bytes = fittable_tile_bytes; + } + } + } + backend->tdm_resolved_tile_bytes = constmem_values.tdm_tile_bytes; +#endif + constmem_values.backend_type = backend->get_type(); #if defined(USE_GDA) if (constmem_values.backend_type == BackendType::GDA_BACKEND) { diff --git a/projects/rocshmem/src/constmem.hpp b/projects/rocshmem/src/constmem.hpp index a95e8a21f0c0..3128fe10d319 100644 --- a/projects/rocshmem/src/constmem.hpp +++ b/projects/rocshmem/src/constmem.hpp @@ -30,7 +30,6 @@ #include #include "gda/gda_enums.hpp" #include "rocshmem/rocshmem_common.hpp" -#include "util.hpp" namespace rocshmem { @@ -46,6 +45,8 @@ struct constmem_t { size_t ipc_sdma_threshold; // SDMA_THRESHOLD_DISABLED = SDMA disabled uintptr_t heap_base; // Local symmetric heap base size_t heap_size; // Local symmetric heap size in bytes + uint32_t tdm_tile_bytes; + size_t tdm_threshold; // transfers smaller than this skip the TDM path } __attribute__ ((aligned (16))); extern __constant__ constmem_t constmem; diff --git a/projects/rocshmem/src/envvar.cpp b/projects/rocshmem/src/envvar.cpp index 875f3a18f26f..19ac19c2f487 100644 --- a/projects/rocshmem/src/envvar.cpp +++ b/projects/rocshmem/src/envvar.cpp @@ -235,6 +235,15 @@ namespace envvar { false); } // namespace sdma + namespace tdm { + const var tile_bytes("TILE_BYTES", + "TDM tile size in bytes for LDS-staged IPC transfers. 0 (default) " + "auto-sets the tile to use all LDS available per block", 0); + const var threshold("THRESHOLD", + "TDM transfer size threshold in bytes; transfers smaller than this use " + "the normal (non-TDM) copy path instead", 16384); + } // namespace tdm + namespace _detail { std::tuple get_var_map() { // construct on first use idiom diff --git a/projects/rocshmem/src/envvar.hpp b/projects/rocshmem/src/envvar.hpp index a7dd9368e027..9a23bd1ae9fc 100644 --- a/projects/rocshmem/src/envvar.hpp +++ b/projects/rocshmem/src/envvar.hpp @@ -130,6 +130,7 @@ namespace envvar { REVERSE_OFFLOAD, GDA, SDMA, + TDM, }; // env var string prefixes @@ -141,6 +142,7 @@ namespace envvar { template <> inline constexpr const char* prefix = "ROCSHMEM_RO"; template <> inline constexpr const char* prefix = "ROCSHMEM_GDA"; template <> inline constexpr const char* prefix = "ROCSHMEM_SDMA"; + template <> inline constexpr const char* prefix = "ROCSHMEM_TDM"; } // namespace category namespace parser { @@ -580,6 +582,12 @@ namespace envvar { extern const var spread_channels; } // namespace sdma + namespace tdm { + template using var = var; + extern const var tile_bytes; + extern const var threshold; + } // namespace tdm + /** * @brief Print mode for environment variables */ diff --git a/projects/rocshmem/src/host/host.cpp b/projects/rocshmem/src/host/host.cpp index 98b865ea8e4e..918228304294 100644 --- a/projects/rocshmem/src/host/host.cpp +++ b/projects/rocshmem/src/host/host.cpp @@ -423,11 +423,12 @@ __host__ void HostInterface::alltoallmem_on_stream(rocshmem_team_t team, // - Query optimal block size using occupancy API // - Limit block size to nelems (number of bytes) to avoid over-subscription // - Always use 1 block (single workgroup collective) + const size_t tdm_lds_bytes = rocshmem_query_tdm_lds_bytes(); int optimal_block_size = 0; int grid_size = 0; CHECK_HIP(hipOccupancyMaxPotentialBlockSize(&grid_size, &optimal_block_size, - rocshmem_alltoallmem_kernel, 0, - 0)); + rocshmem_alltoallmem_kernel, + tdm_lds_bytes, 0)); // Limit block size to nelems (bytes) to avoid over-subscription int num_threads_per_block = (optimal_block_size > static_cast(nelems)) @@ -437,8 +438,8 @@ __host__ void HostInterface::alltoallmem_on_stream(rocshmem_team_t team, // Launch kernel to do alltoall with given stream dim3 gridSize(1); dim3 blockSize(num_threads_per_block); - rocshmem_alltoallmem_kernel<<>>(team, dest, - source, nelems); + rocshmem_alltoallmem_kernel<<>>( + team, dest, source, nelems); } __host__ void HostInterface::broadcastmem_on_stream(rocshmem_team_t team, @@ -450,12 +451,13 @@ __host__ void HostInterface::broadcastmem_on_stream(rocshmem_team_t team, // - Query optimal block size using occupancy API // - Limit block size to nelems (number of bytes) to avoid over-subscription // - Always use 1 block (single workgroup collective) + const size_t tdm_lds_bytes = rocshmem_query_tdm_lds_bytes(); int optimal_block_size = 0; int grid_size = 0; CHECK_HIP(hipOccupancyMaxPotentialBlockSize(&grid_size, &optimal_block_size, rocshmem_broadcastmem_kernel, - 0, + tdm_lds_bytes, 0)); // Limit block size to nelems (bytes) to avoid over-subscription @@ -466,20 +468,19 @@ __host__ void HostInterface::broadcastmem_on_stream(rocshmem_team_t team, // Launch kernel to do broadcast with given stream dim3 gridSize(1); dim3 blockSize(num_threads_per_block); - rocshmem_broadcastmem_kernel<<>>(team, - dest, - source, - nelems, - pe_root); + rocshmem_broadcastmem_kernel<<>>( + team, dest, source, nelems, pe_root); } __host__ void HostInterface::getmem_on_stream(void *dest, const void *source, size_t nelems, int pe, hipStream_t stream) { + const size_t tdm_lds_bytes = rocshmem_query_tdm_lds_bytes(); int optimal_block_size = 0; int grid_size = 0; CHECK_HIP(hipOccupancyMaxPotentialBlockSize(&grid_size, &optimal_block_size, - rocshmem_getmem_kernel, 0, 0)); + rocshmem_getmem_kernel, + tdm_lds_bytes, 0)); // Limit block size to nelems to avoid over-subscription int num_threads_per_block = (optimal_block_size > static_cast(nelems)) @@ -489,17 +490,19 @@ __host__ void HostInterface::getmem_on_stream(void *dest, const void *source, // Launch kernel to do getmem with given stream dim3 gridSize(1); dim3 blockSize(num_threads_per_block); - rocshmem_getmem_kernel<<>>(dest, source, - nelems, pe); + rocshmem_getmem_kernel<<>>( + dest, source, nelems, pe); } __host__ void HostInterface::putmem_on_stream(void *dest, const void *source, size_t nelems, int pe, hipStream_t stream) { + const size_t tdm_lds_bytes = rocshmem_query_tdm_lds_bytes(); int optimal_block_size = 0; int grid_size = 0; CHECK_HIP(hipOccupancyMaxPotentialBlockSize(&grid_size, &optimal_block_size, - rocshmem_putmem_kernel, 0, 0)); + rocshmem_putmem_kernel, + tdm_lds_bytes, 0)); // Limit block size to nelems to avoid over-subscription int num_threads_per_block = (optimal_block_size > static_cast(nelems)) @@ -509,17 +512,19 @@ __host__ void HostInterface::putmem_on_stream(void *dest, const void *source, // Launch kernel to do putmem with given stream dim3 gridSize(1); dim3 blockSize(num_threads_per_block); - rocshmem_putmem_kernel<<>>(dest, source, - nelems, pe); + rocshmem_putmem_kernel<<>>( + dest, source, nelems, pe); } __host__ void HostInterface::putmem_signal_on_stream( void *dest, const void *source, size_t nelems, uint64_t *sig_addr, uint64_t signal, int sig_op, int pe, hipStream_t stream) { + const size_t tdm_lds_bytes = rocshmem_query_tdm_lds_bytes(); int optimal_block_size = 0; int grid_size = 0; CHECK_HIP(hipOccupancyMaxPotentialBlockSize( - &grid_size, &optimal_block_size, rocshmem_putmem_signal_kernel, 0, 0)); + &grid_size, &optimal_block_size, rocshmem_putmem_signal_kernel, + tdm_lds_bytes, 0)); // Limit block size to nelems to avoid over-subscription int num_threads_per_block = (optimal_block_size > static_cast(nelems)) @@ -529,7 +534,7 @@ __host__ void HostInterface::putmem_signal_on_stream( // Launch kernel to do putmem_signal with given stream dim3 gridSize(1); dim3 blockSize(num_threads_per_block); - rocshmem_putmem_signal_kernel<<>>( + rocshmem_putmem_signal_kernel<<>>( dest, source, nelems, sig_addr, signal, sig_op, pe); } diff --git a/projects/rocshmem/src/ipc/context_ipc_device.cpp b/projects/rocshmem/src/ipc/context_ipc_device.cpp index 938eff179a56..5871580ff1cd 100644 --- a/projects/rocshmem/src/ipc/context_ipc_device.cpp +++ b/projects/rocshmem/src/ipc/context_ipc_device.cpp @@ -173,6 +173,21 @@ __device__ void IPCContext::internal_getmem_wg(void *dest, const void *source, __builtin_amdgcn_s_barrier(); } +__device__ void IPCContext::internal_putmem_wg_tdm(void *dest, const void *source, + size_t nelems, int pe) { + uint64_t L_offset = reinterpret_cast(dest) - wrk_sync_pool_bases_[constmem.my_pe]; + memcpy_wg_tdm(wrk_sync_pool_bases_[pe] + L_offset, const_cast(source), nelems); + __builtin_amdgcn_s_barrier(); +} + +__device__ void IPCContext::internal_getmem_wg_tdm(void *dest, const void *source, + size_t nelems, int pe) { + const char *src_typed = reinterpret_cast(source); + uint64_t L_offset = const_cast(src_typed) - wrk_sync_pool_bases_[constmem.my_pe]; + memcpy_wg_tdm(dest, wrk_sync_pool_bases_[pe] + L_offset, nelems); + __builtin_amdgcn_s_barrier(); +} + __device__ void IPCContext::internal_putmem_wave(void *dest, const void *source, size_t nelems, int pe) { uint64_t L_offset = reinterpret_cast(dest) - wrk_sync_pool_bases_[constmem.my_pe]; diff --git a/projects/rocshmem/src/ipc/context_ipc_device.hpp b/projects/rocshmem/src/ipc/context_ipc_device.hpp index 810a6baebab6..53721ed46159 100644 --- a/projects/rocshmem/src/ipc/context_ipc_device.hpp +++ b/projects/rocshmem/src/ipc/context_ipc_device.hpp @@ -535,22 +535,28 @@ class IPCContext : public Context { size_t nelems, int pe); __device__ void internal_putmem_wg(void *dest, const void *source, - size_t nelems, int pe); + size_t nelems, int pe); __device__ void internal_getmem_wg(void *dest, const void *source, - size_t nelems, int pe); + size_t nelems, int pe); + + __device__ void internal_putmem_wg_tdm(void *dest, const void *source, + size_t nelems, int pe); + + __device__ void internal_getmem_wg_tdm(void *dest, const void *source, + size_t nelems, int pe); __device__ void internal_putmem_wave(void *dest, const void *source, - size_t nelems, int pe); + size_t nelems, int pe); __device__ void internal_getmem_wave(void *dest, const void *source, - size_t nelems, int pe); + size_t nelems, int pe); __device__ void internal_alltoallmem_wave(rocshmem_team_t team, void *dst, const void *src, int nelems); __device__ void alltoallmem_linear_wave(rocshmem_team_t team, void *dst, - const void *src, int nelems); + const void *src, int nelems); __device__ void alltoallmem_linear_thread_puts_wave(rocshmem_team_t team, void *dst, const void *src, int nelems); diff --git a/projects/rocshmem/src/ipc_policy.hpp b/projects/rocshmem/src/ipc_policy.hpp index 0ce6364e6ee7..1b9ee060ef85 100644 --- a/projects/rocshmem/src/ipc_policy.hpp +++ b/projects/rocshmem/src/ipc_policy.hpp @@ -39,6 +39,7 @@ #include "bootstrap/bootstrap.hpp" #include "atomic.hpp" #include "sdma_policy.hpp" +#include "tdm_policy.hpp" namespace rocshmem { @@ -263,12 +264,12 @@ class IpcOnImpl { template __device__ void ipcCopy_wg(void *dst, void *src, size_t size, [[maybe_unused]] int local_pe) { - memcpy_wg(dst, src, size); + memcpy_wg_tdm(dst, src, size); } template __device__ void ipcCopy_wave(void *dst, void *src, size_t size, [[maybe_unused]] int local_pe) { - memcpy_wave(dst, src, size); + memcpy_wave_tdm(dst, src, size); } template (dst, src, size); + memcpy_wg_tdm(dst, src, size); } template @@ -451,7 +452,7 @@ class IpcSdmaImpl : public IpcOnImpl { } return; } - memcpy_wave(dst, src, size); + memcpy_wave_tdm(dst, src, size); } template tdm_resolved_tile_bytes); +#else + return 0; +#endif +} + __host__ void rocshmem_global_exit(int status) { VERIFY_BACKEND(); backend->global_exit(status); diff --git a/projects/rocshmem/src/rocshmem_gpu.cpp b/projects/rocshmem/src/rocshmem_gpu.cpp index 0b9623ee583b..2ddf6ff5bbc1 100644 --- a/projects/rocshmem/src/rocshmem_gpu.cpp +++ b/projects/rocshmem/src/rocshmem_gpu.cpp @@ -142,6 +142,21 @@ __device__ void rocshmem_query_thread(int *provided) { #endif } +__device__ size_t rocshmem_query_tdm_lds_bytes() { +#if ROCSHMEM_TDM_ENABLED + return tdm::lds_bytes_for_tile(constmem.tdm_tile_bytes); +#else + return 0; +#endif +} + +__device__ void rocshmem_set_tdm_lds([[maybe_unused]] void *lds, + [[maybe_unused]] size_t lds_bytes) { +#if ROCSHMEM_TDM_ENABLED + tdm::set_lds(lds, lds_bytes); +#endif +} + __device__ void rocshmem_wg_finalize() {} @@ -1307,6 +1322,9 @@ __device__ __forceinline__ void direct_destroy_ctx(rocshmem_ctx_t *ctx) { __device__ int rocshmem_wg_ctx_create(long options, rocshmem_ctx_t *ctx) { LOGD_API("device::wg_ctx_create (options=%ld)", options); +#if ROCSHMEM_TDM_ENABLED + tdm::set_lds(nullptr, 0); +#endif bool result{true}; if (get_flat_block_id() == 0) { ctx->team_opaque = reinterpret_cast(ROCSHMEM_CTX_DEFAULT.team_opaque); @@ -1763,6 +1781,10 @@ __global__ ATTR_NO_INLINE void rocshmem_alltoallmem_kernel(rocshmem_team_t team, void *dest, const void *source, size_t size) { + // rocSHMEM owns this launch (see alltoallmem_on_stream in host.cpp), so it + // can safely claim the kernel's whole dynamic LDS for the TDM fast path + extern __shared__ uint8_t lds[]; + rocshmem_set_tdm_lds(lds, rocshmem_query_tdm_lds_bytes()); rocshmem_ctx_alltoall_wg(ROCSHMEM_CTX_DEFAULT, team, (char *) dest, (const char *) source, (int) size); } @@ -1804,6 +1826,10 @@ __global__ ATTR_NO_INLINE void rocshmem_reduce_on_stream_kernel(rocshmem_team_t __global__ ATTR_NO_INLINE void rocshmem_broadcastmem_kernel( rocshmem_team_t team, void *dest, const void *source, size_t nelems, int pe_root) { + // rocSHMEM owns this launch (see broadcastmem_on_stream in host.cpp), so it + // can safely claim the kernel's whole dynamic LDS for the TDM fast path -- + extern __shared__ uint8_t lds[]; + rocshmem_set_tdm_lds(lds, rocshmem_query_tdm_lds_bytes()); rocshmem_broadcast_wg(ROCSHMEM_CTX_DEFAULT, team, (char *) dest, (const char *) source, (int) nelems, pe_root); } @@ -1811,6 +1837,12 @@ __global__ ATTR_NO_INLINE void rocshmem_broadcastmem_kernel( __global__ ATTR_NO_INLINE void rocshmem_getmem_kernel(void *dest, const void *source, size_t nelems, int pe) { + // rocSHMEM owns this launch (see putmem_on_stream/getmem_on_stream in + // host.cpp), so it can safely claim the kernel's whole dynamic LDS for the + // TDM fast path -- this is the same opt-in rocshmem_set_tdm_lds() API any + // user kernel would call, just wired up internally here. + extern __shared__ uint8_t lds[]; + rocshmem_set_tdm_lds(lds, rocshmem_query_tdm_lds_bytes()); // Use work-group collective getmem with default context rocshmem_getmem_wg(dest, source, nelems, pe); } @@ -1818,6 +1850,8 @@ __global__ ATTR_NO_INLINE void rocshmem_getmem_kernel(void *dest, __global__ ATTR_NO_INLINE void rocshmem_putmem_kernel(void *dest, const void *source, size_t nelems, int pe) { + extern __shared__ uint8_t lds[]; + rocshmem_set_tdm_lds(lds, rocshmem_query_tdm_lds_bytes()); // Use work-group collective putmem with default context rocshmem_putmem_wg(dest, source, nelems, pe); } @@ -1825,6 +1859,8 @@ __global__ ATTR_NO_INLINE void rocshmem_putmem_kernel(void *dest, __global__ ATTR_NO_INLINE void rocshmem_putmem_signal_kernel( void *dest, const void *source, size_t nelems, uint64_t *sig_addr, uint64_t signal, int sig_op, int pe) { + extern __shared__ uint8_t lds[]; + rocshmem_set_tdm_lds(lds, rocshmem_query_tdm_lds_bytes()); // Use work-group collective putmem_signal with default context rocshmem_putmem_signal_wg(dest, source, nelems, sig_addr, signal, sig_op, pe); } diff --git a/projects/rocshmem/src/tdm.hpp b/projects/rocshmem/src/tdm.hpp new file mode 100644 index 000000000000..a953a6a38d10 --- /dev/null +++ b/projects/rocshmem/src/tdm.hpp @@ -0,0 +1,441 @@ +/****************************************************************************** + * Copyright (c) Advanced Micro Devices, Inc. All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + *****************************************************************************/ + +// Tensor Data Mover (TDM) API for gfx1250 (MI450): wraps the +// __builtin_amdgcn_tensor_load_to_lds/store_from_lds intrinsics (descriptor +// construction, cache-policy encoding, bulk tile-copy helpers) plus the +// per-workgroup LDS registration backing rocshmem_set_tdm_lds(). + +#ifndef LIBRARY_SRC_TDM_HPP_ +#define LIBRARY_SRC_TDM_HPP_ + +#include + +#include +#include + +// ============================================================================== +// FEATURE DETECTION +// ============================================================================== + +#ifndef __has_builtin +#define __has_builtin(x) 0 +#endif + +#if defined(__HIP_PLATFORM_AMD__) && \ + (__has_builtin(__builtin_amdgcn_tensor_load_to_lds) && \ + __has_builtin(__builtin_amdgcn_tensor_store_from_lds)) +#define HIP_HAVE_TDM_INTRINSICS 1 +#else +#define HIP_HAVE_TDM_INTRINSICS 0 +#endif + +// Single switch for "TDM is actually usable here" (build opted in, arch +// supports it, compiler exposes the intrinsics). Everything below that's +// TDM-specific is gated on this. Strictly gfx1250 (regular and -strict +// build variants) -- not a wildcard match for other gfx12xx parts. +#if defined(USE_TDM) && (defined(__gfx1250__) || defined(__gfx1250_strict__)) && \ + HIP_HAVE_TDM_INTRINSICS +#define ROCSHMEM_TDM_ENABLED 1 +#else +#define ROCSHMEM_TDM_ENABLED 0 +#endif + +namespace rocshmem { +namespace tdm { + +// LDS bytes for a double-buffered copy_region call with this tile size. +// Host code (constmem.cpp, rocshmem.cpp) needs this under plain USE_TDM, +// not ROCSHMEM_TDM_ENABLED -- its arch macros aren't defined when +// compiling host-only code, which would make this disappear there. +#if defined(USE_TDM) +constexpr uint32_t FlatCopyElementLog2 = 3; // 8 bytes/element + +__host__ __device__ constexpr size_t lds_bytes_for_tile(uint32_t tile_bytes) { + return 2 * static_cast(tile_bytes); +} +#endif // USE_TDM + +#if ROCSHMEM_TDM_ENABLED + +// ============================================================================== +// VECTOR TYPES +// ============================================================================== + +using u32x4 = uint32_t __attribute__((ext_vector_type(4))); +using i32x4 = int32_t __attribute__((ext_vector_type(4))); +using i32x8 = int32_t __attribute__((ext_vector_type(8))); + +// ============================================================================== +// ELEMENT SIZE (GROUP1 data_size field) +// ============================================================================== + +enum class DataSize : uint32_t { Bits8 = 0, Bits16 = 1, Bits32 = 2, Bits64 = 3 }; + +// ============================================================================== +// CACHE POLICY (the `cachepolicy` immediate argument of the TDM intrinsics) +// cachepolicy = th_bits[2:0] | (scope_bits[1:0] << 3); +// ============================================================================== + +// Memory scope (cachepolicy bits[4:3]), pre-shifted into place here. +enum class Scope : unsigned { + Wgp = 0, // workgroup processor: L1-level coherence (default) + ShaderEngine = 8, // GL2/L2-level coherence within one shader engine + Device = 16, // all shader engines on this GPU + System = 24, // CPU + all GPUs on the node +}; + +// Temporal hint (cachepolicy bits[2:0]) for tensor_load_to_lds. th=7 is +// undefined for loads, so it's intentionally not a member here. +enum class LoadTemporalHint : unsigned { + Default = 0, // write-allocate, cached in L1 and L2 + NonTemporal = 1, // TH_LOAD_NT: no L1 allocation + HwManaged = 2, // TH_LOAD_HT: hardware-managed temporal hint + LastUse = 3, // TH_LOAD_LU: evict from L1 after this read + NonTemporalReadThrough = 4, // TH_LOAD_NT_RT + ReadThroughNonTemporal = 5, // TH_LOAD_RT_NT + NonTemporalHwManaged = 6, // TH_LOAD_NT_HT +}; + +// Temporal hint for tensor_store_from_lds. Same bit positions as +// LoadTemporalHint, but th=3/th=7 mean something different for stores. +enum class StoreTemporalHint : unsigned { + Default = 0, // write-back + NonTemporal = 1, // TH_STORE_NT: non-temporal write-through + HwManaged = 2, // TH_STORE_HT: hardware-managed temporal hint + WriteBack = 3, // TH_STORE_WB: explicit write-back + NonTemporalReadThrough = 4, // TH_STORE_NT_RT + ReadThroughNonTemporal = 5, // TH_STORE_RT_NT + NonTemporalHwManaged = 6, // TH_STORE_NT_HT + NonTemporalWriteBack = 7, // TH_STORE_NT_WB +}; + +constexpr int make_cache_policy(LoadTemporalHint th, Scope scope = Scope::Wgp) { + return static_cast(th) | static_cast(scope); +} +constexpr int make_cache_policy(StoreTemporalHint th, Scope scope = Scope::Wgp) { + return static_cast(th) | static_cast(scope); +} + +// ============================================================================== +// DESCRIPTOR CONSTRUCTION +// ============================================================================== + +// GROUP0 (addresses + control) and GROUP1 (tensor layout) for a 2D transfer. +// GROUP2/GROUP3 (3D/gather) aren't modeled; every call site here is 2D. +struct Descriptor2D { + u32x4 group0{}; + i32x8 group1{}; +}; + +// GROUP0 only (LDS address + 48-bit global VA), split out so callers can +// cheaply rebuild just the addresses of an otherwise-invariant descriptor. +__device__ __forceinline__ u32x4 make_group0(uintptr_t lds_addr, uintptr_t global_addr) { + return u32x4{ + 1u, // m_count = 1 + static_cast(lds_addr), + static_cast(global_addr & 0xFFFFFFFFu), + // bits[24:0] = global_addr_hi, bits[31:30] = m_type = 2 (non-gather). + // m_th is left 0: it's overridden by the intrinsic's cachepolicy + // argument, and m_is_store is set automatically by the compiler. + (2u << 30) | static_cast((global_addr >> 32) & 0x1FFFFFFu), + }; +} + +// GROUP1: element size, full-tensor/tile dims, and strides (in elements). +// Strides are genuinely 48-bit here, unlike AMD's own amd_gfx1250_TDM.h +// header, whose tensorDimNStride() setters truncate to 32 bits. +__device__ __forceinline__ i32x8 make_group1(DataSize data_size, uint32_t tensor_dim0, + uint32_t tensor_dim1, uint32_t tile_dim0, + uint32_t tile_dim1, uint64_t tensor_dim0_stride, + uint64_t tensor_dim1_stride, + uint32_t workgroup_mask = 1) { + uint32_t s[8]{}; + s[0] = (workgroup_mask & 0xFFFFu) | (static_cast(data_size) << 16); + s[1] = (tensor_dim0 & 0xFFFFu) << 16; // atomic_barrier_address left 0 + s[2] = ((tensor_dim0 >> 16) & 0xFFFFu) | ((tensor_dim1 & 0xFFFFu) << 16); + s[3] = ((tensor_dim1 >> 16) & 0xFFFFu) | ((tile_dim0 & 0xFFFFu) << 16); + s[4] = (tile_dim1 & 0xFFFFu); // tile_dim2 left 0 (2D only) + s[5] = static_cast(tensor_dim0_stride & 0xFFFFFFFFull); + s[6] = static_cast((tensor_dim0_stride >> 32) & 0xFFFFull) | + (static_cast(tensor_dim1_stride & 0xFFFFull) << 16); + s[7] = static_cast((tensor_dim1_stride >> 16) & 0xFFFFFFFFull); + + i32x8 g1{}; + for (int i = 0; i < 8; ++i) g1[i] = static_cast(s[i]); + return g1; +} + +__device__ __forceinline__ Descriptor2D make_descriptor_2d( + uintptr_t lds_addr, uintptr_t global_addr, DataSize data_size, uint32_t tensor_dim0, + uint32_t tensor_dim1, uint32_t tile_dim0, uint32_t tile_dim1, uint64_t tensor_dim0_stride, + uint64_t tensor_dim1_stride, uint32_t workgroup_mask = 1) { + Descriptor2D d; + d.group0 = make_group0(lds_addr, global_addr); + d.group1 = make_group1(data_size, tensor_dim0, tensor_dim1, tile_dim0, tile_dim1, + tensor_dim0_stride, tensor_dim1_stride, workgroup_mask); + return d; +} + +// ============================================================================== +// RAW INTRINSIC WRAPPERS +// ============================================================================== + +// Stalls until fewer than Cnt+1 of this wave's TDM ops remain in flight; +// wait<0>() drains all of them. Cnt (and CachePolicy below) must be a +// template parameter: the builtins require an immarg (compile-time literal). +template +__device__ __forceinline__ void wait() { + __builtin_amdgcn_s_wait_tensorcnt(Cnt); +} + +template +__device__ __forceinline__ void load_to_lds(u32x4 group0, i32x8 group1, + i32x4 group2 = i32x4{0, 0, 0, 0}, + i32x4 group3 = i32x4{0, 0, 0, 0}) { + const i32x8 group4{0, 0, 0, 0, 0, 0, 0, 0}; // reserved; always zero + __builtin_amdgcn_tensor_load_to_lds(group0, group1, group2, group3, group4, CachePolicy); +} + +template +__device__ __forceinline__ void store_from_lds(u32x4 group0, i32x8 group1, + i32x4 group2 = i32x4{0, 0, 0, 0}, + i32x4 group3 = i32x4{0, 0, 0, 0}) { + const i32x8 group4{0, 0, 0, 0, 0, 0, 0, 0}; // reserved; always zero + __builtin_amdgcn_tensor_store_from_lds(group0, group1, group2, group3, group4, CachePolicy); +} + +template +__device__ __forceinline__ void load_to_lds(const Descriptor2D &d) { + load_to_lds(d.group0, d.group1); +} + +template +__device__ __forceinline__ void store_from_lds(const Descriptor2D &d) { + store_from_lds(d.group0, d.group1); +} + +// ============================================================================== +// BULK TILE-COPY HELPERS: flat global<->global byte copies staged through +// LDS, tiled as 2D descriptors with 8-byte elements. Must be called from +// exactly one lane; callers own that gating and any tail bytes left +// uncovered by covered_bytes(). +// ============================================================================== + +// Bytes of [0, total_bytes) actually moved for a given tile_bytes_cap: the +// largest multiple of tile_bytes_cap not exceeding total_bytes, or a single +// undersized tile if total_bytes is smaller than one. +__device__ __forceinline__ size_t covered_bytes(size_t total_bytes, uint32_t tile_bytes_cap) { + if (total_bytes == 0) return 0; + if (total_bytes < tile_bytes_cap) { + constexpr size_t elem_bytes = size_t{1} << FlatCopyElementLog2; + return total_bytes - (total_bytes % elem_bytes); + } + return (total_bytes / tile_bytes_cap) * static_cast(tile_bytes_cap); +} + +// Copies covered_bytes(bytes, tile_bytes_cap) bytes tile by tile through +// lds_scratch. double_buffered overlaps load(tile N+1) with store(tile N) +// (needs 2*tile_bytes_cap bytes); otherwise each tile is a serialized +// load-wait-store-wait. tile_bytes_cap must be a multiple of 8 bytes. +// +// LoadCachePolicy/StoreCachePolicy are template parameters (see load_to_lds +// above), built with make_cache_policy(hint, scope). +template +__device__ inline void copy_region(void *dst, const void *src, size_t bytes, void *lds_scratch, + uint32_t tile_bytes_cap, bool double_buffered) { + const size_t covered = covered_bytes(bytes, tile_bytes_cap); + if (covered == 0) return; + const uint32_t tile_bytes = + (bytes < tile_bytes_cap) ? static_cast(covered) : tile_bytes_cap; + const size_t full_tiles = covered / tile_bytes; + const uint32_t tile_dim0 = tile_bytes >> FlatCopyElementLog2; + + // Tensor layout is invariant across tiles — only the addresses (GROUP0) + // change per iteration. + Descriptor2D desc; + desc.group1 = make_group1(DataSize::Bits64, /*tensor_dim0=*/tile_dim0, /*tensor_dim1=*/1, + /*tile_dim0=*/tile_dim0, /*tile_dim1=*/1, + /*tensor_dim0_stride=*/tile_dim0, /*tensor_dim1_stride=*/tile_dim0); + + auto *dst_bytes = static_cast(dst); + const auto *src_bytes = static_cast(src); + auto *lds_base = static_cast(lds_scratch); + + if (double_buffered) { + uintptr_t ping = reinterpret_cast(lds_base); + uintptr_t pong = reinterpret_cast(lds_base + tile_bytes); + + // Prologue: start loading tile 0 into ping. + desc.group0 = make_group0(ping, reinterpret_cast(src_bytes)); + load_to_lds(desc); + + for (size_t i = 0; i < full_tiles; ++i) { + const bool has_next = (i + 1) < full_tiles; + + // Wait for the current tile's load to land before storing from it. + wait<0>(); + + // Prefetch tile i+1 into pong and store tile i from ping — both + // enqueued before the next wait, so the tensor engine runs them + // concurrently. + if (has_next) { + desc.group0 = + make_group0(pong, reinterpret_cast(src_bytes + (i + 1) * tile_bytes)); + load_to_lds(desc); + } + desc.group0 = make_group0(ping, reinterpret_cast(dst_bytes + i * tile_bytes)); + store_from_lds(desc); + + uintptr_t tmp = ping; + ping = pong; + pong = tmp; + } + // Drain the last in-flight load+store pair. + wait<0>(); + } else { + const uintptr_t lds_addr = reinterpret_cast(lds_base); + for (size_t i = 0; i < full_tiles; ++i) { + const size_t offset = i * tile_bytes; + + desc.group0 = make_group0(lds_addr, reinterpret_cast(src_bytes + offset)); + load_to_lds(desc); + wait<0>(); + + desc.group0 = make_group0(lds_addr, reinterpret_cast(dst_bytes + offset)); + store_from_lds(desc); + wait<0>(); + } + } +} + +// Same contract as copy_region, but always double-buffered and more +// tightly pipelined: waits only until <=1 op is in flight (wait<1>()) +// instead of draining every iteration, so a load and a store are usually +// both in flight. Tiles issue in load/store pairs (parity picks the +// buffer); each op is issued exactly two slots after the op it depends on, +// so wait<1>() before op N guarantees op N-2 already completed. +template +__device__ inline void copy_region_pipelined(void *dst, const void *src, size_t bytes, + void *lds_scratch, uint32_t tile_bytes_cap) { + const size_t covered = covered_bytes(bytes, tile_bytes_cap); + if (covered == 0) return; + const uint32_t tile_bytes = + (bytes < tile_bytes_cap) ? static_cast(covered) : tile_bytes_cap; + const size_t full_tiles = covered / tile_bytes; + const uint32_t tile_dim0 = tile_bytes >> FlatCopyElementLog2; + + Descriptor2D desc; + desc.group1 = make_group1(DataSize::Bits64, tile_dim0, 1, tile_dim0, 1, tile_dim0, tile_dim0); + + auto *dst_bytes = static_cast(dst); + const auto *src_bytes = static_cast(src); + const uintptr_t buf[2] = { + reinterpret_cast(lds_scratch), + reinterpret_cast(static_cast(lds_scratch) + tile_bytes)}; + + auto issue_load = [&](uintptr_t lds_addr, uintptr_t global_addr) { + desc.group0 = make_group0(lds_addr, global_addr); + load_to_lds(desc); + }; + auto issue_store = [&](uintptr_t lds_addr, uintptr_t global_addr) { + desc.group0 = make_group0(lds_addr, global_addr); + store_from_lds(desc); + }; + + // Prologue: load tile 0 into buf[0]. + issue_load(buf[0], reinterpret_cast(src_bytes)); + + for (size_t tile_even = 0; tile_even < full_tiles; tile_even += 2) { + const size_t tile_odd = tile_even + 1; + const bool has_odd = tile_odd < full_tiles; + + if (has_odd) { + wait<1>(); + issue_load(buf[tile_odd & 1u], reinterpret_cast(src_bytes + tile_odd * tile_bytes)); + } + + // No odd tile: load(tile_even) was only one op ago, not two, so + // wait<1>() would race ahead of it -- drain fully instead. + if (has_odd) { + wait<1>(); + } else { + wait<0>(); + } + issue_store(buf[tile_even & 1u], reinterpret_cast(dst_bytes + tile_even * tile_bytes)); + + if (has_odd) { + wait<1>(); + issue_store(buf[tile_odd & 1u], reinterpret_cast(dst_bytes + tile_odd * tile_bytes)); + + const size_t next_even = tile_even + 2; + if (next_even < full_tiles) { + wait<1>(); + issue_load(buf[next_even & 1u], + reinterpret_cast(src_bytes + next_even * tile_bytes)); + } + } + } + + // Drain everything still in flight before returning. + wait<0>(); +} + +// ============================================================================== +// LDS REGISTRATION: backs rocshmem_set_tdm_lds()/rocshmem_query_tdm_lds_bytes(). +// Every thread in the block must call set_lds() -- only thread (0,0,0)'s +// arguments are used, but its internal barrier needs every thread to reach +// it -- before any work-group or wave put/get in that kernel that should use +// TDM. One registration backs both: work-group put/get (memcpy_wg_tdm) uses +// the whole buffer; wave put/get (memcpy_wave_tdm) splits it one slice per +// wave so concurrent waves don't tile through the same LDS bytes. +// ============================================================================== + +// Each workgroup gets its own private LDS, and set_lds() always runs earlier +// in the same kernel dispatch than any get_lds() that reads it back, so no +// staleness check is needed here. +struct LdsRegistration { + void* ptr; + size_t bytes; +}; + +inline __shared__ LdsRegistration lds_registration; + +__device__ __forceinline__ void set_lds(void* ptr, size_t bytes) { + if (hipThreadIdx_x == 0 && hipThreadIdx_y == 0 && hipThreadIdx_z == 0) { + lds_registration = LdsRegistration{ptr, bytes}; + } + __builtin_amdgcn_s_barrier(); // publish to the rest of the block +} + +__device__ __forceinline__ LdsRegistration get_lds() { + return lds_registration; +} + +#endif // ROCSHMEM_TDM_ENABLED + +} // namespace tdm +} // namespace rocshmem + +#endif // LIBRARY_SRC_TDM_HPP_ diff --git a/projects/rocshmem/src/tdm_policy.hpp b/projects/rocshmem/src/tdm_policy.hpp new file mode 100644 index 000000000000..475583da3b19 --- /dev/null +++ b/projects/rocshmem/src/tdm_policy.hpp @@ -0,0 +1,140 @@ +/****************************************************************************** + * Copyright (c) Advanced Micro Devices, Inc. All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + *****************************************************************************/ + +#ifndef LIBRARY_SRC_TDM_POLICY_HPP_ +#define LIBRARY_SRC_TDM_POLICY_HPP_ + +#include + +#include "rocshmem/rocshmem_config.h" // NOLINT(build/include_subdir) +#include "constmem.hpp" +#include "util.hpp" +#include "tdm.hpp" + +namespace rocshmem { + +// Thread 0 drives TDM into the registered LDS buffer; the rest of the block +// covers what TDM doesn't. Falls back to memcpy_wg if TDM isn't available. +template +[[maybe_unused]] __device__ __forceinline__ void memcpy_wg_tdm(void* dst, void* src, + size_t size) { +#if ROCSHMEM_TDM_ENABLED + // Non-temporal + system scope: correct either direction (peer may be remote). + constexpr int TdmLoadPolicy = tdm::make_cache_policy( + tdm::LoadTemporalHint::NonTemporal, tdm::Scope::System); + constexpr int TdmStorePolicy = tdm::make_cache_policy( + tdm::StoreTemporalHint::NonTemporal, tdm::Scope::System); + + // Below threshold, TDM doesn't pay off (like the SDMA threshold); size + // is small enough here that Unroll=4 is plenty. + if (size < constmem.tdm_threshold) { + memcpy_wg(dst, src, size); + return; + } + const tdm::LdsRegistration reg = tdm::get_lds(); + // Below 16 bytes there's no room for a double-buffered tile. + if (reg.ptr != nullptr && reg.bytes >= 16) { + constexpr uint32_t element_bytes = 1u << tdm::FlatCopyElementLog2; + uint32_t tile_bytes = static_cast( + min(static_cast(constmem.tdm_tile_bytes), reg.bytes / 2)); + // Round down to a whole element -- an unaligned stride would leave + // gaps between tiles. + tile_bytes -= tile_bytes % element_bytes; + if (tile_bytes > 0) { + if (is_thread_zero_in_block()) { + tdm::copy_region(dst, src, size, reg.ptr, tile_bytes, + /*double_buffered=*/true); + } + const size_t covered = tdm::covered_bytes(size, tile_bytes); + if (covered < size) { + // Runs concurrently with thread 0's TDM copy; Unroll=4 trims + // register pressure here for free (too small to need 8 anyway). + memcpy_wg(static_cast(dst) + covered, + static_cast(src) + covered, size - covered); + } + return; + } + } +#endif + memcpy_wg(dst, src, size); +} + +// Wave counterpart of memcpy_wg_tdm: lane 0 of the calling wave drives TDM, +// the rest covers what TDM doesn't. The registered buffer is split into +// one fixed slot per wave the block could launch (not just the ones that +// call this), keeping slots disjoint without runtime coordination -- more +// waves per block means a smaller slot (and tile) per wave. +template +[[maybe_unused]] __device__ __forceinline__ void memcpy_wave_tdm(void* dst, void* src, + size_t size) { +#if ROCSHMEM_TDM_ENABLED + constexpr int TdmLoadPolicy = tdm::make_cache_policy( + tdm::LoadTemporalHint::NonTemporal, tdm::Scope::System); + constexpr int TdmStorePolicy = tdm::make_cache_policy( + tdm::StoreTemporalHint::NonTemporal, tdm::Scope::System); + + if (size < constmem.tdm_threshold) { + memcpy_wave(dst, src, size); + return; + } + const tdm::LdsRegistration reg = tdm::get_lds(); + // Round up: a partial trailing wavefront still needs its own slot -- + // undercounting would let its wave_id index past its slice. + const int num_waves = max((get_flat_block_size() + WF_SIZE - 1) / WF_SIZE, 1); + constexpr uint32_t element_bytes = 1u << tdm::FlatCopyElementLog2; + // Round down to a whole element -- reg.ptr is only guaranteed 8-byte + // aligned, so an unaligned stride could start a wave's slot (and its + // tile) on an unaligned address. + size_t slot_bytes = reg.bytes / static_cast(num_waves); + slot_bytes -= slot_bytes % element_bytes; + // Below 16 bytes there's no room for a double-buffered tile in this slot. + if (reg.ptr != nullptr && slot_bytes >= 16) { + uint32_t tile_bytes = static_cast( + min(static_cast(constmem.tdm_tile_bytes), slot_bytes / 2)); + // Round down to a whole element -- an unaligned stride would leave + // gaps between tiles. + tile_bytes -= tile_bytes % element_bytes; + if (tile_bytes > 0) { + const int wave_id = get_flat_block_id() / WF_SIZE; + void* wave_lds = static_cast(reg.ptr) + + static_cast(wave_id) * slot_bytes; + if (is_thread_zero_in_wave()) { + tdm::copy_region(dst, src, size, wave_lds, tile_bytes, + /*double_buffered=*/true); + } + const size_t covered = tdm::covered_bytes(size, tile_bytes); + if (covered < size) { + memcpy_wave(static_cast(dst) + covered, + static_cast(src) + covered, size - covered); + } + return; + } + } +#endif + memcpy_wave(dst, src, size); +} + +} // namespace rocshmem + +#endif // LIBRARY_SRC_TDM_POLICY_HPP_ diff --git a/projects/rocshmem/src/util.hpp b/projects/rocshmem/src/util.hpp index 7bbf5716449b..405e981a348f 100644 --- a/projects/rocshmem/src/util.hpp +++ b/projects/rocshmem/src/util.hpp @@ -558,13 +558,12 @@ template } } -template +template [[maybe_unused]] __device__ __forceinline__ void memcpy_wg(void* dst, void* src, size_t size) { if (size == 0) return; constexpr int ChunkSize = 16; - constexpr int Unroll = 8; constexpr CachePolicy LP = is_put(Kind) ? CachePolicy::Standard : CachePolicy::SystemScope; diff --git a/projects/rocshmem/tests/functional_tests/broadcast_wave_tester.cpp b/projects/rocshmem/tests/functional_tests/broadcast_wave_tester.cpp index 1c044a4d4680..8d93b4a6a89c 100644 --- a/projects/rocshmem/tests/functional_tests/broadcast_wave_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/broadcast_wave_tester.cpp @@ -72,6 +72,12 @@ __global__ void BroadcastWaveTestKernel(int loop, int skip, long long int *start __syncthreads(); } + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. The TDM LDS + // region sits right after the per-wave ctx slots in the dynamic shared + // buffer; memcpy_wave_tdm splits it evenly across the block's waves. + uint8_t *tdm_lds = reinterpret_cast(&ctx_array[num_waves_per_wg]); + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + // Each wave uses its own buffer slice T1 *my_source = source_buf + flat_wf_id * size; T1 *my_dest = dest_buf + flat_wf_id * size; @@ -155,7 +161,8 @@ void BroadcastWaveTester::preLaunchKernel() { template void BroadcastWaveTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, size_t size) { - size_t shared_bytes = num_warps * sizeof(rocshmem_ctx_t); + size_t shared_bytes = num_warps * sizeof(rocshmem_ctx_t) + + rocshmem_query_tdm_lds_bytes(); int num_elems = size / sizeof(T1); diff --git a/projects/rocshmem/tests/functional_tests/fcollect_wave_tester.cpp b/projects/rocshmem/tests/functional_tests/fcollect_wave_tester.cpp index 44694e430e0a..7c3caef5c206 100644 --- a/projects/rocshmem/tests/functional_tests/fcollect_wave_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/fcollect_wave_tester.cpp @@ -78,6 +78,12 @@ __global__ void FcollectWaveTest(int loop, int skip, long long int *start_time, __syncthreads(); } + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. The TDM LDS + // region sits right after the per-wave ctx slots in the dynamic shared + // buffer; memcpy_wave_tdm splits it evenly across the block's waves. + uint8_t *tdm_lds = reinterpret_cast(&ctx_array[num_waves_per_wg]); + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + int n_pes = rocshmem_ctx_n_pes(ctx_array[wf_id]); // Each wave uses its own buffer slice @@ -165,7 +171,8 @@ void FcollectWaveTester::preLaunchKernel() { template void FcollectWaveTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, size_t size) { - size_t shared_bytes = num_warps * sizeof(rocshmem_ctx_t); + size_t shared_bytes = num_warps * sizeof(rocshmem_ctx_t) + + rocshmem_query_tdm_lds_bytes(); int num_elems = size / sizeof(T1); diff --git a/projects/rocshmem/tests/functional_tests/signaling_operations_tester.cpp b/projects/rocshmem/tests/functional_tests/signaling_operations_tester.cpp index 75ffbae8ed97..b5e21cd84b3a 100644 --- a/projects/rocshmem/tests/functional_tests/signaling_operations_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/signaling_operations_tester.cpp @@ -41,6 +41,10 @@ __global__ void PutmemSignalTest(int loop, int skip, long long int *start_time, int wg_id = get_flat_grid_id(); rocshmem_wg_ctx_create(ctx_type, &ctx); + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + uint64_t signal = 1; for (int i = 0; i < loop + skip; i++) { @@ -214,7 +218,7 @@ void SignalingOperationsTester::resetBuffers([[maybe_unused]] size_t size) { void SignalingOperationsTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, size_t size) { - size_t shared_bytes = 0; + size_t shared_bytes = rocshmem_query_tdm_lds_bytes(); switch (_type) { case SignalAddTestType: diff --git a/projects/rocshmem/tests/functional_tests/team_alltoall_tester.cpp b/projects/rocshmem/tests/functional_tests/team_alltoall_tester.cpp index 103cc482ace3..961770902fb5 100644 --- a/projects/rocshmem/tests/functional_tests/team_alltoall_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/team_alltoall_tester.cpp @@ -66,6 +66,10 @@ __global__ void TeamAlltoallTest(int loop, int skip, long long int *start_time, rocshmem_wg_team_create_ctx(teams[wg_id], ctx_type, &ctx); + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + int n_pes = rocshmem_ctx_n_pes(ctx); source_buf += wg_id * n_pes * num_elems; @@ -145,7 +149,7 @@ void TeamAlltoallTester::preLaunchKernel() { template void TeamAlltoallTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, size_t size) { - size_t shared_bytes = 0; + size_t shared_bytes = rocshmem_query_tdm_lds_bytes(); int num_elems = size / sizeof(T1); diff --git a/projects/rocshmem/tests/functional_tests/team_alltoallv_tester.cpp b/projects/rocshmem/tests/functional_tests/team_alltoallv_tester.cpp index 1c230fb01c6d..187641d0df62 100644 --- a/projects/rocshmem/tests/functional_tests/team_alltoallv_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/team_alltoallv_tester.cpp @@ -86,6 +86,10 @@ __global__ void TeamAlltoallvTest(int loop, int skip, [[maybe_unused]] ShmemContextType ctx_type, rocshmem_team_t *teams) { + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + __syncthreads(); for (int i = 0; i < loop + skip; i++) { @@ -172,7 +176,7 @@ void TeamAlltoallvTester::preLaunchKernel() { template void TeamAlltoallvTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, size_t size) { - size_t shared_bytes = 0; + size_t shared_bytes = rocshmem_query_tdm_lds_bytes(); int num_elems = size / sizeof(T1); size_t disp = 0; diff --git a/projects/rocshmem/tests/functional_tests/team_broadcast_tester.cpp b/projects/rocshmem/tests/functional_tests/team_broadcast_tester.cpp index fd00541ee13b..1cb76e41b3e8 100644 --- a/projects/rocshmem/tests/functional_tests/team_broadcast_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/team_broadcast_tester.cpp @@ -71,6 +71,10 @@ __global__ void TeamBroadcastTest(int loop, int skip, long long int *start_time, rocshmem_wg_team_create_ctx(teams[wg_id], ctx_type, &ctx); + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + [[maybe_unused]] int n_pes = rocshmem_ctx_n_pes(ctx); source_buf += wg_id * size; dest_buf += wg_id * size; @@ -148,7 +152,7 @@ void TeamBroadcastTester::preLaunchKernel() { template void TeamBroadcastTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, size_t size) { - size_t shared_bytes = 0; + size_t shared_bytes = rocshmem_query_tdm_lds_bytes(); int num_elems = size / sizeof(T1); diff --git a/projects/rocshmem/tests/functional_tests/team_fcollect_tester.cpp b/projects/rocshmem/tests/functional_tests/team_fcollect_tester.cpp index ca23f073c7ea..bdddad776bb1 100644 --- a/projects/rocshmem/tests/functional_tests/team_fcollect_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/team_fcollect_tester.cpp @@ -66,6 +66,10 @@ __global__ void TeamFcollectTest(int loop, int skip, long long int *start_time, rocshmem_wg_team_create_ctx(teams[wg_id], ctx_type, &ctx); + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + int n_pes = rocshmem_ctx_n_pes(ctx); source_buf += wg_id * num_elems; dest_buf += wg_id * num_elems * n_pes; @@ -159,7 +163,7 @@ void TeamFcollectTester::preLaunchKernel() { template void TeamFcollectTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, size_t size) { - size_t shared_bytes = 0; + size_t shared_bytes = rocshmem_query_tdm_lds_bytes(); int num_elems = size / sizeof(T1); diff --git a/projects/rocshmem/tests/functional_tests/team_reduce_scatter_tester.cpp b/projects/rocshmem/tests/functional_tests/team_reduce_scatter_tester.cpp index 89f4bbf69e87..768157b94bc1 100644 --- a/projects/rocshmem/tests/functional_tests/team_reduce_scatter_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/team_reduce_scatter_tester.cpp @@ -93,6 +93,10 @@ __global__ void TeamReduceScatterTest(int loop, int skip, rocshmem_wg_team_create_ctx(team, ctx_type, &ctx); + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + __syncthreads(); for (int i = 0; i < loop + skip; i++) { @@ -148,7 +152,7 @@ template void TeamReduceScatterTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, uint64_t size) { - size_t shared_bytes = 0; + size_t shared_bytes = rocshmem_query_tdm_lds_bytes(); hipLaunchKernelGGL(HIP_KERNEL_NAME(TeamReduceScatterTest), gridSize, blockSize, shared_bytes, stream, loop, args.skip, diff --git a/projects/rocshmem/tests/functional_tests/team_reduction_tester.cpp b/projects/rocshmem/tests/functional_tests/team_reduction_tester.cpp index 9d89ae44352d..14bfebc22e4a 100644 --- a/projects/rocshmem/tests/functional_tests/team_reduction_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/team_reduction_tester.cpp @@ -85,6 +85,10 @@ __global__ void TeamReductionTest(int loop, int skip, long long int *start_time, rocshmem_wg_team_create_ctx(team, ctx_type, &ctx); + // In no-TDM build,rocshmem_query_tdm_lds_bytes() return 0. + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + __syncthreads(); for (int i = 0; i < loop + skip; i++) { @@ -136,7 +140,8 @@ void TeamReductionTester::preLaunchKernel() { template void TeamReductionTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, size_t size) { - size_t shared_bytes = 0; + // Dynamic LDS for the TDM (0 when not built with USE_TDM) + size_t shared_bytes = rocshmem_query_tdm_lds_bytes(); hipLaunchKernelGGL(HIP_KERNEL_NAME(TeamReductionTest), gridSize, blockSize, shared_bytes, stream, loop, args.skip, diff --git a/projects/rocshmem/tests/functional_tests/tile_allgather_tester.cpp b/projects/rocshmem/tests/functional_tests/tile_allgather_tester.cpp index c4a11c53fc73..6b94ec848a50 100644 --- a/projects/rocshmem/tests/functional_tests/tile_allgather_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/tile_allgather_tester.cpp @@ -198,6 +198,11 @@ __global__ void TileAllgatherTest(rocshmem_team_t *teams, int num_teams, rocshmem_team_t my_team = teams[wg_id % num_teams]; rocshmem_wg_team_create_ctx(my_team, ctx_type, &ctx_array[0]); + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. The TDM LDS + // region sits right after the one ctx slot in the dynamic shared buffer. + uint8_t *tdm_lds = reinterpret_cast(&ctx_array[1]); + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + int tile_size = tile_extent_0 * tile_extent_1; int src_offset = tile_size * wg_id; int dst_offset = tile_size * n_pes * wg_id; @@ -393,7 +398,7 @@ void TileAllgatherTester::launchKernel(dim3 gridSize, dim3 blockSize, } case TileAllgatherWGTestType: { - size_t wg_shared = sizeof(rocshmem_ctx_t); + size_t wg_shared = sizeof(rocshmem_ctx_t) + rocshmem_query_tdm_lds_bytes(); hipLaunchKernelGGL(TileAllgatherTest, dim3(args.num_wgs), blockSize, wg_shared, stream, teams, args.num_wgs, source, dest, diff --git a/projects/rocshmem/tests/functional_tests/tile_broadcast_tester.cpp b/projects/rocshmem/tests/functional_tests/tile_broadcast_tester.cpp index 339679566610..743e16a6d51b 100644 --- a/projects/rocshmem/tests/functional_tests/tile_broadcast_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/tile_broadcast_tester.cpp @@ -197,6 +197,11 @@ __global__ void TileBroadcastTest(rocshmem_team_t *teams, int num_teams, rocshmem_team_t my_team = teams[wg_id % num_teams]; rocshmem_wg_team_create_ctx(my_team, ctx_type, &ctx_array[0]); + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. The TDM LDS + // region sits right after the one ctx slot in the dynamic shared buffer. + uint8_t *tdm_lds = reinterpret_cast(&ctx_array[1]); + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + int matrix_size = tile_extent_0 * tile_extent_1; int offset = matrix_size * wg_id; @@ -393,7 +398,7 @@ void TileBroadcastTester::launchKernel(dim3 gridSize, dim3 blockSize, } case TileBroadcastWGTestType: { - size_t wg_shared = sizeof(rocshmem_ctx_t); + size_t wg_shared = sizeof(rocshmem_ctx_t) + rocshmem_query_tdm_lds_bytes(); hipLaunchKernelGGL(TileBroadcastTest, dim3(args.num_wgs), blockSize, wg_shared, stream, teams, args.num_wgs, source, dest, diff --git a/projects/rocshmem/tests/functional_tests/tile_reduce_tester.cpp b/projects/rocshmem/tests/functional_tests/tile_reduce_tester.cpp index 08feb024a00e..451fd73c0ca3 100644 --- a/projects/rocshmem/tests/functional_tests/tile_reduce_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/tile_reduce_tester.cpp @@ -237,6 +237,11 @@ __global__ void TileReduceWGTest(rocshmem_team_t *teams, int num_teams, rocshmem_team_t my_team = teams[wg_id % num_teams]; rocshmem_wg_team_create_ctx(my_team, ctx_type, &ctx_array[0]); + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. The TDM LDS + // region sits right after the one ctx slot in the dynamic shared buffer. + uint8_t *tdm_lds = reinterpret_cast(&ctx_array[1]); + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + int matrix_size = tile_extent_0 * tile_extent_1; int offset = matrix_size * wg_id; @@ -400,7 +405,7 @@ void TileReduceTester::launchKernel(dim3 gridSize, dim3 blockSize, } case TileReduceWGTestType: { - size_t wg_shared = sizeof(rocshmem_ctx_t); + size_t wg_shared = sizeof(rocshmem_ctx_t) + rocshmem_query_tdm_lds_bytes(); hipLaunchKernelGGL( HIP_KERNEL_NAME(TileReduceWGTest), dim3(args.num_wgs), blockSize, wg_shared, stream, diff --git a/projects/rocshmem/tests/functional_tests/tile_rma_tester.cpp b/projects/rocshmem/tests/functional_tests/tile_rma_tester.cpp index 02fbac91247a..610890042773 100644 --- a/projects/rocshmem/tests/functional_tests/tile_rma_tester.cpp +++ b/projects/rocshmem/tests/functional_tests/tile_rma_tester.cpp @@ -132,6 +132,10 @@ __global__ void TileRMATest(int loop, int skip, long long int *start_time, int wf_id = t_id / wf_size; rocshmem_wg_ctx_create(ctx_type, &ctx); + // In no-TDM build, rocshmem_query_tdm_lds_bytes() returns 0. + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + __shared__ long long int wf_start_time[32]; // Calculate base offset for this thread/wave/wg's data region. @@ -766,7 +770,7 @@ void TileRMATester::resetBuffers(uint64_t size) { void TileRMATester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, uint64_t size) { - size_t shared_bytes = 0; + size_t shared_bytes = rocshmem_query_tdm_lds_bytes(); // Derive tile dimensions for this iteration from the current message size so // that the reported bandwidth matches the actual bytes transferred. diff --git a/projects/rocshmem/tests/functional_tests/wavefront_primitives.cpp b/projects/rocshmem/tests/functional_tests/wavefront_primitives.cpp index 9f085229c531..bd9ab5c4e960 100644 --- a/projects/rocshmem/tests/functional_tests/wavefront_primitives.cpp +++ b/projects/rocshmem/tests/functional_tests/wavefront_primitives.cpp @@ -46,6 +46,11 @@ __global__ void WaveFrontPrimitiveTest(int loop, int skip, rocshmem_wg_ctx_create(ctx_type, &ctx); + // Opt into TDM for wave put/get on gfx1250 builds with USE_TDM; + // rocshmem_query_tdm_lds_bytes() returns 0 otherwise. + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + // Calculate start index for each wavefront int wf_id = get_flat_block_id() / wf_size; int wg_offset = wg_id * ((get_flat_block_size() - 1 ) / wf_size + 1); @@ -103,11 +108,12 @@ WaveFrontPrimitiveTester::WaveFrontPrimitiveTester(TesterArguments args) char *local = (char *) alloc_test_buffer(buff_size, args.local_buf_type); char *remote = (char *) alloc_test_buffer(buff_size); CHECK_HIP(hipMalloc(&grid_psync, sizeof(int))); + size_t lds_bytes = rocshmem_query_tdm_lds_bytes(); int max_co_resident_wgs_per_cu = 0; CHECK_HIP(hipOccupancyMaxActiveBlocksPerMultiprocessor( &max_co_resident_wgs_per_cu, WaveFrontPrimitiveTest, - args.wg_size, 0)); + args.wg_size, lds_bytes)); const int max_sustainable_wgs = max_co_resident_wgs_per_cu * deviceProps.multiProcessorCount; if (args.num_wgs > static_cast(max_sustainable_wgs)) { @@ -167,7 +173,9 @@ void WaveFrontPrimitiveTester::resetBuffers(size_t size) { void WaveFrontPrimitiveTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, size_t size) { - size_t shared_bytes = 0; + // Dynamic LDS for the TDM fast path (0 when not built with USE_TDM); see + // rocshmem_query_tdm_lds_bytes() / rocshmem_set_tdm_lds() in the kernel. + size_t shared_bytes = rocshmem_query_tdm_lds_bytes(); switch (_type) { case WAVEGetTestType: diff --git a/projects/rocshmem/tests/functional_tests/workgroup_primitives.cpp b/projects/rocshmem/tests/functional_tests/workgroup_primitives.cpp index 3681650082b1..63afea1715c2 100644 --- a/projects/rocshmem/tests/functional_tests/workgroup_primitives.cpp +++ b/projects/rocshmem/tests/functional_tests/workgroup_primitives.cpp @@ -42,6 +42,11 @@ __global__ void WorkGroupPrimitiveTest( int wg_id = get_flat_grid_id(); rocshmem_wg_ctx_create(ctx_type, &ctx); + // Opt into the TDM for work-group put/get on gfx1250 builds with + // USE_TDM; rocshmem_query_tdm_lds_bytes() returns 0 otherwise + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + // Calculate start index for each work group // Each workgroup owns `batch` contiguous slots of `size` bytes. source += size * batch * wg_id; @@ -97,11 +102,12 @@ WorkGroupPrimitiveTester::WorkGroupPrimitiveTester(TesterArguments args) char *local = (char *) alloc_test_buffer(buff_size, args.local_buf_type); char *remote = (char *) alloc_test_buffer(buff_size); CHECK_HIP(hipMalloc(&grid_psync, sizeof(int))); + size_t lds_bytes = rocshmem_query_tdm_lds_bytes(); int max_co_resident_wgs_per_cu = 0; CHECK_HIP(hipOccupancyMaxActiveBlocksPerMultiprocessor( &max_co_resident_wgs_per_cu, WorkGroupPrimitiveTest, - args.wg_size, 0)); + args.wg_size, lds_bytes)); const int max_sustainable_wgs = max_co_resident_wgs_per_cu * deviceProps.multiProcessorCount; if (args.num_wgs > static_cast(max_sustainable_wgs)) { @@ -161,7 +167,9 @@ void WorkGroupPrimitiveTester::resetBuffers(size_t size) { void WorkGroupPrimitiveTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop, size_t size) { - size_t shared_bytes = 0; + // Dynamic LDS for the TDM fast path (0 when not built with USE_TDM); see + // rocshmem_query_tdm_lds_bytes() / rocshmem_set_tdm_lds() in the kernel. + size_t shared_bytes = rocshmem_query_tdm_lds_bytes(); switch (_type) { case WGGetTestType: diff --git a/projects/rocshmem/tests/unit_tests/CMakeLists.txt b/projects/rocshmem/tests/unit_tests/CMakeLists.txt index 8d6ffc2ca7ce..d832b572f3b3 100644 --- a/projects/rocshmem/tests/unit_tests/CMakeLists.txt +++ b/projects/rocshmem/tests/unit_tests/CMakeLists.txt @@ -77,6 +77,16 @@ if (USE_IPC) ) endif() +if (USE_IPC AND USE_TDM) + target_sources( + ${PROJECT_NAME} + PRIVATE + tdm_raw_put_gtest.cpp + tdm_struct_put_gtest.cpp + tdm_wave_put_gtest.cpp + ) +endif() + if (USE_GDA) target_sources( diff --git a/projects/rocshmem/tests/unit_tests/tdm_raw_put_gtest.cpp b/projects/rocshmem/tests/unit_tests/tdm_raw_put_gtest.cpp new file mode 100644 index 000000000000..c83cd28808f5 --- /dev/null +++ b/projects/rocshmem/tests/unit_tests/tdm_raw_put_gtest.cpp @@ -0,0 +1,30 @@ +/****************************************************************************** + * Copyright (c) Advanced Micro Devices, Inc. All rights reserved. + * + * SPDX-License-Identifier: MIT + *****************************************************************************/ + +#include "tdm_raw_put_gtest.hpp" + +using namespace rocshmem; + +TEST_P(TdmRawPutFixture, put_bytes) { + run_put(GetParam()); +} + +static constexpr size_t TDM_TILE_BYTES = TEST_TDM_TILE_BYTES; + +INSTANTIATE_TEST_SUITE_P( + TdmRawPut, + TdmRawPutFixture, + ::testing::Values( + // Sub-tile: exercises memcpy_wg fallback path + static_cast(1024), + // Exactly one TDM tile + static_cast(TDM_TILE_BYTES), + // Multiple tiles: exercises the double-buffered loop + static_cast(TDM_TILE_BYTES * 4), + // Non-power-of-two: partial remainder handled by memcpy_wg + static_cast(TDM_TILE_BYTES * 3 + 1337) + ) +); diff --git a/projects/rocshmem/tests/unit_tests/tdm_raw_put_gtest.hpp b/projects/rocshmem/tests/unit_tests/tdm_raw_put_gtest.hpp new file mode 100644 index 000000000000..f600044e9401 --- /dev/null +++ b/projects/rocshmem/tests/unit_tests/tdm_raw_put_gtest.hpp @@ -0,0 +1,195 @@ +/****************************************************************************** + * Copyright (c) Advanced Micro Devices, Inc. All rights reserved. + * + * SPDX-License-Identifier: MIT + *****************************************************************************/ + +// Phase 1: Raw byte-array put tests exercising memcpy_wg_tdm via ipcCopy_wg. +// +// Requires 2 MPI ranks. PE 0 fills a symmetric buffer with an iota byte +// pattern and calls ipcCopy_wg to transfer it to PE 1. When USE_TDM is +// defined and the hardware is gfx1250, ipcCopy_wg internally calls +// memcpy_wg_tdm; otherwise it falls back to memcpy_wg. +// +// The kernel registers its dynamic LDS buffer via rocshmem_set_tdm_lds(), +// sized at launch via rocshmem_query_tdm_lds_bytes() -- the same public API +// any user kernel would use to opt into the TDM fast path. +// +// Build: cmake -DUSE_TDM=ON -DUSE_IPC=ON -DAMDGPU_TARGETS=gfx1250 ... +// Run: mpirun -n 2 ./tests/unit_tests/rocshmem_unit_tests \ +// --gtest_filter="TdmRaw*" + +#ifndef ROCSHMEM_TDM_RAW_PUT_GTEST_HPP +#define ROCSHMEM_TDM_RAW_PUT_GTEST_HPP + +#include "gtest/gtest.h" + +#include +#include +#include +#include +#include + +#include "rocshmem/rocshmem.hpp" +#include "../src/atomic.hpp" +#include "../src/constmem.hpp" +#include "../src/ipc_policy.hpp" +#include "../src/memory/notifier.hpp" +#include "../src/memory/symmetric_heap.hpp" +#include "../src/util.hpp" +#include "ipc_test_config.hpp" + +namespace rocshmem { + +static constexpr uint32_t TDM_RAW_SIGNAL_OFFSET = 67108864u; // 64 MB + +// Fixed TDM tile size for these tests +static constexpr uint32_t TEST_TDM_TILE_BYTES = 16384; + +// --------------------------------------------------------------------------- +// Kernels +// --------------------------------------------------------------------------- + +// PE 1 validator: spins on a signal word, then checks data. +template +__global__ +void tdm_raw_validator(bool* error, const uint8_t* golden, const uint8_t* dest, + size_t bytes, NotifierT* notifier) { + if (!get_flat_id()) { + uint32_t* sig = reinterpret_cast( + const_cast(dest) + TDM_RAW_SIGNAL_OFFSET); + while (detail::atomic::load(sig) == 0) {} + } + notifier->sync(); + for (size_t i = get_flat_id(); i < bytes; i += get_flat_grid_size()) { + if (golden[i] != dest[i]) { + printf("mismatch@%zu exp=%u got=%u\n", i, (unsigned)golden[i], + (unsigned)dest[i]); + *error = true; + } + } +} + +// PE 0 sender: registers its dynamic LDS buffer for TDM (a no-op on builds/ +// archs where TDM isn't available) and calls the real library path. +template +__global__ +void tdm_raw_put_kernel(IpcImplT* ipc_impl, const uint8_t* src, + uint8_t* remote_dst, size_t bytes) { + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + + ipc_impl->template ipcCopy_wg( + remote_dst, const_cast(src), bytes, /*local_pe=*/0); + + if (is_thread_zero_in_block()) { + ipc_impl->ipcFence(); + uint32_t* sig = reinterpret_cast(remote_dst + TDM_RAW_SIGNAL_OFFSET); + ipc_impl->ipcAMOFetchAdd(sig, uint32_t{1}); + } +} + +// --------------------------------------------------------------------------- +// Fixture +// --------------------------------------------------------------------------- + +class TdmRawPutFixture : public ::testing::TestWithParam { + using IpcImplT = IpcImpl; // whatever IpcImpl is (IpcOnImpl when USE_IPC) + using MPI_T = RemoteHeapInfo; + using NotifierT = Notifier; + using NotifierProxyT = NotifierProxy; + + public: + TdmRawPutFixture() { + MPIInstance::mpilib_dl_init(); + hip_allocator_ = get_default_allocator(); + heap_mem_ = new HeapMemoryType(*hip_allocator_, envvar::heap_size.get_value()); + assert(heap_mem_); + mpi_ = new MPI_T(heap_mem_->get_ptr(), heap_mem_->get_size(), MPI_COMM_WORLD); + ipc_impl_.ipcHostInit(mpi_->my_pe(), mpi_->get_heap_bases(), MPI_COMM_WORLD); + // Force a known tile size so the parameterized byte counts below + // actually land on the tile boundaries their comments claim. + uint32_t tile_bytes = TEST_TDM_TILE_BYTES; + CHECK_HIP(hipMemcpyToSymbol(HIP_SYMBOL(constmem), &tile_bytes, sizeof(tile_bytes), + offsetof(constmem_t, tdm_tile_bytes))); + hip_allocator_->allocate((void**)&ipc_impl_dptr_, sizeof(IpcImplT)); + CHECK_HIP(hipMemcpy(ipc_impl_dptr_, &ipc_impl_, sizeof(IpcImplT), + hipMemcpyHostToDevice)); + hip_allocator_->allocate((void**)&error_dptr_, sizeof(bool)); + CHECK_HIP(hipMemset(error_dptr_, 0, sizeof(bool))); + } + + ~TdmRawPutFixture() override { + if (ipc_impl_dptr_) hip_allocator_->deallocate(ipc_impl_dptr_); + if (error_dptr_) hip_allocator_->deallocate(error_dptr_); + if (golden_dptr_) hip_allocator_->deallocate(golden_dptr_); + ipc_impl_.ipcHostStop(); + delete mpi_; + delete heap_mem_; + MPIInstance::mpilib_dl_close(); + } + + void run_put(size_t bytes) { + assert(bytes + TDM_RAW_SIGNAL_OFFSET + sizeof(uint32_t) <= + heap_mem_->get_size()); + + // Build iota pattern on host, upload to a device buffer. + std::vector host_gold(bytes); + for (size_t i = 0; i < bytes; i++) host_gold[i] = static_cast(i & 0xFF); + hip_allocator_->allocate((void**)&golden_dptr_, bytes); + CHECK_HIP(hipMemcpy(golden_dptr_, host_gold.data(), bytes, hipMemcpyHostToDevice)); + + const int my_pe = mpi_->my_pe(); + + if (my_pe == 0) { + auto* local_heap = reinterpret_cast(ipc_impl_.ipc_bases[0]); + CHECK_HIP(hipMemcpy(local_heap, golden_dptr_, bytes, hipMemcpyDeviceToDevice)); + auto* remote_heap = reinterpret_cast(ipc_impl_.ipc_bases[1]); + CHECK_HIP(hipMemset(remote_heap + TDM_RAW_SIGNAL_OFFSET, 0, sizeof(uint32_t))); + } + + mpi_->barrier(); + + // Block size: full wavefront for TDM to have at least 1 wave. + const dim3 block(WF_SIZE); + const size_t smem = rocshmem_query_tdm_lds_bytes(); + + if (my_pe == 0) { + auto* src = reinterpret_cast(ipc_impl_.ipc_bases[0]); + auto* rdst = reinterpret_cast(ipc_impl_.ipc_bases[1]); + tdm_raw_put_kernel + <<>>(ipc_impl_dptr_, src, rdst, bytes); + CHECK_HIP(hipStreamSynchronize(nullptr)); + } else { + auto* dest = reinterpret_cast(ipc_impl_.ipc_bases[1]); + tdm_raw_validator + <<>>(error_dptr_, golden_dptr_, dest, bytes, + notifier_.get()); + CHECK_HIP(hipStreamSynchronize(nullptr)); + } + + mpi_->barrier(); + + if (my_pe == 1) { + bool host_error = false; + CHECK_HIP(hipMemcpy(&host_error, error_dptr_, sizeof(bool), + hipMemcpyDeviceToHost)); + ASSERT_EQ(host_error, false); + } + } + + private: + HIPAllocator* hip_allocator_{nullptr}; + NotifierProxyT notifier_{}; + HeapMemoryType* heap_mem_{nullptr}; + MPI_T* mpi_{nullptr}; + IpcImplT ipc_impl_{}; + IpcImplT* ipc_impl_dptr_{nullptr}; + bool* error_dptr_{nullptr}; + uint8_t* golden_dptr_{nullptr}; +}; + +} // namespace rocshmem + +#endif // ROCSHMEM_TDM_RAW_PUT_GTEST_HPP diff --git a/projects/rocshmem/tests/unit_tests/tdm_struct_put_gtest.cpp b/projects/rocshmem/tests/unit_tests/tdm_struct_put_gtest.cpp new file mode 100644 index 000000000000..0f50f20da88e --- /dev/null +++ b/projects/rocshmem/tests/unit_tests/tdm_struct_put_gtest.cpp @@ -0,0 +1,32 @@ +/****************************************************************************** + * Copyright (c) Advanced Micro Devices, Inc. All rights reserved. + * + * SPDX-License-Identifier: MIT + *****************************************************************************/ + +#include "tdm_struct_put_gtest.hpp" + +using namespace rocshmem; + +TEST_P(TdmStructPutFixture, put_records) { + run_put(GetParam()); +} + +// Fixed TDM tile size (see TEST_TDM_TILE_BYTES) expressed in whole +// TdmTestRecords (32 bytes each). +static constexpr size_t TDM_TILE_RECORDS = TEST_TDM_TILE_BYTES / sizeof(TdmTestRecord); // 128 + +INSTANTIATE_TEST_SUITE_P( + TdmStructPut, + TdmStructPutFixture, + ::testing::Values( + // Sub-tile: below TDM_TILE_RECORDS, exercises memcpy_wg fallback + static_cast(64), + // Exactly one tile + static_cast(TDM_TILE_RECORDS), + // Four tiles + static_cast(4 * TDM_TILE_RECORDS), + // Non-aligned: partial remainder + static_cast(3 * TDM_TILE_RECORDS + 17) + ) +); diff --git a/projects/rocshmem/tests/unit_tests/tdm_struct_put_gtest.hpp b/projects/rocshmem/tests/unit_tests/tdm_struct_put_gtest.hpp new file mode 100644 index 000000000000..3820b0b12563 --- /dev/null +++ b/projects/rocshmem/tests/unit_tests/tdm_struct_put_gtest.hpp @@ -0,0 +1,208 @@ +/****************************************************************************** + * Copyright (c) Advanced Micro Devices, Inc. All rights reserved. + * + * SPDX-License-Identifier: MIT + *****************************************************************************/ + +// Phase 2: Structured-data put via ipcCopy_wg / memcpy_wg_tdm. +// Transfers arrays of TdmTestRecord (32 bytes each) and validates field-by-field. +// The kernel registers its dynamic LDS buffer via rocshmem_set_tdm_lds(), +// sized at launch via rocshmem_query_tdm_lds_bytes() -- the same public API +// any user kernel would use to opt into the TDM fast path. + +#ifndef ROCSHMEM_TDM_STRUCT_PUT_GTEST_HPP +#define ROCSHMEM_TDM_STRUCT_PUT_GTEST_HPP + +#include "gtest/gtest.h" + +#include +#include +#include +#include +#include + +#include "rocshmem/rocshmem.hpp" +#include "../src/atomic.hpp" +#include "../src/constmem.hpp" +#include "../src/ipc_policy.hpp" +#include "../src/memory/notifier.hpp" +#include "../src/memory/symmetric_heap.hpp" +#include "../src/util.hpp" +#include "ipc_test_config.hpp" + +namespace rocshmem { + +struct __attribute__((packed)) TdmTestRecord { + int id; + float val[7]; +}; +static_assert(sizeof(TdmTestRecord) == 32, "TdmTestRecord must be 32 bytes"); + +static constexpr uint32_t TDM_STRUCT_SIGNAL_OFFSET = 67108864u; + +// Fixed TDM tile size for these tests +static constexpr uint32_t TEST_TDM_TILE_BYTES = 4096; + +// --------------------------------------------------------------------------- +// Kernels +// --------------------------------------------------------------------------- + +template +__global__ +void tdm_struct_validator(bool* error, const TdmTestRecord* golden, + const TdmTestRecord* dest, size_t count, + NotifierT* notifier) { + if (!get_flat_id()) { + uint32_t* sig = reinterpret_cast( + const_cast(reinterpret_cast(dest)) + + TDM_STRUCT_SIGNAL_OFFSET); + while (detail::atomic::load(sig) == 0) {} + } + notifier->sync(); + for (size_t i = get_flat_id(); i < count; i += get_flat_grid_size()) { + if (dest[i].id != golden[i].id) { + printf("record[%zu].id exp=%d got=%d\n", i, golden[i].id, dest[i].id); + *error = true; + } + for (int f = 0; f < 7; f++) { + if (dest[i].val[f] != golden[i].val[f]) { + printf("record[%zu].val[%d] exp=%f got=%f\n", + i, f, (double)golden[i].val[f], (double)dest[i].val[f]); + *error = true; + } + } + } +} + +template +__global__ +void tdm_struct_put_kernel(IpcImplT* ipc_impl, const TdmTestRecord* src, + TdmTestRecord* remote_dst, size_t count) { + const size_t bytes = count * sizeof(TdmTestRecord); + + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + + ipc_impl->template ipcCopy_wg( + remote_dst, const_cast(src), bytes, /*local_pe=*/0); + + if (is_thread_zero_in_block()) { + ipc_impl->ipcFence(); + uint32_t* sig = reinterpret_cast( + reinterpret_cast(remote_dst) + TDM_STRUCT_SIGNAL_OFFSET); + ipc_impl->ipcAMOFetchAdd(sig, uint32_t{1}); + } +} + +// --------------------------------------------------------------------------- +// Fixture +// --------------------------------------------------------------------------- + +class TdmStructPutFixture : public ::testing::TestWithParam { + using IpcImplT = IpcImpl; + using MPI_T = RemoteHeapInfo; + using NotifierT = Notifier; + using NotifierProxyT = NotifierProxy; + + public: + TdmStructPutFixture() { + MPIInstance::mpilib_dl_init(); + hip_allocator_ = get_default_allocator(); + heap_mem_ = new HeapMemoryType(*hip_allocator_, envvar::heap_size.get_value()); + assert(heap_mem_); + mpi_ = new MPI_T(heap_mem_->get_ptr(), heap_mem_->get_size(), MPI_COMM_WORLD); + ipc_impl_.ipcHostInit(mpi_->my_pe(), mpi_->get_heap_bases(), MPI_COMM_WORLD); + // Force a known tile size so the parameterized record counts below + // actually land on the tile boundaries their comments claim. + uint32_t tile_bytes = TEST_TDM_TILE_BYTES; + CHECK_HIP(hipMemcpyToSymbol(HIP_SYMBOL(constmem), &tile_bytes, sizeof(tile_bytes), + offsetof(constmem_t, tdm_tile_bytes))); + hip_allocator_->allocate((void**)&ipc_impl_dptr_, sizeof(IpcImplT)); + CHECK_HIP(hipMemcpy(ipc_impl_dptr_, &ipc_impl_, sizeof(IpcImplT), + hipMemcpyHostToDevice)); + hip_allocator_->allocate((void**)&error_dptr_, sizeof(bool)); + CHECK_HIP(hipMemset(error_dptr_, 0, sizeof(bool))); + } + + ~TdmStructPutFixture() override { + if (ipc_impl_dptr_) hip_allocator_->deallocate(ipc_impl_dptr_); + if (error_dptr_) hip_allocator_->deallocate(error_dptr_); + if (golden_dptr_) hip_allocator_->deallocate(golden_dptr_); + ipc_impl_.ipcHostStop(); + delete mpi_; + delete heap_mem_; + MPIInstance::mpilib_dl_close(); + } + + void run_put(size_t record_count) { + const size_t bytes = record_count * sizeof(TdmTestRecord); + assert(bytes + TDM_STRUCT_SIGNAL_OFFSET + sizeof(uint32_t) <= + heap_mem_->get_size()); + + std::vector host_gold(record_count); + for (size_t i = 0; i < record_count; i++) { + host_gold[i].id = static_cast(i); + for (int f = 0; f < 7; f++) + host_gold[i].val[f] = static_cast(i * 7 + f) * 0.1f; + } + hip_allocator_->allocate((void**)&golden_dptr_, + record_count * sizeof(TdmTestRecord)); + CHECK_HIP(hipMemcpy(golden_dptr_, host_gold.data(), + record_count * sizeof(TdmTestRecord), + hipMemcpyHostToDevice)); + + const int my_pe = mpi_->my_pe(); + + if (my_pe == 0) { + auto* local_heap = reinterpret_cast(ipc_impl_.ipc_bases[0]); + CHECK_HIP(hipMemcpy(local_heap, golden_dptr_, + record_count * sizeof(TdmTestRecord), + hipMemcpyDeviceToDevice)); + auto* remote_heap = reinterpret_cast(ipc_impl_.ipc_bases[1]); + CHECK_HIP(hipMemset(remote_heap + TDM_STRUCT_SIGNAL_OFFSET, 0, sizeof(uint32_t))); + } + + mpi_->barrier(); + + const dim3 block(WF_SIZE); + const size_t smem = rocshmem_query_tdm_lds_bytes(); + + if (my_pe == 0) { + auto* src = reinterpret_cast(ipc_impl_.ipc_bases[0]); + auto* rdst = reinterpret_cast(ipc_impl_.ipc_bases[1]); + tdm_struct_put_kernel + <<>>(ipc_impl_dptr_, src, rdst, record_count); + CHECK_HIP(hipStreamSynchronize(nullptr)); + } else { + auto* dest = reinterpret_cast(ipc_impl_.ipc_bases[1]); + tdm_struct_validator + <<>>(error_dptr_, golden_dptr_, dest, + record_count, notifier_.get()); + CHECK_HIP(hipStreamSynchronize(nullptr)); + } + + mpi_->barrier(); + + if (my_pe == 1) { + bool host_error = false; + CHECK_HIP(hipMemcpy(&host_error, error_dptr_, sizeof(bool), + hipMemcpyDeviceToHost)); + ASSERT_EQ(host_error, false); + } + } + + private: + HIPAllocator* hip_allocator_{nullptr}; + NotifierProxyT notifier_{}; + HeapMemoryType* heap_mem_{nullptr}; + MPI_T* mpi_{nullptr}; + IpcImplT ipc_impl_{}; + IpcImplT* ipc_impl_dptr_{nullptr}; + bool* error_dptr_{nullptr}; + TdmTestRecord* golden_dptr_{nullptr}; +}; + +} // namespace rocshmem + +#endif // ROCSHMEM_TDM_STRUCT_PUT_GTEST_HPP diff --git a/projects/rocshmem/tests/unit_tests/tdm_wave_put_gtest.cpp b/projects/rocshmem/tests/unit_tests/tdm_wave_put_gtest.cpp new file mode 100644 index 000000000000..c8df294af433 --- /dev/null +++ b/projects/rocshmem/tests/unit_tests/tdm_wave_put_gtest.cpp @@ -0,0 +1,36 @@ +/****************************************************************************** + * Copyright (c) Advanced Micro Devices, Inc. All rights reserved. + * + * SPDX-License-Identifier: MIT + *****************************************************************************/ + +#include "tdm_wave_put_gtest.hpp" + +using namespace rocshmem; + +TEST_P(TdmWavePutFixture, put_bytes) { + auto [bytes_per_wave, num_waves] = GetParam(); + run_put(bytes_per_wave, num_waves); +} + +static constexpr size_t TDM_TILE_BYTES = TEST_TDM_TILE_BYTES; + +INSTANTIATE_TEST_SUITE_P( + TdmWavePut, + TdmWavePutFixture, + ::testing::Combine( + ::testing::Values( + // Sub-tile regardless of split: exercises memcpy_wave fallback + static_cast(1024), + // At least one tile per wave even when evenly split several ways + static_cast(TDM_TILE_BYTES), + // Several tiles per wave: exercises the double-buffered loop + static_cast(TDM_TILE_BYTES * 4), + // Non-power-of-two: partial remainder handled by memcpy_wave + static_cast(TDM_TILE_BYTES * 3 + 1337) + ), + // 1 wave (degenerate: no splitting), a power-of-two split, and an + // uneven split (reg.bytes doesn't divide evenly by 3). + ::testing::Values(1, 4, 3) + ) +); diff --git a/projects/rocshmem/tests/unit_tests/tdm_wave_put_gtest.hpp b/projects/rocshmem/tests/unit_tests/tdm_wave_put_gtest.hpp new file mode 100644 index 000000000000..d3b30bba21b1 --- /dev/null +++ b/projects/rocshmem/tests/unit_tests/tdm_wave_put_gtest.hpp @@ -0,0 +1,205 @@ +/****************************************************************************** + * Copyright (c) Advanced Micro Devices, Inc. All rights reserved. + * + * SPDX-License-Identifier: MIT + *****************************************************************************/ + +// Wave-granularity put tests exercising memcpy_wave_tdm via ipcCopy_wave. +// +// Requires 2 MPI ranks. PE 0 launches one block whose waves each call +// ipcCopy_wave concurrently on disjoint slices of a shared symmetric +// buffer, exercising the per-wave LDS slicing that keeps concurrent waves +// from tiling through each other's bytes. The kernel opts in the same way +// any user kernel would: rocshmem_set_tdm_lds() sized via +// rocshmem_query_tdm_lds_bytes(). +// +// Build: cmake -DUSE_TDM=ON -DUSE_IPC=ON -DAMDGPU_TARGETS=gfx1250 ... +// Run: mpirun -n 2 ./tests/unit_tests/rocshmem_unit_tests \ +// --gtest_filter="TdmWave*" + +#ifndef ROCSHMEM_TDM_WAVE_PUT_GTEST_HPP +#define ROCSHMEM_TDM_WAVE_PUT_GTEST_HPP + +#include "gtest/gtest.h" + +#include +#include +#include +#include +#include +#include + +#include "rocshmem/rocshmem.hpp" +#include "../src/atomic.hpp" +#include "../src/constmem.hpp" +#include "../src/ipc_policy.hpp" +#include "../src/memory/notifier.hpp" +#include "../src/memory/symmetric_heap.hpp" +#include "../src/util.hpp" +#include "ipc_test_config.hpp" + +namespace rocshmem { + +static constexpr uint32_t TDM_WAVE_SIGNAL_OFFSET = 67108864u; // 64 MB + +// Fixed TDM tile size for these tests +static constexpr uint32_t TEST_TDM_TILE_BYTES = 16384; + +// --------------------------------------------------------------------------- +// Kernels +// --------------------------------------------------------------------------- + +// PE 1 validator: spins on a signal word, then checks data. +template +__global__ +void tdm_wave_validator(bool* error, const uint8_t* golden, const uint8_t* dest, + size_t bytes, NotifierT* notifier) { + if (!get_flat_id()) { + uint32_t* sig = reinterpret_cast( + const_cast(dest) + TDM_WAVE_SIGNAL_OFFSET); + while (detail::atomic::load(sig) == 0) {} + } + notifier->sync(); + for (size_t i = get_flat_id(); i < bytes; i += get_flat_grid_size()) { + if (golden[i] != dest[i]) { + printf("mismatch@%zu exp=%u got=%u\n", i, (unsigned)golden[i], + (unsigned)dest[i]); + *error = true; + } + } +} + +// PE 0 sender: every wave in the block puts its own disjoint +// [wave_id * bytes_per_wave, (wave_id+1) * bytes_per_wave) slice, +// concurrently with every other wave. +template +__global__ +void tdm_wave_put_kernel(IpcImplT* ipc_impl, const uint8_t* src, + uint8_t* remote_dst, size_t bytes_per_wave) { + extern __shared__ uint8_t tdm_lds[]; + rocshmem_set_tdm_lds(tdm_lds, rocshmem_query_tdm_lds_bytes()); + + const int wave_id = get_flat_block_id() / WF_SIZE; + const size_t offset = static_cast(wave_id) * bytes_per_wave; + + ipc_impl->template ipcCopy_wave( + remote_dst + offset, const_cast(src) + offset, bytes_per_wave, + /*local_pe=*/0); + + // Waves finish their (independently timed) TDM copies at different + // points; every wave must land before any lane signals PE 1. + __builtin_amdgcn_s_barrier(); + + if (is_thread_zero_in_block()) { + ipc_impl->ipcFence(); + uint32_t* sig = reinterpret_cast(remote_dst + TDM_WAVE_SIGNAL_OFFSET); + ipc_impl->ipcAMOFetchAdd(sig, uint32_t{1}); + } +} + +// --------------------------------------------------------------------------- +// Fixture +// --------------------------------------------------------------------------- + +// Params: (bytes_per_wave, num_waves). +class TdmWavePutFixture + : public ::testing::TestWithParam> { + using IpcImplT = IpcImpl; // whatever IpcImpl is (IpcOnImpl when USE_IPC) + using MPI_T = RemoteHeapInfo; + using NotifierT = Notifier; + using NotifierProxyT = NotifierProxy; + + public: + TdmWavePutFixture() { + MPIInstance::mpilib_dl_init(); + hip_allocator_ = get_default_allocator(); + heap_mem_ = new HeapMemoryType(*hip_allocator_, envvar::heap_size.get_value()); + assert(heap_mem_); + mpi_ = new MPI_T(heap_mem_->get_ptr(), heap_mem_->get_size(), MPI_COMM_WORLD); + ipc_impl_.ipcHostInit(mpi_->my_pe(), mpi_->get_heap_bases(), MPI_COMM_WORLD); + // Force a known tile size so the parameterized byte counts below + // actually land on tile/slice boundaries their comments claim. + uint32_t tile_bytes = TEST_TDM_TILE_BYTES; + CHECK_HIP(hipMemcpyToSymbol(HIP_SYMBOL(constmem), &tile_bytes, sizeof(tile_bytes), + offsetof(constmem_t, tdm_tile_bytes))); + hip_allocator_->allocate((void**)&ipc_impl_dptr_, sizeof(IpcImplT)); + CHECK_HIP(hipMemcpy(ipc_impl_dptr_, &ipc_impl_, sizeof(IpcImplT), + hipMemcpyHostToDevice)); + hip_allocator_->allocate((void**)&error_dptr_, sizeof(bool)); + CHECK_HIP(hipMemset(error_dptr_, 0, sizeof(bool))); + } + + ~TdmWavePutFixture() override { + if (ipc_impl_dptr_) hip_allocator_->deallocate(ipc_impl_dptr_); + if (error_dptr_) hip_allocator_->deallocate(error_dptr_); + if (golden_dptr_) hip_allocator_->deallocate(golden_dptr_); + ipc_impl_.ipcHostStop(); + delete mpi_; + delete heap_mem_; + MPIInstance::mpilib_dl_close(); + } + + void run_put(size_t bytes_per_wave, int num_waves) { + const size_t bytes = bytes_per_wave * static_cast(num_waves); + assert(bytes + TDM_WAVE_SIGNAL_OFFSET + sizeof(uint32_t) <= + heap_mem_->get_size()); + + // Build iota pattern on host, upload to a device buffer. + std::vector host_gold(bytes); + for (size_t i = 0; i < bytes; i++) host_gold[i] = static_cast(i & 0xFF); + hip_allocator_->allocate((void**)&golden_dptr_, bytes); + CHECK_HIP(hipMemcpy(golden_dptr_, host_gold.data(), bytes, hipMemcpyHostToDevice)); + + const int my_pe = mpi_->my_pe(); + + if (my_pe == 0) { + auto* local_heap = reinterpret_cast(ipc_impl_.ipc_bases[0]); + CHECK_HIP(hipMemcpy(local_heap, golden_dptr_, bytes, hipMemcpyDeviceToDevice)); + auto* remote_heap = reinterpret_cast(ipc_impl_.ipc_bases[1]); + CHECK_HIP(hipMemset(remote_heap + TDM_WAVE_SIGNAL_OFFSET, 0, sizeof(uint32_t))); + } + + mpi_->barrier(); + + const dim3 block(static_cast(WF_SIZE * num_waves)); + const size_t smem = rocshmem_query_tdm_lds_bytes(); + + if (my_pe == 0) { + auto* src = reinterpret_cast(ipc_impl_.ipc_bases[0]); + auto* rdst = reinterpret_cast(ipc_impl_.ipc_bases[1]); + tdm_wave_put_kernel + <<>>(ipc_impl_dptr_, src, rdst, bytes_per_wave); + CHECK_HIP(hipStreamSynchronize(nullptr)); + } else { + auto* dest = reinterpret_cast(ipc_impl_.ipc_bases[1]); + tdm_wave_validator + <<>>(error_dptr_, golden_dptr_, dest, bytes, + notifier_.get()); + CHECK_HIP(hipStreamSynchronize(nullptr)); + } + + mpi_->barrier(); + + if (my_pe == 1) { + bool host_error = false; + CHECK_HIP(hipMemcpy(&host_error, error_dptr_, sizeof(bool), + hipMemcpyDeviceToHost)); + ASSERT_EQ(host_error, false); + } + } + + private: + HIPAllocator* hip_allocator_{nullptr}; + NotifierProxyT notifier_{}; + HeapMemoryType* heap_mem_{nullptr}; + MPI_T* mpi_{nullptr}; + IpcImplT ipc_impl_{}; + IpcImplT* ipc_impl_dptr_{nullptr}; + bool* error_dptr_{nullptr}; + uint8_t* golden_dptr_{nullptr}; +}; + +} // namespace rocshmem + +#endif // ROCSHMEM_TDM_WAVE_PUT_GTEST_HPP