feat: add automatic transport routing

scripts/common_env.sh

@@ -79,6 +79,19 @@ if [ ! -r "${ASCEND_DRIVER_PATH}/kernel/inc" ] && [ -d "${ASCEND_HOME_PATH}/${TI
     export ASCEND_DRIVER_PATH=${TILEXR_DRIVER_SHIM_HOME}
 fi
 
+_tilexr_prepend_path_if_dir() {
+    if [ -d "$1" ]; then
+        case ":${PATH}:" in
+            *":$1:"*) ;;
+            *) export PATH="$1:${PATH}" ;;
+        esac
+    fi
+}
+
+_tilexr_prepend_path_if_dir "${ASCEND_HOME_PATH}/tools/bisheng_compiler/bin"
+_tilexr_prepend_path_if_dir "${ASCEND_HOME_PATH}/${TILEXR_OS_ARCH}-linux/bin"
+_tilexr_prepend_path_if_dir "/usr/local/Ascend/cann-${TILEXR_CANN_VER}/tools/bisheng_compiler/bin"
+
 export PATH=${MPI_HOME}/bin:${PATH}
 export PATH=${TILEXR_UTIL_HOME}/cmake/bin:${PATH}
 export PATH=${TILEXR_UTIL_HOME}/ccache:${TILEXR_UTIL_HOME}/ripgrep:${TILEXR_UTIL_HOME}/sshpass/bin:${PATH}

src/comm/CMakeLists.txt

@@ -174,6 +174,7 @@ install(FILES
         ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_sync.h
         ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_data_as_flag.h
         ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_perf_trace.h
+        ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_transport.h
         ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_udma.h
         ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_udma_reg.h
         ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_udma_types.h

src/comm/tilexr_comm.cpp

@@ -137,6 +137,7 @@ int TileXRComm::InitUDMA()
     TileXRUDMAContextOptions options {};
     options.rank = rank_;
     options.rankSize = rankSize_;
+    options.localRankSize = static_cast<int>(localRankSize_);
     options.devId = devId_;
     options.exchange = socketExchange_;
     options.threadMode = !uid_.empty();

src/comm/tools/socket/tilexr_sock_exchange.h

@@ -10,6 +10,7 @@
 #ifndef TILEXR_SOCK_EXCHANGE_H
 #define TILEXR_SOCK_EXCHANGE_H
 
+#include <cstdint>
 #include <vector>
 #include <string>
 #include <memory>
@@ -97,8 +98,14 @@ class TileXRSockExchange {
 
     template <typename T> int Send(int fd, const T *sendBuf, size_t sendSize, int flag) const
     {
-        do {
-            auto ret = send(fd, sendBuf, sendSize, flag);
+        const auto *sendBytes = reinterpret_cast<const uint8_t *>(sendBuf);
+        size_t sentBytes = 0;
+        while (sentBytes < sendSize) {
+            auto ret = send(fd, sendBytes + sentBytes, sendSize - sentBytes, flag);
+            if (ret > 0) {
+                sentBytes += static_cast<size_t>(ret);
+                continue;
+            }
             if (ret < 0) {
                 if (CheckErrno(errno)) {
                     TILEXR_LOG(ERROR) << "send failed: " << strerror(errno);
@@ -107,13 +114,20 @@ class TileXRSockExchange {
                 TILEXR_LOG(DEBUG) << "Send failed: " << strerror(errno);
             }
             return ret;
-        } while (true);
+        }
+        return static_cast<int>(sentBytes);
     }
 
     template <typename T> int Recv(int fd, T *recvBuf, size_t recvSize, int flag) const
     {
-        do {
-            auto ret = recv(fd, recvBuf, recvSize, flag);
+        auto *recvBytes = reinterpret_cast<uint8_t *>(recvBuf);
+        size_t receivedBytes = 0;
+        while (receivedBytes < recvSize) {
+            auto ret = recv(fd, recvBytes + receivedBytes, recvSize - receivedBytes, flag);
+            if (ret > 0) {
+                receivedBytes += static_cast<size_t>(ret);
+                continue;
+            }
             if (ret < 0) {
                 if (CheckErrno(errno)) {
                     TILEXR_LOG(ERROR) << "recv failed: " << strerror(errno);
@@ -122,7 +136,8 @@ class TileXRSockExchange {
                 TILEXR_LOG(DEBUG) << "recv failed: " << strerror(errno);
             }
             return ret;
-        } while (true);
+        }
+        return static_cast<int>(receivedBytes);
     }
 
     template <typename T> int ClientSendRecv(const T *sendBuf, size_t sendSize, T *recvBuf)

src/comm/udma/tilexr_udma_context.cpp

@@ -58,6 +58,7 @@ int TileXRUDMAContext::Init(const TileXRUDMAContextOptions& options)
     TileXRUDMATransportOptions transportOptions {};
     transportOptions.rank = options_.rank;
     transportOptions.rankSize = options_.rankSize;
+    transportOptions.localRankSize = options_.localRankSize;
     transportOptions.devId = options_.devId;
     transportOptions.exchange = options_.exchange;
     int ret = transport_->Init(transportOptions);

src/comm/udma/tilexr_udma_context.h

@@ -30,6 +30,7 @@ using TileXRUDMACommArgsUpdateFn = int (*)(const TileXRUDMACommArgsState& state,
 struct TileXRUDMAContextOptions {
     int rank = 0;
     int rankSize = 0;
+    int localRankSize = 1;
     int devId = 0;
     bool threadMode = false;
     TileXRSockExchange* exchange = nullptr;

src/comm/udma/tilexr_udma_transport.cpp

@@ -324,7 +324,8 @@ int TileXRUDMATransport::Init(const TileXRUDMATransportOptions& options)
     if (options.rankSize <= 1) {
         return TILEXR_SUCCESS;
     }
-    if (options.rank < 0 || options.rank >= options.rankSize || options.exchange == nullptr) {
+    if (options.rank < 0 || options.rank >= options.rankSize || options.localRankSize <= 0 ||
+        options.localRankSize > options.rankSize || options.exchange == nullptr) {
         return TILEXR_ERROR_PARA_CHECK_FAIL;
     }
     options_ = options;
@@ -459,7 +460,7 @@ int TileXRUDMATransport::BuildRoutes()
 
     std::vector<int32_t> localRouteByPeer(options_.rankSize, -1);
     for (int peer = 0; peer < options_.rankSize; ++peer) {
-        if (peer == options_.rank) {
+        if (!UsesUDMAPeer(peer)) {
             continue;
         }
         uint32_t localEid = devEids[0].eidIndex;
@@ -480,7 +481,7 @@ int TileXRUDMATransport::BuildRoutes()
     }
 
     for (int peer = 0; peer < options_.rankSize; ++peer) {
-        if (peer == options_.rank) {
+        if (!UsesUDMAPeer(peer)) {
             continue;
         }
         int32_t remoteEid = allRouteByPeer[peer * options_.rankSize + options_.rank];
@@ -720,6 +721,17 @@ uint32_t TileXRUDMATransport::FallbackLocalEid() const
     return 0;
 }
 
+bool TileXRUDMATransport::UsesUDMAPeer(int peer) const
+{
+    if (peer < 0 || peer >= options_.rankSize || peer == options_.rank) {
+        return false;
+    }
+    if (options_.localRankSize >= options_.rankSize) {
+        return true;
+    }
+    return peer / options_.localRankSize != options_.rank / options_.localRankSize;
+}
+
 int TileXRUDMATransport::RefreshUDMAInfo()
 {
     if (eidCount_ == 0 || states_.empty()) {
@@ -777,7 +789,8 @@ int TileXRUDMATransport::RefreshUDMAInfo()
     for (int rank = 0; rank < options_.rankSize; ++rank) {
         uint32_t localEid = fallbackEid;
         uint32_t remoteEid = fallbackEid;
-        if (rank != options_.rank) {
+        const bool usesUDMA = UsesUDMAPeer(rank);
+        if (usesUDMA) {
             localEid = peerLocalEid_[rank];
             remoteEid = peerRemoteEid_[rank];
         }
@@ -795,7 +808,7 @@ int TileXRUDMATransport::RefreshUDMAInfo()
             if (localMemIt != localMemInfoByEid_.end()) {
                 mem[rank] = localMemIt->second;
             }
-        } else {
+        } else if (usesUDMA) {
             mem[rank] = allMem[rank * eidCount_ + remoteEid];
             mem[rank].tpn = state.tpnList[rank];
         }
@@ -832,14 +845,23 @@ int TileXRUDMATransport::RegisterMemory(GM_ADDR localPtr, size_t bytes)
     }
     int ret = RegisterMemoryOnContexts(localPtr, bytes);
     if (ret != TILEXR_SUCCESS) {
+        TILEXR_LOG(ERROR) << "TileXR UDMA local memory registration failed, rank=" << options_.rank
+                          << ", bytes=" << bytes << ", ret=" << ret;
         return ret;
     }
     registeredPtr_ = localPtr;
     ret = ExchangeAndImportMemory();
     if (ret != TILEXR_SUCCESS) {
+        TILEXR_LOG(ERROR) << "TileXR UDMA remote memory exchange/import failed, rank=" << options_.rank
+                          << ", bytes=" << bytes << ", ret=" << ret;
         return ret;
     }
-    return RefreshUDMAInfo();
+    ret = RefreshUDMAInfo();
+    if (ret != TILEXR_SUCCESS) {
+        TILEXR_LOG(ERROR) << "TileXR UDMA info refresh failed after memory registration, rank=" << options_.rank
+                          << ", bytes=" << bytes << ", ret=" << ret;
+    }
+    return ret;
 }
 
 int TileXRUDMATransport::RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes)
@@ -863,6 +885,9 @@ int TileXRUDMATransport::RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes
         void* lmemHandle = nullptr;
         int ret = loader_.RaCtxLmemRegister(ctxEntry.second, &mrInfo, &lmemHandle);
         if (ret != 0 || lmemHandle == nullptr) {
+            TILEXR_LOG(ERROR) << "TileXR UDMA RaCtxLmemRegister failed, rank=" << options_.rank
+                              << ", eid=" << eidIndex << ", bytes=" << bytes << ", ret=" << ret
+                              << ", handle=" << lmemHandle;
             return TILEXR_ERROR_INTERNAL;
         }
 
@@ -903,6 +928,8 @@ int TileXRUDMATransport::ExchangeAndImportMemory()
     std::vector<uint32_t> allCounts(options_.rankSize);
     int ret = options_.exchange->AllGather(&localCount, 1, allCounts.data());
     if (ret != TILEXR_SUCCESS) {
+        TILEXR_LOG(ERROR) << "TileXR UDMA memory-count AllGather failed, rank=" << options_.rank
+                          << ", localCount=" << localCount << ", ret=" << ret;
         return ret;
     }
     const uint32_t maxCount = *std::max_element(allCounts.begin(), allCounts.end());
@@ -927,12 +954,14 @@ int TileXRUDMATransport::ExchangeAndImportMemory()
     std::vector<ExchangedMrInfo> all(options_.rankSize * maxCount);
     ret = options_.exchange->AllGather(local.data(), local.size(), all.data());
     if (ret != TILEXR_SUCCESS) {
+        TILEXR_LOG(ERROR) << "TileXR UDMA memory-info AllGather failed, rank=" << options_.rank
+                          << ", maxCount=" << maxCount << ", ret=" << ret;
         return ret;
     }
 
     remoteMemHandles_.assign(options_.rankSize, nullptr);
     for (int peer = 0; peer < options_.rankSize; ++peer) {
-        if (peer == options_.rank) {
+        if (!UsesUDMAPeer(peer)) {
             continue;
         }
         const uint32_t remoteEid = peerRemoteEid_[peer];
@@ -945,6 +974,9 @@ int TileXRUDMATransport::ExchangeAndImportMemory()
             }
         }
         if (remote == nullptr) {
+            TILEXR_LOG(ERROR) << "TileXR UDMA remote memory info missing, rank=" << options_.rank
+                              << ", peer=" << peer << ", remoteEid=" << remoteEid
+                              << ", peerCount=" << allCounts[peer];
             return TILEXR_ERROR_INTERNAL;
         }
         const uint32_t localEid = peerLocalEid_[peer];
@@ -956,6 +988,10 @@ int TileXRUDMATransport::ExchangeAndImportMemory()
         void* remoteHandle = nullptr;
         ret = loader_.RaCtxRmemImport(ctxHandleByEid_[localEid], &importInfo, &remoteHandle);
         if (ret != 0 || remoteHandle == nullptr) {
+            TILEXR_LOG(ERROR) << "TileXR UDMA RaCtxRmemImport failed, rank=" << options_.rank
+                              << ", peer=" << peer << ", localEid=" << localEid
+                              << ", remoteEid=" << remoteEid << ", ret=" << ret
+                              << ", handle=" << remoteHandle;
             return TILEXR_ERROR_INTERNAL;
         }
         remoteMemHandles_[peer] = remoteHandle;

src/comm/udma/tilexr_udma_transport.h

@@ -27,6 +27,7 @@ class TileXRSockExchange;
 struct TileXRUDMATransportOptions {
     int rank = 0;
     int rankSize = 0;
+    int localRankSize = 1;
     int devId = 0;
     TileXRSockExchange* exchange = nullptr;
 };
@@ -63,6 +64,7 @@ class TileXRUDMATransport {
     void CleanupMemory();
     void CleanupContexts();
     uint32_t FallbackLocalEid() const;
+    bool UsesUDMAPeer(int peer) const;
 
     TileXRHccpLoader loader_;
     TileXRUDMATransportOptions options_ {};

src/ep/CMakeLists.txt

@@ -138,6 +138,7 @@ add_custom_target(tilexr_ep_combine_kernel ALL DEPENDS "${TILEXR_EP_COMBINE_KERN
 add_library(tilexr-ep SHARED
     host/ep_layout.cpp
     host/ep_dispatch_host.cpp
+    host/ep_transport_route.cpp
     host/ep_launch_context.cpp
     host/ep_kernel_launch.cpp
     host/tilexr_ep_dispatch.cpp
@@ -158,7 +159,6 @@ target_link_directories(tilexr-ep
         ${CMAKE_CURRENT_BINARY_DIR}
         ${ASCEND_DRIVER_PATH}/lib64/driver
         ${ASCEND_HOME_PATH}/${ARCH}-linux/lib64
-        ${ASCEND_HOME_PATH}/${ARCH}-linux/devlib
 )
 
 target_link_libraries(tilexr-ep

src/ep/common/ep_window.h

@@ -6,6 +6,7 @@
 namespace TileXREp {
 
 constexpr int64_t kEpWindowAlignmentBytes = 32;
+constexpr int64_t kEpUdmaReadyStrideBytes = 64;
 constexpr int64_t kEpAssistTupleInts = 4;
 constexpr int64_t kEpWindowHeaderBytes = 64;
 constexpr int64_t kEpSrcSlotHeaderBytes = 64;
@@ -18,6 +19,8 @@ constexpr int32_t kEpStepCombineGatewayReady = 76;
 constexpr int32_t kEpStepCombineRelayReady = 77;
 constexpr int64_t kEpStatusOk = 0;
 constexpr int64_t kEpStatusRemoteReadyTimeout = 1;
+constexpr int64_t kEpStatusDispatchReadyTimeout = 2;
+constexpr int64_t kEpStatusDispatchSlotTimeout = 3;
 constexpr uint32_t kEpWindowMagic = 0x54584550U;
 
 struct EpWindowHeader {

src/ep/host/ep_dispatch_host.cpp

@@ -7,11 +7,6 @@
 namespace TileXREp {
 namespace {
 
-bool TileXREpUsesCrossNodeComm(const TileXR::CommArgs &commArgs)
-{
-    return commArgs.localRankSize > 0 && commArgs.localRankSize < commArgs.rankSize;
-}
-
 int64_t TileXREpEffectiveTpWorldSize(int64_t tpWorldSize)
 {
     return tpWorldSize == 0 ? 1 : tpWorldSize;
@@ -118,18 +113,6 @@ int TileXREpValidateDispatchConfig(const EpDispatchParams &params, const TileXR:
         return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL;
     }
 
-    for (int rank = 0; rank < commArgs.rankSize; ++rank) {
-        if (commArgs.peerMems[rank] == nullptr) {
-            return TileXR::TILEXR_ERROR_NOT_INITIALIZED;
-        }
-    }
-
-    if (TileXREpUsesCrossNodeComm(commArgs) &&
-        (params.workspace == nullptr || (commArgs.extraFlag & TileXR::ExtraFlag::UDMA) == 0 ||
-            commArgs.udmaInfoPtr == nullptr || commArgs.udmaRegistryPtr == nullptr)) {
-        return TileXR::TILEXR_ERROR_NOT_INITIALIZED;
-    }
-
     int ret = TileXREpValidateDispatchV2Config(params, commArgs);
     if (ret != TileXR::TILEXR_SUCCESS) {
         return ret;
@@ -169,18 +152,6 @@ int TileXREpValidateCombineConfig(const EpCombineParams &params, const TileXR::C
         commArgs.rank < 0 || commArgs.rank >= commArgs.rankSize) {
         return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL;
     }
-    if (TileXREpUsesCrossNodeComm(commArgs) &&
-        (params.workspace == nullptr || (commArgs.extraFlag & TileXR::ExtraFlag::UDMA) == 0 ||
-            commArgs.udmaInfoPtr == nullptr || commArgs.udmaRegistryPtr == nullptr)) {
-        return TileXR::TILEXR_ERROR_NOT_INITIALIZED;
-    }
-
-    for (int rank = 0; rank < commArgs.rankSize; ++rank) {
-        if (commArgs.peerMems[rank] == nullptr) {
-            return TileXR::TILEXR_ERROR_NOT_INITIALIZED;
-        }
-    }
-
     return TileXREpBuildWindowConfig(commArgs.rankSize, params.bs, params.h, params.topK, params.moeExpertNum,
         params.dtype, window);
 }

src/ep/host/ep_dispatch_host.h

@@ -5,6 +5,7 @@
 
 #include "acl/acl_base.h"
 #include "ep_layout.h"
+#include "ep_transport_route.h"
 #include "tilexr_api.h"
 
 namespace TileXREp {
@@ -64,6 +65,7 @@ struct EpHostLaunchContext {
     TileXR::CommArgs *hostArgs = nullptr;
     GM_ADDR devArgs = nullptr;
     EpWindowConfig window {};
+    TileXR::TileXRTransportKind transport = TileXR::TileXRTransportKind::MEMORY;
 };
 
 int TileXREpValidateBasicDispatchParams(const EpDispatchParams &params);