Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove all recursion limits #15

Merged
merged 1 commit into from
Oct 16, 2024
Merged

Remove all recursion limits #15

merged 1 commit into from
Oct 16, 2024

Conversation

dennisYatunin
Copy link
Member

This PR removes the default recursion limit from every single function defined in UnrolledUtiliites.

I was previously maintaining a list of which functions defined in UnrolledUtilities can be used recursively, and only disabling the default recursion limits of those functions. However, this list keeps needing to get expanded with new types of functions. It previously included all functions that call themselves and all functions whose arguments can be arbitrary functions (including themselves), but CliMA/ClimaCore.jl#1713 has shown in ClimaCore Buildkite Run 4398 that unrolled_product also needs to have its recursion limit disabled because it is used recursively in ClimaCore, even though it doesn't fall into either of those categories. In order to avoid dealing with any more recursion-related type instabilities due to this package, I am introducing the following generic piece of code at the end of the module:

@static if hasfield(Method, :recursion_relation)
    module_names = names(@__MODULE__; all = true)
    module_values = map(Base.Fix1(getproperty, @__MODULE__), module_names)
    module_functions = filter(Base.Fix2(isa, Function), module_values)
    for f in module_functions, method in methods(f)
        method.recursion_relation = Returns(true)
    end
end

Note to anyone who wants to use this code: Since this disables all recursion limits, it can easily lead to increased compilation time of recursive code. Do not use this unless the compilation time of all functions defined in a module is insignificant compared to their type stability.

Copy link
Member

@charleskawczynski charleskawczynski left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a bit dubious. But I think we can make an exception since that is effectively the essence of this package.

Let's first see how https://buildkite.com/clima/climacore-ci/builds/4399 goes.

@dennisYatunin
Copy link
Member Author

dennisYatunin commented Oct 16, 2024

After some experimentation in ClimaCore, I have found that there are often inference failures reported by JET when a Base.Generator is used in a complex chain of unrolled expressions. Inference failures appear to be particularly common when there are multiple Base.Generators nested within each other, but there are also some failures where this is not the case (e.g., the FieldMatrixSolver unit test inspired by prognostic EDMF). I find this result somewhat dubious, as Base.Generator is a pretty fundamental data structure, and I would expect indexing into it to be roughly as stable as indexing into a Tuple. However, JET claims that this is not the case, so I have also replaced the lazy calls to Iterators.map in unrolled_mapreduce and unrolled_flatmap with eager calls to unrolled_map. This change resulted in a successful ClimaCore build (https://buildkite.com/clima/climacore-ci/builds/4544), and it even fixed a unit test that was previously failing (test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_3.jl). I think this PR should now be good to merge in.

@dennisYatunin dennisYatunin linked an issue Oct 16, 2024 that may be closed by this pull request
@dennisYatunin dennisYatunin merged commit 2d471e1 into main Oct 16, 2024
9 checks passed
@dennisYatunin dennisYatunin deleted the dy/recursion_limit branch October 16, 2024 20:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Inference failures
2 participants