Author Archives: Great Lakes Consulting

Enhance Model Accuracy Using Universal Differential Equations

By: Great Lakes Consulting

Re-posted from: https://blog.glcs.io/ude_symbolic_regression

This post was written by Steven Whitaker.

In this Julia for Devs post,we will explore the fascinating world of using universal differential equations (UDEs)for model discovery.A UDE is a differential equationwhere part (or all) of itis defined by a universal approximator(typically a neural network).These UDEs play a pivotal rolein uncovering missing aspectsof established models.

UDEs enable seamless integrationof known mechanistic modelswith data-driven modeling approaches.Rather than omit or guess at unknown aspects of a model,a neural network can be embedded into the model,which can then be trained using observed data.

Once trained,the embedded neural networkcan be approximated by a mathematical expression.This allows us to replace the neural networkwith an interpretable formula,improving transparencywhile filling in the missing aspectsof the original modelwith new structure learned from observed data.

This post will illustratekey ideasand design decisions madewhile helping one of our clientsuncover missing dynamics in their model,a complex system representingthe human body under treatmentof a rare disease.

By replacing just one aspect of our client’s modelwith a neural network,we achieved a significant milestone:a 30% reduction in the error of the model predictions.This success story is a testament to the powerof UDEs in enhancing model performance.

We will not share client-specific code,but will provide similar examplesto illustrate our approach.

Here are some of the key ideas we will focus on:

  • Ensure nonnegativity of state variables (as needed).
  • Appropriately initialize neural network weights.
  • Normalize neural network inputs.
  • Make dynamics modular.
  • Fit symbolic expression.

Note that we will discussfully connected neural networks,so some comments belowmay not apply if other architectures are used.

And with that,let’s dive in!

Ensure Nonnegativity of States

When working with a modelrepresenting a real phenomenon,chances are that at least some of the model’s state variablesare naturally nonnegative.For example,the energy of a systemor the concentration of a chemicalcannot be negative quantities.

However,a differential equation solversees only math,not the physical phenomenonthe math represents.So,we have to tell the solverto respect the natural constraintsthat exist.

One way to enforce nonnegativityis to pass the isoutofdomain optionto solve,where isoutofdomain is a functionthat returns truewhen at least one of the provided state variablesviolates the nonnegativity constraint.For example,if all state variables should be nonnegative,the function can be defined as:

isoutofdomain(u, p, t) = any(<(0), u)

As another example,if only the first and third states should be nonnegative:

isoutofdomain(u, p, t) = u[1] < 0 || u[3] < 0

For more informationand other toolsfor tackling nonnegativity,check out the SciML docs.

The ideaof telling the solverabout the nonnegativity constraintsmay seem obvious.Still, it has important implicationsfor the initialization of neural networksembedded into the model.Care should be takento ensure the initial network weights(before training)do not cause states to become negative.See the next section below for more details.

Appropriately Initialize Neural Network Weights

To learn missing dynamics,the neural network in a UDEneeds to be trained.

Training the neural network can occuronly if the UDE can be solvedusing the initial network weights.Otherwise,the gradient of the loss,which is used by optimization algorithmsto train the neural network,cannot be computed.

As a result,it is necessaryto choose a good initializationfor the neural network weights.Here are some approaches for initializationand how well they work:

  • Random initialization:Typically,neural network weightsare initialized randomly.However,this results in random UDE dynamics,so whether a state variablestays positive(or satisfies any other constraints)is left to chance.And if constraints are violated,but we try to enforce themvia, e.g., isoutofdomain,the solver will be unable to solve the UDE,preventing computation of the loss gradient.So,random initialization might work,but it might not.
  • Zero initialization:Don’t do this.Network weights won’t ever update during trainingbecause the gradients with respect to those weightswill always be zero.
  • Constant initialization:Don’t do this, either.Initializing all the weightsto the same, non-zero valuewill result in some amount of learning,but it severely limitsthe expressivity of the network.
  • Pre-training:If the neural network in the UDEis used to replace an expressionthat is believed to be incorrect(but otherwise gives a usable model),one way to initialize the network weightsfor UDE trainingis to initialize the weights randomlybut then train the network directly(i.e., not the UDE as a whole)to fit the expression the network is to replace.This works because no constraints are involvedduring pre-training(so randomly initializing the weights is okay),but when training the UDE,the network weights have been setto avoid running into issues with constraints.
  • Smoothed Collocation:Smoothed collocation is another approachfor pre-training the neural networkwithout involving constraints.See the SciML docs for more info.

For our client,randomly initialized weightsfailed to give a usable gradientfor training,so we decided on pre-trainingbecause they had a reasonable guessthat we could fit the initial network weights to.

Normalize Neural Network Inputs

One key to ensuring UDE training proceeds nicelyis to normalize the inputsto the neural networkto ensure all inputsare roughly of the same magnitude.

Normalizing is essentialbecause the gradient of the training loss functionis proportional to the network inputs.If one input tends to be 1000 times largerthan another,the gradient will also tend to be that much larger,dominating the learning algorithm.

Effect of input scale imbalance on training

To illustrate how to normalize,suppose we have the following neural network(created using Lux.jl):

nn = Chain(    Dense(2, 5, tanh),    Dense(5, 5, tanh),    Dense(5, 1),)

Then we can add a function to the Chainthat performs the normalization:

nn_normalized = Chain(    x -> x ./ normalization,    # Other layers as before)

In this example,normalization is a Vectorcontaining the valuesto scale each input variable.And of course,we’re not limited to scaling the inputs;we can implement whatever function we wantto process the inputs(such as z-score standardization).

The neural network in our client’s UDEdid not learn without normalization.After normalizing the inputs,we got the neural networkto learn meaningful dynamics.

Make Dynamics Modular

While not strictly necessary for UDEs,making the dynamics modularvastly simplifies comparingbetween different approaches.

For our client,we wanted to comparethree versions of the dynamics:the original dynamics,the UDE(where part of the original dynamicswas replaced with a neural network),and the updated dynamics(where the trained neural networkwas replaced with a symbolic expressionfit to the neural network).

Rather than duplicate the dynamics function three timeswith each slight variation,we allowed the variable part of the dynamicsto be passed in as an optionto the dynamics function.To illustrate:

function f!(du, u, p, t; g = original_g)    # ... some code ...    # `a` and `b` are values computed above    # or pulled from `u`, etc.    val = g(a, b, p)    # `val` is used to update `du` in some way.    # ... some code ...end

However,ODEProblems don’t allow passing optionsto the dynamics function.But that’s not a problem;we can create a closureto set the optionbefore handing the function to an ODEProblem.For example,if we want g to use a Lux.jl neural network:

nn = Chain(...)(p_nn, st) = Lux.setup(...)# Be sure to include `p_nn` in the `ODEProblem` parameters `p`.g_nn = (a, b, p) -> nn([a, b], p.p_nn, st)[1][1]f_nn! = (du, u, p, t) -> f!(du, u, p, t; g = g_nn)

Setting up our client’s code like thisreduced code duplicationwhile allowing easy comparisonsacross different versions of the dynamics.

Fit Symbolic Expression

One of the key stepsto learning interpretable dynamicsfrom a UDEis to fit a symbolic expressionto the trained neural network.Doing so allows us to replacethe black box neural networkwith a mathematical expressionwe can understand.

The first stepis to collect the data to usefor the fitting.If the neural network takes two inputs,we will need a Matrix of values:

X = [    a1 a2 ... aN;    b1 b2 ... bN;]

The (ai, bi) pairsshould cover the expected range of input values.One way to determine this rangeis to solve the UDEand save off the valuesthat end up going into the neural network.

To get the target values,evaluate the neural networkwith the X we just created:

y = nn(X, p_nn, st)[1] |> vec

(Here we assume the neural networkoutputs a single numberper input pair.)

Now we can figure outa mathematical expressionrelating X to y.One way to do thisis to use SymbolicRegression.jl.

After obtaining an Expression object eq(called tree in the README for SymbolicRegression.jl),we can use it in our dynamics functionas described in the previous section:

# `p` needs to be an input even though it's not used.g_eq = (a, b, p) -> eq([a; b;;])[1]f_eq! = (du, u, p, t) -> f!(du, u, p, t; g = g_eq)

Using SymbolicRegression.jlto fit the trained UDE neural network,we could provide our clientwith a precise mathematical expressionrepresenting the dynamics that were missingfrom their original differential equation.

Summary

In this post,we discussed vital aspectsof implementing UDEsto learn missing dynamics.We learned some tipsfor helping UDE training,including appropriately initializing network weightsand normalizing network inputs.We also saw the benefitsof code modularityand how that can easily allow usto insert a learned mathematical expressioninto the dynamics.For our client,we were able to use these ideasto train a UDEto achieve 30% less errorin the model predictions.

Do you want to leverage UDEsto improve your models’ predictive capabilities?Contact us, and we can help you out!

Additional Links

]]>

Optimize Flight Simulations with Improved Type-Stability

By: Great Lakes Consulting

Re-posted from: https://blog.glcs.io/sim-performance-type-stability

This post was written by Steven Whitaker.

In this Julia for Devs post,we will discuss improving the performanceof simulation code written in Juliaby eliminating sources of type-instabilities.

We wrote another postdetailing what type-stability isand how type-instabilities can degrade performance.We also showed how SnoopCompile.jl and Cthulhu.jlcan be used to pinpoint causes of type-instability.

This post will cover some of the type-instabilitieswe helped one of our clients overcome.

Our client is a technology innovator.They are building a first-of-its-kind logistics systemfocused on autonomous electric deliveryto reduce traffic and air pollution.Their aim is to provideefficient delivery services for life-saving productsto people in both urban and rural areas.Julia is helping them power criticalGuidance, Navigation, and Control (GNC) systems.

With this client, we:

  • Eliminated slowdown-related failureson the most important simulation scenario.
  • Decreased the compilation time of the scenario by 30%.
  • Improved the slowest time stepsfrom 300 ms to 10 ms (30x speedup),enabling 2x real-time performance.

We will not share client-specific code,but will provide similar examplesto illustrate root-cause issuesand suggested resolutions.

Here are the root-cause issues and resolutions we will focus on:

  • Help type-inference by unrolling recursion.
  • Standardize the output of different branches.
  • Avoid loops overTuples.
    • Use SnoopCompile.jl to reveal dynamic dispatches.
    • Investigate functionswith Cthulhu.jl.
  • Avoid dictionaries that map to functions.

Let’s dive in!

Help Type-Inference by Unrolling Recursion

One of the interesting problems we sawwas that there was a part of the client’s codethat SnoopCompile.jl reportedwas resulting in calls to inference,but when we inspected the code with Cthulhu.jlthe code looked perfectly type-stable.

This code consisted of a set of functionsthat recursively called each other,traversing the model treeto grab data from all the submodels.

As it turns out,recursion can pose difficulties for Julia’s type-inference.Basically,if type-inference detects recursionbut cannot prove it terminates(based only on the types of inputs—rememberthat type-inference occurs before runtime),inference gives up,resulting in code that runs like it is type-unstable.(See this Discourse post and comments and links thereinfor more information.)

The solution we implementedwas to use a @generated functionto unroll the recursion at compile time,resulting in a flat implementationthat could be correctly inferred.

Here’s an example that illustratesthe essence of the recursive code:

# Grab all the data in the entire model tree beginning at `model`.get_data(model::NamedTuple) = (; data = model.data, submodels = get_submodel_data(model.submodels))# This generated function is necessary for type-stability.# It calls `get_data` on each of the fields of `submodels`# and returns a `NamedTuple` of the results.# (This is not the generated function implemented in the solution.)@generated function get_submodel_data(submodels::NamedTuple)    assignments = map(fieldnames(submodels)) do field        Expr(:kw, field, :(get_data(submodels.$field)))    end    return Expr(:tuple, Expr(:parameters, assignments...))endget_submodel_data(::@NamedTuple{}) = (;)

Note that in this exampleget_data calls get_submodel_data,which in turn calls get_dataon the submodels.

Here’s the code for unrolling the recursion:

function _gen_get_data(T, path)    subT = _typeof_field(T, :submodels)    subpath = :($path.submodels)    return quote        (;            data = $path.data,            submodels = $(_gen_get_submodel_data(subT, subpath)),        )    endend# This function determines the type of `model.field`# given just the type of `model` (so we can't just call `typeof(model.field)`).# This function is necessary because we need to unroll the recursion# using a generated function, which means we have to work in the type domain# (because the generated function is generated before runtime).function _typeof_field(::Type{NamedTuple{names, T}}, field::Symbol) where {names, T}    i = findfirst(n -> n === field, names)    return T.parameters[i]endfunction _gen_get_submodel_data(::Type{@NamedTuple{}}, subpath)    return quote        (;)    endendfunction _gen_get_submodel_data(subT, subpath)    assignments = map(fieldnames(subT)) do field        T = _typeof_field(subT, field)        path = :($subpath.$field)        Expr(:kw, field, _gen_get_data(T, path))    end    return Expr(:tuple, Expr(:parameters, assignments...))end@generated get_data_generated(model::NamedTuple) = _gen_get_data(model, :model)

Unfortunately,this example doesn’t reproduce the issue our client had,but it does show how to use a @generated functionto unroll recursion.Note that there is still recursion:_gen_get_data and _gen_get_submodel_data call each other.The key, though, is that this recursion happens before inference,which means that when get_data_generated is inferred,the recursion has already taken place,resulting in unrolled codewithout any recursionthat might cause inference issues.

When we implemented this solution for our client,we saw the total memory utilization of their simulationdecrease by ~35%.This was enough to allow them to disable garbage collectionduring the simulation,speeding it upto faster than real-time!And this was the first timethis simulation had run faster than real-time!

Standardize Output of Different Branches

The client had different parts of their modelupdate at different frequencies.As a result,at any particular time steponly a subset of all the submodelsactually needed to update.Here’s an example of what this might look like:

function get_output(model, t)    if should_update(model, t)        out = update_model(model, t)    else        out = get_previous_output(model)    end    return outend

Unfortunately,update_model and get_previous_outputreturned values of different types,resulting in type-instability:the output type of get_outputdepended on the runtime result of should_update.

Furthermore,this function was called at every time pointon every submodel (and every sub-submodel, etc.),so the type-instability in this functionaffected the whole simulation.

The issue was that update_modeltypically returned the minimal subset of informationactually needed for the specific model,whereas get_previous_output was genericand returned a wider set of information.For example,maybe update_model would return a NamedTuplelike (x = [1, 2], xdot = [0, 0]),while get_previous_output would returnsomething like (x = [1, 2], xdot = [0, 0], p = nothing, stop_sim = false).

To fix this issue,rather than manually updating the return valuesof all the methods of update_modelfor all the submodels in the system,we created a function standardize_outputthat took whatever NamedTuple returned by update_modeland added the missing fieldsthat get_previous_output included.Then,the only change needed in get_outputwas to call standardize_output:

out = update_model(model, t) |> standardize_output

The result of making this changewas a 30% decrease in compilation timefor their simulation!

Avoid Loops over Tuples

The client stored submodelsof a parent modelas a Tuple or NamedTuple.This makes sense for type-stabilitybecause each submodel was of a unique type,so storing them in this waypreserved the type informationwhen accessing the submodels.In contrast,storing the submodels as a Vector{Any}would lose the type informationof the submodels.

However,type-stability problems arisewhen looping over Tuplesof different types of objects.The problem is that the compiler needsto compile code for the body of the loop,but the body of the loop needsto be able to handleall types included in the Tuple.As a result,the compiler must resort to dynamic dispatchin the loop body(but see the note on union-splitting further below).

President waiting for receptionist to look up his office

Here’s an example of the issue:

module TupleLoopfunction tupleloop(t::Tuple)    for val in t        do_something(val)    endenddo_something(val::Number) = val + 1do_something(val::String) = val * "!"do_something(val::Vector{T}) where {T} = isempty(val) ? zero(T) : val[1]do_something(val::Dict{String,Int}) = get(val, "hello", 0)end

Using SnoopCompile.jl reveals dynamic dispatches to do_something:

julia> using SnoopCompileCorejulia> tinf = @snoop_inference TupleLoop.tupleloop((1, 2.0, "hi", [10.0], Dict{String,Int}()));julia> using SnoopCompilejulia> tinfInferenceTimingNode: 0.019444/0.020361 on Core.Compiler.Timings.ROOT() with 4 direct childrenjulia> itrigs = inference_triggers(tinf); mtrigs = accumulate_by_source(Method, itrigs)2-element Vector{SnoopCompile.TaggedTriggers{Method}}: eval_user_input(ast, backend::REPL.REPLBackend, mod::Module) @ REPL ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:247 (1 callees from 1 callers) tupleloop(t::Tuple) @ Main.TupleLoop /path/to/TupleLoop.jl:3 (3 callees from 1 callers)

Looking at tupleloop with Cthulhu.jl:

julia> using Cthulhujulia> ascend(mtrigs[2].itrigs[1])Choose a call for analysis (q to quit):     do_something(::String) >     tupleloop(::Tuple{Int64, Float64, String, Vector{Float64}, Dict{String, Int64}}) at /path/to/TupleLoop.jltupleloop(t::Tuple) @ Main.TupleLoop /path/to/TupleLoop.jl:3 3 function tupleloop(t::Tuple{Int64, Float64, String, Vector{Float64}, Dict{String, Int64}}::Tuple)::Core.Const(nothing) 5 5     for val::Any in t::Tuple{Int64, Float64, String, Vector{Float64}, Dict{String, Int64}} 6         do_something(val::Any) 7     end 9 9 end

And we see the problem!Even though the Tuple is inferred,the loop variable val is inferred as Any,which means that calling do_something(val)must be a dynamic dispatch.

Note that in some casesJulia can perform union-splitting automaticallyto remove the dynamic dispatchcaused by this type-instability.In this example,union-splitting occurs when the Tuplecontains 4 (by default) or fewer unique types.However,it’s not a general solution.

One way to remove the dynamic dispatchwithout relying on union-splittingis to eliminate the loop:

do_something(t[1])do_something(t[2])

But we can quickly seethat writing this codeis not at all generic;we have to hard-codethe number of calls to do_something,which means the code will only workwith Tuples of a particular length.Fortunately,there’s a way around this issue.We can write a @generated functionto have the compiler unroll the loopfor us in a generic way:

@generated function tupleloop_generated(t::Tuple)    body = [:(do_something(t[$i])) for i in fieldnames(t)]    return quote        $(body...)        return nothing    endend

(Note that this code would also workif we specified t::NamedTuplein the method signature.)

Due to the way @generated functions work,SnoopCompile.jl still detects dynamic dispatches,but note that tupleloop_generateddoes not have any dynamic dispatches reported:

julia> using SnoopCompileCorejulia> tinf = @snoop_inference TupleLoop.tupleloop_generated((1, 2.0, "hi", [10.0], Dict{String,Int}()));julia> using SnoopCompilejulia> tinfInferenceTimingNode: 0.022208/0.050369 on Core.Compiler.Timings.ROOT() with 5 direct childrenjulia> itrigs = inference_triggers(tinf); mtrigs = accumulate_by_source(Method, itrigs)3-element Vector{SnoopCompile.TaggedTriggers{Method}}: (g::Core.GeneratedFunctionStub)(world::UInt64, source::LineNumberNode, args...) @ Core boot.jl:705 (1 callees from 1 callers) eval_user_input(ast, backend::REPL.REPLBackend, mod::Module) @ REPL ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:247 (1 callees from 1 callers) var"#s1#1"(::Any, t) @ Main.TupleLoop none:0 (3 callees from 1 callers)

