From 57c11f4dd072183f6a3d0c2a594fa607f7c39de7 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Sun, 4 Oct 2026 21:53:28 +0200 Subject: [PATCH 1/8] laplace_marginal_density: do not pre-size matrices that are reassigned R was pre-sized for solvers 1 and 2 and LU_solve_covariance for solver 3, but solvers 2 and 3 assign them with `=`, which gives an arena_matrix a new buffer; the pre-sized n x n buffers were never used. Pre-size R only for solver 1, which writes it in place. --- .../math/mix/functor/laplace_marginal_density.hpp | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/stan/math/mix/functor/laplace_marginal_density.hpp b/stan/math/mix/functor/laplace_marginal_density.hpp index cf505417cc4..0cefbfc9d8d 100644 --- a/stan/math/mix/functor/laplace_marginal_density.hpp +++ b/stan/math/mix/functor/laplace_marginal_density.hpp @@ -197,15 +197,12 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, ll_fun, ll_args_vals, covariance_val, options, msgs); auto ll_args_filter = internal::filter_var_scalar_types(ll_args_copy); // tuple of references to var types - // Solver 1, 2 - const bool solver_1_or_2 - = md_est.solver_used == 1 || md_est.solver_used == 2; - arena_t R(md_est.theta.size() * solver_1_or_2, - md_est.theta.size() * solver_1_or_2); - // Solver 3 - arena_t LU_solve_covariance( - covariance.rows() * (md_est.solver_used == 3), - covariance.cols() * (md_est.solver_used == 3)); + // Solver 1 writes R in place; solvers 2 and 3 assign it, which allocates + const bool solver_1 = md_est.solver_used == 1; + arena_t R(md_est.theta.size() * solver_1, + md_est.theta.size() * solver_1); + // Solver 3 assigns it + arena_t LU_solve_covariance; // Solver 1, 2, 3 arena_t s2(md_est.theta.size()); using stan::math::internal::ZeroOut; From 04da9a0df505f02c473b9ffe0cf267f47b6147cd Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Sun, 4 Oct 2026 22:48:34 +0200 Subject: [PATCH 2/8] laplace_marginal_density: test that solvers 2 and 3 allocate no unused buffers --- .../laplace_marginal_density_arena_test.cpp | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 test/unit/math/laplace/laplace_marginal_density_arena_test.cpp diff --git a/test/unit/math/laplace/laplace_marginal_density_arena_test.cpp b/test/unit/math/laplace/laplace_marginal_density_arena_test.cpp new file mode 100644 index 00000000000..40a750a6768 --- /dev/null +++ b/test/unit/math/laplace/laplace_marginal_density_arena_test.cpp @@ -0,0 +1,83 @@ +#include +#include +#include +#include +#include +#include +#include + +namespace laplace_marginal_density_arena_test { + +std::uintptr_t probe_base = 0; +std::intptr_t probe_peak = 0; + +// Poisson log likelihood that records the highest arena position reached +// at its calls, relative to probe_base +struct probing_poisson { + template + auto operator()(const Theta& theta, const std::vector& y, + std::ostream* msgs) const { + if (probe_base != 0) { + auto position = reinterpret_cast( + stan::math::ChainableStack::instance_->memalloc_.alloc(8)); + probe_peak = std::max(probe_peak, + static_cast(position - probe_base)); + } + return stan::math::poisson_log_lpmf(y, theta); + } +}; + +struct se_covariance { + template + auto operator()(const std::vector& x, const T1& alpha, + const T2& rho, std::ostream* msgs) const { + return stan::math::add_diag(stan::math::gp_exp_quad_cov(x, alpha, rho), + 1e-6); + } +}; + +// Peak arena bytes of one gradient with the given solver, with all +// allocations in one large arena block +std::intptr_t peak_arena_bytes(int solver, int n) { + using stan::math::var; + auto& memalloc = stan::math::ChainableStack::instance_->memalloc_; + stan::math::recover_memory(); + memalloc.alloc(16 << 20); + stan::math::recover_memory(); + memalloc.alloc(1 << 20); // skip the small first block + std::vector x(n, Eigen::VectorXd(2)); + std::vector y(n); + for (int i = 0; i < n; ++i) { + x[i] << std::sin(0.37 * i), std::cos(0.11 * i); + y[i] = i % 4; + } + Eigen::VectorXd theta_0 = Eigen::VectorXd::Zero(n); + var alpha = 1.2; + var rho = 0.7; + probe_base = reinterpret_cast(memalloc.alloc(8)); + probe_peak = 0; + var lp = stan::math::laplace_marginal_tol( + probing_poisson{}, std::forward_as_tuple(y), 1, se_covariance{}, + std::forward_as_tuple(x, alpha, rho), + std::make_tuple(theta_0, 1e-8, 200, solver, 0, true), nullptr); + probe_base = 0; + stan::math::recover_memory(); + return probe_peak; +} + +} // namespace laplace_marginal_density_arena_test + +TEST(laplace_marginal_density, solvers_2_and_3_allocate_no_unused_buffers) { + using laplace_marginal_density_arena_test::peak_arena_bytes; + for (int n : {20, 50}) { + std::intptr_t peak_1 = peak_arena_bytes(1, n); + ASSERT_GT(peak_1, 0); + ASSERT_LT(peak_1, 8 << 20); // all probes in the same block + for (int solver : {2, 3}) { + // The solvers' arena use differs by much less than one n x n matrix; + // an unused n x n buffer adds 8 n^2 bytes. + EXPECT_LT(peak_arena_bytes(solver, n) - peak_1, 4 * n * n) + << "solver " << solver << ", n = " << n; + } + } +} From 15a2b49300f5a6f5f48951605d51e90e1c082a90 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Thu, 15 Oct 2026 00:17:16 +0200 Subject: [PATCH 3/8] laplace arena test: measure with approx_bytes_used, run solver 3 for real, bound each solver --- .../laplace_marginal_density_arena_test.cpp | 88 ++++++------------- 1 file changed, 26 insertions(+), 62 deletions(-) diff --git a/test/unit/math/laplace/laplace_marginal_density_arena_test.cpp b/test/unit/math/laplace/laplace_marginal_density_arena_test.cpp index 40a750a6768..f900077ed91 100644 --- a/test/unit/math/laplace/laplace_marginal_density_arena_test.cpp +++ b/test/unit/math/laplace/laplace_marginal_density_arena_test.cpp @@ -1,83 +1,47 @@ -#include #include +#include #include #include -#include -#include +#include #include namespace laplace_marginal_density_arena_test { -std::uintptr_t probe_base = 0; -std::intptr_t probe_peak = 0; - -// Poisson log likelihood that records the highest arena position reached -// at its calls, relative to probe_base -struct probing_poisson { - template - auto operator()(const Theta& theta, const std::vector& y, - std::ostream* msgs) const { - if (probe_base != 0) { - auto position = reinterpret_cast( - stan::math::ChainableStack::instance_->memalloc_.alloc(8)); - probe_peak = std::max(probe_peak, - static_cast(position - probe_base)); - } - return stan::math::poisson_log_lpmf(y, theta); - } -}; - -struct se_covariance { - template - auto operator()(const std::vector& x, const T1& alpha, - const T2& rho, std::ostream* msgs) const { - return stan::math::add_diag(stan::math::gp_exp_quad_cov(x, alpha, rho), - 1e-6); - } -}; - -// Peak arena bytes of one gradient with the given solver, with all -// allocations in one large arena block -std::intptr_t peak_arena_bytes(int solver, int n) { - using stan::math::var; +// Highest arena use reached at the likelihood calls of one gradient of a +// Poisson GP with the given solver, with every allocation in one arena block +size_t peak_arena_bytes(int solver, int n) { auto& memalloc = stan::math::ChainableStack::instance_->memalloc_; - stan::math::recover_memory(); memalloc.alloc(16 << 20); stan::math::recover_memory(); memalloc.alloc(1 << 20); // skip the small first block - std::vector x(n, Eigen::VectorXd(2)); - std::vector y(n); - for (int i = 0; i < n; ++i) { - x[i] << std::sin(0.37 * i), std::cos(0.11 * i); - y[i] = i % 4; - } + size_t peak = 0; + auto ll_fun = [&](auto&& theta, auto&& y, std::ostream*) { + peak = std::max(peak, memalloc.approx_bytes_used()); + return stan::math::poisson_log_lpmf(y, theta); + }; + std::vector x(n); + std::iota(x.begin(), x.end(), 0.0); + std::vector y(n, 2); Eigen::VectorXd theta_0 = Eigen::VectorXd::Zero(n); - var alpha = 1.2; - var rho = 0.7; - probe_base = reinterpret_cast(memalloc.alloc(8)); - probe_peak = 0; - var lp = stan::math::laplace_marginal_tol( - probing_poisson{}, std::forward_as_tuple(y), 1, se_covariance{}, + stan::math::var alpha = 1.2, rho = 0.7; + stan::math::laplace_marginal_tol( + ll_fun, std::forward_as_tuple(y), 1, + stan::math::test::sqr_exp_kernel_functor{}, std::forward_as_tuple(x, alpha, rho), - std::make_tuple(theta_0, 1e-8, 200, solver, 0, true), nullptr); - probe_base = 0; + std::make_tuple(theta_0, 1e-8, 200, solver, 0, false), nullptr); stan::math::recover_memory(); - return probe_peak; + return peak; } } // namespace laplace_marginal_density_arena_test TEST(laplace_marginal_density, solvers_2_and_3_allocate_no_unused_buffers) { using laplace_marginal_density_arena_test::peak_arena_bytes; - for (int n : {20, 50}) { - std::intptr_t peak_1 = peak_arena_bytes(1, n); - ASSERT_GT(peak_1, 0); - ASSERT_LT(peak_1, 8 << 20); // all probes in the same block - for (int solver : {2, 3}) { - // The solvers' arena use differs by much less than one n x n matrix; - // an unused n x n buffer adds 8 n^2 bytes. - EXPECT_LT(peak_arena_bytes(solver, n) - peak_1, 4 * n * n) - << "solver " << solver << ", n = " << n; - } - } + const int n = 50; + const size_t matrix = 8 * n * n; // bytes of one n x n matrix + size_t peak_1 = peak_arena_bytes(1, n); + // Beyond solver 1's n x n arena matrices, solver 2 needs none and solver 3 + // one (A); a pre-sized buffer that is then reassigned adds a whole one. + EXPECT_LT(peak_arena_bytes(2, n), peak_1 + matrix / 2); + EXPECT_LT(peak_arena_bytes(3, n), peak_1 + matrix * 3 / 2); } From 900c7526f6bbbd2832964bc18251a5c1d6c36b19 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Thu, 15 Oct 2026 00:17:16 +0200 Subject: [PATCH 4/8] laplace_marginal_density: size R with a ternary, declare LU_solve_covariance in the solver-3 branch --- stan/math/mix/functor/laplace_marginal_density.hpp | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/stan/math/mix/functor/laplace_marginal_density.hpp b/stan/math/mix/functor/laplace_marginal_density.hpp index 0cefbfc9d8d..bda64e16ad3 100644 --- a/stan/math/mix/functor/laplace_marginal_density.hpp +++ b/stan/math/mix/functor/laplace_marginal_density.hpp @@ -198,11 +198,8 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, auto ll_args_filter = internal::filter_var_scalar_types(ll_args_copy); // tuple of references to var types // Solver 1 writes R in place; solvers 2 and 3 assign it, which allocates - const bool solver_1 = md_est.solver_used == 1; - arena_t R(md_est.theta.size() * solver_1, - md_est.theta.size() * solver_1); - // Solver 3 assigns it - arena_t LU_solve_covariance; + const Eigen::Index n_R = md_est.solver_used == 1 ? md_est.theta.size() : 0; + arena_t R(n_R, n_R); // Solver 1, 2, 3 arena_t s2(md_est.theta.size()); using stan::math::internal::ZeroOut; @@ -261,7 +258,8 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, s2.deep_copy(s2_tmp); internal::copy_compute_s2(partial_parm, ll_args_filter); } else { // options.solver with LU decomposition - LU_solve_covariance = md_est.LU.solve(covariance_val); + arena_t LU_solve_covariance + = md_est.LU.solve(covariance_val); auto I_minus_BinvKW = Eigen::MatrixXd::Identity(md_est.W_r.rows(), md_est.W_r.cols()) - LU_solve_covariance * md_est.W_r; @@ -295,7 +293,7 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, if (md_est.solver_used == 1 || md_est.solver_used == 2) { v = covariance_val * s2 - covariance_val * R * covariance_val * s2; } else { - v = LU_solve_covariance * s2; + v = md_est.LU.solve(covariance_val * s2); } laplace_likelihood::diff_eta_implicit(ll_fun, v, md_est.theta, ll_args_copy, msgs); From d45cfa5d5f670445433ce82ef639887fbd0c67e5 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Thu, 15 Oct 2026 09:43:49 +0200 Subject: [PATCH 5/8] laplace_marginal_density: keep LU_solve_covariance in the enclosing scope, only drop its pre-sizing --- stan/math/mix/functor/laplace_marginal_density.hpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/stan/math/mix/functor/laplace_marginal_density.hpp b/stan/math/mix/functor/laplace_marginal_density.hpp index bda64e16ad3..a0209f8d037 100644 --- a/stan/math/mix/functor/laplace_marginal_density.hpp +++ b/stan/math/mix/functor/laplace_marginal_density.hpp @@ -200,6 +200,8 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, // Solver 1 writes R in place; solvers 2 and 3 assign it, which allocates const Eigen::Index n_R = md_est.solver_used == 1 ? md_est.theta.size() : 0; arena_t R(n_R, n_R); + // Solver 3 assigns it + arena_t LU_solve_covariance; // Solver 1, 2, 3 arena_t s2(md_est.theta.size()); using stan::math::internal::ZeroOut; @@ -258,8 +260,7 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, s2.deep_copy(s2_tmp); internal::copy_compute_s2(partial_parm, ll_args_filter); } else { // options.solver with LU decomposition - arena_t LU_solve_covariance - = md_est.LU.solve(covariance_val); + LU_solve_covariance = md_est.LU.solve(covariance_val); auto I_minus_BinvKW = Eigen::MatrixXd::Identity(md_est.W_r.rows(), md_est.W_r.cols()) - LU_solve_covariance * md_est.W_r; @@ -293,7 +294,7 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, if (md_est.solver_used == 1 || md_est.solver_used == 2) { v = covariance_val * s2 - covariance_val * R * covariance_val * s2; } else { - v = md_est.LU.solve(covariance_val * s2); + v = LU_solve_covariance * s2; } laplace_likelihood::diff_eta_implicit(ll_fun, v, md_est.theta, ll_args_copy, msgs); From 843d77d65363ceca4c683e29601f0ef495616339 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Thu, 15 Oct 2026 21:13:35 +0200 Subject: [PATCH 6/8] laplace_marginal_density: pre-size R for every solver and write it in place with noalias --- .../mix/functor/laplace_marginal_density.hpp | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/stan/math/mix/functor/laplace_marginal_density.hpp b/stan/math/mix/functor/laplace_marginal_density.hpp index a0209f8d037..7cd2f974d86 100644 --- a/stan/math/mix/functor/laplace_marginal_density.hpp +++ b/stan/math/mix/functor/laplace_marginal_density.hpp @@ -197,9 +197,8 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, ll_fun, ll_args_vals, covariance_val, options, msgs); auto ll_args_filter = internal::filter_var_scalar_types(ll_args_copy); // tuple of references to var types - // Solver 1 writes R in place; solvers 2 and 3 assign it, which allocates - const Eigen::Index n_R = md_est.solver_used == 1 ? md_est.theta.size() : 0; - arena_t R(n_R, n_R); + // Every solver writes R in place; no right-hand side below aliases it + arena_t R(md_est.theta.size(), md_est.theta.size()); // Solver 3 assigns it arena_t LU_solve_covariance; // Solver 1, 2, 3 @@ -243,13 +242,14 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, internal::copy_compute_s2(partial_parm, ll_args_filter); } } else if (md_est.solver_used == 2) { - R = md_est.W_r - - md_est.W_r * md_est.K_root - * md_est.L.transpose() - .template triangularView() - .solve( - md_est.L.template triangularView() - .solve(md_est.K_root.transpose() * md_est.W_r)); + R.deep_copy(md_est.W_r); + R.noalias() + -= md_est.W_r * md_est.K_root + * md_est.L.transpose() + .template triangularView() + .solve( + md_est.L.template triangularView().solve( + md_est.K_root.transpose() * md_est.W_r)); arena_t C = md_est.L.template triangularView().solve( @@ -264,7 +264,7 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, auto I_minus_BinvKW = Eigen::MatrixXd::Identity(md_est.W_r.rows(), md_est.W_r.cols()) - LU_solve_covariance * md_est.W_r; - R = md_est.W_r * I_minus_BinvKW; // == W - W B^{-1} K W + R.noalias() = md_est.W_r * I_minus_BinvKW; // == W - W B^{-1} K W arena_t A = covariance_val - covariance_val * md_est.W_r * LU_solve_covariance; auto s2_tmp = laplace_likelihood::compute_s2(ll_fun, md_est.theta, A, From 9f7da5f94d2c12798e2afc3954500522131260c4 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Thu, 15 Oct 2026 21:23:18 +0200 Subject: [PATCH 7/8] laplace_marginal_density: name the latent dimension once --- stan/math/mix/functor/laplace_marginal_density.hpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/stan/math/mix/functor/laplace_marginal_density.hpp b/stan/math/mix/functor/laplace_marginal_density.hpp index 7cd2f974d86..a439c992aa5 100644 --- a/stan/math/mix/functor/laplace_marginal_density.hpp +++ b/stan/math/mix/functor/laplace_marginal_density.hpp @@ -197,12 +197,13 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, ll_fun, ll_args_vals, covariance_val, options, msgs); auto ll_args_filter = internal::filter_var_scalar_types(ll_args_copy); // tuple of references to var types + const Eigen::Index theta_size = md_est.theta.size(); // Every solver writes R in place; no right-hand side below aliases it - arena_t R(md_est.theta.size(), md_est.theta.size()); + arena_t R(theta_size, theta_size); // Solver 3 assigns it arena_t LU_solve_covariance; // Solver 1, 2, 3 - arena_t s2(md_est.theta.size()); + arena_t s2(theta_size); using stan::math::internal::ZeroOut; if (md_est.solver_used == 1) { if (options.hessian_block_size == 1) { From 9a2fbd0886b5ee945c2e13dd4d5fae77d70b6f93 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Sat, 17 Oct 2026 12:24:28 +0200 Subject: [PATCH 8/8] laplace_marginal_density: remove the arena allocation test --- .../laplace_marginal_density_arena_test.cpp | 47 ------------------- 1 file changed, 47 deletions(-) delete mode 100644 test/unit/math/laplace/laplace_marginal_density_arena_test.cpp diff --git a/test/unit/math/laplace/laplace_marginal_density_arena_test.cpp b/test/unit/math/laplace/laplace_marginal_density_arena_test.cpp deleted file mode 100644 index f900077ed91..00000000000 --- a/test/unit/math/laplace/laplace_marginal_density_arena_test.cpp +++ /dev/null @@ -1,47 +0,0 @@ -#include -#include -#include -#include -#include -#include - -namespace laplace_marginal_density_arena_test { - -// Highest arena use reached at the likelihood calls of one gradient of a -// Poisson GP with the given solver, with every allocation in one arena block -size_t peak_arena_bytes(int solver, int n) { - auto& memalloc = stan::math::ChainableStack::instance_->memalloc_; - memalloc.alloc(16 << 20); - stan::math::recover_memory(); - memalloc.alloc(1 << 20); // skip the small first block - size_t peak = 0; - auto ll_fun = [&](auto&& theta, auto&& y, std::ostream*) { - peak = std::max(peak, memalloc.approx_bytes_used()); - return stan::math::poisson_log_lpmf(y, theta); - }; - std::vector x(n); - std::iota(x.begin(), x.end(), 0.0); - std::vector y(n, 2); - Eigen::VectorXd theta_0 = Eigen::VectorXd::Zero(n); - stan::math::var alpha = 1.2, rho = 0.7; - stan::math::laplace_marginal_tol( - ll_fun, std::forward_as_tuple(y), 1, - stan::math::test::sqr_exp_kernel_functor{}, - std::forward_as_tuple(x, alpha, rho), - std::make_tuple(theta_0, 1e-8, 200, solver, 0, false), nullptr); - stan::math::recover_memory(); - return peak; -} - -} // namespace laplace_marginal_density_arena_test - -TEST(laplace_marginal_density, solvers_2_and_3_allocate_no_unused_buffers) { - using laplace_marginal_density_arena_test::peak_arena_bytes; - const int n = 50; - const size_t matrix = 8 * n * n; // bytes of one n x n matrix - size_t peak_1 = peak_arena_bytes(1, n); - // Beyond solver 1's n x n arena matrices, solver 2 needs none and solver 3 - // one (A); a pre-sized buffer that is then reassigned adds a whole one. - EXPECT_LT(peak_arena_bytes(2, n), peak_1 + matrix / 2); - EXPECT_LT(peak_arena_bytes(3, n), peak_1 + matrix * 3 / 2); -}