optimize vole by cache

.vscode/setting.json

@@ -0,0 +1,11 @@
+{
+    "editor.formatOnSave": true,
+    "files.watcherExclude": {
+        "**/.git/**": true,
+        "**/.cache/**": true,
+        "**/bazel-*/**": true,
+        "**/external/**": true
+    },
+    "files.insertFinalNewline": true,
+    "files.trimTrailingWhitespace": true
+}

psi/algorithm/BUILD.bazel

@@ -29,6 +29,7 @@ psi_cc_library(
     hdrs = ["psi_io.h"],
     deps = [
         ":types",
+        "//psi/utils:hash_bucket_cache",
     ],
 )
 

psi/algorithm/psi_io.h

@@ -14,13 +14,15 @@
 
 #pragma once
 
+#include <cstddef>
 #include <memory>
 #include <string>
 #include <vector>
 
 #include "yacl/base/int128.h"
 
 #include "psi/algorithm/types.h"
+#include "psi/utils/hash_bucket_cache.h"
 
 namespace psi {
 
@@ -57,13 +59,27 @@ class IDataProvider {
  public:
   virtual std::vector<PsiItemData> ReadNext(size_t size) = 0;
   virtual std::vector<PsiItemData> ReadAll() = 0;
+  [[nodiscard]] virtual size_t Size() const = 0;
 };
 
 class IDataStore {
  public:
   [[nodiscard]] virtual size_t GetBucketNum() const = 0;
+  [[nodiscard]] virtual size_t GetBucketDatasize(size_t tag) const = 0;
   virtual std::shared_ptr<IDataProvider> Load(size_t tag) = 0;
   virtual ~IDataStore() = default;
 };
 
+class IBucketDataStore {
+ public:
+  virtual std::vector<HashBucketCache::BucketItem> GetBucketItems(
+      size_t bucket_idx) = 0;
+  virtual void WriteIntersetionItems(
+      size_t bucket_idx, const std::vector<HashBucketCache::BucketItem>& items,
+      const std::vector<uint32_t>& intersection_indices,
+      const std::vector<uint32_t>& peer_dup_cnts) = 0;
+  virtual std::pair<size_t, size_t> GetBucketDatasize(size_t bucket_idx) = 0;
+  virtual ~IBucketDataStore() = default;
+};
+
 }  // namespace psi

psi/algorithm/rr22/BUILD.bazel

@@ -102,6 +102,7 @@ psi_cc_library(
         ":rr22_utils",
         "//psi/proto:psi_v2_cc_proto",
         "//psi/utils:bucket",
+        "//psi/utils:simple_channel",
         "//psi/utils:sync",
     ],
 )

psi/algorithm/rr22/common.cc

@@ -16,6 +16,7 @@
 
 #include "omp.h"
 
