Many of the pdf/logpdf/cdf functions are not inlined when called with ForwardDiff.Dual arguments, although they are inlined for Float64. Adding @inline makes them 13-50% faster. Distributions.jl forwards to StatsFuns, so this affects AD users of Normal, Gamma, Chisq etc. as well, and there a call-site annotation is not an option.
Benchmark: loop over a varying Float64 x with Dual parameters, n = 4096, min of 30 reps, normalised within each run by a loop calling the special function that dominates the body so the numbers do not depend on the machine. Medians over 4 runs of master and 2 with the patch. Same numbers at -O2 and -O3. Julia 1.12.6, StatsFuns 2.2.1, ForwardDiff 1.4.3.
With @inline on all six functions in norm.jl (both the z-form and the 3-arg form), on betalogpdf and on gammalogpdf:
|
master |
@inline |
|
normpdf |
3.43 |
1.73 |
-50% |
normlogpdf |
2.06 |
1.50 |
-27% |
gammalogpdf |
2.32 |
1.68 |
-27% |
chisqlogpdf |
2.34 |
1.79 |
-24% |
binomlogpdf |
6.68 |
5.15 |
-23% |
binompdf |
6.63 |
5.35 |
-19% |
normlogcdf |
0.95 |
0.77 |
-18% |
normlogccdf |
0.98 |
0.80 |
-18% |
normcdf |
0.79 |
0.66 |
-17% |
normccdf |
0.78 |
0.66 |
-16% |
betapdf |
4.66 |
4.07 |
-13% |
betalogpdf |
3.90 |
3.65 |
-6% |
gammapdf +2%, chisqpdf +1%, fdist*, tdist* and pois* unchanged. Values are unchanged.
chisqlogpdf and binomlogpdf are not annotated themselves: they call gammalogpdf and betalogpdf, and get the improvement from those being inlined into their bodies, without their own call sites growing.
chisqpdf and gammapdf do not improve, in contrast to chisqlogpdf and gammalogpdf. For gammapdf the annotation only moves the calls to _logabsgamma, _digamma and log out of gammalogpdf and into gammapdf ([gammalogpdf, exp] becomes [_logabsgamma, DomainError, _digamma, log, log, log, exp]), so nothing is saved. chisqpdf is gammapdf(k / 2, 2, x) and just follows gammapdf.
The numbers are the same whether the parameters are loop-invariant or vary per element, so the improvements are not an artefact of hoisting.
I did benchmark @inline on the remaining functions too, but the code size is not free and there it does not pay. Summed over ten representative callers, LLVM lines emitted at Dual argument types:
|
annotations |
improvements >8% |
LLVM lines |
| master |
- |
- |
293 |
norm*, betalogpdf, gammalogpdf |
16 |
11 |
1400 |
| every function that improves |
22 |
13 |
3093 |
Annotating binomlogpdf itself takes it from -23% to -31% but costs 926 LLVM lines per call site instead of 20, fdistlogpdf costs 466 and tdistlogpdf 361, for -10% and -12%. Not worth it IMO.
Two things worth recording:
- For the functions defined as a
promote wrapper plus a where {T} worker, both methods need the annotation. Annotating only the worker moves the barrier rather than removing it: the worker is inlined into the wrapper, the wrapper then exceeds the threshold, and the call site is unchanged. That variant was slower than master for poislogpdf and tdistlogpdf.
pois is left out because @inline does nothing for poislogpdf and costs poispdf about 12%.
(This issue was written with the assistance of generative AI.)
Many of the pdf/logpdf/cdf functions are not inlined when called with
ForwardDiff.Dualarguments, although they are inlined forFloat64. Adding@inlinemakes them 13-50% faster. Distributions.jl forwards to StatsFuns, so this affects AD users ofNormal,Gamma,Chisqetc. as well, and there a call-site annotation is not an option.Benchmark: loop over a varying
Float64xwithDualparameters,n = 4096, min of 30 reps, normalised within each run by a loop calling the special function that dominates the body so the numbers do not depend on the machine. Medians over 4 runs of master and 2 with the patch. Same numbers at-O2and-O3. Julia 1.12.6, StatsFuns 2.2.1, ForwardDiff 1.4.3.With
@inlineon all six functions innorm.jl(both thez-form and the 3-arg form), onbetalogpdfand ongammalogpdf:@inlinenormpdfnormlogpdfgammalogpdfchisqlogpdfbinomlogpdfbinompdfnormlogcdfnormlogccdfnormcdfnormccdfbetapdfbetalogpdfgammapdf+2%,chisqpdf+1%,fdist*,tdist*andpois*unchanged. Values are unchanged.chisqlogpdfandbinomlogpdfare not annotated themselves: they callgammalogpdfandbetalogpdf, and get the improvement from those being inlined into their bodies, without their own call sites growing.chisqpdfandgammapdfdo not improve, in contrast tochisqlogpdfandgammalogpdf. Forgammapdfthe annotation only moves the calls to_logabsgamma,_digammaandlogout ofgammalogpdfand intogammapdf([gammalogpdf, exp]becomes[_logabsgamma, DomainError, _digamma, log, log, log, exp]), so nothing is saved.chisqpdfisgammapdf(k / 2, 2, x)and just followsgammapdf.The numbers are the same whether the parameters are loop-invariant or vary per element, so the improvements are not an artefact of hoisting.
I did benchmark
@inlineon the remaining functions too, but the code size is not free and there it does not pay. Summed over ten representative callers, LLVM lines emitted atDualargument types:norm*,betalogpdf,gammalogpdfAnnotating
binomlogpdfitself takes it from -23% to -31% but costs 926 LLVM lines per call site instead of 20,fdistlogpdfcosts 466 andtdistlogpdf361, for -10% and -12%. Not worth it IMO.Two things worth recording:
promotewrapper plus awhere {T}worker, both methods need the annotation. Annotating only the worker moves the barrier rather than removing it: the worker is inlined into the wrapper, the wrapper then exceeds the threshold, and the call site is unchanged. That variant was slower than master forpoislogpdfandtdistlogpdf.poisis left out because@inlinedoes nothing forpoislogpdfand costspoispdfabout 12%.(This issue was written with the assistance of generative AI.)