[feat] Add configurable scattergatter option for pcstore (#483)

Add configurable scattergatter

Author: qyh111

ucm/store/pcstore/cc/api/pcstore.cc

@@ -41,7 +41,7 @@ class PcStoreImpl : public PcStore {
                 config.transferLocalRankSize, config.transferDeviceId, config.transferStreamNumber,
                 config.kvcacheBlockSize, config.transferIoSize, config.transferIoDirect,
                 config.transferBufferNumber, this->spaceMgr_.GetSpaceLayout(),
-                config.transferTimeoutMs);
+                config.transferTimeoutMs, config.transferScatterGatherEnable);
             if (status.Failure()) { return status.Underlying(); }
         }
         this->ShowConfig(config);
@@ -93,6 +93,7 @@ class PcStoreImpl : public PcStore {
         UC_INFO("Set UC::StreamNumber to {}.", config.transferStreamNumber);
         UC_INFO("Set UC::BufferNumber to {}.", config.transferBufferNumber);
         UC_INFO("Set UC::TimeoutMs to {}.", config.transferTimeoutMs);
+        UC_INFO("Set UC::ScatterGatherEnable to {}.", config.transferScatterGatherEnable);
     }
 
 private:

ucm/store/pcstore/cc/api/pcstore.h

@@ -42,6 +42,7 @@ class PcStore : CCStore<TransTask> {
         size_t transferStreamNumber{8};
         size_t transferBufferNumber{4096};
         size_t transferTimeoutMs{30000};
+        bool transferScatterGatherEnable{false};
 
         Config(const std::vector<std::string>& storageBackends, const size_t kvcacheBlockSize,
                const bool transferEnable)

ucm/store/pcstore/cc/domain/trans/trans_manager.cc

@@ -29,7 +29,7 @@ namespace UC {
 Status TransManager::Setup(const size_t rankSize, const int32_t deviceId, const size_t streamNumber,
                            const size_t blockSize, const size_t ioSize, const bool ioDirect,
                            const size_t bufferNumber, const SpaceLayout* layout,
-                           const size_t timeoutMs)
+                           const size_t timeoutMs, const bool scatterGatherEnable)
 {
     auto s = Status::OK();
     if (rankSize > 1) {
@@ -38,7 +38,7 @@ Status TransManager::Setup(const size_t rankSize, const int32_t deviceId, const
         if (s.Failure()) { return s; }
     }
     s = this->queue_.Setup(deviceId, streamNumber, blockSize, ioSize, ioDirect, bufferNumber,
-                           layout, &this->failureSet_);
+                           layout, &this->failureSet_, scatterGatherEnable);
     if (s.Failure()) { return s; }
     this->rankSize_ = rankSize;
     this->timeoutMs_ = timeoutMs;

ucm/store/pcstore/cc/domain/trans/trans_manager.h

@@ -33,7 +33,8 @@ class TransManager {
 public:
     Status Setup(const size_t rankSize, const int32_t deviceId, const size_t streamNumber,
                  const size_t blockSize, const size_t ioSize, const bool ioDirect,
-                 const size_t bufferNumber, const SpaceLayout* layout, const size_t timeoutMs);
+                 const size_t bufferNumber, const SpaceLayout* layout, const size_t timeoutMs,
+                 const bool scatterGatherEnable);
     Status Submit(TransTask task, size_t& taskId) noexcept;
     Status Wait(const size_t taskId) noexcept;
     Status Check(const size_t taskId, bool& finish) noexcept;

ucm/store/pcstore/cc/domain/trans/trans_queue.cc

@@ -77,7 +77,8 @@ void TransQueue::FileWorker(BlockTask&& task)
 
 Status TransQueue::Setup(const int32_t deviceId, const size_t streamNumber, const size_t blockSize,
                          const size_t ioSize, const bool ioDirect, const size_t bufferNumber,
-                         const SpaceLayout* layout, TaskSet* failureSet_)
+                         const SpaceLayout* layout, TaskSet* failureSet_,
+                         const bool scatterGatherEnable)
 {
     Trans::Device device;
     auto ts = device.Setup(deviceId);
@@ -91,6 +92,14 @@ Status TransQueue::Setup(const int32_t deviceId, const size_t streamNumber, cons
         UC_ERROR("Failed to make buffer and stream on device({}).", deviceId);
         return Status::Error();
     }
+    if (scatterGatherEnable) {
+        devBuffer_ = device.MakeBuffer();
+        smStream_ = device.MakeSMStream();
+        if (!devBuffer_ || !smStream_) {
+            UC_ERROR("Failed to make devBuffer and smStream on device({}).", deviceId);
+            return Status::Error();
+        }
+    }
     ts = buffer_->MakeHostBuffers(blockSize, bufferNumber);
     if (ts.Failure()) {
         UC_ERROR("Failed({}) to make host buffer({},{}).", ts.ToString(), blockSize, bufferNumber);
@@ -108,13 +117,18 @@ Status TransQueue::Setup(const int32_t deviceId, const size_t streamNumber, cons
     this->ioSize_ = ioSize;
     this->ioDirect_ = ioDirect;
     this->failureSet_ = failureSet_;
+    this->scatterGatherEnable_ = scatterGatherEnable;
     return Status::OK();
 }
 
 void TransQueue::Dispatch(TaskPtr task, WaiterPtr waiter)
 {
     if (task->type == TransTask::Type::DUMP) {
-        this->DispatchDump(task, waiter);
+        if (this->scatterGatherEnable_) {
+            this->DispatchSatterGatherDump(task, waiter);
+        } else {
+            this->DispatchDump(task, waiter);
+        }
         return;
     }
     task->ForEachGroup(
@@ -171,4 +185,40 @@ void TransQueue::DispatchDump(TaskPtr task, WaiterPtr waiter)
     }
 }
 
+void TransQueue::DispatchSatterGatherDump(TaskPtr task, WaiterPtr waiter)
+{
+    std::vector<BlockTask> blocks;
+    blocks.reserve(task->GroupNumber());
+    std::vector<std::shared_ptr<void>> addrs;
+    addrs.reserve(task->GroupNumber());
+    task->ForEachGroup(
+        [task, &blocks, &addrs, this](const std::string& block, std::vector<uintptr_t>& shards) {
+            BlockTask blockTask;
+            blockTask.owner = task->id;
+            blockTask.block = block;
+            blockTask.type = task->type;
+            auto number = shards.size();
+            auto bufferSize = this->ioSize_ * number;
+            blockTask.buffer = buffer_->GetHostBuffer(bufferSize);
+            std::swap(blockTask.shards, shards);
+            auto device = (void*)blockTask.shards.data();
+            auto host = blockTask.buffer.get();
+            auto devAddr = this->devBuffer_->MakeDeviceBuffer(sizeof(void*) * number);
+            smStream_->HostToDeviceAsync(device, devAddr.get(), sizeof(void*) * number);
+            smStream_->DeviceToHostAsync((void**)devAddr.get(), host, this->ioSize_, number);
+            addrs.push_back(devAddr);
+            blocks.push_back(std::move(blockTask));
+        });
+    auto s = smStream_->Synchronized();
+    if (s.Failure()) { this->failureSet_->Insert(task->id); }
+    for (auto&& block : blocks) {
+        if (s.Failure()) {
+            waiter->Done(nullptr);
+            return;
+        }
+        this->filePool_.Push(std::move(block));
+        waiter->Done([task, ioSize = this->ioSize_] { UC_DEBUG("{}", task->Epilog(ioSize)); });
+    }
+}
+
 } // namespace UC

ucm/store/pcstore/cc/domain/trans/trans_queue.h

@@ -51,19 +51,23 @@ class TransQueue {
 public:
     Status Setup(const int32_t deviceId, const size_t streamNumber, const size_t blockSize,
                  const size_t ioSize, const bool ioDirect, const size_t bufferNumber,
-                 const SpaceLayout* layout, TaskSet* failureSet_);
+                 const SpaceLayout* layout, TaskSet* failureSet_, const bool scatterGatherEnable);
     void Dispatch(TaskPtr task, WaiterPtr waiter);
     void DispatchDump(TaskPtr task, WaiterPtr waiter);
+    void DispatchSatterGatherDump(TaskPtr task, WaiterPtr waiter);
 
 private:
     std::unique_ptr<Trans::Buffer> buffer_{nullptr};
     std::unique_ptr<Trans::Stream> stream_{nullptr};
+    std::unique_ptr<Trans::Buffer> devBuffer_{nullptr};
+    std::unique_ptr<Trans::Stream> smStream_{nullptr};
     const SpaceLayout* layout_;
     size_t ioSize_;
     bool ioDirect_;
     ThreadPool<BlockTask> devPool_;
     ThreadPool<BlockTask> filePool_;
     TaskSet* failureSet_;
+    bool scatterGatherEnable_;
 };
 
 } // namespace UC

