Avoiding Mutation vs Allocation

Specific Domains Modelling & Simulations
1

I’ve been writing code for dynamic models in Julia for a few years now, and I keep running into a conflict between allocations and mutation. To improve ODE solve performance, I need to write in-place functions and preallocate, then mutate variables as necessary, but to be compatible with reverse mode AD (through Zygote for BLAS reasons), I need to avoid mutation, which seems to involve a lot of allocations. Right now, I’ve got one model for forward solves and a separate model for gradients, with 1-2 orders of magnitude between them in terms of execution time and memory usage depending what I’m doing. This is very annoying though, and I would prefer to write only one model.

Are there any tricks for writing low-memory usage/high performance code that also avoids mutation? I keep falling back on list comprehensions which are not great for allocations and feel very disorganized, so I feel like there’s got to be a better way, especially for higher-dimensional arrays generated through complicated functions.

Thanks!

2

You could write custom AD rules for performance-critical routines that require mutation.

3

I think Mooncake is adding more and more BLAS rules, have you tried it ? Zygote won’t work for 1.12+ too so replacing it is a good idea.

4

This dilemma ultimately led me to switch from Zygote to Enzyme, which works with mutating code and in general works better with the sort of Julia code that also runs fast.

5

I need to read more of the documentation around custom rules, but it does seem like that is probably the “right” way to fix my problems.

I try Mooncake periodically, and so far I’ve struggled to get it working. I wasn’t aware that Zygote was going to drop support for 1.12+; I’m still using 1.11 since the 1.12 update caused so many issues in the packages I use, so I guess I will need to find a different way of doing things anyway.

I’ll take another look at Enzyme. Historically it’s been pretty unstable for me (likely my fault with bad code), but hopefully things are improved since the last time I tried it.

Seems there’s no silver bullet unfortunately. Thanks for the replies!

6

Not your fault, it is often fiddly to get working and the error messages are challenging, but when it works it is usually fast and correct.

Alas no, but you could try DifferentiationInterface.jl if you want to switch between AD packages quickly.

7

If you have any issues let me know, happy to try to help! For code no one has tried before there’s sometimes a bit of a tail of issues to fix up, but once those are done it should be pretty stable.

Also another option you could use Enzyme from within Reactant. Reactant lets the inner enzyme code completely side step all the complications coming from julia (like type instability), including also doing kernel fusion and device offloading.