+#include "psi/algorithm/rr22/rr22_psi.h"
 #include "psi/utils/bucket.h"
 
 namespace psi::rr22 {
@@ -28,4 +29,52 @@ Rr22PsiOptions GenerateRr22PsiOptions(bool low_comm_mode) {
   return options;
 }
 
+BucketDataStoreImpl::BucketDataStoreImpl(
+    std::shared_ptr<yacl::link::Context> lctx,
+    HashBucketCache* input_bucket_store,
+    IndexWriter* intersection_indices_writer, RecoveryManager* recovery_manager)
+    : IBucketDataStore(),
+      input_bucket_store_(input_bucket_store),
+      intersection_indices_writer_(intersection_indices_writer),
+      recovery_manager_(recovery_manager),
+      lctx_(std::move(lctx)) {
+  self_sizes_ = std::vector<uint32_t>(input_bucket_store_->BucketNum());
+  for (size_t i = 0; i < self_sizes_.size(); i++) {
+    self_sizes_[i] = input_bucket_store_->GetBucketSize(i);
+  }
+  peer_sizes_ = std::vector<uint32_t>(input_bucket_store_->BucketNum());
+  yacl::ByteContainerView buffer(self_sizes_.data(),
+                                 self_sizes_.size() * sizeof(uint32_t));
+  auto data = yacl::link::AllGather(lctx_, buffer, "exchange size");
+  std::memcpy(peer_sizes_.data(), data[lctx_->NextRank()].data(),
+              self_sizes_.size() * sizeof(uint32_t));
+}
+
+std::vector<HashBucketCache::BucketItem> BucketDataStoreImpl::GetBucketItems(
+    size_t bucket_idx) {
+  if (bucket_idx >= input_bucket_store_->BucketNum()) {
+    return {};
+  }
+  return input_bucket_store_->LoadBucketItems(bucket_idx);
+}
+
+void BucketDataStoreImpl::WriteIntersetionItems(
+    size_t bucket_idx, const std::vector<HashBucketCache::BucketItem>& items,
+    const std::vector<uint32_t>& intersection_indices,
+    const std::vector<uint32_t>& peer_dup_cnts) {
+  for (size_t i = 0; i != intersection_indices.size(); ++i) {
+    intersection_indices_writer_->WriteCache(
+        items[intersection_indices[i]].index, peer_dup_cnts[i]);
+  }
+  intersection_indices_writer_->Commit();
+  if (recovery_manager_ != nullptr) {
+    recovery_manager_->UpdateParsedBucketCount(bucket_idx + 1);
+  }
+}
+
+std::pair<size_t, size_t> BucketDataStoreImpl::GetBucketDatasize(
+    size_t bucket_idx) {
+  return std::make_pair(self_sizes_[bucket_idx], peer_sizes_[bucket_idx]);
+}
+
 }  // namespace psi::rr22

psi/algorithm/rr22/common.h

@@ -31,4 +31,30 @@ constexpr bool kDefaultCompress = true;
 
 Rr22PsiOptions GenerateRr22PsiOptions(bool low_comm_mode);
 
+class BucketDataStoreImpl : public IBucketDataStore {
+ public:
+  BucketDataStoreImpl(std::shared_ptr<yacl::link::Context> lctx,
+                      HashBucketCache* input_bucket_store,
+                      IndexWriter* intersection_indices_writer,
+                      RecoveryManager* recovery_manager);
+
+  std::vector<HashBucketCache::BucketItem> GetBucketItems(
+      size_t bucket_idx) override;
+
+  void WriteIntersetionItems(
+      size_t bucket_idx, const std::vector<HashBucketCache::BucketItem>& items,
+      const std::vector<uint32_t>& intersection_indices,
+      const std::vector<uint32_t>& peer_dup_cnts) override;
+
+  std::pair<size_t, size_t> GetBucketDatasize(size_t bucket_idx) override;
+
+ private:
+  HashBucketCache* input_bucket_store_;
+  IndexWriter* intersection_indices_writer_;
+  RecoveryManager* recovery_manager_;
+  std::shared_ptr<yacl::link::Context> lctx_;
+  std::vector<uint32_t> peer_sizes_;
+  std::vector<uint32_t> self_sizes_;
+};
+
 }  // namespace psi::rr22

psi/algorithm/rr22/receiver.cc

