From 0afafacfc582509180e3f1eecfd81b0989b0be81 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Tue, 14 Jul 2026 08:37:32 -0700 Subject: [PATCH] Add fixed-size 6x6 dense LU factorization with benchmark. mju_factorLU6/mju_solveLU6: same algorithm as mju_factorLU/mju_solveLU with compile-time size, allowing full unrolling. At n=6, factor+solve is 25% faster than the runtime-sized version (93 vs 124 ns), and fixed-size LU factorization is faster than generic dense Cholesky (55 vs 61 ns): at this size, runtime-n loop overhead outweighs Cholesky's 2x flop advantage. See new lu_benchmark_test. Results agree with the generic version to rounding, not bitwise: the compiler may fuse (FMA) the unrolled version differently. Also add two DenseLU tests: a pivoting-required matrix with zero diagonal, and fixed-vs-generic agreement. PiperOrigin-RevId: 947705056 Change-Id: I24c54c9510964aa376886e9dd721890eda9889d3 --- src/engine/engine_util_solve.c | 79 +++++++++++ src/engine/engine_util_solve.h | 6 + test/benchmark/CMakeLists.txt | 6 + test/benchmark/lu_benchmark_test.cc | 188 ++++++++++++++++++++++++++ test/engine/engine_util_solve_test.cc | 71 ++++++++++ 5 files changed, 350 insertions(+) create mode 100644 test/benchmark/lu_benchmark_test.cc diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index f275b488..3619ba82 100644 --- a/src/engine/engine_util_solve.c +++ b/src/engine/engine_util_solve.c @@ -834,6 +834,85 @@ void mju_solveLU(mjtNum* restrict x, const mjtNum* LU, const mjtNum* b, const in } +// 6x6 specialization of mju_factorLU: same algorithm with compile-time size, +// allowing full unrolling; produces identical results +int mju_factorLU6(mjtNum A[36], int pivot[6]) { + for (int k=0; k < 6; k++) { + // initialize pivot + pivot[k] = k; + + // find pivot: max absolute value in column k, rows k..n-1 + mjtNum maxval = mju_abs(A[k*6+k]); + int maxrow = k; + for (int i=k+1; i < 6; i++) { + mjtNum val = mju_abs(A[i*6+k]); + if (val > maxval) { + maxval = val; + maxrow = i; + } + } + + // check singularity + if (maxval < mjMINVAL) { + return 0; + } + + // swap rows k and maxrow + if (maxrow != k) { + pivot[k] = maxrow; + for (int j=0; j < 6; j++) { + mjtNum tmp = A[k*6+j]; + A[k*6+j] = A[maxrow*6+j]; + A[maxrow*6+j] = tmp; + } + } + + // compute multipliers and update trailing submatrix + mjtNum diaginv = 1.0 / A[k*6+k]; + for (int i=k+1; i < 6; i++) { + A[i*6+k] *= diaginv; + mjtNum Aik = A[i*6+k]; + for (int j=k+1; j < 6; j++) { + A[i*6+j] -= Aik * A[k*6+j]; + } + } + } + + return 1; +} + + +// solve A*x = b given 6x6 LU factorization from mju_factorLU6 +void mju_solveLU6(mjtNum x[6], const mjtNum LU[36], const mjtNum b[6], const int pivot[6]) { + for (int i=0; i < 6; i++) { + x[i] = b[i]; + } + + // apply row permutation and forward substitution: solve L*y = P*b + for (int i=0; i < 6; i++) { + // apply pivot swap + if (pivot[i] != i) { + mjtNum tmp = x[i]; + x[i] = x[pivot[i]]; + x[pivot[i]] = tmp; + } + + // subtract known terms + for (int j=0; j < i; j++) { + x[i] -= LU[i*6+j] * x[j]; + } + } + + // back substitution: solve U*x = y + for (int i=6-1; i >= 0; i--) { + for (int j=i+1; j < 6; j++) { + x[i] -= LU[i*6+j] * x[j]; + } + x[i] /= LU[i*6+i]; + } +} + + //------------------------------ sparse LU factorization ------------------------------------------- // sparse reverse-order LU factorization, no fill-in (assuming tree topology) diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h index 971dfc86..90db2b3b 100644 --- a/src/engine/engine_util_solve.h +++ b/src/engine/engine_util_solve.h @@ -104,6 +104,12 @@ MJAPI int mju_factorLU(mjtNum* A, int n, int* pivot); // solve A*x = b given LU factorization of A, LU and pivot are output of mju_factorLU MJAPI void mju_solveLU(mjtNum* x, const mjtNum* LU, const mjtNum* b, const int* pivot, int n); +// 6x6 specialization of mju_factorLU (identical results, allows full unrolling) +MJAPI int mju_factorLU6(mjtNum A[36], int pivot[6]); + +// solve A*x = b given 6x6 LU factorization from mju_factorLU6 +MJAPI void mju_solveLU6(mjtNum x[6], const mjtNum LU[36], const mjtNum b[6], const int pivot[6]); + // sparse reverse-order LU factorization, assume tree topology (only dofs in index, if given) // LU = L + U; original = (U+I) * L; scratch is size n void mju_factorLUSparse(mjtNum *LU, int n, int* scratch, diff --git a/test/benchmark/CMakeLists.txt b/test/benchmark/CMakeLists.txt index 3f0be08e..b9aff09f 100644 --- a/test/benchmark/CMakeLists.txt +++ b/test/benchmark/CMakeLists.txt @@ -67,6 +67,12 @@ mujoco_test( ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers ) +mujoco_test( + lu_benchmark_test + MAIN_TARGET benchmark::benchmark_main + ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers +) + mujoco_test( crb_benchmark_test MAIN_TARGET benchmark::benchmark_main diff --git a/test/benchmark/lu_benchmark_test.cc b/test/benchmark/lu_benchmark_test.cc new file mode 100644 index 00000000..8ad9a56a --- /dev/null +++ b/test/benchmark/lu_benchmark_test.cc @@ -0,0 +1,188 @@ +// 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. + +// A benchmark comparing generic (runtime-sized) and fixed-size dense LU +// factorization and solve at n = 6, with dense Cholesky as a reference. +// The 6x6 case is the free-body block in the implicitfast integrator. + +#include +#include + +#include "benchmark/benchmark.h" +#include +#include +#include "src/engine/engine_util_blas.h" +#include "src/engine/engine_util_solve.h" + +namespace mujoco { +namespace { + +constexpr int kN = 6; +constexpr int kNumMat = 256; // cycle through matrices to avoid cache pinning + +// generate matrices with the structure of the free-body block M - h*D + h*G: +// symmetric positive definite plus a scaled antisymmetric part +struct LUData { + std::vector mats; // kNumMat x 36 + std::vector rhs; // kNumMat x 6 + + LUData() : mats(kNumMat * kN * kN), rhs(kNumMat * kN) { + std::mt19937_64 rng(42); + std::normal_distribution dist(0, 1); + + for (int m = 0; m < kNumMat; m++) { + mjtNum* A = mats.data() + m * kN * kN; + mjtNum sqrtH[kN * kN]; + for (int i = 0; i < kN * kN; i++) { + sqrtH[i] = dist(rng); + } + + // A = sqrtH' * sqrtH + n*I: SPD + mju_mulMatTMat(A, sqrtH, sqrtH, kN, kN, kN); + for (int i = 0; i < kN; i++) { + A[i * kN + i] += kN; + } + + // add antisymmetric part (gyroscopic-like), scaled to ~10% of diagonal + for (int r = 0; r < kN; r++) { + for (int c = r + 1; c < kN; c++) { + mjtNum g = 0.6 * dist(rng); + A[r * kN + c] += g; + A[c * kN + r] -= g; + } + } + + for (int i = 0; i < kN; i++) { + rhs[m * kN + i] = dist(rng); + } + } + } +}; + +LUData& GetData() { + static LUData* data = new LUData(); + return *data; +} + +enum class Variant { + kGeneric, // mju_factorLU / mju_solveLU with n = 6 + kFixed, // mju_factorLU6 / mju_solveLU6 + kCholesky // mju_cholFactor / mju_cholSolve (symmetric part; reference) +}; + +// factor only (includes the copy, identical across variants) +ABSL_ATTRIBUTE_NO_TAIL_CALL static void FactorBenchmark(benchmark::State& state, + Variant variant) { + LUData& data = GetData(); + mjtNum A[kN * kN]; + int pivot[kN]; + int m = 0; + + for (auto s : state) { + mju_copy(A, data.mats.data() + m * kN * kN, kN * kN); + m = (m + 1) % kNumMat; + + int rank; + switch (variant) { + case Variant::kGeneric: + rank = mju_factorLU(A, kN, pivot); + break; + case Variant::kFixed: + rank = mju_factorLU6(A, pivot); + break; + case Variant::kCholesky: + rank = mju_cholFactor(A, kN, 0); + break; + } + benchmark::DoNotOptimize(rank); + benchmark::DoNotOptimize(A); + } + + state.counters["factor/s"] = + benchmark::Counter(state.iterations(), benchmark::Counter::kIsRate); +} + +// factor and solve +ABSL_ATTRIBUTE_NO_TAIL_CALL static void FactorSolveBenchmark( + benchmark::State& state, Variant variant) { + LUData& data = GetData(); + mjtNum A[kN * kN]; + mjtNum x[kN]; + int pivot[kN]; + int m = 0; + + for (auto s : state) { + mju_copy(A, data.mats.data() + m * kN * kN, kN * kN); + const mjtNum* b = data.rhs.data() + m * kN; + m = (m + 1) % kNumMat; + + switch (variant) { + case Variant::kGeneric: + mju_factorLU(A, kN, pivot); + mju_solveLU(x, A, b, pivot, kN); + break; + case Variant::kFixed: + mju_factorLU6(A, pivot); + mju_solveLU6(x, A, b, pivot); + break; + case Variant::kCholesky: + mju_cholFactor(A, kN, 0); + mju_cholSolve(x, A, b, kN); + break; + } + benchmark::DoNotOptimize(x); + } + + state.counters["solve/s"] = + benchmark::Counter(state.iterations(), benchmark::Counter::kIsRate); +} + +void BM_factorLU_generic6(benchmark::State& state) { + FactorBenchmark(state, Variant::kGeneric); +} +BENCHMARK(BM_factorLU_generic6); + +void BM_factorLU_fixed6(benchmark::State& state) { + FactorBenchmark(state, Variant::kFixed); +} +BENCHMARK(BM_factorLU_fixed6); + +void BM_factorChol6(benchmark::State& state) { + FactorBenchmark(state, Variant::kCholesky); +} +BENCHMARK(BM_factorChol6); + +void BM_factorSolveLU_generic6(benchmark::State& state) { + FactorSolveBenchmark(state, Variant::kGeneric); +} +BENCHMARK(BM_factorSolveLU_generic6); + +void BM_factorSolveLU_fixed6(benchmark::State& state) { + FactorSolveBenchmark(state, Variant::kFixed); +} +BENCHMARK(BM_factorSolveLU_fixed6); + +void BM_factorSolveChol6(benchmark::State& state) { + FactorSolveBenchmark(state, Variant::kCholesky); +} +BENCHMARK(BM_factorSolveChol6); + +} // namespace +} // namespace mujoco + +int main(int argc, char** argv) { + benchmark::Initialize(&argc, argv); + benchmark::RunSpecifiedBenchmarks(); + return 0; +} diff --git a/test/engine/engine_util_solve_test.cc b/test/engine/engine_util_solve_test.cc index 32703efd..d28fb26f 100644 --- a/test/engine/engine_util_solve_test.cc +++ b/test/engine/engine_util_solve_test.cc @@ -1235,5 +1235,76 @@ TEST_F(DenseLUTest, Singular) { EXPECT_EQ(mju_factorLU(A, n, pivot), 0); } +// 6x6 specialization: identical results to the generic version +TEST_F(DenseLUTest, Fixed6MatchesGeneric) { + constexpr int n = 6; + std::mt19937_64 rng; + rng.seed(3); + std::normal_distribution dist(0, 1); + + for (int trial = 0; trial < 100; trial++) { + mjtNum A[n * n], A6[n * n], b[n], x[n], x6[n]; + int pivot[n], pivot6[n]; + + // random asymmetric matrix; first trial zeroes the diagonal to force + // pivoting + for (int i = 0; i < n * n; i++) { + A[i] = dist(rng); + } + if (trial == 0) { + for (int i = 0; i < n; i++) { + A[i * n + i] = 0; + } + } + for (int i = 0; i < n; i++) { + b[i] = dist(rng); + } + mju_copy(A6, A, n * n); + + // same algorithm; results agree to rounding (not bitwise: the compiler may + // contract the unrolled fixed-size version differently) + int rank = mju_factorLU(A, n, pivot); + int rank6 = mju_factorLU6(A6, pivot6); + ASSERT_EQ(rank, rank6); + if (!rank) { + continue; + } + for (int i = 0; i < n; i++) { + EXPECT_EQ(pivot[i], pivot6[i]); + } + for (int i = 0; i < n * n; i++) { + EXPECT_NEAR(A[i], A6[i], MjTol(1e-13, 1e-4)); + } + + mju_solveLU(x, A, b, pivot, n); + mju_solveLU6(x6, A6, b, pivot6); + for (int i = 0; i < n; i++) { + EXPECT_NEAR(x[i], x6[i], MjTol(1e-12, 1e-3)); + } + } +} + +// nonsingular matrix with all-zero diagonal: requires pivoting +TEST_F(DenseLUTest, ZeroDiagonal) { + constexpr int n = 3; + // A = [0 2 1; 1 0 3; 4 1 0], det(A) = 25 + mjtNum A[n * n] = { + 0, 2, 1, 1, 0, 3, 4, 1, 0, + }; + mjtNum A_orig[n * n]; + mju_copy(A_orig, A, n * n); + int pivot[n]; + mjtNum b[n] = {1, 2, 3}; + mjtNum x[n], Ax[n]; + + EXPECT_EQ(mju_factorLU(A, n, pivot), 1); + mju_solveLU(x, A, b, pivot, n); + + // verify A_orig * x == b + mju_mulMatVec(Ax, A_orig, x, n, n); + mjtNum eps = MjTol(1e-14, 1e-6); + EXPECT_THAT(AsVector(Ax, n), Pointwise(MjNear(eps, eps), AsVector(b, n))); +} + } // namespace } // namespace mujoco