From 219bf3dc4982215bd76a131e26cfa1b6d7d3e858 Mon Sep 17 00:00:00 2001 From: Alex Breslow Date: Tue, 21 Jul 2026 15:56:54 -0700 Subject: [PATCH 1/8] Add first pass at async logic --- src/header/TransferBench.hpp | 11 +- src/header/asyncCopy.h | 460 +++++++++++++++++++++++++++++++++++ src/header/cachePolicy.h | 97 ++++++++ 3 files changed, 565 insertions(+), 3 deletions(-) create mode 100644 src/header/asyncCopy.h create mode 100644 src/header/cachePolicy.h diff --git a/src/header/TransferBench.hpp b/src/header/TransferBench.hpp index 00be3a1..9dd1367 100644 --- a/src/header/TransferBench.hpp +++ b/src/header/TransferBench.hpp @@ -80,6 +80,7 @@ THE SOFTWARE. #endif #include "tdmCopy.h" +#include "asyncCopy.h" #endif /// @endcond @@ -5870,10 +5871,14 @@ namespace { size_t const sizeBytes = p.N * sizeof(float); int subIterations = 0; - while (1) { - tdm::tdmCopy(dst, src, sizeBytes, shmem, ldsBytes); + bool useTDM = 0; + for(int subIterations = 0; subIterations < numSubIterations; subIterations++) { + if constexpr(useTDM){ + tdm::tdmCopy(dst, src, sizeBytes, shmem, ldsBytes); + else{ + async::tdmCopy(dst, src, sizeBytes, shmem, ldsBytes); + } __syncthreads(); // Wait for all warps to finish - if (++subIterations == numSubIterations) break; } if (shouldRecordTiming) { diff --git a/src/header/asyncCopy.h b/src/header/asyncCopy.h new file mode 100644 index 0000000..c0b070c --- /dev/null +++ b/src/header/asyncCopy.h @@ -0,0 +1,460 @@ +/* +Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. + +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 __ASYNCCOPY_H +#define __ASYNCCOPY_H + +#include "cachePolicy.h" + +enum struct SyncPolicy : uint32_t { + Async, + Sync, +}; + +constexpr SyncPolicy DEFAULT_SYNC_POLICY = SyncPolicy::Async; + +// Used for setting TDM descriptor fields and arguments to the async load/store builtins +using __rccl_int32x2 = int32_t __attribute__((__vector_size__(8))); +using __rccl_int32x4 = int32_t __attribute__((__vector_size__(16))); +using __rccl_int32x8 = int32_t __attribute__((__vector_size__(32))); + +// The builtins take an immediate byte offset that is folded into the memory instruction, so it must be a +// compile-time constant. We always pass a zero immediate and instead apply the per-lane offset directly to +// the src/dst pointers, which lets every lane target a different chunk of the transfer. +namespace { + constexpr int32_t ZERO_OFFSET = 0; +} + +// Waits until at most WAIT_CNT outstanding async-to/from-LDS transfers remain in flight. The count +// is baked into the hardware instruction, so it must be a compile-time constant. +template +__device__ void asyncWait(){ + __builtin_amdgcn_s_wait_asynccnt(WAIT_CNT); +} + +/* Async load/Store APIs */ +// The async-to/from-LDS builtins move a single b8/b32/b64/b128 access per lane between global memory and LDS. +// A whole warp issues one instruction, so a warp moves (warpSize * accessWidth) bytes at a time, with each +// lane targeting its own slice via a per-lane pointer offset. The hardware only provides b8/b32/b64/b128 +// variants -- there is no 16-bit per lane instruction. +namespace async_detail { + +// Direction of the transfer. In both cases the first pointer is the global-memory side and the second is the +// LDS side, which keeps the per-width helpers below symmetric. +enum struct AsyncDir { Load, Store }; + +using __rccl_global_char = __attribute__((address_space(1))) char; +using __rccl_lds_char = __attribute__((address_space(3))) char; +using __rccl_global_int32 = __attribute__((address_space(1))) int32_t; +using __rccl_lds_int32 = __attribute__((address_space(3))) int32_t; +using __rccl_global_int32x2 = __attribute__((address_space(1))) __rccl_int32x2; +using __rccl_lds_int32x2 = __attribute__((address_space(3))) __rccl_int32x2; +using __rccl_global_int32x4 = __attribute__((address_space(1))) __rccl_int32x4; +using __rccl_lds_int32x4 = __attribute__((address_space(3))) __rccl_int32x4; + +// One b8 access (1 byte) for a single lane. +template +__device__ inline void asyncCopyB8(const uint8_t* global, uint8_t* lds){ + if constexpr (DIR == AsyncDir::Load) { + __builtin_amdgcn_global_load_async_to_lds_b8( + (__rccl_global_char*)global, (__rccl_lds_char*)lds, ZERO_OFFSET, cp); + } else { + __builtin_amdgcn_global_store_async_from_lds_b8( + (__rccl_global_char*)global, (__rccl_lds_char*)lds, ZERO_OFFSET, cp); + } +} + +// One b32 access (4 bytes) for a single lane. +template +__device__ inline void asyncCopyB32(const uint8_t* global, uint8_t* lds){ + if constexpr (DIR == AsyncDir::Load) { + __builtin_amdgcn_global_load_async_to_lds_b32( + (__rccl_global_int32*)global, (__rccl_lds_int32*)lds, ZERO_OFFSET, cp); + } else { + __builtin_amdgcn_global_store_async_from_lds_b32( + (__rccl_global_int32*)global, (__rccl_lds_int32*)lds, ZERO_OFFSET, cp); + } +} + +// One b64 access (8 bytes) for a single lane. +template +__device__ inline void asyncCopyB64(const uint8_t* global, uint8_t* lds){ + if constexpr (DIR == AsyncDir::Load) { + __builtin_amdgcn_global_load_async_to_lds_b64( + (__rccl_global_int32x2*)global, (__rccl_lds_int32x2*)lds, ZERO_OFFSET, cp); + } else { + __builtin_amdgcn_global_store_async_from_lds_b64( + (__rccl_global_int32x2*)global, (__rccl_lds_int32x2*)lds, ZERO_OFFSET, cp); + } +} + +// One b128 access (16 bytes) for a single lane. +template +__device__ inline void asyncCopyB128(const uint8_t* global, uint8_t* lds){ + if constexpr (DIR == AsyncDir::Load) { + __builtin_amdgcn_global_load_async_to_lds_b128( + (__rccl_global_int32x4*)global, (__rccl_lds_int32x4*)lds, ZERO_OFFSET, cp); + } else { + __builtin_amdgcn_global_store_async_from_lds_b128( + (__rccl_global_int32x4*)global, (__rccl_lds_int32x4*)lds, ZERO_OFFSET, cp); + } +} + +// Bytes moved by a single lane in the widest (b128) access. +constexpr size_t BYTES_PER_LANE_B128 = 16; +// The hot loop must start on a 128-byte boundary, so the peel aligns the global pointer to this. +constexpr size_t NATURAL_ALIGNMENT_BYTES = 128; +// How many b128 iterations the hot loop unrolls at a time. +constexpr int HOT_LOOP_UNROLL = 4; + +// Warp-cooperative byte copy between global memory and LDS. The whole warp must call this with identical +// `global`, `lds`, and `sizeInBytes` arguments -- the per-lane work distribution is derived internally. +// +// The copy is split into three phases: +// 1. Peel the leading bytes so the global pointer reaches 128-byte (cache-line) alignment: first with b8 +// accesses up to a 16-byte boundary, then with b128 accesses up to the 128-byte boundary. +// 2. Run a hot, unrolled loop of full-warp b128 transfers (warpSize * 16 bytes per iteration), starting on +// the 128-byte aligned boundary produced by phase 1. +// 3. Drain the tail (< warpSize * 16 bytes) with a descending b128 -> b64 -> b32 -> b8 cascade. +// +// Alignment is peeled against the global-memory pointer; the LDS staging buffer is assumed to share the same +// alignment (RCCL allocates it 128-byte aligned), so applying the same offset keeps the LDS side aligned too. +template +__device__ inline void warpAsyncCopy(const uint8_t* global, uint8_t* lds, size_t sizeInBytes){ + const unsigned lane = __lane_id(); + const size_t laneCount = (size_t)__builtin_amdgcn_wavefrontsize(); + const size_t laneByte = (size_t)lane * BYTES_PER_LANE_B128; + + size_t offset = 0; + + // Phase 1a: peel leading bytes until `global` is 16-byte aligned, one byte per lane. + const size_t misalign = (uintptr_t)global & (BYTES_PER_LANE_B128 - 1); + if (misalign != 0) { + const size_t head = (BYTES_PER_LANE_B128 - misalign) < sizeInBytes + ? (BYTES_PER_LANE_B128 - misalign) + : sizeInBytes; + if (lane < head) { + asyncCopyB8(global + lane, lds + lane); + } + offset += head; + } + + // Phase 1b: peel whole 16-byte chunks until `global + offset` reaches 128-byte alignment, one b128 per + // leading lane. `global + offset` is already 16-byte aligned, so the number of bytes left to the next + // 128-byte boundary is a multiple of 16 (at most 112 -> 7 chunks, well within a warp). + { + const size_t align128 = ((uintptr_t)global + offset) & (NATURAL_ALIGNMENT_BYTES - 1); + const size_t bytesToLine = (NATURAL_ALIGNMENT_BYTES - align128) & (NATURAL_ALIGNMENT_BYTES - 1); + size_t peelChunks = bytesToLine / BYTES_PER_LANE_B128; + const size_t availChunks = (sizeInBytes - offset) / BYTES_PER_LANE_B128; + if (peelChunks > availChunks) peelChunks = availChunks; + if (lane < peelChunks) { + const size_t laneOff = offset + laneByte; + asyncCopyB128(global + laneOff, lds + laneOff); + } + offset += peelChunks * BYTES_PER_LANE_B128; + } + + // Phase 2: hot unrolled loop of full-warp b128 transfers, now starting on a 128-byte boundary. + const size_t bytesPerIter = laneCount * BYTES_PER_LANE_B128; + size_t remaining = sizeInBytes - offset; + #pragma unroll HOT_LOOP_UNROLL + while (remaining >= bytesPerIter) { + const size_t laneOff = offset + laneByte; + asyncCopyB128(global + laneOff, lds + laneOff); + offset += bytesPerIter; + remaining -= bytesPerIter; + } + + // Phase 3a: the remaining full 16-byte chunks, one b128 per leading lane. `remaining < warpSize * 16`, so + // `full128 < warpSize` and every chunk maps to a distinct lane. + const size_t full128 = remaining / BYTES_PER_LANE_B128; + if (lane < full128) { + const size_t laneOff = offset + laneByte; + asyncCopyB128(global + laneOff, lds + laneOff); + } + offset += full128 * BYTES_PER_LANE_B128; + remaining -= full128 * BYTES_PER_LANE_B128; // now < 16 bytes + + // Phase 3b: 0..15 trailing bytes. `offset` and `remaining` are uniform across the warp here, so the + // descending cascade below is taken identically by every lane. Because `offset` is 16-byte aligned, the + // b64 (8-byte) and b32 (4-byte) accesses are naturally aligned. + if (remaining >= 8) { + if (lane == 0) { + asyncCopyB64(global + offset, lds + offset); + } + offset += 8; + remaining -= 8; + } + if (remaining >= 4) { + if (lane == 0) { + asyncCopyB32(global + offset, lds + offset); + } + offset += 4; + remaining -= 4; + } + // Final 1..3 bytes, one b8 per lane. + if (lane < remaining) { + asyncCopyB8(global + offset + lane, lds + offset + lane); + } +} + +} // namespace async_detail + +// Warp-level async copy from global memory into LDS. The entire warp calls this with the same arguments. +template +__device__ void asyncLoadToLDS(const uint8_t* globalSrc, uint8_t* ldsDst, size_t sizeInBytes){ + async_detail::warpAsyncCopy(globalSrc, ldsDst, sizeInBytes); + if constexpr (sp == SyncPolicy::Sync) { + asyncWait<0>(); + } +} + +// Warp-level async copy from LDS into global memory. The entire warp calls this with the same arguments. +template +__device__ void asyncStoreFromLDS(const uint8_t* ldsSrc, uint8_t* globalDst, size_t sizeInBytes){ + // warpAsyncCopy shares one non-const `lds` parameter across load and store; the store path only reads from + // it, so dropping const here is safe. + async_detail::warpAsyncCopy(globalDst, const_cast(ldsSrc), sizeInBytes); + if constexpr (sp == SyncPolicy::Sync) { + asyncWait<0>(); + } +} + +// ============================================================================ +// DROP-IN PARITY API (mirrors tdm/tdmCopy.h) +// ---------------------------------------------------------------------------- +// The entry points below expose the EXACT same public surface as tdm/tdmCopy.h +// -- same names, signatures, defaults, and memcpy-style semantics -- so the two +// libraries are interchangeable. The only difference is internal: this +// implementation stages HBM->LDS->HBM through the async-to/from-LDS builtins +// (see warpAsyncCopy above) instead of the tensor-data-mover load/store +// instructions. +// +// AVAILABILITY: the async-to/from-LDS builtins are a gfx1250 feature. As with +// tdmCopy.h, ASYNC_COPY_SUPPORTED is 1 only for a device pass on a capable arch +// whose compiler exposes the builtin; otherwise each entry point is `= delete`d, +// so including the header is fine but CALLING one on an unsupported target is a +// compile-time error at the call site. async::IsTdmCopySupported() is always +// callable (host and device) and is the runtime/host-side guard. +// ============================================================================ +#ifndef __has_builtin +# define __has_builtin(x) 0 +#endif + +#if defined(__gfx1250__) && __has_builtin(__builtin_amdgcn_global_load_async_to_lds_b128) + /* extend: || (defined(__gfxNNNN__) && ...) */ +# define ASYNC_COPY_SUPPORTED 1 +#else +# define ASYNC_COPY_SUPPORTED 0 +#endif + +#if ASYNC_COPY_SUPPORTED +# define ASYNC_API inline +# define ASYNC_DELETED +#else +# define ASYNC_API +# define ASYNC_DELETED = delete +#endif + +namespace async { + +/// \brief Report whether async-to/from-LDS copies are usable. Always available +/// (never deleted), so it is safe to call on any target as a guard. +/// \see tdm::IsTdmCopySupported for the identical query on the TDM library. +__host__ __device__ inline bool IsTdmCopySupported(int deviceId = 0) { +#if defined(__HIP_DEVICE_COMPILE__) + (void)deviceId; + return ASYNC_COPY_SUPPORTED; // compile-time constant for this arch pass +#else + hipDeviceProp_t prop; + if (hipGetDeviceProperties(&prop, deviceId) != hipSuccess) return false; + // gcnArchName looks like "gfx1250:sramecc+:xnack-"; match the arch prefix. + const char* arch = prop.gcnArchName; + const char* p = "gfx1250"; + while (*p && *arch == *p) { ++arch; ++p; } + return (*p == '\0'); +#endif +} + +#if ASYNC_COPY_SUPPORTED + +namespace detail { + +// LDS is subdivided into per-warp staging windows; each window is a multiple of +// this many bytes (matching the widest async b128 access / 128B cache line). +constexpr uint32_t WINDOW_GRAIN = 128; + +// Direct global->global byte copy by one warp, used only when the LDS staging +// area is too small to hold even a single window (the vector fallback, mirroring +// tdm::detail::warpVecCopy's role). Byte granularity keeps it correct for any +// source/destination alignment. +template +__device__ inline void warpGlobalCopy(const uint8_t* s, uint8_t* d, size_t n, + uint32_t warpThread, uint32_t warpThreads) { + for (size_t i = warpThread; i < n; i += warpThreads) d[i] = s[i]; +} + +// core: partition + issue the whole copy for the team [start, stop). Work + LDS +// partition by rank within the team (warpId - start), so each team indexes its +// own ldsBuffer from zero. Safe to call collectively (warps outside the range +// return immediately) or only from the team's warps. NO final wait -- each async +// tile already drains its own load/store (RAW/WAR) inside the loop, so nothing +// is left in flight on return; async::tdmWait() below is therefore a no-op guard +// kept for API parity with tdm::tdmWait(). +template +__device__ inline void issue(void* dst, const void* src, size_t sizeBytes, + void* ldsBuffer, size_t ldsBufferBytes, + uint32_t startWarpId, uint32_t stopWarpId) { + const uint8_t* s = reinterpret_cast(src); + uint8_t* d = reinterpret_cast(dst); + uint8_t* lds = reinterpret_cast(ldsBuffer); + const uint32_t ldsBytes = static_cast(ldsBufferBytes); // LDS is small + + const uint32_t W = warpSize; + const uint32_t nThreads = blockDim.x * blockDim.y * blockDim.z; + const uint32_t tid = (threadIdx.z * blockDim.y + threadIdx.y) * blockDim.x + + threadIdx.x; + const uint32_t warpThread = tid % W; // thread index within its warp + const uint32_t warpId = tid / W; + const uint32_t nWarps = (nThreads + W - 1) / W; + + // --- team membership: this warp participates iff in [start, stop) -------- + const uint32_t teamStop = (stopWarpId > nWarps) ? nWarps : stopWarpId; + if (startWarpId >= teamStop || warpId < startWarpId || warpId >= teamStop) + return; // not on this team + const uint32_t rank = warpId - startWarpId; // rank within the team + const uint32_t teamWarps = teamStop - startWarpId; // >= 1 + + // active threads in THIS warp (handles partial final warp); stride for the + // byte fallback below. + const uint32_t warpThreads = (nThreads - warpId * W < W) ? (nThreads - warpId * W) : W; + + if (sizeBytes == 0) return; + + // --- tiny-LDS fallback: not enough LDS for a single window --------------- + if (ldsBytes < WINDOW_GRAIN) { + if (rank == 0) warpGlobalCopy(s, d, sizeBytes, warpThread, warpThreads); + return; + } + + // --- give each issuing warp a 128B-multiple LDS window ------------------- + uint32_t maxIssuers = ldsBytes / WINDOW_GRAIN; // #warps we can give a window + uint32_t issuers = teamWarps < maxIssuers ? teamWarps : maxIssuers; + uint32_t window = (ldsBytes / issuers) & ~(WINDOW_GRAIN - 1); // per-warp 128B-multiple + if (rank >= issuers) return; // this warp doesn't issue + + // distribute the byte range across issuers by team rank (contiguous blocks) + size_t base = sizeBytes / issuers; + size_t extra = sizeBytes % issuers; + size_t myBytes = base + (rank < extra ? 1u : 0u); + size_t myStart = rank * base + (rank < extra ? rank : extra); + if (myBytes == 0) return; + + uint8_t* myLds = lds + static_cast(rank) * window; + + // Stage each chunk HBM->LDS->HBM through this warp's single window. The Sync + // policy waits after the load (RAW: store must see the filled LDS) and after + // the store (WAR: the next load must not overwrite LDS still being drained). + for (size_t o = 0; o < myBytes; o += window) { + const size_t chunk = (myBytes - o < window) ? (myBytes - o) : window; + asyncLoadToLDS(s + myStart + o, myLds, chunk); + asyncStoreFromLDS(myLds, d + myStart + o, chunk); + } +} + +} // namespace detail + +#endif // ASYNC_COPY_SUPPORTED + +/// \brief Drain the CALLING WARP's outstanding async-to/from-LDS ops. +/// \see tdm::tdmWait. The blocking copy forms already drain internally. +__device__ ASYNC_API void tdmWait() ASYNC_DELETED; + +/// \brief Non-blocking block-collective copy: issue and return. +/// \see tdm::tdmCopyAsync. +template +__device__ ASYNC_API void tdmCopyAsync(void* dst, const void* src, size_t sizeBytes, + void* ldsBuffer, size_t ldsBufferBytes) ASYNC_DELETED; + +/// \brief Blocking block-collective copy of [0, sizeBytes): dst <- src. +/// \see tdm::tdmCopy. +template +__device__ ASYNC_API void tdmCopy(void* dst, const void* src, size_t sizeBytes, + void* ldsBuffer, size_t ldsBufferBytes) ASYNC_DELETED; + +/// \brief Non-blocking WARP-SPECIALIZED copy by one contiguous warp team. +/// \see tdm::tdmCopyAsyncByTeam. +template +__device__ ASYNC_API void tdmCopyAsyncByTeam(void* dst, const void* src, size_t sizeBytes, + void* ldsBuffer, size_t ldsBufferBytes, + uint32_t startWarpId, uint32_t stopWarpId) ASYNC_DELETED; + +/// \brief Blocking WARP-SPECIALIZED copy by one contiguous warp team. +/// \see tdm::tdmCopyByTeam. +template +__device__ ASYNC_API void tdmCopyByTeam(void* dst, const void* src, size_t sizeBytes, + void* ldsBuffer, size_t ldsBufferBytes, + uint32_t startWarpId, uint32_t stopWarpId) ASYNC_DELETED; + +#if ASYNC_COPY_SUPPORTED + +__device__ inline void tdmWait() { asyncWait<0>(); } + +template +__device__ inline void tdmCopyAsync(void* dst, const void* src, size_t sizeBytes, + void* ldsBuffer, size_t ldsBufferBytes) { + detail::issue(dst, src, sizeBytes, ldsBuffer, ldsBufferBytes, /*start=*/0, /*stop=*/~0u); +} + +template +__device__ inline void tdmCopy(void* dst, const void* src, size_t sizeBytes, + void* ldsBuffer, size_t ldsBufferBytes) { + detail::issue(dst, src, sizeBytes, ldsBuffer, ldsBufferBytes, /*start=*/0, /*stop=*/~0u); + tdmWait(); +} + +template +__device__ inline void tdmCopyAsyncByTeam(void* dst, const void* src, size_t sizeBytes, + void* ldsBuffer, size_t ldsBufferBytes, + uint32_t startWarpId, uint32_t stopWarpId) { + detail::issue(dst, src, sizeBytes, ldsBuffer, ldsBufferBytes, startWarpId, stopWarpId); +} + +template +__device__ inline void tdmCopyByTeam(void* dst, const void* src, size_t sizeBytes, + void* ldsBuffer, size_t ldsBufferBytes, + uint32_t startWarpId, uint32_t stopWarpId) { + detail::issue(dst, src, sizeBytes, ldsBuffer, ldsBufferBytes, startWarpId, stopWarpId); + tdmWait(); // no-op on any warp that issued nothing / is off-team +} + +#endif // ASYNC_COPY_SUPPORTED + +} // namespace async + +#undef ASYNC_API +#undef ASYNC_DELETED + + +#endif // __ASYNCCOPY_H + diff --git a/src/header/cachePolicy.h b/src/header/cachePolicy.h new file mode 100644 index 0000000..d733cef --- /dev/null +++ b/src/header/cachePolicy.h @@ -0,0 +1,97 @@ +* +Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +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. +*/ + +/// \file cachePolicy.h +/// \brief Shared cache-policy encoding for the TDM copy libraries. +/// +/// Both the tensor-data-mover implementation (tdmCopy.h) and the async-to/from-LDS +/// implementation (asyncCopy.h) accept the same compile-time cache policy so their +/// public APIs stay identical. A CachePolicy packs a temporal hint together with a +/// memory scope into the integer immediate the memory instructions expect. + +#ifndef __TDM_CACHE_POLICY_H +#define __TDM_CACHE_POLICY_H + +#include + +#ifndef __host__ +#define __host__ +#endif +#ifndef __device__ +#define __device__ +#endif + +using CachePolicy = uint32_t; + +enum struct MemScope : uint32_t { + WGP = 0, // Workgroup processor scope - warps running on the same WGP should be able to see the effect of the operation + SE, // Shader engine a.k.a cluster scope + DEV, // Device scope + SYS, // System scope +}; + +enum struct TemporalHint : uint32_t { + RT = 0, // Regular temporal (nothing special) + NT, // Not temporal + HT, // High temporal + LU, // Last use + NT_RT, + RT_NT, + NT_HT, +}; + +__host__ __device__ constexpr CachePolicy createCachePolicy(TemporalHint temporal, MemScope scope) noexcept { + return static_cast(scope) << 3 | static_cast(temporal); +} + +static_assert(createCachePolicy(TemporalHint::RT, MemScope::WGP) == 0); +static_assert(createCachePolicy(TemporalHint::NT, MemScope::WGP) == 1); +static_assert(createCachePolicy(TemporalHint::HT, MemScope::WGP) == 2); +static_assert(createCachePolicy(TemporalHint::LU, MemScope::WGP) == 3); +static_assert(createCachePolicy(TemporalHint::NT_RT, MemScope::WGP) == 4); +static_assert(createCachePolicy(TemporalHint::RT_NT, MemScope::WGP) == 5); +static_assert(createCachePolicy(TemporalHint::NT_HT, MemScope::WGP) == 6); +static_assert(createCachePolicy(TemporalHint::RT, MemScope::SE) == 8); +static_assert(createCachePolicy(TemporalHint::NT, MemScope::SE) == 9); +static_assert(createCachePolicy(TemporalHint::HT, MemScope::SE) == 10); +static_assert(createCachePolicy(TemporalHint::LU, MemScope::SE) == 11); +static_assert(createCachePolicy(TemporalHint::NT_RT, MemScope::SE) == 12); +static_assert(createCachePolicy(TemporalHint::RT_NT, MemScope::SE) == 13); +static_assert(createCachePolicy(TemporalHint::NT_HT, MemScope::SE) == 14); +static_assert(createCachePolicy(TemporalHint::RT, MemScope::DEV) == 16); +static_assert(createCachePolicy(TemporalHint::NT, MemScope::DEV) == 17); +static_assert(createCachePolicy(TemporalHint::HT, MemScope::DEV) == 18); +static_assert(createCachePolicy(TemporalHint::LU, MemScope::DEV) == 19); +static_assert(createCachePolicy(TemporalHint::NT_RT, MemScope::DEV) == 20); +static_assert(createCachePolicy(TemporalHint::RT_NT, MemScope::DEV) == 21); +static_assert(createCachePolicy(TemporalHint::NT_HT, MemScope::DEV) == 22); +static_assert(createCachePolicy(TemporalHint::RT, MemScope::SYS) == 24); +static_assert(createCachePolicy(TemporalHint::NT, MemScope::SYS) == 25); +static_assert(createCachePolicy(TemporalHint::HT, MemScope::SYS) == 26); +static_assert(createCachePolicy(TemporalHint::LU, MemScope::SYS) == 27); +static_assert(createCachePolicy(TemporalHint::NT_RT, MemScope::SYS) == 28); +static_assert(createCachePolicy(TemporalHint::RT_NT, MemScope::SYS) == 29); +static_assert(createCachePolicy(TemporalHint::NT_HT, MemScope::SYS) == 30); + +constexpr CachePolicy DEFAULT_CACHE_POLICY = createCachePolicy(TemporalHint::RT, MemScope::SYS); + +#endif // __TDM_CACHE_POLICY_H From 015863b3a1c92fd9c67e2edb8462387bec34bab5 Mon Sep 17 00:00:00 2001 From: Alex Breslow Date: Tue, 21 Jul 2026 16:41:15 -0700 Subject: [PATCH 2/8] Add workaround for simulator --- Makefile | 12 ++++++++++++ src/header/TransferBench.hpp | 13 ++++++------- src/header/cachePolicy.h | 2 +- 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/Makefile b/Makefile index 1e52f02..aca2795 100644 --- a/Makefile +++ b/Makefile @@ -78,6 +78,18 @@ ifeq ($(filter clean,$(MAKECMDGOALS)),) ifneq ($(strip $(ROCM_DEVICE_LIB_PATH)),) HIPFLAGS += --rocm-device-lib-path=$(ROCM_DEVICE_LIB_PATH) endif + + # gfx1250 TDM descriptor header fallback. src/header/tdmCopy.h keys TDM support + # on __has_include(); some ROCm SDKs omit that + # descriptor header, in which case every tdm:: entry point compiles to a deleted + # no-op stub and IsTdmCopySupported() reports the GPU as unsupported. When the SDK + # lacks the header we add a vendored copy (src/tdm_fallback) to the include search + # path. The path is added only when the SDK is missing it, so a real SDK header is + # never shadowed. + ifeq ($(wildcard $(ROCM_PATH)/include/hip/amd_detail/amd_gfx1250_TDM.h),) + $(info - SDK lacks hip/amd_detail/amd_gfx1250_TDM.h; using vendored fallback in ./src/tdm_fallback) + CXXFLAGS += -I./src/tdm_fallback + endif endif ifeq ($(SINGLE_KERNEL), 1) diff --git a/src/header/TransferBench.hpp b/src/header/TransferBench.hpp index 9dd1367..f11e20f 100644 --- a/src/header/TransferBench.hpp +++ b/src/header/TransferBench.hpp @@ -5870,13 +5870,12 @@ namespace { float* __restrict__ dst = (float*)p.dst[0]; size_t const sizeBytes = p.N * sizeof(float); - int subIterations = 0; - bool useTDM = 0; - for(int subIterations = 0; subIterations < numSubIterations; subIterations++) { - if constexpr(useTDM){ - tdm::tdmCopy(dst, src, sizeBytes, shmem, ldsBytes); - else{ - async::tdmCopy(dst, src, sizeBytes, shmem, ldsBytes); + constexpr bool useTDM = false; + for (int subIterations = 0; subIterations < numSubIterations; subIterations++) { + if constexpr (useTDM) { + tdm::tdmCopy(dst, src, sizeBytes, shmem, ldsBytes); + } else { + async::tdmCopy(dst, src, sizeBytes, shmem, ldsBytes); } __syncthreads(); // Wait for all warps to finish } diff --git a/src/header/cachePolicy.h b/src/header/cachePolicy.h index d733cef..1c6e856 100644 --- a/src/header/cachePolicy.h +++ b/src/header/cachePolicy.h @@ -1,4 +1,4 @@ -* +/* Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. Permission is hereby granted, free of charge, to any person obtaining a copy From 7c243838a8b9340a616b32132962af3c79b247a9 Mon Sep 17 00:00:00 2001 From: Alex Breslow Date: Tue, 21 Jul 2026 16:53:17 -0700 Subject: [PATCH 3/8] Fix issue relating to simulator not implementing the XCC instruction --- Makefile | 5 +++++ src/header/TransferBench.hpp | 13 +++++++++++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index aca2795..3cf08d5 100644 --- a/Makefile +++ b/Makefile @@ -103,6 +103,11 @@ ifeq ($(filter clean,$(MAKECMDGOALS)),) endif COMMON_FLAGS += -I./src/header -I./src/client -I./src/client/Presets -I./third-party/ibverbs + # Extra user-supplied preprocessor/compiler flags, e.g. + # make EXTRA_CXXFLAGS=-DTB_DISABLE_HWID_QUERY ... + # to stub the hardware-ID reads for functional emulators. + COMMON_FLAGS += $(EXTRA_CXXFLAGS) + # libibverbs is loaded dynamically at runtime via dlopen/dlsym (see # third-party/ibverbs/IbvDynLoad.hpp), so the build never links against -libverbs # and does not require libibverbs-dev to be installed. We only need -ldl so diff --git a/src/header/TransferBench.hpp b/src/header/TransferBench.hpp index f11e20f..5be78e2 100644 --- a/src/header/TransferBench.hpp +++ b/src/header/TransferBench.hpp @@ -714,8 +714,15 @@ namespace TransferBench // Helper macro functions //========================================================================================== +// TB_DISABLE_HWID_QUERY: stub the hardware-ID reads (CU/SM and XCC) to 0. Some +// targets (e.g. functional emulators) cannot decode the s_getreg/s_sendmsg_rtn +// instructions these macros emit; define this at build time to run the kernels +// there. Real-hardware profiling is unaffected unless this is defined. + // Macro for collecting CU/SM GFX kernel is running on -#if defined(__GFX9__) +#if defined(TB_DISABLE_HWID_QUERY) + #define GetHwId(hwId) hwId = 0 +#elif defined(__GFX9__) #define GetHwId(hwId) asm volatile ("s_getreg_b32 %0, hwreg(HW_REG_HW_ID)" : "=s" (hwId)) #elif defined(__GFX10__) || defined(__GFX11__) || defined(__GFX12__) #define GetHwId(hwId) asm volatile ("s_getreg_b32 %0, hwreg(HW_REG_HW_ID1)" : "=s" (hwId)) @@ -726,7 +733,9 @@ namespace TransferBench #endif // Macro for collecting XCC GFX kernel is running on -#if defined(__gfx942__) || defined(__gfx950__) +#if defined(TB_DISABLE_HWID_QUERY) +#define GetXccId(val) val = 0 +#elif defined(__gfx942__) || defined(__gfx950__) #define GetXccId(val) asm volatile ("s_getreg_b32 %0, hwreg(HW_REG_XCC_ID)" : "=s" (val)) #elif defined(__GFX12__) #define GetXccId(val) \ From d566e31118460aaa3fea4c266e390f7fd30b20d4 Mon Sep 17 00:00:00 2001 From: Alex Breslow Date: Wed, 22 Jul 2026 10:29:43 -0700 Subject: [PATCH 4/8] A2A preset for TDM --- src/client/Presets/AllToAll.hpp | 17 ++++++-- src/client/Presets/AllToAllSweep.hpp | 58 ++++++++++++++++++++-------- 2 files changed, 56 insertions(+), 19 deletions(-) diff --git a/src/client/Presets/AllToAll.hpp b/src/client/Presets/AllToAll.hpp index b9717a5..e532739 100644 --- a/src/client/Presets/AllToAll.hpp +++ b/src/client/Presets/AllToAll.hpp @@ -55,6 +55,7 @@ int AllToAllPreset(EnvVars& ev, int numSubExecs = EnvVars::GetEnvVar("NUM_SUB_EXEC" , 8); int showDetails = EnvVars::GetEnvVar("SHOW_DETAILS" , 0); int useDmaExec = EnvVars::GetEnvVar("USE_DMA_EXEC" , 0); + int useTdmExec = EnvVars::GetEnvVar("USE_TDM_EXEC" , 0); int useRemoteRead = EnvVars::GetEnvVar("USE_REMOTE_READ", 0); // Check that all ranks have at least the number of GPUs requested @@ -108,6 +109,7 @@ int AllToAllPreset(EnvVars& ev, ev.Print("NUM_SUB_EXEC" , numSubExecs , "Using %d subexecutors/CUs per Transfer", numSubExecs); ev.Print("SHOW_DETAILS" , showDetails , "%s full Test details", showDetails ? "Showing" : "Hiding"); ev.Print("USE_DMA_EXEC" , useDmaExec , "Using %s executor", useDmaExec ? "DMA" : "GFX"); + ev.Print("USE_TDM_EXEC" , useTdmExec , "%s TDM executor", useTdmExec ? "Using" : "Not using"); ev.Print("USE_REMOTE_READ", useRemoteRead, "Using %s as executor", useRemoteRead ? "DST" : "SRC"); printf("\n"); } @@ -121,13 +123,22 @@ int AllToAllPreset(EnvVars& ev, Utils::Print("[ERROR] DMA execution can only be used for copies (A2A_MODE=0)\n"); return ERR_FATAL; } + if (useTdmExec && useDmaExec) { + Utils::Print("[ERROR] USE_TDM_EXEC and USE_DMA_EXEC are mutually exclusive\n"); + return ERR_FATAL; + } + if (useTdmExec && (numSrcs != 1 || numDsts != 1)) { + Utils::Print("[ERROR] TDM execution can only be used for copies (A2A_MODE=0)\n"); + return ERR_FATAL; + } if (numResults * 2 > numRanks) { Utils::Print("[ERROR] Number of extrema results requested exceeds number of ranks. NUM_RESULTS should be at most half the number of ranks\n"); return ERR_FATAL; } // Collect the number of GPU devices to use - ExeType exeType = useDmaExec ? EXE_GPU_DMA : EXE_GPU_GFX; + ExeType exeType = useTdmExec ? EXE_GPU_TDM : (useDmaExec ? EXE_GPU_DMA : EXE_GPU_GFX); + char const* exeName = useTdmExec ? "TDM" : (useDmaExec ? "DMA" : "GFX"); std::vector, int>> reIndex(numRanks); std::vector transfers; @@ -184,10 +195,10 @@ int AllToAllPreset(EnvVars& ev, } } - Utils::Print("GPU-%s All-To-All benchmark:\n", useDmaExec ? "DMA" : "GFX"); + Utils::Print("GPU-%s All-To-All benchmark:\n", exeName); Utils::Print("==============================\n"); Utils::Print("[%lu bytes per Transfer] [%s:%d] [%d Read(s) %d Write(s)] [MemType:%s] [NIC QueuePairs:%d] [#Ranks:%d]\n", - numBytesPerTransfer, useDmaExec ? "DMA" : "GFX", numSubExecs, numSrcs, numDsts, + numBytesPerTransfer, exeName, numSubExecs, numSrcs, numDsts, devMemTypeStr.c_str(), numQueuePairs, numRanks); if (transfers.size() == 0) { diff --git a/src/client/Presets/AllToAllSweep.hpp b/src/client/Presets/AllToAllSweep.hpp index 36e571d..61fa3b5 100644 --- a/src/client/Presets/AllToAllSweep.hpp +++ b/src/client/Presets/AllToAllSweep.hpp @@ -56,11 +56,20 @@ int AllToAllSweepPreset(EnvVars& ev, int showMinOnly = EnvVars::GetEnvVar("SHOW_MIN_ONLY", 1); int useRemoteRead = EnvVars::GetEnvVar("USE_REMOTE_READ", 0); int useSpray = EnvVars::GetEnvVar("USE_SPRAY", 0); + int useTdmExec = EnvVars::GetEnvVar("USE_TDM_EXEC", 0); int verbose = EnvVars::GetEnvVar("VERBOSE", 0); std::vector blockList = EnvVars::GetEnvVarArray("BLOCKSIZES", {256,512,768,1024}); std::vector unrollList = EnvVars::GetEnvVarArray("UNROLLS", {1,2,3,4,6,8}); std::vector numSesList = EnvVars::GetEnvVarArray("NUM_SUB_EXECS", {4,8,12,16,24,32}); + std::vector ldsList = EnvVars::GetEnvVarArray("LDSBYTES", {0}); + + // GFX and TDM expose different "middle" sweep dimensions shown in the second column: + // GFX -> UNROLLS (cfg.gfx.unrollFactor) + // TDM -> LDSBYTES (cfg.tdm.ldsBytes, per-block LDS staging bytes; 0 = device max) + // Each executor lacks the other's knob, so only the relevant list is swept. + std::vector& midList = useTdmExec ? ldsList : unrollList; + char const* midLabel = useTdmExec ? "LDS" : "UnR"; // A2A_MODE may be 0,1,2 or else custom numSrcs:numDsts int numSrcs, numDsts; @@ -94,9 +103,13 @@ int AllToAllSweepPreset(EnvVars& ev, ev.Print("NUM_GPU_DEVICES", numGpus , "Using %d GPUs", numGpus); ev.Print("NUM_SUB_EXECS" , numSesList.size(), EnvVars::ToStr(numSesList).c_str()); ev.Print("SHOW_MIN_ONLY" , showMinOnly , showMinOnly ? "Showing only slowest GPU results" : "Showing slowest and fastest GPU results"); - ev.Print("UNROLLS" , unrollList.size(), EnvVars::ToStr(unrollList).c_str()); + if (useTdmExec) + ev.Print("LDSBYTES" , ldsList.size() , "%s (0 = device max)", EnvVars::ToStr(ldsList).c_str()); + else + ev.Print("UNROLLS" , unrollList.size(), EnvVars::ToStr(unrollList).c_str()); ev.Print("USE_REMOTE_READ", useRemoteRead , "Using %s as executor", useRemoteRead ? "DST" : "SRC"); ev.Print("USE_SPRAY" , useSpray , "%s per SubExecutor", useSpray ? "All targets" : "One target"); + ev.Print("USE_TDM_EXEC" , useTdmExec , "Using %s executor", useTdmExec ? "TDM" : "GFX"); ev.Print("VERBOSE" , verbose , verbose ? "Display test results" : "Display summary only"); printf("\n"); } @@ -112,8 +125,18 @@ int AllToAllSweepPreset(EnvVars& ev, exit(1); } + if (useTdmExec && (numSrcs != 1 || numDsts != 1)) { + printf("[ERROR] TDM execution can only be used for copies (A2A_MODE=0)\n"); + exit(1); + } + if (useTdmExec && useSpray) { + printf("[ERROR] Cannot use USE_TDM_EXEC with USE_SPRAY (TDM copies one src/dst pair)\n"); + exit(1); + } + // Collect the number of GPU devices to use - ExeType exeType = EXE_GPU_GFX; + ExeType exeType = useTdmExec ? EXE_GPU_TDM : EXE_GPU_GFX; + char const* exeName = useTdmExec ? "TDM" : "GFX"; std::vector transfers; @@ -180,7 +203,8 @@ int AllToAllSweepPreset(EnvVars& ev, } } - Utils::Print("GPU-GFX All-To-All Sweep benchmark (%lu bytes, local=%s). All values are %s GB/s\n", + Utils::Print("GPU-%s All-To-All Sweep benchmark (%lu bytes, local=%s). All values are %s GB/s\n", + exeName, numBytesPerTransfer, a2aLocal ? "yes" : "no", ev.useHipEvents ? "GPU-Event-Timed (min over GPUs)": "CPU-Timed"); @@ -196,10 +220,10 @@ int AllToAllSweepPreset(EnvVars& ev, char sep = ev.outputToCsv ? ',' : ' '; double bestMinBw = 0.0; - int bestBlock = -1, bestUnroll = -1, bestNumSes = -1; + int bestBlock = -1, bestMid = -1, bestNumSes = -1; // Print header once - Utils::Print(" BlkS %c UnR ", sep); + Utils::Print(" BlkS %c %3s ", sep, midLabel); for (int c : numSesList) { Utils::Print("%c SE %03d", sep, c); if (ev.useHipEvents && !showMinOnly) { @@ -208,15 +232,17 @@ int AllToAllSweepPreset(EnvVars& ev, } Utils::Print("\n"); - // Results keyed by (blockSize, numSes, unroll) for verbose output + // Results keyed by (blockSize, numSes, midValue) for verbose output std::map, TransferBench::TestResults> results; for (int blockSize : blockList) { - cfg.gfx.blockSize = blockSize; + if (useTdmExec) cfg.tdm.blockSize = blockSize; + else cfg.gfx.blockSize = blockSize; - for (int u : unrollList) { - cfg.gfx.unrollFactor = u; - Utils::Print("%5d %c %3d ", blockSize, sep, u); + for (int mid : midList) { + if (useTdmExec) cfg.tdm.ldsBytes = mid; // per-block LDS staging bytes + else cfg.gfx.unrollFactor = mid; // GFX-kernel unroll factor + Utils::Print("%5d %c %3d ", blockSize, sep, mid); fflush(stdout); for (int c : numSesList) { @@ -247,11 +273,11 @@ int AllToAllSweepPreset(EnvVars& ev, if (minBw > bestMinBw) { bestMinBw = minBw; bestBlock = blockSize; - bestUnroll = u; + bestMid = mid; bestNumSes = c; } if (verbose) { - results[std::make_tuple(blockSize, c, u)] = result; + results[std::make_tuple(blockSize, c, mid)] = result; } } Utils::Print("%c%8.2f", sep, minBw); @@ -270,13 +296,13 @@ int AllToAllSweepPreset(EnvVars& ev, int testNum = 0; for (int blockSize : blockList) { for (int c : numSesList) { - for (int u : unrollList) { + for (int mid : midList) { auto verboseTransfers = transfers; for (auto& t : verboseTransfers) { t.numSubExecs = useSpray ? (c * targetCount) : c; } - Utils::Print("BlockSize: %d SubExecs: %d Unroll: %d\n", blockSize, c, u); - Utils::PrintResults(ev, ++testNum, verboseTransfers, results[std::make_tuple(blockSize, c, u)]); + Utils::Print("BlockSize: %d SubExecs: %d %s: %d\n", blockSize, c, midLabel, mid); + Utils::PrintResults(ev, ++testNum, verboseTransfers, results[std::make_tuple(blockSize, c, mid)]); } } } @@ -287,7 +313,7 @@ int AllToAllSweepPreset(EnvVars& ev, Utils::Print("Highest %s bandwidth found: %7.2f GB/s\n", ev.useHipEvents ? "GPU-event-timed (min)" : "CPU-timed", bestMinBw); Utils::Print(" BlockSize : %7d\n", bestBlock); - Utils::Print(" Unroll : %7d\n", bestUnroll); + Utils::Print(" %-10s : %7d\n", midLabel, bestMid); Utils::Print(" NumSubExec : %7d\n", bestNumSes); } From c6f2f20b3dbf9602390e1ac15a7b92b1022e37cd Mon Sep 17 00:00:00 2001 From: Alex Breslow Date: Wed, 22 Jul 2026 11:56:55 -0700 Subject: [PATCH 5/8] Add logic for skipping misalignment logic where it's not necessary --- src/header/asyncCopy.h | 106 ++++++++++++++++++++++++++++------------- 1 file changed, 74 insertions(+), 32 deletions(-) diff --git a/src/header/asyncCopy.h b/src/header/asyncCopy.h index c0b070c..c2c67ed 100644 --- a/src/header/asyncCopy.h +++ b/src/header/asyncCopy.h @@ -138,7 +138,12 @@ constexpr int HOT_LOOP_UNROLL = 4; // // Alignment is peeled against the global-memory pointer; the LDS staging buffer is assumed to share the same // alignment (RCCL allocates it 128-byte aligned), so applying the same offset keeps the LDS side aligned too. -template +// +// `Aligned` fast path: when the caller has already proven that both `global` and `lds` start on a 128-byte +// boundary, the phase-1 peel is provably a no-op. Compiling it out removes the peel branches and lets the +// hot loop start at offset 0 with folded, compile-time-known strides. Callers must only pass Aligned=true +// after checking the runtime pointers (see async::detail::issue). +template __device__ inline void warpAsyncCopy(const uint8_t* global, uint8_t* lds, size_t sizeInBytes){ const unsigned lane = __lane_id(); const size_t laneCount = (size_t)__builtin_amdgcn_wavefrontsize(); @@ -146,22 +151,22 @@ __device__ inline void warpAsyncCopy(const uint8_t* global, uint8_t* lds, size_t size_t offset = 0; - // Phase 1a: peel leading bytes until `global` is 16-byte aligned, one byte per lane. - const size_t misalign = (uintptr_t)global & (BYTES_PER_LANE_B128 - 1); - if (misalign != 0) { - const size_t head = (BYTES_PER_LANE_B128 - misalign) < sizeInBytes - ? (BYTES_PER_LANE_B128 - misalign) - : sizeInBytes; - if (lane < head) { - asyncCopyB8(global + lane, lds + lane); + if constexpr (!Aligned) { + // Phase 1a: peel leading bytes until `global` is 16-byte aligned, one byte per lane. + const size_t misalign = (uintptr_t)global & (BYTES_PER_LANE_B128 - 1); + if (misalign != 0) { + const size_t head = (BYTES_PER_LANE_B128 - misalign) < sizeInBytes + ? (BYTES_PER_LANE_B128 - misalign) + : sizeInBytes; + if (lane < head) { + asyncCopyB8(global + lane, lds + lane); + } + offset += head; } - offset += head; - } - // Phase 1b: peel whole 16-byte chunks until `global + offset` reaches 128-byte alignment, one b128 per - // leading lane. `global + offset` is already 16-byte aligned, so the number of bytes left to the next - // 128-byte boundary is a multiple of 16 (at most 112 -> 7 chunks, well within a warp). - { + // Phase 1b: peel whole 16-byte chunks until `global + offset` reaches 128-byte alignment, one b128 per + // leading lane. `global + offset` is already 16-byte aligned, so the number of bytes left to the next + // 128-byte boundary is a multiple of 16 (at most 112 -> 7 chunks, well within a warp). const size_t align128 = ((uintptr_t)global + offset) & (NATURAL_ALIGNMENT_BYTES - 1); const size_t bytesToLine = (NATURAL_ALIGNMENT_BYTES - align128) & (NATURAL_ALIGNMENT_BYTES - 1); size_t peelChunks = bytesToLine / BYTES_PER_LANE_B128; @@ -221,20 +226,22 @@ __device__ inline void warpAsyncCopy(const uint8_t* global, uint8_t* lds, size_t } // namespace async_detail // Warp-level async copy from global memory into LDS. The entire warp calls this with the same arguments. -template +// Set Aligned=true only when globalSrc and ldsDst are known to be 128-byte aligned (skips the peel). +template __device__ void asyncLoadToLDS(const uint8_t* globalSrc, uint8_t* ldsDst, size_t sizeInBytes){ - async_detail::warpAsyncCopy(globalSrc, ldsDst, sizeInBytes); + async_detail::warpAsyncCopy(globalSrc, ldsDst, sizeInBytes); if constexpr (sp == SyncPolicy::Sync) { asyncWait<0>(); } } // Warp-level async copy from LDS into global memory. The entire warp calls this with the same arguments. -template +// Set Aligned=true only when globalDst and ldsSrc are known to be 128-byte aligned (skips the peel). +template __device__ void asyncStoreFromLDS(const uint8_t* ldsSrc, uint8_t* globalDst, size_t sizeInBytes){ // warpAsyncCopy shares one non-const `lds` parameter across load and store; the store path only reads from // it, so dropping const here is safe. - async_detail::warpAsyncCopy(globalDst, const_cast(ldsSrc), sizeInBytes); + async_detail::warpAsyncCopy(globalDst, const_cast(ldsSrc), sizeInBytes); if constexpr (sp == SyncPolicy::Sync) { asyncWait<0>(); } @@ -314,6 +321,22 @@ __device__ inline void warpGlobalCopy(const uint8_t* s, uint8_t* d, size_t n, for (size_t i = warpThread; i < n; i += warpThreads) d[i] = s[i]; } +// Stage [myStart, myStart+myBytes) HBM->LDS->HBM through this warp's single `window`-byte LDS buffer. +// LoadAligned/StoreAligned pick the peel-free warpAsyncCopy fast path for the src/dst side respectively; +// they are loop-invariant (myStart and window are both 128B multiples), so the choice is made once by the +// caller and templated in here rather than branched per iteration. The Sync policy waits after the load +// (RAW: the store must see the filled LDS) and after the store (WAR: the next load must not overwrite LDS +// that is still draining). +template +__device__ inline void stageLoop(const uint8_t* s, uint8_t* d, uint8_t* myLds, + size_t myStart, size_t myBytes, uint32_t window) { + for (size_t o = 0; o < myBytes; o += window) { + const size_t chunk = (myBytes - o < window) ? (myBytes - o) : window; + asyncLoadToLDS (s + myStart + o, myLds, chunk); + asyncStoreFromLDS(myLds, d + myStart + o, chunk); + } +} + // core: partition + issue the whole copy for the team [start, stop). Work + LDS // partition by rank within the team (warpId - start), so each team indexes its // own ldsBuffer from zero. Safe to call collectively (warps outside the range @@ -363,23 +386,42 @@ __device__ inline void issue(void* dst, const void* src, size_t sizeBytes, uint32_t window = (ldsBytes / issuers) & ~(WINDOW_GRAIN - 1); // per-warp 128B-multiple if (rank >= issuers) return; // this warp doesn't issue - // distribute the byte range across issuers by team rank (contiguous blocks) - size_t base = sizeBytes / issuers; - size_t extra = sizeBytes % issuers; - size_t myBytes = base + (rank < extra ? 1u : 0u); - size_t myStart = rank * base + (rank < extra ? rank : extra); + // Distribute the byte range across issuers in contiguous, 128B-aligned blocks. Splitting on whole + // WINDOW_GRAIN (128B) grains -- rather than raw sizeBytes/issuers -- keeps every issuing warp's start a + // multiple of 128B, so `s + myStart` and `d + myStart` inherit the base pointers' alignment and stay + // aligned across chunks (the stride `window` is a 128B multiple). The sub-128B remainder is appended to + // the last issuing warp so coverage remains contiguous and gap-free. + const size_t grain = WINDOW_GRAIN; + const size_t nGrains = sizeBytes / grain; + const size_t tail = sizeBytes - nGrains * grain; // 0..127 trailing bytes + + const size_t baseG = nGrains / issuers; + const size_t extraG = nGrains % issuers; + const size_t myG = baseG + (rank < extraG ? 1u : 0u); + const size_t startG = rank * baseG + (rank < extraG ? rank : extraG); + + size_t myStart = startG * grain; // 128B-aligned + size_t myBytes = myG * grain; + if (rank == issuers - 1) myBytes += tail; // last warp mops up the < 128B tail if (myBytes == 0) return; uint8_t* myLds = lds + static_cast(rank) * window; - // Stage each chunk HBM->LDS->HBM through this warp's single window. The Sync - // policy waits after the load (RAW: store must see the filled LDS) and after - // the store (WAR: the next load must not overwrite LDS still being drained). - for (size_t o = 0; o < myBytes; o += window) { - const size_t chunk = (myBytes - o < window) ? (myBytes - o) : window; - asyncLoadToLDS(s + myStart + o, myLds, chunk); - asyncStoreFromLDS(myLds, d + myStart + o, chunk); - } + // Pick the peel-free fast path per direction (load reads src, store writes dst). myStart and window are + // both 128B multiples, so each side's alignment is loop-invariant and computed once here. + const uintptr_t mask = WINDOW_GRAIN - 1; + const bool ldsAligned = ((uintptr_t)myLds & mask) == 0; + const bool loadAligned = ldsAligned && (((uintptr_t)(s + myStart) & mask) == 0); + const bool storeAligned = ldsAligned && (((uintptr_t)(d + myStart) & mask) == 0); + + if (loadAligned && storeAligned) + stageLoop(s, d, myLds, myStart, myBytes, window); + else if (loadAligned) + stageLoop(s, d, myLds, myStart, myBytes, window); + else if (storeAligned) + stageLoop(s, d, myLds, myStart, myBytes, window); + else + stageLoop(s, d, myLds, myStart, myBytes, window); } } // namespace detail From cedaee627750c24de3ab7d97b1a93ce22e17a784 Mon Sep 17 00:00:00 2001 From: Alex Breslow Date: Wed, 22 Jul 2026 14:30:20 -0700 Subject: [PATCH 6/8] Add quick branch to skip tail checks if possible. --- src/header/asyncCopy.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/header/asyncCopy.h b/src/header/asyncCopy.h index c2c67ed..1bfecac 100644 --- a/src/header/asyncCopy.h +++ b/src/header/asyncCopy.h @@ -200,6 +200,8 @@ __device__ inline void warpAsyncCopy(const uint8_t* global, uint8_t* lds, size_t offset += full128 * BYTES_PER_LANE_B128; remaining -= full128 * BYTES_PER_LANE_B128; // now < 16 bytes + if (remaining == 0) return; // common case after a 16B-multiple chunk: nothing left to drain + // Phase 3b: 0..15 trailing bytes. `offset` and `remaining` are uniform across the warp here, so the // descending cascade below is taken identically by every lane. Because `offset` is 16-byte aligned, the // b64 (8-byte) and b32 (4-byte) accesses are naturally aligned. From 97cb090e5de2c189ef11bf211c57ab71518e5945 Mon Sep 17 00:00:00 2001 From: Alex Breslow Date: Thu, 23 Jul 2026 09:56:05 -0700 Subject: [PATCH 7/8] Unroll of 2 was the best --- src/header/asyncCopy.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/header/asyncCopy.h b/src/header/asyncCopy.h index 1bfecac..c8f86c2 100644 --- a/src/header/asyncCopy.h +++ b/src/header/asyncCopy.h @@ -124,7 +124,7 @@ constexpr size_t BYTES_PER_LANE_B128 = 16; // The hot loop must start on a 128-byte boundary, so the peel aligns the global pointer to this. constexpr size_t NATURAL_ALIGNMENT_BYTES = 128; // How many b128 iterations the hot loop unrolls at a time. -constexpr int HOT_LOOP_UNROLL = 4; +constexpr int HOT_LOOP_UNROLL = 2; // Warp-cooperative byte copy between global memory and LDS. The whole warp must call this with identical // `global`, `lds`, and `sizeInBytes` arguments -- the per-lane work distribution is derived internally. From 9bf828d3abe3f98d39b30dcd4e7eb4ae4ded4d69 Mon Sep 17 00:00:00 2001 From: Alex Breslow Date: Thu, 23 Jul 2026 10:08:52 -0700 Subject: [PATCH 8/8] Make -DTB_USE_ASYNC_COPY=1/0 the source of truth regarding whether to use the async copy backend for the TDM path --- src/header/TransferBench.hpp | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/header/TransferBench.hpp b/src/header/TransferBench.hpp index 5be78e2..81278ab 100644 --- a/src/header/TransferBench.hpp +++ b/src/header/TransferBench.hpp @@ -5860,6 +5860,14 @@ namespace { // TDM Executor-related functions //======================================================================================== +// Compile-time selection of the LDS-staged copy backend used by GpuTdmKernel: +// TB_USE_ASYNC_COPY=1 (default) -> async-to/from-LDS path (async::tdmCopy, asyncCopy.h) +// TB_USE_ASYNC_COPY=0 -> tensor-data-mover path (tdm::tdmCopy, tdmCopy.h) +// Override without touching source, e.g.: make ... EXTRA_CXXFLAGS=-DTB_USE_ASYNC_COPY=0 +#ifndef TB_USE_ASYNC_COPY +#define TB_USE_ASYNC_COPY 1 +#endif + #if TDM_SUPPORTED __global__ void GpuTdmKernel(SubExecParam* params, uint32_t ldsBytes, @@ -5879,7 +5887,7 @@ namespace { float* __restrict__ dst = (float*)p.dst[0]; size_t const sizeBytes = p.N * sizeof(float); - constexpr bool useTDM = false; + constexpr bool useTDM = (TB_USE_ASYNC_COPY == 0); for (int subIterations = 0; subIterations < numSubIterations; subIterations++) { if constexpr (useTDM) { tdm::tdmCopy(dst, src, sizeBytes, shmem, ldsBytes);