WIP

Author: Binyang2014

src/include/utils_internal.hpp

@@ -74,7 +74,7 @@ class TokenPool : public std::enable_shared_from_this<TokenPool> {
   uint64_t* baseAddr_;
   uint64_t tailMask_;
   std::shared_ptr<uint64_t> tokens_;
-  std::vector<uint64_t> allocationMap_;
+  std::vector<std::bitset<UINT64_WIDTH>> allocationMap_;
 };
 
 }  // namespace mscclpp

src/semaphore.cc

@@ -29,7 +29,6 @@ struct SemaphoreStub::Impl {
 };
 
 std::shared_ptr<uint64_t> SemaphoreStub::Impl::gpuCallocToken(std::shared_ptr<Context> context) {
-  // NVLS is not supported in ROCm/HIP
 #if (CUDA_NVLS_API_AVAILABLE)
   if (isNvlsSupported()) {
     return context->pimpl_->getToken();

src/utils_internal.cc

@@ -251,22 +251,22 @@ std::shared_ptr<uint64_t> TokenPool::getToken() {
     size_t index = (token - self->baseAddr_) / UINT64_WIDTH;
     size_t bit = (token - self->baseAddr_) % UINT64_WIDTH;
     uint64_t mask = 1UL << bit;
-    if ((self->allocationMap_[index] & mask) == 0) {
-      WARN("TokenPool tried to free a token that was not allocated");
-      return;
-    }
     self->allocationMap_[index] &= ~mask;
   };
 
   size_t size = allocationMap_.size();
   for (size_t i = 0; i < size; i++) {
+    uint64_t ullong = allocationMap_[i].to_ullong();
     uint64_t mask = (i + 1 == size) ? tailMask_ : ~0ULL;
-    uint64_t holes = (~allocationMap_[i]) & mask;
+    uint64_t holes = (~ullong) & mask;
     if (!holes) continue;
-    size_t bit = __builtin_ctzll(holes);
-    allocationMap_[i] |= (1UL << bit);
-    INFO(MSCCLPP_ALLOC, "TokenPool allocated token at addr %p", baseAddr_ + i * UINT64_WIDTH + bit);
-    return std::shared_ptr<uint64_t>(baseAddr_ + i * UINT64_WIDTH + bit, deleter);
+    for (int bit = 0; bit < UINT64_WIDTH; bit++) {
+      if (holes & (1UL << bit)) {
+        allocationMap_[i].set(bit);
+        INFO(MSCCLPP_ALLOC, "TokenPool allocated token at addr %p", baseAddr_ + i * UINT64_WIDTH + bit);
+        return std::shared_ptr<uint64_t>(baseAddr_ + i * UINT64_WIDTH + bit, deleter);
+      }
+    }
   }
   throw Error("TokenPool is exhausted", ErrorCode::InternalError);
 }