Lazy neighborhood filter kernel compilation - #1708
Open
cmdupuis3 wants to merge 3 commits into
Open
Conversation
cmdupuis3
added
bug
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
Benchmarks that have got worse:
|
Sevans711
requested changes
Aug 26, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
Thank you looking into this! Assuming you wanted me to give an initial review now because you requested review, but not 100% sure since it is still marked as draft.
The regression test is clever, and I confirmed it crashes on main. This should be one good way to guard against import uxarray slowdowns in the future!
I left inline comments about the caching. Basically, it feels strange to me to use functools.cache in this situation, and I think it leads to some more-complicated-than-necessary code logic at the kernel call sites. I think it would be easier to maintain/understand quickly in the future if using a single cache which directly contains the cached kernels themselves, instead of caching on a per-function basis. That might also help with introspection/debugging later if you ever want to quickly check which kernels have been compiled so far. For example, there could even be a test which asserts _NUMBA_KERNELS is empty after import uxarray.
| def mean(self, uxda): | ||
| """Mean of each neighborhood.""" | ||
| return self._apply_kernel(uxda, _MEAN_KERNEL, 0.0) | ||
| return self._apply_kernel(uxda, _mean_kernel, 0.0) |
Collaborator
There was a problem hiding this comment.
I'm pretty sure this isn't equivalent anymore; the equivalent syntax if sticking with functools here would be _mean_kernel(), right?
I would have a slight preference for a solution that doesn't use functools.cache, to avoid this confusion. Something like this would be more readable and less likely to cause typo, in my opinion:
_NUMBA_KERNELS = {}
def _numba_kernel(kernel_name):
"""returns numba kernel object, such as _make_kernel(lambda window, _: np.mean(window)).
kernels are cached after initially being created.
Supported kernels are: mean, sum, ...
"""
if kernel_name in _NUMBA_KERNELS:
return _NUMBA_KERNELS[kernel_name]
# else, need to make kernel for the first time and cache it:
if kernel_name=='mean':
kernel = _make_kernel(lambda window, _: np.mean(window))
elif kernel_name=='sum':
kernel = _make_kernel(lambda window, _: np.sum(window))
# ... similar pattern for all supported names
else:
raise ValueError(f'unrecognized kernel_name: {kernel_name!r}')
_NUMBA_KERNELS[kernel_name] = kernel
return kernelSidenote: upon initially reading the code I actually had the feeling that something like _NUMBA_KERNELS probably should actually be attached to the Neighborhood class for now, since these kernels are specifically used by the Neighborhood class and nowhere else. If you think these will always just be specific to the Neighborhood class, I'd recommend something like:
class Neighborhood():
_NUMBA_KERNELS = {}
@classmethod
def _numba_kernel(cls, kernel_name):
# same as above, but use cls._NUMBA_KERNELS instead.
Collaborator
Author
There was a problem hiding this comment.
So, the first commit on this branch took an OO approach that should be equivalent to the current approach, but I thought it was messier than the functools.cache implementation here. I think this approach has some of the same clarity issues, where we need a lot of OO and control flow to get a relatively simple result.
I'll move the kernels inside Neighborhood, and go from there.
| if block.dtype not in (np.float64, np.float32): | ||
| block = block.astype(np.float64) | ||
| return kernel(block, *arrays, param) | ||
| return kernel()(block, *arrays, param) |
Collaborator
There was a problem hiding this comment.
Ah, I see now that you added an extra call here… probably to deal with the issue I noted in my previous comment? I think this is a confusing abstraction though; basically, with this syntax it means that kernel isn't actually a compiled kernel at all, it is a "function factory which returns a compiled kernel function". This feels like it is returning to functional programming abstractions which I recall you mentioning you wanted to avoid.
If you dislike my previous suggestion and really want to keep the functools solution, I would really want to see the calls used above, e.g. self._apply_kernel(uxda, _mean_kernel(), 0.0) instead of here.
(If you really like the code as-is, the minimal change I would want to see would be to add clearer comments / docstrings to clarify for future developers that kernel isn't actually the kernel, but rather a function which returns a kernel function.)
Collaborator
Author
There was a problem hiding this comment.
That would be okay, either way is fine.
To be specific, I was talking about getting away from the functional-style API, rather than necessarily the underlying mechanics. I think considering how numba and @guvectorize are interacting here, a functional approach at some level is unavoidable.
cmdupuis3
marked this pull request as ready for review
August 26, 2026 19:16
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #1706
Overview
This PR should solve an issue introduced with PR #941, wherein the neighborhood filter gufunc-based kernels would compile eagerly, forcing a major slowdown of basic
import uxarrayimports, as well as creating threadpools that would eventually block safe forking that is needed in PR #1700.The idea is to have lazy compilation of gufuncs introduced with neighborhood filters for performant reductions. This prevents the otherwise eager compilation of every one of these kernels upon
import uxarray.Combined, this PR and PR #1700 can bring the total benchmark suit runtime down to about half of the pre-neighborhood filters runtime of ~40min, yielding about 20min without neighborhood filter benchmarks and 26-40ish minutes with them.
On its own, this PR can at least mitigate some of the benchmark performance regression, on the order of about 30-50% or so.
PR Checklist
General
Testing & Benchmarking
Documentation
AI Disclosure
AI Usage: Claude Opus 5