Apply callbacks with a type-stable generated function.

[meson800]

Summary

Currently, as referenced in SciML/DifferentialEquations.jl#971, the old implementation of handle_callbacks! directly calls apply_callback! on continuous_callbacks[idx], which is inherently type-unstable because apply_callback! is specialized on the callback type.

This commit adds a generated function apply_ith_callback! which generates type-stable code to do the same thing, where for each callback tuple type, the generated function unrolls the tuple by checking the callback index against static indices.

Efficiency

As a nice bonus, this generated function seems to often be converted into a switch statement at the LLVM level, e.g. calling apply_ith_callback! on a nine-length callback tuple just branches with a switch:

   switch i64 %4, label %L46 [
    i64 9, label %L3
    i64 8, label %L8
    i64 7, label %L13
    i64 6, label %L18
    i64 5, label %L23
    i64 4, label %L28
    i64 3, label %L33
    i64 2, label %L38
    i64 1, label %L43
  ]

Interestingly, this only happens if I structured the generated function as a nested if/else clause instead of a series of if statements, even though there are returns in each if branch.

Testing

For testing, I added an allocation test which sets up a simple ODE problem, steps the integrator manually to before the first callback, then manipulates integrator state past the first callback point. This way, we can directly call handle_callbacks! and write a test on the allocation count.

I confirm that (at least testing against commit SciML/DiffEqBase.jl@1799fc3, the current master branch tip in DiffEqBase.jl), the new method does not allocate, whereas the old one allocates. This may not be the case until a new release is cut of DiffEqBase.jl, because the old version of find_first_continuous_callback might allocate.

I am still running the entire 37-min test suite to check that no other test fails, but a run earlier today on very similar code didn't turn up any failures.

Formatting

I ran JuliaFormatter as requested, but it wanted to make a lot of other whitespace changes within integrator_utils.jl. I just committed the formatting changes just in the section I changed to not generate a lot of spam whitespace edits attributed to me.

Extra changes

Cthulhu'ing around, I noticed that handle_callbacks! is actually type unstable, as it returns nothing or bool. Looking at the call sites, it seems like nothing is intended, so I added an explicit nothing return. I can change this back if this is not correct.

[ai-maintainer]

AI-Maintainer Review for PR - Apply callbacks with a type-stable generated function

Title and Description 👍

The Title and Description are clear, concise, and helpful

The title and description of the pull request are clear and concise. They effectively communicate the purpose of the changes, which is to apply callbacks with a type-stable generated function. The description provides a detailed summary of the problem being addressed, the approach taken, and the expected efficiency improvements. It also mentions the testing performed and any additional changes made.

Scope of Changes 👍

The changes are narrowly focused

The changes appear to be narrowly focused on resolving a specific issue related to type stability and efficiency in handling callbacks. The author introduces a generated function `apply_ith_callback!` to address the type instability in the original implementation. The changes also include a new test to verify the absence of allocations and some formatting adjustments. While there are additional changes mentioned, such as adding an explicit `nothing` return and running the test suite, these seem to be related to the main issue being addressed.

Testing 👍

Appropriate testing has been performed

The description provides details on how the author tested the changes. It mentions a specific allocation test that sets up a simple ODE problem with callbacks and manually triggers the first callback to check its allocations. The author confirms that the new method does not allocate, while the old one does. Additionally, the author mentions running the entire test suite to ensure that no other tests fail.

Docstrings 👍

All new functions have appropriate docstrings

All newly added functions, classes, or methods have docstrings describing their behavior, arguments, and return values. There are no functions, classes, or methods in the diff that need docstrings.

Suggested Changes

No changes are suggested as the pull request appears to be well-structured and addresses the issue at hand effectively. The author has also taken care to test the changes thoroughly and ensure that they do not introduce any new issues.

Keep up the good work!

Reviewed with AI Maintainer

[ChrisRackauckas]

Great!

Cthulhu'ing around, I noticed that handle_callbacks! is actually type unstable, as it returns nothing or bool. Looking at the call sites, it seems like nothing is intended, so I added an explicit nothing return. I can change this back if this is not correct.

Yeah since the return isn't used that should just get cleaned up in DCE, but it is indeed always best to just clean such things up, thanks.

As a nice bonus, this generated function seems to often be converted into a switch statement at the LLVM level

Oh nice, didn't know it would do that.

I ran JuliaFormatter as requested, but it wanted to make a lot of other whitespace changes within integrator_utils.jl. I just committed the formatting changes just in the section I changed to not generate a lot of spam whitespace edits attributed to me.

The formatter got an update, so some of the formats are out of whack. Don't worry about that.