Gibbs Compilation Performance

[joelkandiah]

The runtime of the compilation of the Gibbs sampler appears to be O(N_component_samplers), however some tests show that the cost to compile any given call to sample is high, and the cost of each individual component sampler is high.

I have attached a test file, which allows a user to pass in an argument for the number of components to use (between 3 and 9) where 17 Normal(0,1) variables are sampled. These are sampled using a default Metropolis Hastings sampler (i.e. from the prior) and no data is provided. We run this for two samples to allow the code to compile the sample call with and without varinfo being passed.

(Bad practice: On my local machine) This estimates a runtime of ~25-30 seconds per extra component sampler, where for 3 components the time taken is ~150 seconds and for 8 components the time taken is ~290 seconds.
This is significantly longer than using just an MH sampler without Gibbs ~10 seconds.

using Turing # v0.37.0
using Distributions

# --- Probabilistic Model Definition ---
@model function test_gibbs_individual()
    # Define the model parameters with Normal priors
    p1 ~ Normal(0, 1)
    p2 ~ Normal(0, 1)
    p3 ~ Normal(0, 1)
    p4 ~ Normal(0, 1)
    p5 ~ Normal(0, 1)
    p6 ~ Normal(0, 1)
    p7 ~ Normal(0, 1)
    p8 ~ Normal(0, 1)

    vp1 ~ Normal(0, 1)
    vp2 ~ Normal(0, 1)
    vp3 ~ Normal(0, 1)
    vp4 ~ Normal(0, 1)
    vp5 ~ Normal(0, 1)
    vp6 ~ Normal(0, 1)
    vp7 ~ Normal(0, 1)
    vp8 ~ Normal(0, 1)

    global_p1 ~ Normal(0, 1)

    return
end


# --- Gibbs Construction --- #
# Define a simple gibbs construction for 9 compartments (dynamic construction with the same number of parameters can be found in the attached .txt file
gibbs_MH = Gibbs(
    (@varname(global_p1),) => MH(),
    (@varname(p1), @varname(vp1)) => MH(),
    (@varname(p2), @varname(vp2)) => MH(),
    (@varname(p3), @varname(vp3)) => MH(),
    (@varname(p4), @varname(vp4)) => MH(),
    (@varname(p5), @varname(vp5)) => MH(),
    (@varname(p6), @varname(vp6)) => MH(),
    (@varname(p7), @varname(vp7)) => MH(),
    (@varname(p8), @varname(vp8)) => MH()
)

# --- Time compilation --- #
# Call sample and use this to estimate the compilation time (model evals should be fast
time1 = time_ns()
sample(test_gibbs_individual(), gibbs_MH, MCMCSerial(), 2, 1; progress=false) # progress=false to avoid output interfering with timing
time2 = time_ns()

# --- Output Results --- #
println("n_components: 9")
println("Approximate Compilation Time: $(time2 - time1) ns")
println("Approximate Compilation Time: $((time2 - time1)/1e9) s")

The code has been attached as a .txt file

main.txt

[mhauru]

We had a chat about this with @joelkandiah, @penelopeysm and myself last week. It's not obvious to me how to best fix this, but I do think we should fix it, because having your sample call first do type inference for 10 minutes straight is not a good user experience.

I may be overlooking something, but I suspect we do want tilde_obssume!! calls to specialise on each component sampler, for performance. This would mean that the linear scaling remains. The question is then, why does it compile/infer for so long for each component sampler. Could we shift the specialisation burden to some other level of the call stack where it goes quicker, or could we somehow simplify the types in play (GibbsContext) to give the compiler an easier time?

Hopefully we would fix this together with #2403 and with doing the same changes that Penny did for ConditionContext (TuringLang/DynamicPPL.jl#857) for GibbsContext.

[mhauru]

@joelkandiah, @penelopeysm, and I have been looking into this and discussing on Slack. I'm moving some of the findings and future conversation here for broader visibility and more permanence.

As Joel's above snippet shows, Gibbs results in very long inference times when there are multiple components. Some clues as to what might be going on:

