fix(rocshmem): make Anvil SDMA usable on MI300X CPX and QPX, DPX up to two OAMs

projects/rccl/src/CMakeLists.txt

@@ -1199,7 +1199,7 @@ if (ENABLE_ROCSHMEM_GIN)
     # SDMA anvil host code: copied into build tree (no sed — host code needs
     # unqualified log.hpp to find the host shim via include path ordering).
     set(_GIN_SDMA_DIR "${PROJECT_BINARY_DIR}/gin/sdma")
-    foreach(_f anvil.cpp anvil.hpp)
+    foreach(_f anvil.cpp anvil.hpp anvil_engine_map.hpp)
       add_custom_command(
         OUTPUT "${_GIN_SDMA_DIR}/${_f}"
         COMMAND ${CMAKE_COMMAND} -E make_directory "${_GIN_SDMA_DIR}"

projects/rocshmem/src/sdma/anvil.hpp

@@ -33,33 +33,64 @@
 #include <iostream>
 #include <memory>
 #include <mutex>
+#include <string>
 #include <unordered_map>
 #include <vector>
 
 #include "anvil_device.hpp"
+#include "anvil_engine_map.hpp"
 #include "hsa/hsa_ext_amd.h"
 #include "hsakmt/hsakmt.h"
 #include "hsakmt/hsakmttypes.h"
 
 namespace sdma_anvil {
 
+// How an engine was chosen for one peer, and the values that explain the choice. The two failure
+// logs that report it -- the queue-create error and the missing-xGMI-id error -- need the same
+// fields, so they travel as one record: a second copy means the next diagnostic field has to be
+// remembered in two places.
+struct EngineSelection {
+  int engineId{-1};  // -1 when no engine could be mapped for this pair
+  int srcDeviceId{-1};
+  int dstDeviceId{-1};
+  // Result of hsa_amd_memory_get_preferred_copy_engine. queried stays false when the agents were
+  // not valid enough to ask, which is different from asking and being refused.
+  hsa_status_t preferredStatus{HSA_STATUS_ERROR};
+  uint32_t preferredMask{0};
+  bool preferredQueried{false};
+  bool usedPreferred{false};
+  // Engine counts this node reports, which bound every id above.
+  uint32_t numSdmaEngines{0};
+  uint32_t numSdmaXgmiEngines{0};
+  uint32_t numSdmaEnginesTotal{0};
+};
 
 class SdmaQueue {
  public:
-  SdmaQueue(int localDeviceId, int remoteDeviceId, const hsa_agent_t& localAgent, uint32_t engineId);
+  SdmaQueue(int localDeviceId, int remoteDeviceId, const hsa_agent_t& localAgent, uint32_t engineId,
+            const EngineSelection& selection);
   ~SdmaQueue();
 
   SdmaQueueDeviceHandle* deviceHandle() const;
   SdmaQueueSingleProducerDeviceHandle* singleProducerDeviceHandle() const;
   void dump(std::ofstream& logFile);
+  // False when the constructor could not create an SDMA queue. The constructor releases whatever it
+  // acquired before giving up, so an invalid queue must be dropped without being used or destroyed
+  // any further.
+  bool valid() const;
+  // Status of the queue-creation attempt, so the caller can tell an exhausted queue budget
+  // (NO_MEMORY) apart from other failures.
+  HSAKMT_STATUS createStatus() const;
 
  private:
-  int remoteDeviceId_;
-  uint64_t* cachedWptr_;
-  uint64_t* committedWptr_;
-  void* queueBuffer_;
-  HsaQueueResource queue_;
-  SdmaQueueDeviceHandle* deviceHandle_;
+  bool valid_{false};
+  HSAKMT_STATUS createStatus_{HSAKMT_STATUS_ERROR};
+  int remoteDeviceId_{-1};
+  uint64_t* cachedWptr_{nullptr};
+  uint64_t* committedWptr_{nullptr};
+  void* queueBuffer_{nullptr};
+  HsaQueueResource queue_{};
+  SdmaQueueDeviceHandle* deviceHandle_{nullptr};
   SdmaQueueSingleProducerDeviceHandle* singleProducerDeviceHandle_{nullptr};
 };
 
@@ -80,9 +111,18 @@ class AnvilLib {
   void init();
   bool connect(int srcDeviceId, int dstDeviceId, int numChannels = 1);
   void disconnect();
+  // Drop only the queues held for one destination. The channel map is process-global and a
+  // second communicator reuses a previous create's queues, so a failed create must not take
+  // disconnect()'s whole-process path and destroy handles already live on the GPU.
+  void disconnectDevice(int dstDeviceId);
+  // Queues this process currently holds, counted across all engines. The budget is otherwise only
+  // visible in a log line, which leaves the refusal in connect() and the teardown in
+  // disconnectDevice() with nothing a test can assert: a refused connect must leave this unchanged
+  // and a fully abandoned create must bring it back to where it started.
+  uint32_t queuesUsed() const { return queuesUsedTotal_; }
   SdmaQueue* getSdmaQueue(int srcDeviceId, int dstDeviceId, int channel_idx = 0);
   SdmaQueue* createSdmaQueue(int srcDeviceId, int dstDeviceId, uint32_t engineId,
-                             int* channelIdx = nullptr);
+                             const EngineSelection& selection, int* channelIdx = nullptr);
 
  private:
   /*
@@ -111,13 +151,30 @@ class AnvilLib {
   uint32_t numSdmaEngines_{0};
   uint32_t numSdmaXgmiEngines_{0};
   uint32_t numSdmaEnginesTotal_{0};
+  // KFD caps user SDMA queues per engine, so a partition with few engines also has a small total
+  // queue budget. Track usage to refuse a mesh that cannot fit before KFD returns NO_MEMORY part
+  // way through building it.
+  uint32_t numSdmaQueuesPerEngine_{0};
+  // Queues already taken by this process, counted across the usable even engines rather than per
+  // engine. connect() compares this against usableSdmaEngines(numSdmaEnginesTotal_) *
+  // numSdmaQueuesPerEngine_; an engine-pinned rejection fails instead of retrying on an
+  // unconstrained engine.
+  uint32_t queuesUsedTotal_{0};
+  HSAKMT_STATUS lastQueueStatus_{HSAKMT_STATUS_SUCCESS};
+  // Selection in progress, so getOamId can report why the map was consulted without the caller
+  // threading the same values back down.
+  EngineSelection selection_;
 
   void buildGpuAgentMap();
   hsa_agent_t getHipGpuAgent(int hipDeviceId) const;
   void querySdmaEngineCounts();
+  // Pop `vec` back to `keep` entries, crediting each popped queue to the budget.
+  void releaseQueues(std::vector<std::unique_ptr<SdmaQueue>>& vec, size_t keep);
   int getOamId(int deviceId);
   int getSdmaEngineIdFromOamMap(int srcDeviceId, int dstDeviceId);
-  int getSdmaEngineId(int srcDeviceId, int dstDeviceId);
+  // Chooses an engine for the pair and records why, so the caller can hand the same record to
+  // every queue it then creates.
+  EngineSelection getSdmaEngineId(int srcDeviceId, int dstDeviceId);
 
   std::once_flag init_flag;
   std::vector<hsa_agent_t> gpuAgentsByHipDev_;

projects/rocshmem/src/sdma/anvil_engine_map.hpp

@@ -0,0 +1,122 @@
+/******************************************************************************
+ * 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_SDMA_ANVIL_ENGINE_MAP_HPP_
+#define LIBRARY_SRC_SDMA_ANVIL_ENGINE_MAP_HPP_
+
+// Pure engine-selection and queue-budget helpers for AnvilLib. Kept free of HSA/KFD headers so the
+// unit tests can build them without USE_SDMA.
+
+#include <cstdint>
+#include <string>
+
+namespace sdma_anvil {
+
+// Every queue slot on an even-numbered MI300 SDMA engine is reachable by a GPU-written doorbell.
+// Slots 0-3 on odd engines are not (ROCM-32598), and this code cannot choose a slot, so use only
+// even engines until the slot-aware workaround can inspect and park unusable queues.
+inline bool isUsableSdmaEngine(int engineId) { return engineId >= 0 && engineId % 2 == 0; }
+
+// Number of even engines in a node-local engine range. A zero total means the count is unknown.
+inline uint32_t usableSdmaEngines(uint32_t numSdmaEnginesTotal) {
+  return (numSdmaEnginesTotal + 1) / 2;
+}
+
+// Select an even engine from an HSA preferred-copy-engine mask. Prefer an even bit from the mask;
+// if it contains only odd engines, use the paired even engine immediately below the lowest one.
+// Return -1 for an empty mask.
+inline int preferredUsableEngine(uint32_t preferredMask) {
+  if (preferredMask == 0) return -1;
+  const uint32_t evenMask = preferredMask & 0x55555555u;
+  return evenMask != 0 ? __builtin_ctz(evenMask) : (__builtin_ctz(preferredMask) & ~1);
+}
+
+// A CPX/DPX partition owns host SDMA engines but no xGMI engines. Callers pass a total of 0 when
+// the engine counts could not be queried, so a failed query is never read as a partition.
+inline bool isSdmaPartition(uint32_t numSdmaXgmiEngines, uint32_t numSdmaEnginesTotal) {
+  return numSdmaXgmiEngines == 0 && numSdmaEnginesTotal > 0;
+}
+
+// True when the doubled OAM-map id cannot be used as-is. A partition (no xGMI
+// engines) folds even when the doubled id is in range: same-device peers share
+// the map diagonal and would otherwise all land on engine 0. An id past the
+// engines this node reports folds too. numSdmaEnginesTotal == 0 never folds,
+// so the modulo below is not asked to divide by zero.
+inline bool oamMapEngineNeedsFold(uint32_t numSdmaXgmiEngines, uint32_t numSdmaEnginesTotal,
+                                  int doubledEngineId) {
+  const bool partition = isSdmaPartition(numSdmaXgmiEngines, numSdmaEnginesTotal);
+  const bool outOfRange =
+      numSdmaEnginesTotal > 0 && static_cast<uint32_t>(doubledEngineId) >= numSdmaEnginesTotal;
+  return partition || outOfRange;
+}
+
+// Fold the undoubled OAM-map value and both PCI functions onto an even engine this node reports.
+// numEnginesTotal must be > 0. An unreadable function (-1) counts as 0 and therefore collides with
+// function 0.
+inline int foldOamMapEngine(int oamEngine, int srcFn, int dstFn, uint32_t numEnginesTotal) {
+  const int src = srcFn < 0 ? 0 : srcFn;
+  const int dst = dstFn < 0 ? 0 : dstFn;
+  return 2 * ((oamEngine + src + dst) %
+              static_cast<int>(usableSdmaEngines(numEnginesTotal)));
+}
+
+// PCI function and the function-0 BDF of the same device. function is -1 when the tail is not a
+// function digit, so an unreadable id is not treated as function 0. The physical BDF is rewritten
+// only in that case; a bad tail is left unchanged.
+struct PciFunctionBus {
+  std::string busId;
+  std::string physBusId;
+  int function;
+};
+
+// Split a BDF into its function digit and the function-0 BDF of the same device. Pure, and beside
+// the two helpers above because it decides the same fallback they do: a tail outside '0'-'7'
+// leaves function at -1 and physBusId equal to busId, which collapses getOamId's candidate list to
+// one entry so the physical-BDF read never runs. PCI function numbers are three bits, so '0'-'7'
+// is the whole range.
+inline PciFunctionBus pciFunctionBus(const std::string& busId) {
+  PciFunctionBus loc;
+  loc.busId = busId;
+  loc.physBusId = busId;
+  loc.function = -1;
+  if (busId.empty()) return loc;
+  const char c = busId.back();
+  if (c >= '0' && c <= '7') {
+    loc.function = c - '0';
+    loc.physBusId.back() = '0';
+  }
+  return loc;
+}
+
+// True when taking `requested` more queues would exceed the queue budget, counted across the
+// engines the caller can safely use. Either count being 0 means the budget is unknown, which is
+// never reported as exhausted.
+inline bool queueBudgetExceeded(uint32_t used, int requested, uint32_t numEnginesTotal,
+                                uint32_t queuesPerEngine) {
+  const uint32_t budget = numEnginesTotal * queuesPerEngine;
+  return requested > 0 && budget > 0 && used + static_cast<uint32_t>(requested) > budget;
+}
+
+}  // namespace sdma_anvil
+
+#endif  // LIBRARY_SRC_SDMA_ANVIL_ENGINE_MAP_HPP_

projects/rocshmem/src/sdma_policy.cpp

@@ -64,12 +64,18 @@ __host__ void SdmaImpl::sdmaHostInit(int pe, int num_pes, int rank) {
   int deviceId;
   CHECK_HIP(hipGetDevice(&deviceId));
 
-  // Create SDMA connections to all local PEs including self
+  // Create SDMA connections to all local PEs including self. A failed connect aborts rather than
+  // clearing sdmaEnabled on this rank alone: sdmaEnabled also selects the alltoall algorithm, and
+  // ranks that disagree on it run incompatible synchronization and hang.
   for (int i = 0; i < shm_size; i++) {
     if (i != deviceId) {
       sdma_anvil::EnablePeerAccess(deviceId, i);
     }
-    sdma_anvil::anvil.connect(deviceId, i, numChannels);
+    if (!sdma_anvil::anvil.connect(deviceId, i, numChannels)) {
+      LOG_ERROR_ABORT("SDMA: connect failed from device %d to %d with %d channel(s); set "
+                      "ROCSHMEM_SDMA_ENABLED=0 or lower ROCSHMEM_SDMA_NUM_CHANNELS",
+                      deviceId, i, numChannels);
+    }
   }
 
   // Total number of handles: shm_size * numChannels
@@ -103,7 +109,11 @@ __host__ void SdmaImpl::sdmaHostStop() {
     CHECK_HIP(hipFree(deviceHandles_d));
     deviceHandles_d = nullptr;
   }
-  sdma_anvil::anvil.disconnect();
+  // Release only the peers sdmaHostInit wired: the channel map is process-global, so disconnect()
+  // would also destroy queues another user still holds. A disabled init wired none.
+  if (sdmaEnabled) {
+    for (int i = 0; i < shm_size; i++) sdma_anvil::anvil.disconnectDevice(i);
+  }
 }
 
 #endif  // USE_SDMA

projects/rocshmem/tests/unit_tests/CMakeLists.txt

@@ -64,6 +64,7 @@ target_sources(
     atomic_wf_queue_gtest.cpp
     envvar_gtest.cpp
     backend_bc_buffer_register_gtest.cpp
+    anvil_engine_map_gtest.cpp
 )
 
 if (USE_IPC)