supoort reset cuda error state

Author: leoda1

flagcx/adaptor/device/cuda_adaptor.cc

@@ -930,6 +930,11 @@ flagcxResult_t cudaAdaptorSymMulticastFree(void *) { return flagcxSuccess; }
 
 #endif // CUDART_VERSION >= 12010
 
+flagcxResult_t cudaAdaptorGetLastError() {
+  cudaError_t err = cudaGetLastError();
+  return err == cudaSuccess ? flagcxSuccess : flagcxSystemError;
+}
+
 struct flagcxDeviceAdaptor cudaAdaptor {
   "CUDA",
       // Basic functions
@@ -996,6 +1001,7 @@ struct flagcxDeviceAdaptor cudaAdaptor {
       cudaAdaptorSymFlatUnmap, cudaAdaptorSymMulticastSupported,
       cudaAdaptorSymMulticastCreate, cudaAdaptorSymMulticastBind,
       cudaAdaptorSymMulticastTeardown, cudaAdaptorSymMulticastFree,
+      cudaAdaptorGetLastError,
 };
 
 #endif // USE_NVIDIA_ADAPTOR

flagcx/adaptor/include/flagcx_device_adaptor.h

@@ -243,6 +243,8 @@ struct flagcxDeviceAdaptor_latest {
   // Release the multicast object handle returned by symMulticastCreate.
   // Must be called after all ranks have torn down their mappings.
   flagcxResult_t (*symMulticastFree)(void *mcHandle);
+
+  flagcxResult_t (*getLastError)();
 };
 
 #define flagcxDeviceAdaptor flagcxDeviceAdaptor_latest

flagcx/core/flagcx_p2p.cc

@@ -1477,7 +1477,7 @@ static int detectPtrTypeAndMaybeCacheIpc(void *ptr, char *ipcHandleBuf,
     deviceAdaptor->ipcMemHandleFree(handle);
     return FLAGCX_PTR_CUDA;
   }
-
+  if (deviceAdaptor->getLastError) deviceAdaptor->getLastError();
   deviceAdaptor->ipcMemHandleFree(handle);
   return FLAGCX_PTR_HOST;
 }
@@ -2632,11 +2632,22 @@ bool flagcxP2pEngineConnIsLocal(FlagcxP2pConn *conn) {
   return conn != NULL && conn->isLocal;
 }
 
