amirsojoodi · Oct 26, 2026 · Sep 30, 2026 · Sep 30, 2026 · Oct 1, 2026 · Oct 22, 2026
@@ -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)

@@ -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

@@ -257,3 +257,24 @@ 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
+        | put/get on gfx1250 (requires ``USE_TDM`` build option). Doubled
+        | internally for double buffering; see ``rocshmem_query_tdm_lds_bytes()``.
+      - ``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
+        | 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.
@@ -196,6 +196,18 @@ __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 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). 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 +830,40 @@ __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 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). 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.
+ * 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.
+ *
+ * @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.

@@ -259,6 +259,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.
@@ -506,6 +511,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.

@@ -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
    */

@@ -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<uintptr_t>(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<size_t>(max_shared_mem_per_block) > reserve_bytes)
+            ? static_cast<size_t>(max_shared_mem_per_block) - reserve_bytes
+            : 0;
+    const uint32_t avail_tile_bytes =
+        static_cast<uint32_t>((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<size_t>(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) {

@@ -30,7 +30,6 @@
 #include <cstdint>
 #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;

@@ -235,6 +235,15 @@ namespace envvar {
         false);
   }  // namespace sdma
 
+  namespace tdm {
+    const var<uint32_t> 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<size_t> 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<var_map_t&, std::mutex&> get_var_map() {
       // construct on first use idiom

@@ -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<tag::REVERSE_OFFLOAD> = "ROCSHMEM_RO";
     template <> inline constexpr const char* prefix<tag::GDA> = "ROCSHMEM_GDA";
     template <> inline constexpr const char* prefix<tag::SDMA> = "ROCSHMEM_SDMA";
+    template <> inline constexpr const char* prefix<tag::TDM> = "ROCSHMEM_TDM";
   }  // namespace category
 
   namespace parser {
@@ -580,6 +582,12 @@ namespace envvar {
     extern const var<bool> spread_channels;
   }  // namespace sdma
 
+  namespace tdm {
+    template <typename T> using var = var<T, category::tag::TDM>;
+    extern const var<uint32_t> tile_bytes;
+    extern const var<size_t> threshold;
+  }  // namespace tdm
+
   /**
    * @brief Print mode for environment variables
    */

@@ -476,10 +476,12 @@ __host__ void HostInterface::broadcastmem_on_stream(rocshmem_team_t team,
 __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<int>(nelems))
@@ -489,17 +491,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<<<gridSize, blockSize, 0, stream>>>(dest, source,
-                                                             nelems, pe);
+  rocshmem_getmem_kernel<<<gridSize, blockSize, tdm_lds_bytes, stream>>>(
+      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<int>(nelems))
@@ -509,17 +513,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<<<gridSize, blockSize, 0, stream>>>(dest, source,
-                                                             nelems, pe);
+  rocshmem_putmem_kernel<<<gridSize, blockSize, tdm_lds_bytes, stream>>>(
+      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<int>(nelems))
@@ -529,7 +535,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<<<gridSize, blockSize, 0, stream>>>(
+  rocshmem_putmem_signal_kernel<<<gridSize, blockSize, tdm_lds_bytes, stream>>>(
       dest, source, nelems, sig_addr, signal, sig_op, pe);
 }
 

@@ -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<char *>(dest) - wrk_sync_pool_bases_[constmem.my_pe];
+  memcpy_wg_tdm<MemcpyKind::Put>(wrk_sync_pool_bases_[pe] + L_offset, const_cast<void *>(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<const char *>(source);
+  uint64_t L_offset = const_cast<char *>(src_typed) - wrk_sync_pool_bases_[constmem.my_pe];
+  memcpy_wg_tdm<MemcpyKind::Get>(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<char *>(dest) - wrk_sync_pool_bases_[constmem.my_pe];