Benchmark for tuple generators (#83)

Summary: Pull Request resolved: #83 Adding benchmarks for the tuple generator and two-party tuple generator. Reviewed By: RuiyuZhu Differential Revision: D34833460 fbshipit-source-id: 700ef5ad205aab55eb6a99e699bbca21ceb58d0d
facebookresearch · Mar 16, 2022 · 4c87e42 · 4c87e42
1 parent 4861c1d
commit 4c87e42
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 9 deletions.
diff --git a/fbpcf/engine/tuple_generator/test/TupleGeneratorTest.cpp b/fbpcf/engine/tuple_generator/test/TupleGeneratorTest.cpp
@@ -84,17 +84,14 @@ TEST(TupleGeneratorTest, testWithDummyProductShareGenerator) {
   int numberOfParty = 4;
 
   testTupleGenerator(
-      numberOfParty,
-
-      createInMemoryTupleGeneratorFactoryWithDummyProductShareGenerator);
+      numberOfParty, createTupleGeneratorFactoryWithDummyProductShareGenerator);
 }
 
 TEST(TupleGeneratorTest, testWithSecureProductShareGenerator) {
   int numberOfParty = 4;
 
   testTupleGenerator(
-      numberOfParty,
-      createInMemoryTupleGeneratorFactoryWithRealProductShareGenerator);
+      numberOfParty, createTupleGeneratorFactoryWithRealProductShareGenerator);
 }
 
 TEST(TupleGeneratorTest, testTwoPartyTupleGeneratorWithDummyRcot) {

diff --git a/fbpcf/engine/tuple_generator/test/TupleGeneratorTestHelper.h b/fbpcf/engine/tuple_generator/test/TupleGeneratorTestHelper.h
@@ -12,8 +12,6 @@
 #include <memory>
 
 #include "fbpcf/engine/communication/IPartyCommunicationAgentFactory.h"
-#include "fbpcf/engine/communication/InMemoryPartyCommunicationAgentFactory.h"
-#include "fbpcf/engine/communication/InMemoryPartyCommunicationAgentHost.h"
 #include "fbpcf/engine/communication/test/AgentFactoryCreationHelper.h"
 #include "fbpcf/engine/tuple_generator/DummyProductShareGeneratorFactory.h"
 #include "fbpcf/engine/tuple_generator/DummyTupleGeneratorFactory.h"
@@ -54,7 +52,7 @@ inline std::unique_ptr<ITupleGeneratorFactory> createDummyTupleGeneratorFactory(
 }
 
 inline std::unique_ptr<ITupleGeneratorFactory>
-createInMemoryTupleGeneratorFactoryWithDummyProductShareGenerator(
+createTupleGeneratorFactoryWithDummyProductShareGenerator(
     int numberOfParty,
     int myId,
     communication::IPartyCommunicationAgentFactory& agentFactory) {
@@ -69,7 +67,7 @@ createInMemoryTupleGeneratorFactoryWithDummyProductShareGenerator(
 }
 
 inline std::unique_ptr<ITupleGeneratorFactory>
-createInMemoryTupleGeneratorFactoryWithRealProductShareGenerator(
+createTupleGeneratorFactoryWithRealProductShareGenerator(
     int numberOfParty,
     int myId,
     communication::IPartyCommunicationAgentFactory& agentFactory) {

diff --git a/fbpcf/engine/tuple_generator/test/benchmarks/TupleGeneratorBenchmark.cpp b/fbpcf/engine/tuple_generator/test/benchmarks/TupleGeneratorBenchmark.cpp
@@ -87,6 +87,97 @@ BENCHMARK_COUNTERS(ProductShareGenerator, counters) {
   benchmark.runBenchmark(counters);
 }
 
+class BaseTupleGeneratorBenchmark : public util::NetworkedBenchmark {
+ public:
+  void setup() override {
+    auto [agentFactory0, agentFactory1] = util::getSocketAgentFactories();
+    agentFactory0_ = std::move(agentFactory0);
+    agentFactory1_ = std::move(agentFactory1);
+
+    senderFactory_ = getTupleGeneratorFactory(0, *agentFactory0_);
+    receiverFactory_ = getTupleGeneratorFactory(1, *agentFactory1_);
+  }
+
+  void runSender() override {
+    sender_ = senderFactory_->create();
+    sender_->getBooleanTuple(size_);
+  }
+
+  void runReceiver() override {
+    receiver_ = receiverFactory_->create();
+    receiver_->getBooleanTuple(size_);
+  }
+
+  std::pair<uint64_t, uint64_t> getTrafficStatistics() override {
+    return sender_->getTrafficStatistics();
+  }
+
+ protected:
+  virtual std::unique_ptr<ITupleGeneratorFactory> getTupleGeneratorFactory(
+      int myId,
+      communication::IPartyCommunicationAgentFactory& agentFactory) = 0;
+
+  size_t bufferSize_ = 1600000;
+
+ private:
+  size_t size_ = 1000000;
+
+  std::unique_ptr<communication::IPartyCommunicationAgentFactory>
+      agentFactory0_;
+  std::unique_ptr<communication::IPartyCommunicationAgentFactory>
+      agentFactory1_;
+
+  std::unique_ptr<ITupleGeneratorFactory> senderFactory_;
+  std::unique_ptr<ITupleGeneratorFactory> receiverFactory_;
+
+  std::unique_ptr<ITupleGenerator> sender_;
+  std::unique_ptr<ITupleGenerator> receiver_;
+};
+
+class TupleGeneratorBenchmark final : public BaseTupleGeneratorBenchmark {
+ protected:
+  std::unique_ptr<ITupleGeneratorFactory> getTupleGeneratorFactory(
+      int myId,
+      communication::IPartyCommunicationAgentFactory& agentFactory) override {
+    auto otFactory = std::make_unique<
+        oblivious_transfer::RcotBasedBidirectionObliviousTransferFactory<bool>>(
+        myId, agentFactory, oblivious_transfer::createFerretRcotFactory());
+    auto productShareGeneratorFactory =
+        std::make_unique<ProductShareGeneratorFactory<bool>>(
+            std::make_unique<util::AesPrgFactory>(bufferSize_),
+            std::move(otFactory));
+    return std::make_unique<TupleGeneratorFactory>(
+        std::move(productShareGeneratorFactory),
+        std::make_unique<util::AesPrgFactory>(),
+        bufferSize_,
+        myId,
+        2);
+  }
+};
+
+BENCHMARK_COUNTERS(TupleGenerator, counters) {
+  TupleGeneratorBenchmark benchmark;
+  benchmark.runBenchmark(counters);
+}
+
+class TwoPartyTupleGeneratorBenchmark final
+    : public BaseTupleGeneratorBenchmark {
+ protected:
+  std::unique_ptr<ITupleGeneratorFactory> getTupleGeneratorFactory(
+      int myId,
+      communication::IPartyCommunicationAgentFactory& agentFactory) override {
+    return std::make_unique<TwoPartyTupleGeneratorFactory>(
+        oblivious_transfer::createFerretRcotFactory(),
+        agentFactory,
+        myId,
+        bufferSize_);
+  }
+};
+
+BENCHMARK_COUNTERS(TwoPartyTupleGenerator, counters) {
+  TwoPartyTupleGeneratorBenchmark benchmark;
+  benchmark.runBenchmark(counters);
+}
 } // namespace fbpcf::engine::tuple_generator
 
 int main(int argc, char* argv[]) {