fix(reduce_sum_static): wrap Args with ref_type_t to materialize Eigen expressions - #3305
fix(reduce_sum_static): wrap Args with ref_type_t to materialize Eigen expressions#3305tmchow wants to merge 6 commits into
Conversation
Sorry, something went wrong.
…n expressions
reduce_sum has been wrapping shared Args with ref_type_t<Args&&>... since
v4.8+ so that Eigen expression-template temporaries (e.g. M.row(0),
matrix.col(i)) materialize before being handed off to TBB's parallel
reducer. reduce_sum_static's STAN_THREADS branch was not updated in the
same pass and still forwarded plain Args..., so callers passing a
temporary Eigen expression as a shared argument could get a dangling
reference once the temporary went out of scope.
Mirror the reduce_sum wrapping exactly:
internal::reduce_sum_impl<..., ref_type_t<Args&&>...>()(...)
No other changes. ref_type_t is already in scope via
stan/math/prim/meta.hpp (already included).
Closes stan-dev#3304
Jenkins Console Log Machine information No LSB modules are available. Distributor ID: Ubuntu Description: Ubuntu 20.04.3 LTS Release: 20.04 Codename: focal CPU: G++: Clang: |
Sorry, something went wrong.
|
Thanks! Though I think I would prefer if we used ref_type_t in the class and not in the template instantiation directly. So we would pass Args&& and then the class would handle any ownership logic |
Sorry, something went wrong.
Move the ref_type_t<Args> wrapping out of the call sites in reduce_sum and reduce_sum_static and into the recursive_reducer's args_tuple_ member. Both prim arithmetic and rev var specializations now own the materialized Eigen expression by typing args_tuple_ as std::tuple<ref_type_t<Args>...>; the call sites pass plain Args... and the implicit conversion materializes the temporary at construction. Per @SteveBronder review on stan-dev#3305.
|
@SteveBronder Pushed f73f114.
Verified locally with STAN_THREADS=true: test/unit/math/prim/functor/reduce_sum_test and test/unit/math/rev/functor/reduce_sum_test all pass. Jenkins i think will exercise the rest including the mix tests and the perf benchmarks. |
Sorry, something went wrong.
The recursive_reducer's args_tuple_ stores ref_type_t<Args>... so Eigen expression-template args (e.g. M.row(0)) materialize at storage time. The rev specialization's outer operator() was still calling save_varis(args...) on the raw expressions before constructing the worker, then materializing in the args_tuple_ initialization, walking the underlying expression twice. ExpressionTestRev.reduce_sum1 caught this: matrix2_expr3_counter == 2 vs expected <= 1. Materialize args... once at the function entry into a local std::tuple<ref_type_t<Args>...> args_refs, then route count_vars, save_varis, and worker construction through stan::math::apply on args_refs. Subsequent calls see plain matrices, so the counter increments once (in args_refs initialization) and not again. The prim specialization needs no change: its outer operator() does not call count_vars or save_varis on the shared args, so args_tuple_ initialization is the only evaluation. Closes stan-dev#3304.
|
@SteveBronder Pushed 4efff0c72f. CI failure on the previous push was a real regression in the rev specialization, not an infra issue. Locally the failing expression test passes and the broader unit suite stays green: test/expressions/tests*_test (Prim/Rev/Fwd) -> 3/3 PASS test/unit/math/rev/functor/reduce_sum_test -> 9/9 PASS (also with STAN_THREADS=true) test/unit/math/prim/functor/reduce_sum_test -> 32/32 PASS |
Sorry, something went wrong.
Jenkins Console Log Machine information No LSB modules are available. Distributor ID: Ubuntu Description: Ubuntu 20.04.3 LTS Release: 20.04 Codename: focal CPU: G++: Clang: |
Sorry, something went wrong.
|
@SteveBronder anything else I should do for this? Lmk |
Sorry, something went wrong.
Summary
Closes #3304. reduce_sum wraps its shared arguments with ref_type_t<Args&&>... before handing them to TBB so that Eigen expression-template temporaries (e.g. M.row(0), matrix.col(i)) materialize instead of being held as references to a destroyed temporary. reduce_sum_static's STAN_THREADS branch was not updated in the same pass and still forwarded plain Args....
A caller that passes a temporary Eigen expression as a shared argument to reduce_sum_static can therefore end up with a dangling reference once the temporary falls out of scope at the call site.
Fix
stan/math/prim/functor/reduce_sum_static.hpp:
This exactly mirrors the wrapping already in place in reduce_sum.hpp (see line 208).
Notes
Closes #3304