Plumb pinned core through benchmarking abstractions

This patch does the necessary plumbing to enable setting the
benchmarking core from the Python wrapper around exegesis that is used
for benchmarking.

Pull Request: google#272

gematria/datasets/exegesis_benchmark.cc

@@ -113,7 +113,8 @@ int main(int Argc, char *Argv[]) {
         Benchmark->parseJSONBlock(*AnnotatedBlockObject, BlockIndex));
 
     double Throughput = exitOnFileError(
-        AnnotatedBlocksJson, Benchmark->benchmarkBasicBlock(BenchCode));
+        AnnotatedBlocksJson,
+        Benchmark->benchmarkBasicBlock(BenchCode, std::nullopt));
 
     std::optional<StringRef> HexValue = AnnotatedBlockObject->getString("Hex");
     // The block has already been parsed previously, and thus should have thrown

gematria/datasets/exegesis_benchmark_lib.cc

@@ -341,7 +341,7 @@ Expected<BenchmarkCode> ExegesisBenchmark::processAnnotatedBlock(
 }
 
 Expected<double> ExegesisBenchmark::benchmarkBasicBlock(
-    const BenchmarkCode &BenchCode) {
+    const BenchmarkCode &BenchCode, std::optional<int> BenchmarkProcessCPU) {
   std::unique_ptr<const SnippetRepetitor> SnipRepetitor =
       SnippetRepetitor::Create(Benchmark::RepetitionModeE::MiddleHalfLoop,
                                ExegesisState, BenchCode.Key.LoopRegister);
@@ -355,15 +355,15 @@ Expected<double> ExegesisBenchmark::benchmarkBasicBlock(
   if (!RC2) return RC2.takeError();
 
   std::pair<Error, Benchmark> BenchmarkResultAOrErr =
-      BenchRunner->runConfiguration(std::move(*RC1), {}, std::nullopt);
+      BenchRunner->runConfiguration(std::move(*RC1), {}, BenchmarkProcessCPU);
 
   if (std::get<0>(BenchmarkResultAOrErr))
     return std::move(std::get<0>(BenchmarkResultAOrErr));
 
   AllResults.push_back(std::move(std::get<1>(BenchmarkResultAOrErr)));
 
   std::pair<Error, Benchmark> BenchmarkResultBOrErr =
-      BenchRunner->runConfiguration(std::move(*RC2), {}, std::nullopt);
+      BenchRunner->runConfiguration(std::move(*RC2), {}, BenchmarkProcessCPU);
 
   if (std::get<0>(BenchmarkResultBOrErr))
     return std::move(std::get<0>(BenchmarkResultBOrErr));

gematria/datasets/exegesis_benchmark_lib.h

@@ -14,6 +14,7 @@
 
 #include <cstddef>
 #include <memory>
+#include <optional>
 #include <string_view>
 
 #include "gematria/proto/execution_annotation.pb.h"
@@ -46,7 +47,8 @@ class ExegesisBenchmark {
       std::string_view BlockHex, const ExecutionAnnotations &Annotations);
 
   llvm::Expected<double> benchmarkBasicBlock(
-      const llvm::exegesis::BenchmarkCode &BenchCode);
+      const llvm::exegesis::BenchmarkCode &BenchCode,
+      std::optional<int> BenchmarkProcessCPU);
 
  private:
   // This is a simple wrapper around functionality in ExegesisState that maps

gematria/datasets/exegesis_benchmark_lib_test.cc

@@ -14,7 +14,10 @@
 
 #include "gematria/datasets/exegesis_benchmark_lib.h"
 
+#include <sched.h>
+
 #include <memory>
+#include <optional>
 #include <string>
 
 #include "gematria/proto/execution_annotation.pb.h"
@@ -92,7 +95,7 @@ class ExegesisBenchmarkTest : public testing::Test {
         Benchmark->parseJSONBlock(*BlockValue->getAsObject(), 0);
     if (!BenchCode) return BenchCode.takeError();
 
-    return Benchmark->benchmarkBasicBlock(*BenchCode);
+    return Benchmark->benchmarkBasicBlock(*BenchCode, std::nullopt);
   }
 };
 
@@ -441,7 +444,45 @@ TEST_F(ExegesisBenchmarkTest, TestBenchmarkFromAnnotatedBlock) {
   ASSERT_TRUE(static_cast<bool>(BenchmarkConfiguration));
 
   Expected<double> BenchmarkResult =
-      Benchmark->benchmarkBasicBlock(*BenchmarkConfiguration);
+      Benchmark->benchmarkBasicBlock(*BenchmarkConfiguration, std::nullopt);
+  EXPECT_LT(*BenchmarkResult, 10);
+}
+
+TEST_F(ExegesisBenchmarkTest, TestBenchmarkCorePinning) {
+  // clang-format off
+  const ExecutionAnnotations Annotations = ParseTextProto(R"pb(
+    code_start_address: 0
+    block_size: 4096
+    block_contents: 34359738376
+    accessed_blocks: 86016
+    initial_registers: [
+      {
+        register_name: "RCX"
+        register_value: 86016
+      },
+      {
+        register_name: "RSI"
+        register_value: 86016
+      }
+    ]
+    loop_register: "RAX"
+  )pb");
+  // clang-format on
+
+  Expected<BenchmarkCode> BenchmarkConfiguration =
+      Benchmark->processAnnotatedBlock("3b31", Annotations);
+  ASSERT_TRUE(static_cast<bool>(BenchmarkConfiguration));
+
+  cpu_set_t process_cpu_mask;
+  ASSERT_FALSE(
+      sched_getaffinity(0, sizeof(process_cpu_mask), &process_cpu_mask));
+  if (!CPU_ISSET(0, &process_cpu_mask)) {
+    GTEST_SKIP()
+        << "The CPU (0) used for testing core pinning is not available.";
+  }
+
+  Expected<double> BenchmarkResult =
+      Benchmark->benchmarkBasicBlock(*BenchmarkConfiguration, 0);
   EXPECT_LT(*BenchmarkResult, 10);
 }
 

gematria/datasets/python/exegesis_benchmark.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include <memory>
+#include <optional>
 
 #include "absl/status/statusor.h"
 #include "gematria/datasets/bhive_to_exegesis.h"
@@ -26,6 +27,8 @@
 #include "pybind11/cast.h"
 #include "pybind11/detail/common.h"
 #include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
+#include "pybind11/stl.h"  // IWYU pragma: keep
 #include "pybind11_abseil/import_status_module.h"
 #include "pybind11_abseil/status_casters.h"         // IWYU pragma: keep
 #include "pybind11_protobuf/native_proto_caster.h"  // IWYU pragma: keep
@@ -116,11 +119,13 @@ PYBIND11_MODULE(exegesis_benchmark, m) {
           )")
       .def(
           "benchmark_basic_block",
-          [](ExegesisBenchmark& Self, const BenchmarkCode& InputBenchmark) {
+          [](ExegesisBenchmark& Self, const BenchmarkCode& InputBenchmark,
+             std::optional<int> BenchmarkProcessCPU) {
             return LlvmExpectedToStatusOr(
-                Self.benchmarkBasicBlock(InputBenchmark));
+                Self.benchmarkBasicBlock(InputBenchmark, BenchmarkProcessCPU));
           },
           py::arg("input_benchmark"),
+          py::arg("benchmark_process_cpu") = py::none(),
           R"(Benchmarks a block in the form of a BenchmarkCode instance.
 
           Takes a BenchmarkCode instance and then executes it, collecting
@@ -129,6 +134,8 @@ PYBIND11_MODULE(exegesis_benchmark, m) {
           Args:
             input_benchmark: The BenchmarkCode instance formed from the block and
               annotations of interest that should be benchmarked.
+            benchmark_process_cpu: An optional integer specifying the CPU ID
+              that the benchmarking subprocess should be pinned to.
           
           Returns:
             A floating point value representing the inverse throughput (i.e.,

gematria/datasets/python/exegesis_benchmark_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
 from absl.testing import absltest
 
 from gematria.datasets.python import bhive_to_exegesis
@@ -31,7 +33,7 @@ def setUp(self):
     )
     self.exegesis_benchmark = exegesis_benchmark.ExegesisBenchmark.create()
 
-  def test_benchmarking(self):
+  def _get_block_for_benchmarking(self) -> exegesis_benchmark.BenchmarkCode:
     execution_annotations = self.bhive_to_exegesis.annotate_basic_block(
         "4829d38b44246c8b54246848c1fb034829d04839c3",
         bhive_to_exegesis.AnnotatorType.fast,
@@ -49,12 +51,27 @@ def test_benchmarking(self):
         block_with_annotations
     )
 
+    return benchmark_code
+
+  def test_benchmarking(self):
+    benchmark_code = self._get_block_for_benchmarking()
+
     block_measurement = self.exegesis_benchmark.benchmark_basic_block(
         benchmark_code
     )
 
     self.assertLess(block_measurement, 10)
 
+  def test_benchmarking_pinned_core(self):
+    benchmark_code = self._get_block_for_benchmarking()
+    benchmark_core = os.sched_getaffinity(0).pop()
+
+    block_measurement = self.exegesis_benchmark.benchmark_basic_block(
+        benchmark_code, benchmark_core
+    )
+
+    self.assertLess(block_measurement, 10)
+
 
 if __name__ == "__main__":
   absltest.main()