@@ -112,32 +112,17 @@ void Rr22PsiReceiver::Online() {
   Rr22PsiOptions rr22_options = GenerateRr22PsiOptions(
       config_.protocol_config().rr22_config().low_comm_mode());
 
-  PreProcessFunc pre_f =
-      [&](size_t idx) -> std::vector<HashBucketCache::BucketItem> {
-    if (idx >= input_bucket_store_->BucketNum()) {
-      return {};
-    }
-    return input_bucket_store_->LoadBucketItems(idx);
-  };
-  PostProcessFunc post_f =
-      [&](size_t bucket_idx,
-          const std::vector<HashBucketCache::BucketItem>& bucket_items,
-          const std::vector<uint32_t>& indices,
-          const std::vector<uint32_t>& peer_cnt) {
-        for (size_t i = 0; i != indices.size(); ++i) {
-          intersection_indices_writer_->WriteCache(
-              bucket_items[indices[i]].index, peer_cnt[i]);
-        }
-        intersection_indices_writer_->Commit();
-        if (recovery_manager_) {
-          recovery_manager_->UpdateParsedBucketCount(bucket_idx + 1);
-        }
-      };
-
+  BucketDataStoreImpl data_processor(lctx_, input_bucket_store_.get(),
+                                     intersection_indices_writer_.get(),
+                                     recovery_manager_.get());
   Rr22Runner runner(lctx_, rr22_options, input_bucket_store_->BucketNum(),
-                    config_.protocol_config().broadcast_result(), pre_f,
-                    post_f);
-  SyncWait(lctx_, [&] { runner.AsyncRun(bucket_idx, false); });
+                    config_.protocol_config().broadcast_result(),
+                    &data_processor);
+  auto scoped_temp_dir = std::make_unique<ScopedTempDir>();
+  scoped_temp_dir->CreateUniqueTempDirUnderPath(GetTaskDir());
+  SyncWait(lctx_, [&] {
+    runner.AsyncRun(bucket_idx, false, scoped_temp_dir->path());
+  });
   SPDLOG_INFO("[Rr22PsiReceiver::Online] end");
 }
 

psi/algorithm/rr22/rr22_operator.cc

@@ -14,10 +14,90 @@
 
 #include "psi/algorithm/rr22/rr22_operator.h"
 
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
 #include <memory>