@snoop_inference claims that the vast majority of time is not spent doing type inference. For a slightly different benchmarking snippet than Joel's above, I got InferenceTimingNode: 184.275454/193.283657 on Core.Compiler.Timings.ROOT() with 734 direct children, so 184s out of 193s on something other than inference.

The above is somewhat supported by doing an @profile run that includes compilation. I get flame graphs like this:

The towers on the left all look like type inference. I don't know what the vast emptiness is, but I think LLVM codegen stuff would not show up in profiling, so that tracks.

When I tried @snoop_llvm on one test snippet it claimed that 96s out of a total of 184s were spent in code generation. I need to investigate the @snoop_llvm results in more detail.

[mhauru]

Here's another hint that LLVM is the problem: One test snippet I have takes 197s to run, but if I run it with julia -O0, turning off (most?) compiler optimisations, that comes down to 29s. The flame graph continues to have the massive blank though 🤷. (I'm using Profile.init(n = 10^9, delay = 0.01) which should allow 1000 times the default number of instruction pointers to be stored.)

[mhauru]

Looking at the @snoop_llvm results for the example case that I've been using, they capture a total of 96 seconds of LLVM/codegen time. Here are the top contributor functions to that time:

         1.316336624 => "DynamicPPL.evaluate_threadsafe!!"
         1.554801575 => ""
         1.921519249 => "DynamicPPL.assume"
  3.6262722049999994 => "DynamicPPL._evaluate!!"
   6.574277672000004 => "DynamicPPL.tilde_assume!!"
         7.555176872 => "Main.test_gibbs_individual"
   66.84144883400002 => "DynamicPPL.tilde_assume"

test_gibbs_individual is the name of my model function. But tilde_assume is the dominant cost.

For posterity, the snippet to produce those times:

module Analysis

using DataFrames, SnoopCompile

function get_name(pair)
    time, e = pair
    s = first(e)
    m = match(r"^Tuple\{typeof\(([^\)]*)\)", s)
    if m === nothing
        println("No match for $s")
        return ""
    else
        return m.captures[1]
    end
end

times, info = SnoopCompile.read_snoop_llvm("func_names.csv", "llvm_timings.yaml", tmin_secs = 0.000001);

times_names = [first(p) => get_name(p) for p in times]

df_names = DataFrame(; times=first.(times_names), names=last.(times_names))
total_names_by_times = [sum(g[!, :times]) => g[1,:names] for g in groupby(df_names, :names)]

@show sort(total_names_by_times)

end

Digging into the tilde_assume LLVM times a bit: We have compiled 1121 different specialisations for tilde_assume (!). The total time of 66s is not dominated by a few specialisations, but rather by the bulk of those 1121, each of which only takes something like 10-100ms.

Contributions to the tilde_assume LLVM time by context:

3-element Vector{Pair{Float64, SubString{String}}}:
 0.7451583359999998 => "DynamicPPL.DefaultContext"
       26.095196426 => "DynamicPPL.ValuesAsInModelContext"
  40.00109407199997 => "Turing.Inference.GibbsContext"

In trying to understand why we have the 1121 tilde_assume specialisations, after some further analysis, a picture starts to emerge. In my case, I have 17 variables and 9 component samplers. The different specialisations can be counted roughly as

9 for the number of component samplers
* 15 for the number of variables that are conditioned for each component
* 2 for the sampler being either the component sampler or `SampleFromPrior`
* 2 for having just plain GibbsContext or GibbsContext wrapped in ValuesAsInModelContext
* 2 for the VarInfo + ThreadSafeVarInfo
= 1080

And then there's some extras on top because one of the components only has one variable, and there's DefaultContext, and blahblah.

I think I'll have to take a break and mull on this a bit. Hopefully the combinatorial explosion here could be slowed down, maybe by rejigging some of the tilde pipeline, changing where the function boundaries are. I also wonder how accumulators are going to change this picture, we could at least get rid of ValuesAsInModelContext with TuringLang/DynamicPPL.jl#908.