HACK: Implement Numba VM with caching of individual nodes #1604

ricardoV94 · 2025-09-02T15:52:04Z

Disclaimer: this is still 100% on hack status, and I don't understand half of the things I did

When we tried #811 it was obvious that numba compile times were prohibitive.

This PR tries a different approach (still at the hacking stage), of using a mode more like the CVM, where individual nodes are jit_compiled, but the whole graph (i.e., the "VM") is not. This allows reusing pre-compiled/cached nodes across different functions, bringing the compilation cost down.

It requires interfacing with the numba cache locator to direct it to cached objects, which requires defining our own cache keys. Numba usually uses the file line position and contents as the cache key, but this doesn't work for dynamically generated files (at least not if stored in a random temp file) nor really for nested functions like those built for Elemwise, Also some Ops are string-generated, and others are regular python functions with globals which numba can usually cache. All this has to be re-examined.

We are also not calling njit on inner Ops (the store_core_outputs / ScalarOp) of Elemwise, but instead doing register_jittable. This was needed for caching to work, but I also don't have a good mental picture of why yet.
If we need this approach, we must move the control of when does an Op get jitted. Sounds like the right place would be in the numba_funcify_FunctionGraph.

Results:

Second pass over tests/tensor/rewriting/test_basic.py (to allow compiling everything first):
2s with C_VM backend
54s with Numba backend
34s with Numba VM without cache
4s with Numba VM with cache

We're finally approaching the speed of the previous backend (at least for single function compilation + eval). Probably could get it there with more optimizing, but a small slowdown is acceptable.

TODO:

We are still writing python strings to the filesystem to compile them, this is probably not needed as explored in Cache numba stuff #1326 (last commit?)
We have to compile some functions that don't really need so we can cache it, such as with Elemwise. This is related to https://numba.discourse.group/t/caching-redefined-functions/3057 but I don't yet have a clear picture.
Proper cache keys, I just hacked some quick things. Perhaps use the source code of the generated functions?
- Composite key is certainly broken
- Cache whole FunctionGraph, this would avoid recompiling identical graphs in the regular Numba mode, not just NumbaCVM (it's also needed for correct cache of Composite/Blockwise/Scan,OpFromGraph (i.e., anything with inner Ops)).
Figure out what happens with Ops that run with object mode?
Handle functions with pointers / large constants that can't traditionally be cached (not sure what's happening now). Related to cache=True failures with locally defined functions numba/numba#10098
Benchmark slowdown from the "VM" approach in realistic functions. Consider using/adapting CVM to orchestrate the calls to the individuals nodes (would need to use the thunk approach). Right now the VM is the python source code generated by the outermost unjitted FunctionGraph

review-notebook-app · 2025-09-02T15:52:10Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

jorenham

left some comments; hope you don't mind

jorenham · 2025-09-05T11:04:36Z

pytensor/link/numba/dispatch/basic.py

-    kwargs.setdefault("cache", config.numba__cache)
-    kwargs.setdefault("no_cpython_wrapper", True)
-    kwargs.setdefault("no_cfunc_wrapper", True)
+    kwargs.setdefault("cache", True)


Note that the numba currently can't detect changes in other modules, which could lead to outdated caches. I'm guessing that's why they've decided to make it opt-in

Not sure what you mean? So far numba caching has been fine, to the extent that it was actually used (not much). It plays a much bigger role in this PR approach, but we are also pretty-much customizing it's behavior completely.

jorenham · 2025-09-05T11:07:23Z

pytensor/link/numba/linker.py

+        else:
+            from pytensor.link.numba.dispatch.basic import numba_njit

-        jitted_fn = numba_njit(fn, no_cpython_wrapper=False, no_cfunc_wrapper=False)
-        return jitted_fn
+            jitted_fn = numba_njit(fn, no_cpython_wrapper=False, no_cfunc_wrapper=False)
+            return jitted_fn


I suppose there's no need for this else branch 🤷🏻