+#include <utility>
+#include <vector>
+
+#include "psi/algorithm/rr22/rr22_psi.h"
 
 namespace psi::rr22 {
 
+namespace {
+class BucketDataStoreImpl : public IBucketDataStore {
+ public:
+  BucketDataStoreImpl(std::shared_ptr<yacl::link::Context> lctx,
+                      IDataStore* input_store, IResultStore* output_store,
+                      RecoveryManager* recovery_manager)
+      : input_store_(input_store),
+        output_store_(output_store),
+        recovery_manager_(recovery_manager),
+        lctx_(std::move(lctx)) {
+    self_sizes_ = std::vector<uint32_t>(input_store_->GetBucketNum());
+    for (size_t i = 0; i < self_sizes_.size(); i++) {
+      self_sizes_[i] = input_store_->GetBucketDatasize(i);
+    }
+    peer_sizes_ = std::vector<uint32_t>(input_store_->GetBucketNum());
+    yacl::ByteContainerView buffer(self_sizes_.data(),
+                                   self_sizes_.size() * sizeof(uint32_t));
+    auto data = yacl::link::AllGather(lctx_, buffer, "exchange size");
+    std::memcpy(peer_sizes_.data(), data[lctx_->NextRank()].data(),
+                self_sizes_.size() * sizeof(uint32_t));
+  };
+
+  std::vector<HashBucketCache::BucketItem> GetBucketItems(
+      size_t bucket_idx) override {
+    if (bucket_idx >= input_store_->GetBucketNum()) {
+      return {};
+    }
+    std::vector<HashBucketCache::BucketItem> bucket_items;
+    auto provider = input_store_->Load(bucket_idx);
+    auto item_datas = provider->ReadAll();
+    bucket_items.reserve(item_datas.size());
+    for (size_t i = 0; i < item_datas.size(); ++i) {
+      bucket_items.emplace_back(HashBucketCache::BucketItem{
+          .index = i,
+          .extra_dup_cnt = item_datas[i].cnt - 1,
+          .base64_data = std::move(item_datas[i].buf)});
+    }
+    return bucket_items;
+  };
+
+  void WriteIntersetionItems(
+      size_t bucket_idx, const std::vector<HashBucketCache::BucketItem>& items,
+      const std::vector<uint32_t>& intersection_indices,
+      const std::vector<uint32_t>& peer_dup_cnts) override {
+    std::vector<PsiResultIndex> indices;
+    indices.reserve(intersection_indices.size());
+
+    for (size_t i = 0; i < intersection_indices.size(); ++i) {
+      indices.emplace_back(
+          PsiResultIndex{.data = items[intersection_indices[i]].index,
+                         .peer_item_cnt = peer_dup_cnts[i] + 1});
+    }
+    auto recevier = output_store_->GetReceiver(bucket_idx);
+    recevier->Add(std::move(indices));
+    if (recovery_manager_ != nullptr) {
+      recovery_manager_->UpdateParsedBucketCount(bucket_idx + 1);
+    }
+  };
+
+  std::pair<size_t, size_t> GetBucketDatasize(size_t bucket_idx) override {
+    return std::make_pair(self_sizes_[bucket_idx], peer_sizes_[bucket_idx]);
+  };
+
+ private:
+  IDataStore* input_store_;
+  IResultStore* output_store_;
+  RecoveryManager* recovery_manager_;
+  std::shared_ptr<yacl::link::Context> lctx_;
+  std::vector<uint32_t> peer_sizes_;
+  std::vector<uint32_t> self_sizes_;
+};
+}  // namespace
+
 Rr22Operator::Rr22Operator(Options opts,
                            std::shared_ptr<IDataStore> input_store,
                            std::shared_ptr<IResultStore> output_store)
@@ -32,36 +112,9 @@ bool Rr22Operator::ReceiveResult() {
 void Rr22Operator::OnInit() {}
 
 void Rr22Operator::OnRun() {
-  PreProcessFunc pre_process_func =
-      [this](size_t bucket_idx) -> std::vector<HashBucketCache::BucketItem> {
-    std::vector<HashBucketCache::BucketItem> bucket_items;
-    auto provider = input_store_->Load(bucket_idx);
-    auto item_datas = provider->ReadAll();
-    bucket_items.reserve(item_datas.size());
-    for (size_t i = 0; i < item_datas.size(); ++i) {
-      bucket_items.emplace_back(HashBucketCache::BucketItem{
-          .index = i,
-          .extra_dup_cnt = item_datas[i].cnt - 1,
-          .base64_data = std::move(item_datas[i].buf)});
-    }
-    return bucket_items;
-  };
-  PostProcessFunc post_process_func =
-      [this](size_t bucket_idx,
-             const std::vector<HashBucketCache::BucketItem>& bucket_items,
-             const std::vector<uint32_t>& intersection_indices,
-             const std::vector<uint32_t>& peer_dup_cnts) {
-        std::vector<PsiResultIndex> indices;
-        indices.reserve(intersection_indices.size());
-
-        for (size_t i = 0; i < intersection_indices.size(); ++i) {
-          indices.emplace_back(PsiResultIndex{
-              .data = bucket_items[intersection_indices[i]].index,
-              .peer_item_cnt = peer_dup_cnts[i] + 1});
-        }
-        auto recevier = output_store_->GetReceiver(bucket_idx);
-        recevier->Add(std::move(indices));
-      };
+  BucketDataStoreImpl data_processor(link_ctx_, input_store_.get(),
+                                     output_store_.get(),
+                                     recovery_manager_.get());
 
   size_t bucket_idx =
       recovery_manager_
@@ -70,11 +123,11 @@ void Rr22Operator::OnRun() {
           : 0;
 
   Rr22Runner runner(link_ctx_, opts_.rr22_opts, input_store_->GetBucketNum(),
-                    opts_.broadcast_result, pre_process_func,
-                    post_process_func);
+                    opts_.broadcast_result, &data_processor);
 
   if (opts_.pipeline_mode) {
-    runner.AsyncRun(bucket_idx, opts_.lctx->Rank() != opts_.receiver_rank);
+    runner.AsyncRun(bucket_idx, opts_.lctx->Rank() != opts_.receiver_rank,
+                    opts_.cache_dir);
   } else {
     runner.ParallelRun(bucket_idx, opts_.lctx->Rank() != opts_.receiver_rank,
                        opts_.parallel_level);

psi/algorithm/rr22/rr22_operator.h

@@ -31,6 +31,8 @@ class Rr22Operator : public PsiOperator {
     size_t parallel_level = 6;
 
     std::shared_ptr<RecoveryManager> recovery_manager = nullptr;
+    std::string cache_dir =
+        std::filesystem::temp_directory_path() / GetRandomString();
   };
 
  public: