[avehtari]

In reverse mode with vector arguments, with this PR normal_lpdf takes 0.17 to 0.99 of the develop time in the cases that I measured (see Speed). Values and gradients agree with develop up to rounding. Exceptions have the same types and texts as on develop, unless sigma and y or mu are both invalid; then the text is about sigma.

The same speedup could be made for other distributions, but I start with this one to make this easier to review. Some of the ideas have already been used in normal_id_glm_lpdf and with this kind of speedup + O1, we would get most of the benefits of hand written glm functions without need to actually add them to Stan math and language (I have done those speed experiments with other distributions, but not showing the timings here as this is about normal distribution).

The scripts behind the numbers in this text are in https://github.com/avehtari/stan-pr-checks/tree/724274d2d3ea10b1e2ce11d08ac71d941118b14d/pr3498. The README.md there gives the commands and the expected output for each claim (sections 1 to 7). Section 8 shows the to_ref copy mentioned in the Summary.

This PR was Claude-assisted.

Summary

On develop, two parts of a vectorized call cost 56 to 76 % of its time in case A below (measured with a prototype):

• The checks check_not_nan(y) and check_finite(mu) run before the arithmetic. Each is a loop with a test per element.
• y_scaled is a heap array of N doubles. So is y_scaled_sq when sigma is an autodiff argument, and scaled_diff when two or more arguments are autodiff arguments. In the prototype, removing y_scaled_sq or scaled_diff saved 1.7 to 4.4 ns per element from N = 1e4. to_ref also copies the values of a Matrix<var> argument into a heap array: its return type ref_type_t is a plain array when Eigen::Ref cannot bind the value view.

The changes:

• The function checks sigma first, right after check_consistent_sizes.
• In the arithmetic path, the checks of y and mu run only if the sum of squares is not finite. A NaN y or a non-finite mu always makes the sum NaN or infinite, also when sigma is +inf. An infinite y with a finite mu passes the checks and gives -inf, as on develop. normal_id_glm_lpdf has used this approach since Glm optimizations #1185.
• The checks of y and mu still run before the size-zero return and before the return for all-data arguments with propto = true.
• If mu is an autodiff vector, its partials first hold (y - mu) / sigma; otherwise the partials of an autodiff y vector do. The function takes the sum of squares from them and then scales them in place. The partial of y is the negated partial of mu.
• For an autodiff vector sigma, its partials first hold 1 / sigma. If mu or y is also an autodiff vector, one loop then writes both partials; else an expression writes the partials of sigma in place.
• The function reads the values of autodiff containers through views, without a copy.
• The scalar sigma partial is inv_sigma * (sum_sq - N), not a sum over N elements.

Apart from the scalar sigma partial, the value and the partials use the develop formulas.

Speed

Triton, Xeon Gold 6248, g++ 14.3.0, -O3, one core. One evaluation creates the autodiff arguments, calls normal_lpdf, runs grad() and recover_memory(). Each case and version is a separate binary; the develop and branch binaries alternate over 5 rounds. Each cell is the time of the branch divided by the time of develop (median over the rounds). D is data, V is an autodiff argument, varmat is var_value<VectorXd>.

case	y	mu	sigma	propto	N = 100	1e3	1e4	1e5
A	D vector	varmat	V scalar	true	0.51	0.50	0.17	0.26
A	D vector	varmat	V scalar	false	0.52	0.50	0.17	0.26
B	D vector	V Matrix	V scalar	true	0.70	0.71	0.36	0.51
C	D vector	V Matrix	V Matrix	true	0.83	0.88	0.54	0.56
C	D vector	V Matrix	V Matrix	false	0.84	0.88	0.54	0.57
D	V Matrix	V scalar	V scalar	true	0.76	0.77	0.48	0.43
E	D vector	V Matrix	D scalar	true	0.75	0.77	0.60	0.67
F	D vector	D vector	V scalar	true	0.34	0.25	0.25	0.48
G	D std::vector	V std::vector	V scalar	true	0.72	0.71	0.40	0.60
H	D vector	D vector	D scalar	false	0.36	0.30	0.38	0.80
Y	D vector	V Matrix	D vector	true	0.82	0.84	0.44	0.55

In case A, develop takes 4.6 ns per element at N = 100 and 11 to 15 ns from N = 1e4. The branch takes 1.9 to 4.0 ns. The largest ratio of a single round in the table is 0.89 (case C, N = 1e3).

The smallest speedup I measured is for a data vector sigma with scalar data y and an autodiff mu vector: 0.96 to 0.99 at N <= 1e3 and 0.62 to 0.72 from N = 1e4. There the branch computes 1 / sigma twice per element, where develop computes it once into a heap array.

