perf(crypto): CRP-2416 add a pin_cpu flag to rust_bench and set to Tr…

…ue in crypto benches
dfinity · Feb 19, 2024 · e6ccca2 · e6ccca2
1 parent d747b2b
commit e6ccca2
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 4 deletions.
diff --git a/bazel/defs.bzl b/bazel/defs.bzl
@@ -254,13 +254,14 @@ def rust_ic_test(env = {}, data = [], **kwargs):
         **kwargs
     )
 
-def rust_bench(name, env = {}, data = [], **kwargs):
+def rust_bench(name, env = {}, data = [], pin_cpu = False, **kwargs):
     """A rule for defining a rust benchmark.
 
     Args:
       name: the name of the executable target.
       env: additional environment variables to pass to the benchmark binary.
       data: data dependencies required to run the benchmark.
+      pin_cpu: pins the benchmark process to a single CPU if set `True`.
       **kwargs: see docs for `rust_binary`.
     """
 
@@ -277,14 +278,18 @@ def rust_bench(name, env = {}, data = [], **kwargs):
         testonly = kwargs.get("testonly", False),
     )
 
+    bench_prefix = "taskset -c 0 " if pin_cpu else ""
+
     # The benchmark binary is a shell script that runs the binary
     # (similar to how `cargo bench` runs the benchmark binary).
     native.sh_binary(
         srcs = ["//bazel:generic_rust_bench.sh"],
         name = name,
         # Allow benchmark targets to use test-only libraries.
         testonly = kwargs.get("testonly", False),
-        env = dict(env.items() + {"BAZEL_DEFS_BENCH_BIN": "$(location :%s)" % binary_name_publish}.items()),
+        env = dict(env.items() +
+                   [("BAZEL_DEFS_BENCH_PREFIX", bench_prefix)] +
+                   {"BAZEL_DEFS_BENCH_BIN": "$(location :%s)" % binary_name_publish}.items()),
         data = data + [":" + binary_name_publish],
         tags = kwargs.get("tags", []) + ["rust_bench"],
     )

diff --git a/bazel/generic_rust_bench.sh b/bazel/generic_rust_bench.sh
@@ -2,5 +2,7 @@
 
 set -euo pipefail
 
-echo "running ${BAZEL_DEFS_BENCH_BIN}"
-"${BAZEL_DEFS_BENCH_BIN}" --bench "$@"
+CMD="${BAZEL_DEFS_BENCH_PREFIX}${BAZEL_DEFS_BENCH_BIN} --bench $@"
+
+echo "running ${CMD}"
+${CMD}
diff --git a/rs/crypto/BUILD.bazel b/rs/crypto/BUILD.bazel
@@ -170,6 +170,7 @@ rust_bench(
     testonly = True,
     srcs = ["benches/basic_sig.rs"],
     aliases = ALIASES,
+    pin_cpu = True,
     proc_macro_deps = MACRO_DEPENDENCIES,
     deps = DEPENDENCIES + DEV_DEPENDENCIES,
 )
@@ -188,6 +189,7 @@ rust_bench(
     testonly = True,
     srcs = ["benches/threshold_sig.rs"],
     aliases = ALIASES,
+    pin_cpu = True,
     proc_macro_deps = MACRO_DEPENDENCIES,
     deps = DEPENDENCIES + DEV_DEPENDENCIES + [":crypto"],
 )
@@ -197,6 +199,7 @@ rust_bench(
     testonly = True,
     srcs = ["benches/ni_dkg.rs"],
     aliases = ALIASES,
+    pin_cpu = True,
     proc_macro_deps = MACRO_DEPENDENCIES,
     deps = DEPENDENCIES + DEV_DEPENDENCIES,
 )
@@ -206,6 +209,7 @@ rust_bench(
     testonly = True,
     srcs = ["benches/idkg.rs"],
     aliases = ALIASES,
+    pin_cpu = True,
     proc_macro_deps = MACRO_DEPENDENCIES,
     deps = DEPENDENCIES + DEV_DEPENDENCIES,
 )
@@ -215,6 +219,7 @@ rust_bench(
     testonly = True,
     srcs = ["benches/tecdsa.rs"],
     aliases = ALIASES,
+    pin_cpu = True,
     proc_macro_deps = MACRO_DEPENDENCIES,
     deps = DEPENDENCIES + DEV_DEPENDENCIES,
 )
@@ -224,6 +229,7 @@ rust_bench(
     testonly = True,
     srcs = ["benches/multi_sig.rs"],
     aliases = ALIASES,
+    pin_cpu = True,
     proc_macro_deps = MACRO_DEPENDENCIES,
     deps = DEPENDENCIES + DEV_DEPENDENCIES,
 )
@@ -233,6 +239,7 @@ rust_bench(
     testonly = True,
     srcs = ["benches/canister_sig.rs"],
     aliases = ALIASES,
+    pin_cpu = True,
     proc_macro_deps = MACRO_DEPENDENCIES,
     deps = DEPENDENCIES + DEV_DEPENDENCIES,
 )