staticfloat
commented
5 months ago @ChrisRackauckas I don't know how common it is for callbacks or other pieces of code in the SciML universe to have memory requirements that cannot be known beforehand (such as is the case here, where the linearizer is recursive and you need a few u vectors at each level of the recursion). This CachePool technique (especially the one where you have a single thread-safe cache at the top level, and each ensemble solve can then get a thread-unsafe cache pool for its own use, and those can be re-used from solution to solution) seems to work quite well. There's still a little bit more blood to be squeezed from this stone (the allocation profiler is still showing some things in my code, and in IDA, but much, much less) but this takes a large step forward. If this CachePool stuff would be useful elsewhere, feel free to move it somewhere more general.
One implementation note; I've avoided adding a dependency on Sundials, but I have Sundials-specific workarounds for the fact that you need to allocate NVector's instead of Vector{Float64}'s for your u vectors if your solver is IDA. I think a cleaner implementation would use a package extension on Sundials, but I haven't had the chance to put that in yet. If you think I should do that, feel free to hold off on merging this until I get around to it.
ChrisRackauckas
This adds a new type,
LinearizingSavingCallbackCacheand some sub-types to allow for efficient re-use of memory as the callback executes over the course of a solve, as well as re-use of that memory in future solves when operating on a large ensemble simulation.The top-level
LinearizingSavingCallbackCachecreates thread-safe cache pool objects that are then used to acquire thread-unsafe cache pool objects to be used within a single solve. Those thread-unsafe cache pool objects can then be released and acquired anew by the next solve. The thread-unsafe pool objects allow for acquisition of pieces of memory such as temporaryuvectors (the recusrive nature of theLinearizingSavingCallbackmeans that we must allocate unknown numbers of temporaryuvectors) and chunks ofublocks that are then compacted into a single large matrix in the finalize method of the callback. All these pieces of memory are stored within that set of thread-unsafe caches, and these are released back to the top-level thread-safe cache pool, for the next solve to acquire and make use of those pieces of memory in the cache pool.Using these techniques, the solve time of a large ensemble simulation with low per-simulation computation has reduced dramatically. The simulation solves a butterworth 3rd-order filter circuit over a certain timespan, swept across different simulus frequencies and circuit parameters. The parameter sweep results in a 13500-element ensemble simulation, that when run with 8 threads on a M1 Pro takes:
Now, after these caching optimizations, we solve the same ensemble in:
As a side note, the size requirements of the raw linearized solution data itself is
1.04 GB. In general, we expect to allocate somewhere between 2-3x the final output data to account for temporaries and inefficient sharing, so while there is still some more work to be done, this gets us significantly closer to minimal overhead.Checklist