Fix: stabilize HCCL communicator initialization

- pass device_id through comm_init so communicator setup binds the
  correct physical device for each rank
- align a2a3 HCCL init with the pto-isa sequence by linking hcomm,
  running rank-0 rtSetDevice, and synchronizing rootinfo exchange
- move distributed worker device setup after comm init and keep the
  A5/sim backends signature-compatible with the new API

Made-with: Cursor

Author: echo_stone

examples/scripts/distributed_worker.py

@@ -105,14 +105,12 @@ def buf_bytes(b):
 
     lib_path = artifact_dir / "libhost_runtime.so"
     Runtime = bind_host_binary(str(lib_path))
-    set_device(args.device_id)
-
-    sys.stderr.write(f"[rank {args.rank}] Library loaded, device {args.device_id} set\n")
+    sys.stderr.write(f"[rank {args.rank}] Library loaded\n")
 
     # ----------------------------------------------------------------
     # 2. Comm init + alloc windows
     # ----------------------------------------------------------------
-    comm = comm_init(args.rank, args.nranks, args.rootinfo_file)
+    comm = comm_init(args.rank, args.nranks, args.device_id, args.rootinfo_file)
 
     total_win = args.win_sync_prefix
     for b in buffers:
@@ -124,6 +122,9 @@ def buf_bytes(b):
 
     sys.stderr.write(f"[rank {args.rank}] Comm initialized, local_base=0x{local_base:x}\n")
 
+    set_device(args.device_id)
+    sys.stderr.write(f"[rank {args.rank}] Device {args.device_id} set for runtime\n")
+
     # ----------------------------------------------------------------
     # 3. Allocate buffers
     # ----------------------------------------------------------------

python/bindings.py

@@ -181,7 +181,7 @@ def _setup_functions(self):
         self.lib.enable_runtime_profiling.restype = c_int
 
         # --- Distributed communication API (comm_*) ---
-        self.lib.comm_init.argtypes = [c_int, c_int, c_char_p]
+        self.lib.comm_init.argtypes = [c_int, c_int, c_int, c_char_p]
         self.lib.comm_init.restype = c_void_p
 
         self.lib.comm_alloc_windows.argtypes = [c_void_p, c_size_t, POINTER(c_uint64)]
@@ -543,13 +543,14 @@ def launch_runtime(
 # ============================================================================
 
 
-def comm_init(rank: int, nranks: int, rootinfo_path: str) -> int:
+def comm_init(rank: int, nranks: int, device_id: int, rootinfo_path: str) -> int:
     """
     Initialize a distributed communicator for the given rank.
 
     Args:
         rank: This process's rank (0-based)
         nranks: Total number of ranks
+        device_id: Physical device ID used by this process
         rootinfo_path: Filesystem path for root info exchange
 
     Returns:
@@ -562,7 +563,7 @@ def comm_init(rank: int, nranks: int, rootinfo_path: str) -> int:
     if _lib is None:
         raise RuntimeError("Runtime not loaded. Call bind_host_binary() first.")
 
-    handle = _lib.comm_init(rank, nranks, rootinfo_path.encode('utf-8'))
+    handle = _lib.comm_init(rank, nranks, device_id, rootinfo_path.encode('utf-8'))
     if not handle:
         raise RuntimeError(f"comm_init failed for rank {rank}")
     return handle

src/a2a3/platform/include/host/comm.h

@@ -32,11 +32,12 @@ typedef struct CommHandle_* CommHandle;
  *
  * @param rank           This process's rank (0-based).
  * @param nranks         Total number of ranks.
+ * @param device_id      Physical device ID used by this process.
  * @param rootinfo_path  Filesystem path used to exchange root info between
  *                       ranks (rank 0 writes, others read).
  * @return Opaque handle, or NULL on failure.
  */
-CommHandle comm_init(int rank, int nranks, const char* rootinfo_path);
+CommHandle comm_init(int rank, int nranks, int device_id, const char* rootinfo_path);
 
 /**
  * Allocate RDMA / shared-memory windows and populate the device context.

src/a2a3/platform/onboard/host/CMakeLists.txt

@@ -81,22 +81,15 @@ target_link_directories(host_runtime
         ${ASCEND_HOME_PATH}/runtime/lib64
 )
 
-# Detect HCCL library version: CANN 9.x ships hccl_v2 instead of hccl
-# Prioritize hccl_v2 since CANN 9.x may have both hccl.so and hccl_v2.so,
-# and hccl_v2 is the complete, actively-maintained library.
-find_library(HCCL_V2_LIB NAMES hccl_v2 PATHS ${ASCEND_HOME_PATH}/lib64 NO_DEFAULT_PATH)
-find_library(HCCL_LIB NAMES hccl PATHS ${ASCEND_HOME_PATH}/lib64 NO_DEFAULT_PATH)
-
-if(HCCL_V2_LIB)
-    set(HCCL_LINK_TARGETS hccl_v2 hccl_plf)
-    target_compile_definitions(host_runtime PRIVATE HCCL_USE_V2_API=1)
-    message(STATUS "Using HCCL library: hccl_v2 + hccl_plf (CANN 9.x)")
-elseif(HCCL_LIB)
-    set(HCCL_LINK_TARGETS hccl)
-    message(STATUS "Using HCCL library: hccl")
+# CANN 9.x exposes the working non-V2 HCCL entry points through libhcomm.
+# Link it explicitly so comm_hccl.cpp can follow the same initialization path
+# as the pto-isa communication tests.
+find_library(HCOMM_LIB NAMES hcomm PATHS ${ASCEND_HOME_PATH}/lib64 NO_DEFAULT_PATH)
+if(HCOMM_LIB)
+    set(HCCL_LINK_TARGETS hcomm)
+    message(STATUS "Using HCCL library: hcomm")
 else()
-    message(WARNING "Neither hccl_v2 nor hccl found, linking against hccl and hoping for the best")
-    set(HCCL_LINK_TARGETS hccl)
+    message(FATAL_ERROR "libhcomm not found under ${ASCEND_HOME_PATH}/lib64")
 endif()
 
 # Optionally link nnopbase (provides aclCreateTensor/aclDestroyTensor for SdmaWorkspaceManager)
@@ -115,7 +108,6 @@ target_link_libraries(host_runtime
         runtime
         ascendcl
         ${HCCL_LINK_TARGETS}
-        hccl_fwk
         ${NNOPBASE_LINK}
         dl
 )

src/a2a3/platform/onboard/host/comm_hccl.cpp

@@ -24,45 +24,14 @@
 
 using CommTopo = uint32_t;
 
-// Internal HCCL APIs (not in public headers).
-// CANN 9.x renames all HCCL symbols with a V2 suffix and ships libhccl_v2.so.
-// The public header still declares weak non-V2 symbols, so we declare V2 variants
-// separately and dispatch via inline wrappers.
-#ifdef HCCL_USE_V2_API
-extern "C" HcclResult HcclGetRootInfoV2(HcclRootInfo* rootInfo);
-extern "C" HcclResult HcclCommInitRootInfoV2(uint32_t nRanks, const HcclRootInfo* rootInfo,
-                                              uint32_t rank, HcclComm* comm);
-extern "C" HcclResult HcclGetCommNameV2(HcclComm comm, char* commName);
-extern "C" HcclResult HcclBarrierV2(HcclComm comm, aclrtStream stream);
-extern "C" HcclResult HcclCommDestroyV2(HcclComm comm);
-extern "C" HcclResult HcclAllocComResourceByTilingV2(HcclComm comm, void* stream,
-                                                      void* mc2Tiling, void** commContext);
-extern "C" HcclResult HcomGetCommHandleByGroupV2(const char* group, HcclComm* commHandle);
-extern "C" HcclResult HcomGetL0TopoTypeExV2(const char* group, CommTopo* topoType,
-                                             uint32_t isSetDevice);
-
-static inline HcclResult hccl_get_root_info(HcclRootInfo* ri)
-    { return HcclGetRootInfoV2(ri); }
-static inline HcclResult hccl_comm_init_root_info(uint32_t n, const HcclRootInfo* ri, uint32_t r, HcclComm* c)
-    { return HcclCommInitRootInfoV2(n, ri, r, c); }
-static inline HcclResult hccl_get_comm_name(HcclComm c, char* name)
-    { return HcclGetCommNameV2(c, name); }
-static inline HcclResult hccl_barrier(HcclComm c, aclrtStream s)
-    { return HcclBarrierV2(c, s); }
-static inline HcclResult hccl_comm_destroy(HcclComm c)
-    { return HcclCommDestroyV2(c); }
-static inline HcclResult hccl_alloc_com_resource(HcclComm c, void* s, void* t, void** ctx)
-    { return HcclAllocComResourceByTilingV2(c, s, t, ctx); }
-static inline HcclResult hccl_get_comm_handle_by_group(const char* g, HcclComm* c)
-    { return HcomGetCommHandleByGroupV2(g, c); }
-static inline HcclResult hccl_get_l0_topo_type_ex(const char* g, CommTopo* t, uint32_t f)
-    { return HcomGetL0TopoTypeExV2(g, t, f); }
-#else
+// Internal HCCL helpers are exported by libhcomm on CANN 9.x.  The public
+// HCCL APIs below intentionally use the standard, non-V2 entry points to match
+// the working pto-isa initialization sequence.
 extern "C" HcclResult HcclAllocComResourceByTiling(HcclComm comm, void* stream,
-                                                    void* mc2Tiling, void** commContext);
+                                                   void* mc2Tiling, void** commContext);
 extern "C" HcclResult HcomGetCommHandleByGroup(const char* group, HcclComm* commHandle);
 extern "C" HcclResult HcomGetL0TopoTypeEx(const char* group, CommTopo* topoType,
-                                           uint32_t isSetDevice);
+                                          uint32_t isSetDevice);
 
 static inline HcclResult hccl_get_root_info(HcclRootInfo* ri)
     { return HcclGetRootInfo(ri); }
@@ -80,13 +49,13 @@ static inline HcclResult hccl_get_comm_handle_by_group(const char* g, HcclComm*
     { return HcomGetCommHandleByGroup(g, c); }
 static inline HcclResult hccl_get_l0_topo_type_ex(const char* g, CommTopo* t, uint32_t f)
     { return HcomGetL0TopoTypeEx(g, t, f); }
-#endif
 
 static constexpr uint32_t COMM_IS_NOT_SET_DEVICE = 0;
 static constexpr uint32_t COMM_TOPO_MESH = 0b1u;
 
 using rtStream_t = void*;
 static constexpr int32_t RT_STREAM_PRIORITY_DEFAULT = 0;
+extern "C" int32_t rtSetDevice(int32_t device);
 extern "C" int32_t rtStreamCreate(rtStream_t* stream, int32_t priority);
 extern "C" int32_t rtStreamDestroy(rtStream_t stream);
 
@@ -325,7 +294,7 @@ static void file_barrier(const std::string& dir, int rank, int nranks, const std
 // API implementation
 // ============================================================================
 
-extern "C" CommHandle comm_init(int rank, int nranks, const char* rootinfo_path) {
+extern "C" CommHandle comm_init(int rank, int nranks, int device_id, const char* rootinfo_path) {
     auto* h = new (std::nothrow) CommHandle_{};
     if (!h) return nullptr;
 
@@ -342,10 +311,26 @@ extern "C" CommHandle comm_init(int rank, int nranks, const char* rootinfo_path)
         return nullptr;
     }
 
-    // NOTE: Do NOT call aclrtSetDevice here — the caller (distributed_worker)
-    // already set the correct physical device via set_device(device_id).
-    // Calling aclrtSetDevice(rank) would override the context when
-    // rank != device_id (e.g. devices=[2,4,5,7]).
+    if (rank == 0) {
+        int32_t rtRet = rtSetDevice(device_id);
+        if (rtRet != 0) {
+            fprintf(stderr, "[comm rank %d] rtSetDevice(%d) failed: %d\n",
+                    rank, device_id, rtRet);
+            delete h;
+            return nullptr;
+        }
+    }
+
+    // HCCL requires an ACL runtime context bound to the physical device.
+    // This cannot be inferred from rank because distributed runs may map
+    // ranks to arbitrary device lists (for example devices=[2,4,5,7]).
+    aRet = aclrtSetDevice(device_id);
+    if (aRet != ACL_SUCCESS) {
+        fprintf(stderr, "[comm rank %d] aclrtSetDevice(%d) failed: %d\n",
+                rank, device_id, (int)aRet);
+        delete h;
+        return nullptr;
+    }
 
     // RootInfo exchange
     HcclRootInfo rootInfo{};
@@ -369,6 +354,13 @@ extern "C" CommHandle comm_init(int rank, int nranks, const char* rootinfo_path)
         fin.read(rootInfo.internal, HCCL_ROOT_INFO_BYTES);
     }
 
+    std::string barrier_dir = h->rootinfo_path;
+    auto last_slash = barrier_dir.rfind('/');
+    if (last_slash != std::string::npos) {
+        barrier_dir = barrier_dir.substr(0, last_slash);
+    }
+    file_barrier(barrier_dir, h->rank, h->nranks, "rootinfo_ready");
+
     // Create stream for HCCL operations
     rtStreamCreate(&h->stream, RT_STREAM_PRIORITY_DEFAULT);
 
@@ -403,8 +395,9 @@ extern "C" int comm_alloc_windows(CommHandle h, size_t /*win_size*/, uint64_t* d
     // File barrier so all ranks have completed HcclCommInitRootInfo
     std::string barrier_dir = h->rootinfo_path;
     auto last_slash = barrier_dir.rfind('/');
-    if (last_slash != std::string::npos)
+    if (last_slash != std::string::npos) {
         barrier_dir = barrier_dir.substr(0, last_slash);
+    }
     file_barrier(barrier_dir, h->rank, h->nranks, "hccl_init");
 
     // Tiling configuration for HcclAllocComResourceByTiling

src/a2a3/platform/sim/host/comm_sim.cpp

@@ -68,9 +68,10 @@ struct CommHandle_ {
 // API implementation
 // ============================================================================
 
-extern "C" CommHandle comm_init(int rank, int nranks, const char* rootinfo_path) {
+extern "C" CommHandle comm_init(int rank, int nranks, int device_id, const char* rootinfo_path) {
     auto* h = new (std::nothrow) CommHandle_{};
     if (!h) return nullptr;
+    (void)device_id;
 
     h->rank = rank;
     h->nranks = nranks;

src/a5/platform/include/host/comm.h

@@ -32,11 +32,12 @@ typedef struct CommHandle_* CommHandle;
  *
  * @param rank           This process's rank (0-based).
  * @param nranks         Total number of ranks.
+ * @param device_id      Physical device ID used by this process.
  * @param rootinfo_path  Filesystem path used to exchange root info between
  *                       ranks (rank 0 writes, others read).
  * @return Opaque handle, or NULL on failure.
  */
-CommHandle comm_init(int rank, int nranks, const char* rootinfo_path);
+CommHandle comm_init(int rank, int nranks, int device_id, const char* rootinfo_path);
 
 /**
  * Allocate RDMA / shared-memory windows and populate the device context.

src/a5/platform/onboard/host/comm_hccl.cpp

@@ -274,7 +274,7 @@ static void file_barrier(const std::string& dir, int rank, int nranks, const std
 // API implementation
 // ============================================================================
 
-extern "C" CommHandle comm_init(int rank, int nranks, const char* rootinfo_path) {
+extern "C" CommHandle comm_init(int rank, int nranks, int device_id, const char* rootinfo_path) {
     auto* h = new (std::nothrow) CommHandle_{};
     if (!h) return nullptr;
 
@@ -291,10 +291,16 @@ extern "C" CommHandle comm_init(int rank, int nranks, const char* rootinfo_path)
         return nullptr;
     }
 
-    // NOTE: Do NOT call aclrtSetDevice here — the caller (distributed_worker)
-    // already set the correct physical device via set_device(device_id).
-    // Calling aclrtSetDevice(rank) would override the context when
-    // rank != device_id (e.g. devices=[2,4,5,7]).
+    // HCCL requires an ACL runtime context bound to the physical device.
+    // This cannot be inferred from rank because distributed runs may map
+    // ranks to arbitrary device lists (for example devices=[2,4,5,7]).
+    aRet = aclrtSetDevice(device_id);
+    if (aRet != ACL_SUCCESS) {
+        fprintf(stderr, "[comm rank %d] aclrtSetDevice(%d) failed: %d\n",
+                rank, device_id, (int)aRet);
+        delete h;
+        return nullptr;
+    }
 
     // RootInfo exchange
     HcclRootInfo rootInfo{};

src/a5/platform/sim/host/comm_sim.cpp

@@ -68,9 +68,10 @@ struct CommHandle_ {
 // API implementation
 // ============================================================================
 
-extern "C" CommHandle comm_init(int rank, int nranks, const char* rootinfo_path) {
+extern "C" CommHandle comm_init(int rank, int nranks, int device_id, const char* rootinfo_path) {
     auto* h = new (std::nothrow) CommHandle_{};
     if (!h) return nullptr;
+    (void)device_id;
 
     h->rank = rank;
     h->nranks = nranks;