Skip to content

pdf/logpdf/cdf functions are not inlined for ForwardDiff.Dual arguments #223

Description

@devmotion

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.)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions