Implementation of benchmark with grouped queries

Author: pgmpablo157321

loadgen/bindings/c_api.cc

@@ -88,6 +88,9 @@ class QuerySampleLibraryTrampoline : public QuerySampleLibrary {
   const std::string& Name() override { return name_; }
   size_t TotalSampleCount() override { return total_sample_count_; }
   size_t PerformanceSampleCount() override { return performance_sample_count_; }
+  size_t GroupSize(size_t i) override { return 1; }
+  size_t GroupOf(size_t i) override { return i; }
+  size_t NumberOfGroups() override { return total_sample_count_; }
 
   void LoadSamplesToRam(const std::vector<QuerySampleIndex>& samples) override {
     (*load_samples_to_ram_cb_)(client_data_, samples.data(), samples.size());

loadgen/bindings/python_api.cc

@@ -28,6 +28,7 @@ limitations under the License.
 #include "pybind11/pybind11.h"
 #include "pybind11/stl.h"
 #include "pybind11/stl_bind.h"
+#include "pybind11/numpy.h"
 
 namespace mlperf {
 
@@ -109,8 +110,10 @@ class QuerySampleLibraryTrampoline : public QuerySampleLibrary {
   ~QuerySampleLibraryTrampoline() override = default;
 
   const std::string& Name() override { return name_; }
-  size_t TotalSampleCount() { return total_sample_count_; }
-  size_t PerformanceSampleCount() { return performance_sample_count_; }
+  size_t TotalSampleCount() override { return total_sample_count_; }
+  size_t PerformanceSampleCount() override { return performance_sample_count_; }
+  size_t GroupSize(size_t i) override { return 1; }
+  size_t NumberOfGroups() override { return total_sample_count_; }
 
   void LoadSamplesToRam(const std::vector<QuerySampleIndex>& samples) override {
     pybind11::gil_scoped_acquire gil_acquirer;
@@ -130,6 +133,65 @@ class QuerySampleLibraryTrampoline : public QuerySampleLibrary {
   UnloadSamplesFromRamCallback unload_samples_from_ram_cb_;
 };
 
+// Forwards QuerySampleLibrary calls to relevant callbacks.
+class GroupedQuerySampleLibraryTrampoline : public QuerySampleLibrary {
+ public:
+  GroupedQuerySampleLibraryTrampoline(
+      std::string name,
+      size_t performance_sample_count,
+      LoadSamplesToRamCallback load_samples_to_ram_cb,
+      UnloadSamplesFromRamCallback unload_samples_from_ram_cb,
+      pybind11::array_t<size_t>& group_sizes)
+      : name_(std::move(name)),
+        performance_sample_count_(performance_sample_count),
+        load_samples_to_ram_cb_(load_samples_to_ram_cb),
+        unload_samples_from_ram_cb_(unload_samples_from_ram_cb) {
+
+      total_sample_count_ = 0;
+      if(group_sizes.ndim() != 1){
+        throw std::runtime_error("Group sizes should be a 1D Numpy array");
+      }
+      auto buffer = group_sizes.request();
+      size_t* ptr = (size_t*)buffer.ptr;
+
+      for(ssize_t i = 0; i < group_sizes.shape()[0]; i++){
+        group_sizes_.push_back(ptr[i]);
+        total_sample_count_ += ptr[i];
+        for(ssize_t j = 0; j < ptr[i]; j++){
+          group_idx_.push_back(i);
+        }
+      }
+    }
+  ~GroupedQuerySampleLibraryTrampoline() override = default;
+
+  const std::string& Name() override { return name_; }
+  size_t TotalSampleCount() override { return total_sample_count_; }
+  size_t PerformanceSampleCount() override { return performance_sample_count_; }
+  size_t GroupSize(size_t i) override { return group_sizes_[i]; }
+  size_t GroupOf(size_t i) override { return group_idx_[i]; }
+  size_t NumberOfGroups() override { return group_sizes_.size(); }
+
+  void LoadSamplesToRam(const std::vector<QuerySampleIndex>& samples) override {
+    pybind11::gil_scoped_acquire gil_acquirer;
+    load_samples_to_ram_cb_(samples);
+  }
+  void UnloadSamplesFromRam(
+      const std::vector<QuerySampleIndex>& samples) override {
+    pybind11::gil_scoped_acquire gil_acquirer;
+    unload_samples_from_ram_cb_(samples);
+  }
+
+ private:
+  std::string name_;
+  std::vector<size_t> group_sizes_;
+  std::vector<size_t> group_idx_;
+  size_t total_sample_count_;
+  size_t performance_sample_count_;
+  LoadSamplesToRamCallback load_samples_to_ram_cb_;
+  UnloadSamplesFromRamCallback unload_samples_from_ram_cb_;
+};
+
+
 // A QDL that allows defining callbacks for
 // IssueQuery, FlushQueries, and Name methods.
 class QueryDispatchLibraryTrampoline : public QueryDispatchLibrary {
@@ -226,6 +288,23 @@ void DestroyQDL(uintptr_t qdl) {
   delete qdl_cast;
 }
 
+uintptr_t ConstructGroupedQSL(
+    pybind11::array_t<size_t>& group_sizes,
+    size_t performance_sample_count,
+    LoadSamplesToRamCallback load_samples_to_ram_cb,
+    UnloadSamplesFromRamCallback unload_samples_from_ram_cb) {
+  GroupedQuerySampleLibraryTrampoline* qsl = new GroupedQuerySampleLibraryTrampoline(
+      "PyQSL", performance_sample_count,
+      load_samples_to_ram_cb, unload_samples_from_ram_cb, group_sizes);
+  return reinterpret_cast<uintptr_t>(qsl);
+}
+
+void DestroyGroupedQSL(uintptr_t qdl) {
+  QueryDispatchLibraryTrampoline* qdl_cast =
+      reinterpret_cast<QueryDispatchLibraryTrampoline*>(qdl);
+  delete qdl_cast;
+}
+
 void StartTest(uintptr_t sut, uintptr_t qsl, mlperf::TestSettings test_settings,
                const std::string& audit_config_filename) {
   pybind11::gil_scoped_release gil_releaser;
@@ -251,6 +330,20 @@ void StartTestWithLogSettings(uintptr_t sut, uintptr_t qsl,
                     audit_config_filename);
 }
 
+void StartTestWithGroupedTest(
+  uintptr_t sut, uintptr_t qsl, mlperf::TestSettings test_settings,
+               const std::string& audit_config_filename){
+  pybind11::gil_scoped_release gil_releaser;
+  SystemUnderTestTrampoline* sut_cast =
+      reinterpret_cast<SystemUnderTestTrampoline*>(sut);
+  GroupedQuerySampleLibraryTrampoline* qsl_cast =
+      reinterpret_cast<GroupedQuerySampleLibraryTrampoline*>(qsl);
+  LogSettings default_log_settings;
+  assert(TestSettings.use_grouped_qsl);
+  mlperf::StartTest(sut_cast, qsl_cast, test_settings, default_log_settings,
+                    audit_config_filename);
+}
+
 using ResponseCallback = std::function<void(QuerySampleResponse*)>;
 
 /// TODO: Get rid of copies.
@@ -342,6 +435,7 @@ PYBIND11_MODULE(mlperf_loadgen, m) {
       .def_readwrite("test05_schedule_rng_seed",
                      &TestSettings::test05_schedule_rng_seed)
       .def_readwrite("use_token_latencies", &TestSettings::use_token_latencies)
+      .def_readwrite("use_grouped_qsl", &TestSettings::use_grouped_qsl)
       .def_readwrite("ttft_latency", &TestSettings::server_ttft_latency)
       .def_readwrite("tpot_latency", &TestSettings::server_tpot_latency)
       .def_readwrite("infer_token_latencies",
@@ -456,6 +550,11 @@ PYBIND11_MODULE(mlperf_loadgen, m) {
   m.def("DestroyQDL", &py::DestroyQDL,
         "Destroy the object created by ConstructQDL.");
 
+  m.def("ConstructGroupedQSL", &py::ConstructGroupedQSL,
+        "Construct the query sample library.");
+  m.def("DestroyGroupedQSL", &py::DestroyQSL,
+        "Destroy the object created by ConstructGroupedQSL.");
+
   m.def("StartTest", &py::StartTest,
         "Run tests on a SUT created by ConstructSUT() with the provided QSL. "
         "Uses default log settings.",

loadgen/loadgen.cc

@@ -264,6 +264,7 @@ auto SampleDistributionEqualIssue(size_t sample_count, size_t set_size,
 /// the QPS as scheduled is equal to the QPS as requested.
 template <TestScenario scenario, TestMode mode>
 std::vector<QueryMetadata> GenerateQueries(
+    QuerySampleLibrary* qsl,
     const TestSettingsInternal& settings,
     const LoadableSampleSet& loaded_sample_set, SequenceGen* sequence_gen,
     ResponseDelegate* response_delegate) {
@@ -347,6 +348,29 @@ std::vector<QueryMetadata> GenerateQueries(
   // Choose a single sample to repeat when in performance_issue_same mode
   QuerySampleIndex same_sample = settings.performance_issue_same_index;
 
+  // Variables for handling group test
+  QuerySampleIndex global_idx = 0;
+  std::vector<size_t> groups;
+  std::vector<size_t> groups_first;
+  size_t number_of_groups = 0;
+  size_t g, group_size;
+
+  if (settings.use_grouped_qsl) {
+    size_t current_idx = 0;
+    while (current_idx < loaded_samples.size())
+    {
+      size_t current_group = qsl->GroupOf(loaded_samples[current_idx]);
+      groups.push_back(current_group);
+      groups_first.push_back(current_idx);
+      current_idx += qsl->GroupSize(loaded_samples[current_idx]);
+      number_of_groups++;
+    }
+    
+  }
+
+  auto grouped_sample_distribution = SampleDistribution<mode>(
+      number_of_groups, sample_stride, &sample_rng);
+
   while (prev_timestamp < gen_duration || queries.size() < min_queries) {
     if (kIsMultiStream) {
       QuerySampleIndex sample_i = settings.performance_issue_unique
@@ -393,6 +417,9 @@ std::vector<QueryMetadata> GenerateQueries(
           assert(remainder == 0);
         }
       }
+    } else if (settings.use_grouped_qsl) {
+      g = grouped_sample_distribution(sample_rng);
+      group_size = qsl->GroupSize(loaded_samples[groups_first[g]]);
     } else {
       for (auto& s : samples) {
         s = loaded_samples[settings.performance_issue_unique
@@ -403,10 +430,21 @@ std::vector<QueryMetadata> GenerateQueries(
                                : sample_distribution(sample_rng)];
       }
     }
-    queries.emplace_back(samples, timestamp, response_delegate, sequence_gen);
+    if (!settings.use_grouped_qsl) {
+      queries.emplace_back(samples, timestamp, response_delegate, sequence_gen);
+    } else {
+      for (size_t i = 0; i < group_size; i++){
+        samples[0] = loaded_samples[groups_first[g]+i];
+        queries.emplace_back(samples, timestamp, response_delegate, sequence_gen);
+      }
+    }
     prev_timestamp = timestamp;
     if (settings.server_constant_gen && (scenario == TestScenario::Server)){
-      timestamp += schedule_constant_distribution(schedule_rng);
+      if(!settings.use_grouped_qsl){
+        timestamp += schedule_constant_distribution(schedule_rng);
+      } else {
+        timestamp += group_size * schedule_constant_distribution(schedule_rng);
+      }
     } else {
       timestamp += schedule_distribution(schedule_rng);
     }
@@ -453,7 +491,8 @@ std::vector<QueryMetadata> GenerateQueries(
 //       no longer generates queries on the fly. Should we reduce the
 //       use of templates?
 template <TestScenario scenario, TestMode mode>
-PerformanceResult IssueQueries(SystemUnderTest* sut,
+PerformanceResult IssueQueries(SystemUnderTest* sut, 
+                               QuerySampleLibrary* qsl,
                                const TestSettingsInternal& settings,
                                const LoadableSampleSet& loaded_sample_set,
                                SequenceGen* sequence_gen) {
@@ -469,7 +508,7 @@ PerformanceResult IssueQueries(SystemUnderTest* sut,
   // Generate queries.
   auto sequence_id_start = sequence_gen->CurrentSampleId();
   std::vector<QueryMetadata> queries = GenerateQueries<scenario, mode>(
-      settings, loaded_sample_set, sequence_gen, &response_logger);
+      qsl, settings, loaded_sample_set, sequence_gen, &response_logger);
 
   // Calculated expected number of queries
   uint64_t expected_queries =
@@ -644,12 +683,35 @@ std::vector<LoadableSampleSet> GenerateLoadableSets(
   // Generate indices for all available samples in the QSL.
   const size_t qsl_total_count = qsl->TotalSampleCount();
   std::vector<QuerySampleIndex> samples(qsl_total_count);
-  for (size_t i = 0; i < qsl_total_count; i++) {
-    samples[i] = static_cast<QuerySampleIndex>(i);
-  }
+  std::vector<QuerySampleIndex> groupIdx(qsl_total_count);
+  if (!settings.use_grouped_qsl){
+    for (size_t i = 0; i < qsl_total_count; i++) {
+      samples[i] = static_cast<QuerySampleIndex>(i);
+    }
 
-  // Randomize the order of the samples.
-  std::shuffle(samples.begin(), samples.end(), qsl_rng);
+    // Randomize the order of the samples.
+    std::shuffle(samples.begin(), samples.end(), qsl_rng);
+  } else {
+    // If using grouped qsl, we randomized the groups.
+    // The samples within a group mantain their order.
+    size_t number_of_groups = qsl->NumberOfGroups();
+    size_t acumCount = 0, idx = 0;
+    std::vector<QuerySampleIndex> groups(number_of_groups);
+    std::vector<QuerySampleIndex> acumSizes(number_of_groups);
+    for (size_t i = 0; i < number_of_groups; i++) {
+      groups[i] = static_cast<QuerySampleIndex>(i);
+      acumSizes[i] = acumCount;
+      acumCount += qsl->GroupSize(i);
+    }
+    std::shuffle(groups.begin(), groups.end(), qsl_rng);
+    for (size_t i = 0; i < number_of_groups; i++) {
+      for (size_t j = 0; j < qsl->GroupSize(groups[i]); j++) {
+        samples[idx] = acumSizes[groups[i]] + j;
+        groupIdx[idx] = groups[i];
+        idx++;
+      }
+    }
+  }
 
   // Partition the samples into loadable sets.
   const size_t set_size = settings.performance_sample_count;
@@ -659,12 +721,30 @@ std::vector<LoadableSampleSet> GenerateLoadableSets(
   std::vector<QuerySampleIndex> loadable_set;
   loadable_set.reserve(set_size + set_padding);
 
-  for (auto s : samples) {
-    loadable_set.push_back(s);
-    if (loadable_set.size() == set_size) {
-      result.push_back({std::move(loadable_set), set_size});
-      loadable_set.clear();
-      loadable_set.reserve(set_size + set_padding);
+  if (!settings.use_grouped_qsl){
+    for (auto s : samples) {
+      loadable_set.push_back(s);
+      if (loadable_set.size() == set_size) {
+        result.push_back({std::move(loadable_set), set_size});
+        loadable_set.clear();
+        loadable_set.reserve(set_size + set_padding);
+      }
+    }
+  } else {
+    size_t idx = 0;
+    size_t number_of_groups = qsl->NumberOfGroups();
+    for (size_t i = 0; i < number_of_groups; i++) {
+      size_t group_size = qsl->GroupSize(groupIdx[idx]);
+      if (loadable_set.size() + group_size < set_size) {
+        for (size_t j = 0; j < group_size; j++) {
+          loadable_set.push_back(samples[idx]);
+          idx++;
+        }
+      } else {
+        result.push_back({std::move(loadable_set), loadable_set.size()});
+        loadable_set.clear();
+        loadable_set.reserve(set_size + set_padding);
+      }
     }
   }
 
@@ -771,7 +851,7 @@ std::pair<PerformanceSummary, PerformanceSummary> FindBoundaries(
   LoadSamplesToRam(qsl, performance_set.set);
 
   PerformanceResult u_pr(IssueQueries<scenario, TestMode::PerformanceOnly>(
-      sut, u_settings, performance_set, sequence_gen));
+      sut, qsl, u_settings, performance_set, sequence_gen));
   PerformanceSummary u_perf_summary{sut->Name(), u_settings, std::move(u_pr)};
 
   qsl->UnloadSamplesFromRam(performance_set.set);
@@ -823,7 +903,7 @@ PerformanceSummary FindPeakPerformanceBinarySearch(
   });
 
   PerformanceResult m_pr(IssueQueries<scenario, TestMode::PerformanceOnly>(
-      sut, m_settings, performance_set, sequence_gen));
+      sut, qsl, m_settings, performance_set, sequence_gen));
   PerformanceSummary m_perf_summary{sut->Name(), m_settings, std::move(m_pr)};
 
   std::string tmp;
@@ -868,7 +948,7 @@ void RunPerformanceMode(SystemUnderTest* sut, QuerySampleLibrary* qsl,
   }
 
   PerformanceResult pr(IssueQueries<scenario, TestMode::PerformanceOnly>(
-      sut, settings, performance_set, sequence_gen));
+      sut, qsl, settings, performance_set, sequence_gen));
 
   // Measure PerfClock/system_clock timer durations for comparison vs
   // external timer.
@@ -991,7 +1071,7 @@ void FindPeakPerformanceMode(SystemUnderTest* sut, QuerySampleLibrary* qsl,
   LoadSamplesToRam(qsl, base_performance_set.set);
 
   PerformanceResult base_pr(IssueQueries<scenario, TestMode::PerformanceOnly>(
-      sut, base_settings, base_performance_set, sequence_gen));
+      sut, qsl, base_settings, base_performance_set, sequence_gen));
   PerformanceSummary base_perf_summary{sut->Name(), base_settings,
                                        std::move(base_pr)};
 
@@ -1113,7 +1193,7 @@ void RunAccuracyMode(SystemUnderTest* sut, QuerySampleLibrary* qsl,
     }
 
     PerformanceResult pr(IssueQueries<scenario, TestMode::AccuracyOnly>(
-        sut, settings, loadable_set, sequence_gen));
+        sut, qsl, settings, loadable_set, sequence_gen));
 
     {
       auto tracer = MakeScopedTracer(

loadgen/query_sample_library.h

@@ -66,6 +66,9 @@ class QuerySampleLibrary {
   ///   * A previously unloaded sample will not be unloaded again.
   virtual void UnloadSamplesFromRam(
       const std::vector<QuerySampleIndex>& samples) = 0;
+  virtual size_t GroupSize(size_t i) = 0;
+  virtual size_t GroupOf(size_t i) = 0;
+  virtual size_t NumberOfGroups() = 0;
 };
 
 /// @}

loadgen/test_settings.h

@@ -239,6 +239,9 @@ struct TestSettings {
   uint64_t test05_sample_index_rng_seed = 0;
   uint64_t test05_schedule_rng_seed = 0;
 
+  /// \brief Flag to use a grouped qsl when running the benchmark test
+  bool use_grouped_qsl = false;
+
   /// \brief Load mlperf parameter config from file.
   int FromConfig(const std::string &path, const std::string &model,
                  const std::string &scenario, int conf_type = 1);