It's a stylistic choice. Should be enforced by ruff I guess

Feel free to open an issue to add https://docs.astral.sh/ruff/rules/superfluous-else-return/ to our rules.

jorenham · 2025-09-05T11:08:16Z

pytensor/link/utils.py

+        with open(filename, "wb") as f:
+            f.write(src.encode())


Suggested change

with open(filename, "wb") as f:

f.write(src.encode())

filename.write_bytes(src.encode())

pytensor/link/utils.py

ricardoV94 · 2025-09-05T17:11:28Z

left some comments; hope you don't mind

I don't, but it's still too early for that sort of feedback. I'm just thinkering around at this point.

And make that the default backend

ricardoV94 · 2025-09-14T10:58:43Z

pytensor/link/numba/dispatch/scalar.py

    )

-    signature = create_numba_signature(node, force_scalar=True)
+    # signature = create_numba_signature(node, force_scalar=True)


This was causing eager compilation during dispatch, nullifying any caching benefits

ricardoV94 · 2025-09-14T10:59:39Z

pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -239,11 +241,21 @@ def codegen(
            output_core_shapes,
        )

+        core_signature = typingctx.resolve_function_type(


This causes compilation of the inner function, so I moved it to the codegen, otherwise we had to pay the cost, even if we could cache it? Not sure. But even if not, I guess we generally want to lazy compile by default?

jorenham reviewed Sep 5, 2025

View reviewed changes

ricardoV94 force-pushed the numba_cache branch from 037e889 to aff2a9a Compare September 14, 2025 10:12

ricardoV94 added 4 commits September 14, 2025 12:15

Add error message in Numba implementation of SpecifyShape

aefe989

Make Numba the default backend

fb60edb

Implement Numba VM with caching

40e6dad

And make that the default backend

Cache more Ops

978b701

ricardoV94 force-pushed the numba_cache branch from aff2a9a to 978b701 Compare September 14, 2025 10:30

ricardoV94 changed the title ~~Implement Numba VM~~ Implement Numba VM with caching of individual nodes Sep 14, 2025

ricardoV94 changed the title ~~Implement Numba VM with caching of individual nodes~~ HACK: Implement Numba VM with caching of individual nodes Sep 14, 2025

ricardoV94 commented Sep 14, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

HACK: Implement Numba VM with caching of individual nodes #1604

HACK: Implement Numba VM with caching of individual nodes #1604

Uh oh!

ricardoV94 commented Sep 2, 2025 •

edited

Loading

Uh oh!

review-notebook-app bot commented Sep 2, 2025

Uh oh!

jorenham left a comment

Uh oh!

jorenham Sep 5, 2025

Uh oh!

ricardoV94 Sep 14, 2025

Uh oh!

jorenham Sep 5, 2025

Uh oh!

ricardoV94 Sep 14, 2025 •

edited

Loading

Uh oh!

ricardoV94 Sep 14, 2025

Uh oh!

jorenham Sep 5, 2025

Uh oh!

ricardoV94 Sep 14, 2025

Uh oh!

Uh oh!

ricardoV94 commented Sep 5, 2025

Uh oh!

ricardoV94 Sep 14, 2025

Uh oh!

ricardoV94 Sep 14, 2025

Uh oh!

Uh oh!

	with open(filename, "wb") as f:
	f.write(src.encode())
	filename.write_bytes(src.encode())

HACK: Implement Numba VM with caching of individual nodes #1604

Are you sure you want to change the base?

HACK: Implement Numba VM with caching of individual nodes #1604

Uh oh!

Conversation

ricardoV94 commented Sep 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

TODO:

Uh oh!

review-notebook-app bot commented Sep 2, 2025

Uh oh!

jorenham left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Sep 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ricardoV94 commented Sep 5, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ricardoV94 commented Sep 2, 2025 •

edited

Loading

ricardoV94 Sep 14, 2025 •

edited

Loading