And we can verify with Cthulhu.jlthat there are no more dynamic dispatches in tupleloop_generated:

julia> using Cthulhujulia> ascend(mtrigs[2].itrigs[1])Choose a call for analysis (q to quit): >   tupleloop_generated(::Tuple{Int64, Float64, String, Vector{Float64}, Dict{String, Int64}})       eval at ./boot.jl:430 => eval_user_input(::Any, ::REPL.REPLBackend, ::Module) at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-1tupleloop_generated(t::Tuple) @ Main.TupleLoop /path/to/TupleLoop.jl:11Variables  #self#::Core.Const(Main.TupleLoop.tupleloop_generated)  t::Tuple{Int64, Float64, String, Vector{Float64}, Dict{String, Int64}}Body::Core.Const(nothing)    @ /path/to/TupleLoop.jl:11 within `tupleloop_generated`    @ /path/to/TupleLoop.jl:15 within `macro expansion`1  %1  = Main.TupleLoop.do_something::Core.Const(Main.TupleLoop.do_something)   %2  = Base.getindex(t, 1)::Int64         (%1)(%2)   %4  = Main.TupleLoop.do_something::Core.Const(Main.TupleLoop.do_something)   %5  = Base.getindex(t, 2)::Float64         (%4)(%5)   %7  = Main.TupleLoop.do_something::Core.Const(Main.TupleLoop.do_something)   %8  = Base.getindex(t, 3)::String         (%7)(%8)   %10 = Main.TupleLoop.do_something::Core.Const(Main.TupleLoop.do_something)   %11 = Base.getindex(t, 4)::Vector{Float64}         (%10)(%11)   %13 = Main.TupleLoop.do_something::Core.Const(Main.TupleLoop.do_something)   %14 = Base.getindex(t, 5)::Dict{String, Int64}         (%13)(%14)   @ /path/to/TupleLoop.jl:16 within `macro expansion`       return Main.TupleLoop.nothing

Here we have to examine the so-called “Typed” code(since the source code was generated via metaprogramming),but we see that there is no loop in this code.As a result,each call to do_somethingis a static dispatchwith a concretely inferred input.Hooray!

Avoid Dictionaries that Map to Functions

The client registered functionsfor updating their simulation visualizationvia a dictionary that mapped from a String keyto the appropriate update function.

Sometimes it can be convenientto have a dictionary of functions,for example:

d = Dict{String, Function}(    "sum" => sum,    "norm" => norm,    # etc.)x = [1.0, 2.0, 3.0]d["sum"](x) # Compute the sum of the elements of `x`d["norm"](x) # Compute the norm of `x`

This allows you to write generic codethat can call the appropriate intermediate functionbased on a key supplied by the caller.

You could use multiple dispatch to achieve similar results,but it requires a bit more thoughtto organize the code in such a waythat ensures the caller has access to the types to dispatch on.

As another alternative,you could also have the callerjust pass in the function to call.But again,it takes a bit more effortto organize the codeto make it work.

Unfortunately,using a dictionary in this wayis type-unstable:Julia can’t figure out what functionwill be calleduntil runtime,when the precise dictionary key is known.And since the function is unknown,the type of the result of the function is also unknown.

One partial solutionis to use type annotations:

d[func_key](x)::Float64

Then at least the output of the functioncan be used in a type-stable way.However,this only works if all the functions in the dictionaryreturn values of the same typegiven the same input.

A slightly less stringent alternativeis to explicitly convertthe result to a common type,but this requires conversion to be possible.

Our client updated a dictionaryusing the output of the registered function,so the full solution we implemented for our clientwas to remove the dictionaryand instead have explicit branches in the code.That is,instead of

updates[key] = d[key](updates[key])

we had

if key == "k1"    updates[key] = f1(updates[key]::OUTPUT_TYPE_F1)elseif key == "k2"    updates[key] = f2(updates[key]::OUTPUT_TYPE_F2)# Additional branches as neededend

Note that we needed the type annotationsOUTPUT_TYPE_F1 and OUTPUT_TYPE_F2because updates had an abstractly typed value type.The key that makes this solution workis recognizing that in the first branchupdates[key] is the output of f1from the previous time step in the simulation(and similarly for the other branches).Therefore,in each branch we know what the type of updates[key] is,so we can give the compiler that type information.

Also note that the previously mentioned ideasof using multiple dispatchor just passing in the functions to usedon’t work in this situationwithout removing the updates dictionary(and refactoring the affected code).

Making the above changecompletely removed type-instabilitiesin that part of the client’s code.

Summary

In this post,we explored a few problemsrelated to type-stabilitythat we helped our client resolve.We were able to diagnose issuesusing SnoopCompile.jl and Cthulhu.jland make code improvementsthat enabled our client’smost important simulation scenarioto pass tests for the first time.This was possiblebecause our solutions enabled the scenarioto run faster than real-timeand reduced compilation time by 30%.

Do you have type-instabilities that plague your Julia code?Contact us, and we can help you out!

Additional Links

The cover image backgroundof a person at the start of a racewas found at freepik.com.

The cartoon about looping over Tupleswas generated with AI.

]]>

