diff --git a/test/benchmark/CMakeLists.txt b/test/benchmark/CMakeLists.txt index 6821aa03..b9aff09f 100644 --- a/test/benchmark/CMakeLists.txt +++ b/test/benchmark/CMakeLists.txt @@ -79,12 +79,6 @@ mujoco_test( ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers ) -mujoco_test( - island_benchmark_test - MAIN_TARGET benchmark::benchmark_main - ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers -) - mujoco_test( engine_util_sparse_benchmark_test MAIN_TARGET benchmark::benchmark_main diff --git a/test/benchmark/island_benchmark_test.cc b/test/benchmark/island_benchmark_test.cc deleted file mode 100644 index 21a922b8..00000000 --- a/test/benchmark/island_benchmark_test.cc +++ /dev/null @@ -1,374 +0,0 @@ -// Copyright 2026 DeepMind Technologies Limited -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Benchmarks island discovery on deterministic geodesic Rips graphs. The -// corpus spans the connectivity transition of points sampled on S^2 and adds -// MuJoCo-relevant static and repeated incidences. Corpus construction and -// validation are deliberately outside the timed region. - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "src/engine/engine_island.h" - -namespace mujoco { -namespace { - -struct Edge { - int first; - int second; -}; - -struct GraphCase { - std::string name; - int node_count; - std::vector incidences; - std::vector expected_partition; - int active_nodes; - int expected_components; - int pre_bridge_components; - bool bridge_added; - std::uint64_t expected_checksum; -}; - -struct Point { - double x; - double y; - double z; -}; - -std::uint64_t SplitMix64(std::uint64_t& state) { - state += 0x9e3779b97f4a7c15ULL; - std::uint64_t value = state; - value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; - value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; - return value ^ (value >> 31); -} - -double Uniform01(std::uint64_t& state) { - return static_cast(SplitMix64(state) >> 11) * 0x1.0p-53; -} - -std::vector SampleSphere(int count, std::uint64_t seed) { - constexpr double kTwoPi = 6.283185307179586476925286766559; - std::vector points; - points.reserve(count); - for (int i = 0; i < count; ++i) { - const double z = 2.0 * Uniform01(seed) - 1.0; - const double angle = kTwoPi * Uniform01(seed); - const double radial = std::sqrt(std::max(0.0, 1.0 - z*z)); - points.push_back({radial * std::cos(angle), radial * std::sin(angle), z}); - } - return points; -} - -double Dot(const Point& a, const Point& b) { - return a.x*b.x + a.y*b.y + a.z*b.z; -} - -std::vector RipsEdges(const std::vector& points, double target_degree) { - const double probability = target_degree / (points.size() - 1); - const double radius = 2.0 * std::asin(std::sqrt(probability)); - const double minimum_dot = std::cos(radius); - std::vector edges; - for (int i = 0; i < static_cast(points.size()); ++i) { - for (int j = i + 1; j < static_cast(points.size()); ++j) { - if (Dot(points[i], points[j]) >= minimum_dot) { - edges.push_back({i, j}); - } - } - } - return edges; -} - -std::vector CanonicalPartition(int node_count, const std::vector& edges) { - std::vector> adjacency(node_count); - std::vector active(node_count, false); - for (const Edge& edge : edges) { - active[edge.first] = true; - active[edge.second] = true; - if (edge.first != edge.second) { - adjacency[edge.first].push_back(edge.second); - adjacency[edge.second].push_back(edge.first); - } - } - - std::vector partition(node_count, -1); - std::queue pending; - for (int start = 0; start < node_count; ++start) { - if (!active[start] || partition[start] != -1) { - continue; - } - partition[start] = start; - pending.push(start); - while (!pending.empty()) { - const int node = pending.front(); - pending.pop(); - for (int neighbor : adjacency[node]) { - if (partition[neighbor] == -1) { - partition[neighbor] = start; - pending.push(neighbor); - } - } - } - } - return partition; -} - -int CountComponents(const std::vector& partition) { - int count = 0; - for (int node = 0; node < static_cast(partition.size()); ++node) { - count += partition[node] == node; - } - return count; -} - -std::uint64_t PartitionChecksum(const std::vector& partition) { - std::uint64_t hash = 1469598103934665603ULL; - for (int value : partition) { - hash ^= static_cast(value); - hash *= 1099511628211ULL; - } - return hash; -} - -std::vector CanonicalizeLabels(const std::vector& labels) { - std::vector minimum(labels.size(), std::numeric_limits::max()); - for (int node = 0; node < static_cast(labels.size()); ++node) { - if (labels[node] >= 0) { - minimum[labels[node]] = std::min(minimum[labels[node]], node); - } - } - std::vector canonical(labels.size(), -1); - for (int node = 0; node < static_cast(labels.size()); ++node) { - if (labels[node] >= 0) { - canonical[node] = minimum[labels[node]]; - } - } - return canonical; -} - -void DeterministicShuffle(std::vector& edges, std::uint64_t seed) { - for (std::size_t i = edges.size(); i > 1; --i) { - const std::size_t j = SplitMix64(seed) % i; - std::swap(edges[i - 1], edges[j]); - } -} - -bool AddCriticalBridge(const std::vector& points, std::vector& edges) { - const std::vector partition = CanonicalPartition(points.size(), edges); - if (CountComponents(partition) < 2) { - return false; - } - - double best_dot = -2.0; - Edge bridge{-1, -1}; - for (int i = 0; i < static_cast(points.size()); ++i) { - for (int j = i + 1; j < static_cast(points.size()); ++j) { - if (partition[i] >= 0 && partition[j] >= 0 && partition[i] != partition[j] && - Dot(points[i], points[j]) > best_dot) { - best_dot = Dot(points[i], points[j]); - bridge = {i, j}; - } - } - } - if (bridge.first >= 0) { - edges.push_back(bridge); - return true; - } - return false; -} - -GraphCase MakeCase(std::string name, int node_count, double target_degree, - std::uint64_t seed, bool critical_bridge, bool static_rows, - bool repeated_rows) { - const std::vector points = SampleSphere(node_count, seed); - std::vector edges = RipsEdges(points, target_degree); - const int pre_bridge_components = - critical_bridge ? CountComponents(CanonicalPartition(node_count, edges)) : -1; - bool bridge_added = false; - if (critical_bridge) { - bridge_added = AddCriticalBridge(points, edges); - } - - const std::vector unique_edges = edges; - if (static_rows) { - for (int node = 0; node < node_count; node += 17) { - edges.push_back({node, node}); - } - } - if (repeated_rows) { - for (std::size_t i = 0; i < unique_edges.size(); i += 11) { - edges.push_back(unique_edges[i]); - edges.push_back({unique_edges[i].second, unique_edges[i].first}); - } - } - DeterministicShuffle(edges, seed ^ 0xd1b54a32d192ed03ULL); - - std::vector expected = CanonicalPartition(node_count, edges); - const int active_nodes = std::count_if(expected.begin(), expected.end(), - [](int component) { return component >= 0; }); - const int components = CountComponents(expected); - const std::uint64_t checksum = PartitionChecksum(expected); - return {std::move(name), node_count, std::move(edges), std::move(expected), active_nodes, - components, pre_bridge_components, bridge_added, checksum}; -} - -const std::vector& Corpus() { - static const std::vector corpus = { - MakeCase("StableSparse_S2Rips_64", 64, 2.0, 0x33960001ULL, false, false, false), - MakeCase("CriticalBridge_S2Rips_256", 256, 0.75 * std::log(256.0), - 0x33960002ULL, true, false, false), - MakeCase("SupercriticalDense_S2Rips_256", 256, 2.0 * std::ceil(std::log(256.0)), - 0x33960003ULL, false, false, false), - MakeCase("GroundedStaticRepeated_S2Rips_256", 256, - 2.0 * std::ceil(std::log(256.0)), 0x33960004ULL, false, true, true), - MakeCase("StableRepeated_S2Rips_1024", 1024, 2.0, 0x33960005ULL, - false, false, true), - MakeCase("CriticalLarge_S2Rips_1024", 1024, std::ceil(std::log(1024.0)), - 0x33960006ULL, true, false, false), - }; - return corpus; -} - -struct FloodFillWorkspace { - explicit FloodFillWorkspace(int node_count) - : adjacency(node_count * node_count), rownnz(node_count), rowadr(node_count), - colind(node_count * node_count), stack(node_count * node_count + node_count), - island(node_count) {} - - std::vector adjacency; - std::vector rownnz; - std::vector rowadr; - std::vector colind; - std::vector stack; - std::vector island; -}; - -int RunFloodFill(const GraphCase& graph, FloodFillWorkspace& work) { - const int n = graph.node_count; - std::fill(work.adjacency.begin(), work.adjacency.end(), 0); - std::fill(work.rownnz.begin(), work.rownnz.end(), 0); - for (const Edge& edge : graph.incidences) { - work.adjacency[edge.first*n + edge.second] = 1; - work.adjacency[edge.second*n + edge.first] = 1; - } - - int address = 0; - for (int row = 0; row < n; ++row) { - work.rowadr[row] = address; - for (int column = 0; column < n; ++column) { - if (work.adjacency[row*n + column]) { - work.colind[address++] = column; - ++work.rownnz[row]; - } - } - } - return mj_floodFill(work.island.data(), n, work.rownnz.data(), work.rowadr.data(), - work.colind.data(), work.stack.data()); -} - -struct DsuWorkspace { - explicit DsuWorkspace(int node_count) - : parent(node_count), island(node_count), dof_count(node_count, 1) {} - - std::vector parent; - std::vector island; - std::vector dof_count; -}; - -int RunDsu(const GraphCase& graph, DsuWorkspace& work) { - _mjPRIVATE_dsuInit(work.parent.data(), graph.node_count); - for (const Edge& edge : graph.incidences) { - _mjPRIVATE_dsuUnion(work.parent.data(), edge.first, edge.second); - } - int dof_count = 0; - return _mjPRIVATE_dsuAssign(work.island.data(), work.parent.data(), work.dof_count.data(), - graph.node_count, &dof_count); -} - -bool Validate(const GraphCase& graph) { - FloodFillWorkspace flood(graph.node_count); - DsuWorkspace dsu(graph.node_count); - const int flood_components = RunFloodFill(graph, flood); - const int dsu_components = RunDsu(graph, dsu); - const bool bridge_valid = graph.pre_bridge_components < 0 || - (graph.bridge_added && - graph.pre_bridge_components == graph.expected_components + 1); - return bridge_valid && flood_components == graph.expected_components && - dsu_components == graph.expected_components && - CanonicalizeLabels(flood.island) == graph.expected_partition && - CanonicalizeLabels(dsu.island) == graph.expected_partition && - PartitionChecksum(graph.expected_partition) == graph.expected_checksum; -} - -void BM_FloodFill(benchmark::State& state, const GraphCase* graph) { - if (!Validate(*graph)) { - state.SkipWithError("invalid S2-Rips graph fixture"); - return; - } - FloodFillWorkspace work(graph->node_count); - state.SetLabel("edges=" + std::to_string(graph->incidences.size()) + - " active=" + std::to_string(graph->active_nodes) + - " components=" + std::to_string(graph->expected_components) + - " pre_bridge=" + std::to_string(graph->pre_bridge_components) + - " bridge_added=" + std::to_string(graph->bridge_added) + - " checksum=" + std::to_string(graph->expected_checksum)); - for (auto _ : state) { - int components = RunFloodFill(*graph, work); - benchmark::DoNotOptimize(components); - benchmark::ClobberMemory(); - } - state.SetItemsProcessed(state.iterations() * graph->incidences.size()); -} - -void BM_Dsu(benchmark::State& state, const GraphCase* graph) { - if (!Validate(*graph)) { - state.SkipWithError("invalid S2-Rips graph fixture"); - return; - } - DsuWorkspace work(graph->node_count); - state.SetLabel("edges=" + std::to_string(graph->incidences.size()) + - " active=" + std::to_string(graph->active_nodes) + - " components=" + std::to_string(graph->expected_components) + - " pre_bridge=" + std::to_string(graph->pre_bridge_components) + - " bridge_added=" + std::to_string(graph->bridge_added) + - " checksum=" + std::to_string(graph->expected_checksum)); - for (auto _ : state) { - int components = RunDsu(*graph, work); - benchmark::DoNotOptimize(components); - benchmark::ClobberMemory(); - } - state.SetItemsProcessed(state.iterations() * graph->incidences.size()); -} - -const bool kRegistered = [] { - for (const GraphCase& graph : Corpus()) { - benchmark::RegisterBenchmark(("Island/FloodFill/" + graph.name).c_str(), BM_FloodFill, &graph); - benchmark::RegisterBenchmark(("Island/DSU/" + graph.name).c_str(), BM_Dsu, &graph); - } - return true; -}(); - -} // namespace -} // namespace mujoco