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
This commit is contained in:
Yuval Tassa
committed by
Copybara-Service
Copybara-Service
parent
b7af32ac87
commit
0afafacfc5
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 <random>
|
||||
#include <vector>
|
||||
|
||||
#include "benchmark/benchmark.h"
|
||||
#include <absl/base/attributes.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
#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<mjtNum> mats; // kNumMat x 36
|
||||
std::vector<mjtNum> rhs; // kNumMat x 6
|
||||
|
||||
LUData() : mats(kNumMat * kN * kN), rhs(kNumMat * kN) {
|
||||
std::mt19937_64 rng(42);
|
||||
std::normal_distribution<double> 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;
|
||||
}
|
||||
@@ -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<double> 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
|
||||
|
||||
Reference in New Issue
Block a user