Type-Stability with SnoopCompile.jl and Cthulhu.jl for High-Performance Julia

By: Great Lakes Consulting

Re-posted from: https://blog.glcs.io/type-stability

This post was written by Steven Whitaker.

The Julia programming languageis a high-level languagethat boasts the ability to achieve C-like speeds.Julia can run fast,despite being a dynamic language,because it is compiledand has smart type-inference.

Type-inference is the processof the Julia compilerreasoning about the types of objects,enabling compilation to create efficient machine codefor the types at hand.

However,Julia code can be written in a waythat prevents type-inference from succeeding—specifically,by writing type-unstable functions.(I’ll explain type-instability later on.)When type-inference fails,Julia has to compile generic machine codethat can handle any type of input,sacrificing the C-like performanceand instead running more like an interpreted languagesuch as Python.

Fortunately,there are tools that Julia developers can useto track down what code causes type-inference to fail.Among the most powerful of these toolsare SnoopCompile.jl and Cthulhu.jl.Using these tools,developers can fix type-inference failuresand restore the C-like performancethey were hoping to achieve.

In this post,we will learn about type-stabilityand how it impacts performance.Then we will see how to useSnoopCompile.jl and Cthulhu.jlto locate and resolve type-instabilities.

Type-Stability

A function is type-stableif the type of the function’s output can be concretely determinedgiven the types of the inputs to the function,without any runtime information.

To illustrate, consider the following function methods:

f(x::Int) = "stable"f(x::Float64) = rand(Bool) ? 1 : 2.0

In this example,if we call f(x) where x is an Int,the compiler can figure outthat the output will be a Stringwithout knowing the value of x,so f(x::Int) is type-stable.In other words,it doesn’t matter whether x is 1, -1, or 176859431;the return value will always be a Stringif x is an Int.

On the other hand,if we call f(x) where x is a Float64,the compiler doesn’t knowwhether the output will be an Int or a Float64because that depends on the result of rand(Bool),which is computed at runtime.Therefore,f(x::Float64) is type-unstable.

Here’s a more subtle example of type-instability:

function g(x)    if x < 0        return 0    else        return x    endend

In this example,g(x) is type-unstablebecause the output will either be an Intor whatever the type of x is,and it all depends on the value of x,which isn’t known at compile time.(Note, however,that g(x) is type-stableif x is an Intbecause then both branches of the if statementreturn the same type of value.)

And sometimes a function that might look type-stablecan be type-unstabledepending on the input.For example:

h(x::Array) = x[1] + 1

In this case,h([1]) is type-stable,but h(Any[1]) is not.Why?Because with h([1]),x is a Vector{Int},so the compiler knowsthat the type of x[1] will be Int.On the other hand,with h(Any[1]),x is a Vector{Any},so the compiler thinksx[1] could be of any type.

To reiterate:a function is type-stableif the compiler can figure outthe concrete type of the return valuegiven only the types of the inputs,without any runtime information.

When Compilation Occurs

Another aspect of type-inferencethat is useful to understandis when compilation (including type-inference) occurs.

In a static language like C,an entire program is compiledbefore any code runs.This is possible because the types of all variablesare known in advance,so machine code specific to those typescan be generated in advance.

In an interpreted language like Python,no code is ever compiledbecause variables are dynamic,meaning their types aren’t really ever knownuntil variables are actually used(i.e., during runtime).

Julia programs can liepretty much anywhere betweenthe extremes of C and Python,and where on that spectrum a program liesdepends on type-stability.

In a just-in-time (JIT) compiled language like Julia,compilation occurs once types are known.

  • If a Julia program is completely type-stable,type-inference can figure out the types of all variablesin the programbefore running any code.As a result,the entire program can be compiledas if it were written in a static language.This is what allows Julia to achieve C-like speeds.
  • If a Julia program is entirely type-unstable,every function has to be compiled individually.In this case,compilation occurs at the momentthe function is calledbecause that’s when the runtime informationof all the input types is finally known.Furthermore,the machine code for a type-unstable functioncannot be efficientbecause it must be able tohandle a wide rangeof potential types.As a result,despite being compiled,the code runs essentially like an interpreted language.

Running a Julia program with type-instabilitiesis like driving down the streetand hitting all the red lights.Julia will compile all the codefor which type-inference succeedsand then start running.But when the program reaches a function callthat could not be inferred,that’s like a car stopping at a red light;Julia stops running the codeto compile the function callnow that it knows the runtime types of the inputs.After the function is compiled,the program can continue execution,like how the car can continue drivingonce the light turns green.