-int flagcxP2pEngineReg(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
-                       FlagcxP2pMr &mrId) {
+int flagcxP2pEngineRegEx(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
+                         int hintType, FlagcxP2pMr &mrId) {
   if (engine == NULL || data == 0)
     return -1;
 
+  auto resolvePtrType = [&](char *ipcHandleBuf,
+                            uint32_t *ipcHandleSize) -> int {
+    if (hintType == FLAGCX_PTR_HOST || hintType == FLAGCX_PTR_CUDA) {
+      if (ipcHandleSize)
+        *ipcHandleSize = 0;
+      return hintType;
+    }
+    return detectPtrTypeAndMaybeCacheIpc(reinterpret_cast<void *>(data),
+                                         ipcHandleBuf, ipcHandleSize);
+  };
+
   if (!flagcxParamMrSortedLookup()) {
     /* Legacy: mutex + hash maps */
     std::lock_guard<std::mutex> lock(gMemMutex);
@@ -2666,8 +2677,7 @@ int flagcxP2pEngineReg(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
     entry.ibDevN = ibDevN;
 
     setEngineDevice(engine);
-    entry.ptrType = detectPtrTypeAndMaybeCacheIpc(
-        reinterpret_cast<void *>(data), entry.ipcHandle, &entry.ipcHandleSize);
+    entry.ptrType = resolvePtrType(entry.ipcHandle, &entry.ipcHandleSize);
     entry.hasIpc = entry.ptrType == FLAGCX_PTR_CUDA && entry.ipcHandleSize > 0;
 
     if (engine->adaptor->regMr(&devCtx, reinterpret_cast<void *>(data), size,
@@ -2721,8 +2731,7 @@ int flagcxP2pEngineReg(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
   memset(ipcHandle, 0, sizeof(ipcHandle));
 
   setEngineDevice(engine);
-  int ptrType = detectPtrTypeAndMaybeCacheIpc(reinterpret_cast<void *>(data),
-                                              ipcHandle, &ipcHandleSize);
+  int ptrType = resolvePtrType(ipcHandle, &ipcHandleSize);
   bool hasIpc = ptrType == FLAGCX_PTR_CUDA && ipcHandleSize > 0;
 
   /* Register with adaptor */
@@ -2766,6 +2775,11 @@ int flagcxP2pEngineReg(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
   return 0;
 }
 
+int flagcxP2pEngineReg(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
+                       FlagcxP2pMr &mrId) {
+  return flagcxP2pEngineRegEx(engine, data, size, 0, mrId);
+}
+
 void flagcxP2pEngineMrDestroy(FlagcxP2pEngine *engine, FlagcxP2pMr mr) {
   if (engine == NULL)
     return;
@@ -3480,6 +3494,21 @@ int flagcxP2pRpcRegister(void *engine, uint64_t addr, uint64_t size,
   return 0;
 }
 
+int flagcxP2pRpcRegisterHost(void *engine, uint64_t addr, uint64_t size,
+                             uint64_t *mrIdOut) {
+  if (mrIdOut == NULL)
+    return -1;
+  FlagcxP2pMr mrId = 0;
+  const int rc = flagcxP2pEngineRegEx(
+      reinterpret_cast<FlagcxP2pEngine *>(engine),
+      static_cast<uintptr_t>(addr), static_cast<size_t>(size),
+      FLAGCX_PTR_HOST, mrId);
+  if (rc != 0)
+    return rc;
+  *mrIdOut = mrId;
+  return 0;
+}
+
 void *flagcxP2pRpcGetConn(void *engine, const char *session) {
   return reinterpret_cast<void *>(flagcxP2pEngineGetConn(
       reinterpret_cast<FlagcxP2pEngine *>(engine), session));

flagcx/include/flagcx_p2p.h

@@ -618,6 +618,12 @@ int flagcxP2pRpcStartServer(void *engine);
 int flagcxP2pRpcRegister(void *engine, uint64_t addr, uint64_t size,
                          uint64_t *mrIdOut);
 
+/* Same interface as flagcxP2pRpcRegister, but on chips/adaptors whose runtime
+   cannot reset the current error state, this host entry avoids producing the 
+   error at all. */
+int flagcxP2pRpcRegisterHost(void *engine, uint64_t addr, uint64_t size,
+                             uint64_t *mrIdOut);
+
 /* Get (or lazily establish) a cached connection to "host:port". */
 void *flagcxP2pRpcGetConn(void *engine, const char *session);
 

plugin/interservice/flagcx_wrapper.py

@@ -419,6 +419,10 @@ class FLAGCXLibrary:
             flagcxP2pEngine_t, ctypes.c_uint64, ctypes.c_uint64,
             ctypes.POINTER(ctypes.c_uint64)
         ]),
+        Function("flagcxP2pRpcRegisterHost", ctypes.c_int, [
+            flagcxP2pEngine_t, ctypes.c_uint64, ctypes.c_uint64,
+            ctypes.POINTER(ctypes.c_uint64)
+        ]),
         Function("flagcxP2pRpcGetConn", flagcxP2pConn_t, [
             flagcxP2pEngine_t, ctypes.c_char_p
         ]),
@@ -447,9 +451,12 @@ def _find_default_library() -> str:
             so_path = os.path.join(flagcx_path, "lib", "libflagcx.so")
             if os.path.isfile(so_path):
                 return so_path
+            build_so_path = os.path.join(flagcx_path, "build", "lib", "libflagcx.so")
+            if os.path.isfile(build_so_path):
+                return build_so_path
             raise FileNotFoundError(
-                f"FLAGCX_PATH is set to '{flagcx_path}' but "
-                f"'{so_path}' does not exist. "
+                f"FLAGCX_PATH is set to '{flagcx_path}' but neither "
+                f"'{so_path}' nor '{build_so_path}' exists. "
                 f"Please build FlagCX or check FLAGCX_PATH."
             )
         # 2. Fall back to <repo_root>/build/lib/libflagcx.so
@@ -818,6 +825,17 @@ def flagcxP2pRegister(self, engine: flagcxP2pEngine_t,
                 f"flagcxP2pRegister failed (addr={hex(addr)}, size={size})")
         return mr_id.value
 
+    def flagcxP2pRegisterHost(self, engine: flagcxP2pEngine_t,
+                              addr: int, size: int) -> int:
+        mr_id = ctypes.c_uint64(0)
+        rc = self._funcs["flagcxP2pRpcRegisterHost"](
+            engine, ctypes.c_uint64(addr), ctypes.c_uint64(size),
+            ctypes.byref(mr_id))
+        if rc != 0:
+            raise RuntimeError(
+                f"flagcxP2pRegisterHost failed (addr={hex(addr)}, size={size})")
+        return mr_id.value
+
     def flagcxP2pGetConn(self, engine: flagcxP2pEngine_t,
                          session: str) -> flagcxP2pConn_t:
         conn = self._funcs["flagcxP2pRpcGetConn"](