Testing

• A check program compares the branch with a copy of the develop function in the same binary. It covers:

  • 25 type cases: data, Matrix<var>, row Matrix<var>, varmat column and row vectors, std::vector<var>, scalars;
  • 16 to 19 input cases per type case: NaN, infinite and overflowing values, invalid sigma together with invalid y or mu, size zero, size mismatch;
  • propto true and false, reverse mode and fvar<double>.

  In 1512 of the 1692 comparisons, the branch gives the same exception text as develop, or values and gradients within 1e-12 times max(1, |reference|). The largest error is 2.4e-15 for the values and 1.4e-15 for the gradients. The other 180 comparisons have an invalid sigma and an invalid y or mu, and the branch reports sigma. A malloc counter finds no allocation of 8 N bytes or more on the branch in any of the 25 type cases; develop makes up to 7 such allocations per evaluation.

• New tests in test/unit/math/prim/prob/normal_test.cpp and test/unit/math/mix/prob/normal_test.cpp:

  • the full exception texts for invalid arguments of each argument type (data, Matrix<var>, varmat, std::vector<var>, fvar);
  • an infinite y, and size zero with an invalid scalar;
  • values and gradients against mpmath references for Matrix<var>, row Matrix<var>, varmat (column and row) and std::vector<var>;
  • expect_ad with vector arguments.

• The new tests also pass with the develop header, since no call has more than one invalid argument. To check that they cover the new code, I removed, in turn, the check after the sum of squares in each of the two branches of the arithmetic (one for an autodiff mu or y vector, one for the other cases). At 760f6a1f65, 8 and 5 assertions failed.

• These tests pass: prim/prob/normal_test, mix/prob/normal_test, rev/prob/normal_log_test, rev/prob/normal_id_glm_lpdf_test and the generated test/prob/normal lpdf tests. rev/prob/normal_id_glm_lpdf_test uses normal_lpdf as reference.

• opencl/rev/normal_lpdf_test passes on a V100. It compares the unchanged OpenCL function with the changed CPU function.

• cpplint, make test-math-dependencies, the header check, make test-headers and clang-format 10 pass. make doxygen with doxygen 1.9.7 gives no warnings.

• The unit tests, the generated tests, the OpenCL test and the checklist items ran at 421386dec5 (the change applied to develop at dc66941969) with g++ 14.3.0.

With fvar<var> arguments and an invalid y or mu, the function now throws after it has made var nodes; on develop it threw before. The nodes stay on the stack until recover_memory().

Related PRs and issues

• [WIP] Allow distributions to return vectors #2751 and Cleanup: Remove unused #includes in stan/math and tests (~minus 3kLOC) #3470 also change normal_lpdf.hpp. Cleanup: Remove unused #includes in stan/math and tests (~minus 3kLOC) #3470 removes three includes there; this PR changes the include block too, so the two conflict in that block only.
• All make_partials_propagator to return its input args as a tuple #3423 proposes a new partials interface and uses normal_lpdf as its example. This PR uses the present partials_propagator.
• Add check_foreach(...) function that does multiple checks at once on an array or matrix #3417 proposes one pass over an argument for several checks. This PR avoids the pass over y and mu when the inputs are valid.

Release notes

Vectorized normal_lpdf is faster: it no longer allocates heap arrays of size N, and it checks y and mu only when the sum of squares is not finite.

Checklist

• Copyright holder: Aalto University

  The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
  - Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
  - Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

• the basic tests are passing

  • new unit tests pass (to run, use: ./runTests.py test/unit)
  • header checks pass, (make test-headers)
  • dependencies checks pass, (make test-math-dependencies)
  • docs build, (make doxygen)
  • code passes the built in C++ standards checks (make cpplint)

• the code is written in idiomatic C++ and changes are documented in the doxygen

• the new changes are tested

[WardBrian]

to_ref copies the values of a Matrix argument into another heap array.

It does? I'm pretty sure to_ref is just a wrapper around std::forward, no?

After the size check, the function checks sigma before the arithmetic. If sigma is not valid, the function runs the three develop checks in the develop order, so the first invalid argument gives the message.

I don't think we need to jump through hoops to keep errors in the same order

[stan-buildbot]