Type-Stability and Performance

As this analogy implies,and as I’ve stated before,type-stability has performance implications.Type-instabilities can cause various performance degradations, including:

  • Dynamic (aka runtime) dispatch.If the compiler knows the input types to a function,the generated machine code can include a callto the specific method determined by those types.But if the compiler doesn’t know those types,the machine code has to include instructionsto perform dynamic dispatch.As a result,rather than jumping directly to the correct method,Julia has to spend runtime CPU cyclesto look up the correct method to call.
  • Increased memory allocations.If the compiler doesn’t know what typea variable will have,it’s impossible to put it in a registeror even allocate stack space for it.As a result,it has to be heap-allocatedand managed by the garbage collector.
  • Suboptimal compiled code.Imagine summing the contents of an array in a loop.If the compiler knows the array contains just Float64s,it can perform optimizations to compute the sumas efficiently as possible,e.g., by using specialized CPU instructions.Such optimizations cannot occurif the compiler doesn’t know what type of datait’s working with.

Here’s an example(inspired by this Stack Overflow answer)that illustrates the impacttype-stability can have on performance:

# Type-unstable because `x` is a non-constant global variable.x = 0f() = [i + x for i = 1:10^6]# Type-stable because `y` is constant and therefore always an `Int`.const y = 0g() = [i + y for i = 1:10^6]using BenchmarkTools@btime f() #  16.868 ms (1998983 allocations: 38.13 MiB)@btime g() # 190.755 s (3 allocations: 7.63 MiB)

Note that the type-unstable versionis two orders of magnitude slower!Also note, however,that this is an extreme examplewhere essentially the entire computationis type-unstable.In practice,some type-instabilities will not impact performance very much.Type-stability mainly matters in “hot loops”,i.e., in parts of the codethat run very frequentlyand contribute to a significant portionof the program’s overall run time.

Detecting Type-Instabilities with SnoopCompile.jl

Now the question is,how do we know if or where our code is type-unstable?One excellent toolfor discovering where type-instabilities occur in codeis SnoopCompile.jl.This package provides functionalityfor reporting how many timesa Julia program needs to stop to compile code.(Remember that a perfectly type-stable programcan compile everything in one go,so every time execution stops for compilationindicates a type-instability was encountered.)

Let’s use an example to illustrate how to use SnoopCompile.jl.First, the code we want to analyze:

module Originalstruct Alg1 endstruct Alg2 endfunction process(alg::String)    if alg == "alg1"        a = Alg1()    elseif alg == "alg2"        a = Alg2()    end    data = get_data(a)    result = _process(a, data)    return resultendget_data(::Alg1) = (1, 1.0, 0x00, 1.0f0, "hi", [0.0], (1, 2.0))function _process(::Alg1, data)    val = data[1]    if val < 0        val = -val    end    result = map(data) do d        process_item(d, val)    end    return resultendprocess_item(d::Int, val) = d + valprocess_item(d::AbstractFloat, val) = d * valprocess_item(d::Unsigned, val) = d - valprocess_item(d::String, val) = d * string(val)process_item(d::Array, val) = d .+ valprocess_item(d::Tuple, val) = d .- valget_data(::Alg2) = rand(5)_process(::Alg2, data) = error("not implemented")end

We’ll use the @snoop_inference macroto analyze this code.Note that this macro should be usedin a fresh Julia session(after loading the code to be analyzed,but before running anything)to get the most accurate analysis results.

julia> using SnoopCompileCorejulia> tinf = @snoop_inference Original.process("alg1");julia> using SnoopCompilejulia> tinfInferenceTimingNode: 0.144601/0.247183 on Core.Compiler.Timings.ROOT() with 8 direct children

You can consult the SnoopCompile.jl docsfor more information about what we just did,but for now,notice that displaying tinf revealed 8 direct children.That means compilation occurred 8 timeswhile running Original.process("alg1").If this function were completely type-stable,@snoop_inference would have reported just 1 direct child,so we know there are type-instabilities somewhere.

Each of the 8 direct childrenis an inference trigger,i.e., calling the specific methodindicated in the inference triggercaused compilation to occur.We can collect the inference triggers:

julia> itrigs = inference_triggers(tinf) Inference triggered to call process(::String) from eval (./boot.jl:430) inlined into REPL.eval_user_input(::Any, ::REPL.REPLBackend, ::Module) (/cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/usr/share/julia/stdlib/v1.11/REPL/src/REPL.jl:261) Inference triggered to call process_item(::Int64, ::Int64) from #1 (./REPL[1]:30) with specialization (::var"#1#2")(::Int64) Inference triggered to call process_item(::Float64, ::Int64) from #1 (./REPL[1]:30) with specialization (::var"#1#2")(::Float64) Inference triggered to call process_item(::UInt8, ::Int64) from #1 (./REPL[1]:30) with specialization (::var"#1#2")(::UInt8) Inference triggered to call process_item(::Float32, ::Int64) from #1 (./REPL[1]:30) with specialization (::var"#1#2")(::Float32) Inference triggered to call process_item(::String, ::Int64) from #1 (./REPL[1]:30) with specialization (::var"#1#2")(::String) Inference triggered to call process_item(::Vector{Float64}, ::Int64) from #1 (./REPL[1]:30) with specialization (::var"#1#2")(::Vector{Float64}) Inference triggered to call process_item(::Tuple{Int64, Float64}, ::Int64) from #1 (./REPL[1]:30) with specialization (::var"#1#2")(::Tuple{Int64, Float64})

The first inference triggercorresponds to compiling the top-level process functionwe called(this is the inference trigger we always expect to see).But then it looks like Julia had to stop runningto compile several different methods of process_item.

Inference triggers tell us that type-instabilities existedwhen calling the given functions,but what we really want to know iswhere these type-instabilities originated.You’ll note that each displayed inference trigger abovealso indicates the calling functionby specifying from <calling function>.(Note that the from #1 in the above exampleindicates process_item was calledfrom an anonymous function.)

We can use accumulate_by_sourceto get an aggregated viewof what functions made calls via dynamic dispatch:

julia> mtrigs = accumulate_by_source(Method, itrigs)2-element Vector{SnoopCompile.TaggedTriggers{Method}}: eval_user_input(ast, backend::REPL.REPLBackend, mod::Module) @ REPL ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:247 (1 callees from 1 callers) (::var"#1#2")(d) @ Main REPL[1]:30 (7 callees from 7 callers)

From this,we can see that the example codereally has only one problematic function:the anonymous function var"#1#2".

Diving in with Cthulhu.jl

Now that we have a rough ideaof where the type-instabilities come from,we can drill down into the codeand pinpoint the precise causeswith Cthulhu.jl.We can use the ascend functionon an inference triggerto start investigating:

julia> using Cthulhujulia> ascend(itrigs[2]) # Skip `itrigs[1]` because that's the top-level compilation that should always occur.

ascend provides a menu that shows process_itemand the anonymous function.Select the anonymous function and press Enter.Here’s a screenshot of the Cthulhu output:

Cthulhu output

Reading the output of Cthulhu.jltakes some time to get used to(especially when it can’t display source code,as in this example),but the main thing to rememberis that red is bad.See the Cthulhu.jl README for more information.

In this example,the source of the type-instabilitywas fairly easy to pinpoint.I annotated the screenshotto indicate from where the type-instability arose,which is this Core.Box thing.These are always bad;they are essentially containersthat can hold values of any type,hence the type-instability that ariseswhen accessing the contents.In this particular case,Core.getfield(#self#, :val) indicatesval is a variablethat was captured by the anonymous function.

Once we determine what caused the type-instability,the solution varies on a case-by-case basis.Some potential solutions may include:

  • Ensure different branches of an if statementreturn data of the same type.
  • Add a type annotation to help out inference.For example, 
    x = Any[1]y = do_something(x[1]::Int)
  • Make sure a container type has a concrete element type.For example, x = Int[], not x = [].
  • Avoid loops over heterogeneous Tuples.
  • Use let blocks to define closures.(See this section of the Julia manual for more details.)

We’ll use this last solution in our example.The anonymous function in questionis defined by the do block in _process.So, let’s fix the issue of the captured variable val:

module Corrected# All other code is the same as in module `Original`.function _process(::Alg1, data)    val = data[1]    if val < 0        val = -val    end    f = let val = val        d -> process_item(d, val)    end    result = map(f, data)    return resultend# All other code is the same as in module `Original`.end

Now let’s see what @snoop_inference says:

julia> using SnoopCompileCorejulia> tinf = @snoop_inference Corrected.process("alg1");julia> using SnoopCompilejulia> tinfInferenceTimingNode: 0.113669/0.183888 on Core.Compiler.Timings.ROOT() with 1 direct children

There’s just one direct child.Hooray, type-stability!

Let’s see how performance compares:

julia> using BenchmarkToolsjulia> @btime Original.process("alg1");  220.506 ns (16 allocations: 496 bytes)julia> @btime Corrected.process("alg1");  51.104 ns (8 allocations: 288 bytes)

Awesome, the improved code is ~4 times faster!

Summary

In this post,we learned about type-stabilityand how type-instabilities affect compilationand runtime performance.We also walked through an examplethat demonstrated how to use SnoopCompile.jl and Cthulhu.jlto pinpoint the sources of type-instability in a program.Even though the example in this postwas a relatively easy fix,the principles discussed apply to more complicated programs as well.And, of course,check out the documentation for SnoopCompile.jl and Cthulhu.jlfor further examples to bolster your understanding.

Do you have type-instabilities that plague your Julia code?Contact us, and we can help you out!

Additional Links

]]>