ucm/store/pcstore/cpy/pcstore.py.cc

@@ -101,6 +101,8 @@ PYBIND11_MODULE(ucmpcstore, module)
     config.def_readwrite("transferIoSize", &UC::PcStorePy::Config::transferIoSize);
     config.def_readwrite("transferBufferNumber", &UC::PcStorePy::Config::transferBufferNumber);
     config.def_readwrite("transferTimeoutMs", &UC::PcStorePy::Config::transferTimeoutMs);
+    config.def_readwrite("transferScatterGatherEnable",
+                         &UC::PcStorePy::Config::transferScatterGatherEnable);
     store.def(py::init<>());
     store.def("CCStoreImpl", &UC::PcStorePy::CCStoreImpl);
     store.def("Setup", &UC::PcStorePy::Setup);

ucm/store/pcstore/pcstore_connector.py

@@ -53,6 +53,7 @@ def __init__(self, config: Dict):
             param.transferStreamNumber = config.get("stream_number", 8)
             param.transferBufferNumber = config.get("buffer_number", 4096)
             param.transferLocalRankSize = config.get("local_rank_size", 8)
+            param.transferScatterGatherEnable = config.get("use_scatter_gatter", False)
         ret = self.store.Setup(param)
         if ret != 0:
             msg = f"Failed to initialize ucmpcstore, errcode: {ret}."