Name	Old Result	New Result	Ratio	Performance change( 1 - new / old )
stat_comp_benchmarks/benchmarks/gp_regr/gp_regr.stan	0.24	0.24	1.01	0.66% faster
stat_comp_benchmarks/benchmarks/gp_regr/gen_gp_data.stan	0.06	0.06	0.99	-0.63% slower
stat_comp_benchmarks/benchmarks/low_dim_corr_gauss/low_dim_corr_gauss.stan	0.02	0.02	1.0	-0.09% slower
stat_comp_benchmarks/benchmarks/arma/arma.stan	0.71	0.64	1.1	8.91% faster
stat_comp_benchmarks/benchmarks/arK/arK.stan	3.41	3.7	0.92	-8.35% slower
stat_comp_benchmarks/benchmarks/pkpd/one_comp_mm_elim_abs.stan	42.66	42.7	1.0	-0.1% slower
stat_comp_benchmarks/benchmarks/pkpd/sim_one_comp_mm_elim_abs.stan	0.58	0.61	0.96	-3.92% slower
stat_comp_benchmarks/benchmarks/gp_pois_regr/gp_pois_regr.stan	4.68	4.72	0.99	-0.8% slower
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix_collapse/low_dim_gauss_mix_collapse.stan	22.06	18.02	1.22	18.35% faster
stat_comp_benchmarks/benchmarks/garch/garch.stan	1.0	1.1	0.91	-9.57% slower
stat_comp_benchmarks/benchmarks/eight_schools/eight_schools.stan	0.11	0.11	1.04	3.83% faster
stat_comp_benchmarks/benchmarks/irt_2pl/irt_2pl.stan	9.26	10.76	0.86	-16.23% slower
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix/low_dim_gauss_mix.stan	6.58	6.91	0.95	-4.91% slower
stat_comp_benchmarks/benchmarks/sir/sir.stan	228.56	229.34	1.0	-0.34% slower
performance.compilation	366.19	363.07	1.01	0.85% faster
Mean result: 0.9979739419992885

Jenkins Console Log
Jenkins Build Stages
Commit hash: e7cab1646d71f360d212c6aecb8537bca67ebe6f

Machine information

Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.3 LTS
Release:	20.04
Codename:	focal

CPU:

Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Byte Order:                              Little Endian
Address sizes:                           43 bits physical, 48 bits virtual
CPU(s):                                  256
On-line CPU(s) list:                     0-255
Thread(s) per core:                      2
Core(s) per socket:                      64
Socket(s):                               2
NUMA node(s):                            2
Vendor ID:                               AuthenticAMD
CPU family:                              23
Model:                                   49
Model name:                              AMD EPYC 7742 64-Core Processor
Stepping:                                0
Frequency boost:                         enabled
CPU MHz:                                 1496.827
CPU max MHz:                             3416.0681
CPU min MHz:                             1500.0000
BogoMIPS:                                4491.56
Virtualization:                          AMD-V
L1d cache:                               4 MiB
L1i cache:                               4 MiB
L2 cache:                                64 MiB
L3 cache:                                512 MiB
NUMA node0 CPU(s):                       0-63,128-191
NUMA node1 CPU(s):                       64-127,192-255
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Old microcode:             Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Mitigation; untrained return thunk; SMT enabled with STIBP protection
Vulnerability Spec rstack overflow:      Mitigation; Safe RET
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; Retpolines; IBPB conditional; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Mitigation; IBPB before exit to userspace
Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sev sev_es

G++:

g++ (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

Clang:

clang version 10.0.0-4ubuntu1 
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin

[avehtari]

to_ref copies the values of a Matrix<var> argument into another heap array.

It does? I'm pretty sure to_ref is just a wrapper around std::forward, no?

I have to admit that I'm one edge of my knowledge. This is what Claude says:
Yes, for a Matrix<var> argument. The body is return std::forward<T>(a);, but the return type is ref_type_t<T&&>, which is plain_type_t<T> when Eigen::Ref cannot bind the expression. On develop, normal_lpdf calls to_ref(as_value_column_array_or_scalar(y_ref)). For a Matrix<var>, that argument is a view that reads the values through the vari pointers, so the return statement evaluates it into a new array. For a VectorXd there is no copy. A probe with a malloc counter finds one allocation of at least 8000 bytes for a Matrix<var> of size 1000, and none for a VectorXd. I reworded the sentence in the description.

I don't think we need to jump through hoops to keep errors in the same order

Agreed, I removed that code in dde60be194. normal_lpdf now calls check_positive(sigma) right after the size check. The y and mu checks stay where they were: after the sum of squares, and before the size-zero and all-data propto returns. The exception text now differs from develop only when sigma and y or mu are both invalid; it is then about sigma. In the check program, that is 180 of the 1692 comparisons (a NaN y or an infinite mu, each with a negative sigma), and the other 1512 still agree with develop (README section 7). I dropped the tests of the check order, and 421386dec5 adds tests for an invalid varmat and std::vector<var> sigma. I updated the description.