Cache keys for numba Op dispatches

Author: ricardoV94

pytensor/link/numba/dispatch/blockwise.py

@@ -1,21 +1,26 @@
+from hashlib import sha256
 from typing import cast
 
 from numba.core.extending import overload
 from numba.np.unsafe.ndarray import to_fixed_tuple
 
-from pytensor.link.numba.dispatch.basic import numba_funcify, numba_njit
+from pytensor.link.numba.cache import compile_numba_function_src
+from pytensor.link.numba.dispatch.basic import (
+    numba_funcify_and_cache_key,
+    numba_njit,
+    register_funcify_and_cache_key,
+)
 from pytensor.link.numba.dispatch.vectorize_codegen import (
     _jit_options,
     _vectorized,
     encode_literals,
     store_core_outputs,
 )
-from pytensor.link.utils import compile_function_src
 from pytensor.tensor import TensorVariable, get_vector_length
 from pytensor.tensor.blockwise import Blockwise, BlockwiseWithCoreShape
 
 
-@numba_funcify.register(BlockwiseWithCoreShape)
+@register_funcify_and_cache_key(BlockwiseWithCoreShape)
 def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
     [blockwise_node] = op.fgraph.apply_nodes
     blockwise_op: Blockwise = blockwise_node.op
@@ -28,7 +33,7 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
         cast(tuple[TensorVariable], node.inputs[:nin]),
         propagate_unbatched_core_inputs=True,
     )
-    core_op_fn = numba_funcify(
+    core_op_fn, core_op_key = numba_funcify_and_cache_key(
         core_op,
         node=core_node,
         parent_node=node,
@@ -56,7 +61,7 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
     src += ")"
 
     to_tuple = numba_njit(
-        compile_function_src(
+        compile_numba_function_src(
             src,
             "to_tuple",
             global_env={"to_fixed_tuple": to_fixed_tuple},
@@ -85,9 +90,27 @@ def blockwise(*inputs_and_core_shapes):
     def ov_blockwise(*inputs_and_core_shapes):
         return blockwise_wrapper
 
-    # The outer caller won't create a wrapper for our overloaded function, so we do it here
-    @numba_njit
-    def blockwise_wrapped(*args):
-        return blockwise(*args)
+    if core_op_key is None:
+        # We were told the core op cannot be cached
+        # The outer caller won't create a wrapper for our overloaded function, so we do it here
+        @numba_njit
+        def blockwise_wrapped(*args):
+            return blockwise(*args)
 
-    return blockwise_wrapped
+        return blockwise_wrapped, None
+    else:
+        blockwise_key = "_".join(
+            map(
+                str,
+                (
+                    type(op),
+                    type(blockwise_op),
+                    tuple(blockwise_op.destroy_map.items()),
+                    blockwise_op.signature,
+                    input_bc_patterns,
+                    core_op_key,
+                ),
+            )
+        )
+        blockwise_key = sha256(blockwise_key.encode()).hexdigest()
+        return blockwise, blockwise_key

pytensor/link/numba/dispatch/compile_ops.py

@@ -1,3 +1,5 @@
+from hashlib import sha256
+
 import numpy as np
 
 from pytensor.compile.builders import OpFromGraph
@@ -8,14 +10,16 @@
 from pytensor.ifelse import IfElse
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import (
-    numba_funcify,
+    numba_funcify_and_cache_key,
     numba_njit,
+    register_funcify_and_cache_key,
+    register_funcify_default_op_cache_key,
 )
 from pytensor.raise_op import CheckAndRaise
 from pytensor.tensor.type import TensorType
 
 
-@numba_funcify.register(OpFromGraph)
+@register_funcify_and_cache_key(OpFromGraph)
 def numba_funcify_OpFromGraph(op, node=None, **kwargs):
     _ = kwargs.pop("storage_map", None)
 
@@ -30,7 +34,7 @@ def numba_funcify_OpFromGraph(op, node=None, **kwargs):
         accept_inplace=True,
     )
     NUMBA.optimizer(fgraph)
-    fgraph_fn = numba_njit(numba_funcify(op.fgraph, **kwargs))
+    fgraph_fn, fgraph_cache_key = numba_funcify_and_cache_key(op.fgraph, **kwargs)
 
     if len(op.fgraph.outputs) == 1:
 
@@ -44,10 +48,23 @@ def opfromgraph(*inputs):
         def opfromgraph(*inputs):
             return fgraph_fn(*inputs)
 
-    return opfromgraph
+    if fgraph_cache_key is None:
+        # Can't cache the inner graph
+        ofg_cache_key = None
+    else:
+        ofg_cache_key = sha256(
+            str(
+                (
+                    type(op),
+                    fgraph_cache_key,
+                )
+            ).encode()
+        ).hexdigest()
+
+    return opfromgraph, ofg_cache_key
 
 
-@numba_funcify.register(TypeCastingOp)
+@register_funcify_default_op_cache_key(TypeCastingOp)
 def numba_funcify_type_casting(op, **kwargs):
     @numba_njit
     def identity(x):
@@ -56,7 +73,7 @@ def identity(x):
     return identity
 
 
-@numba_funcify.register(DeepCopyOp)
+@register_funcify_default_op_cache_key(DeepCopyOp)
 def numba_funcify_DeepCopyOp(op, node, **kwargs):
     if isinstance(node.inputs[0].type, TensorType):
 
@@ -73,7 +90,7 @@ def deepcopy(x):
     return deepcopy
 
 
-@numba_funcify.register(IfElse)
+@register_funcify_default_op_cache_key(IfElse)
 def numba_funcify_IfElse(op, **kwargs):
     n_outs = op.n_outs
 
@@ -102,7 +119,7 @@ def ifelse(cond, *args):
     return ifelse
 
 
-@numba_funcify.register(CheckAndRaise)
+@register_funcify_and_cache_key(CheckAndRaise)
 def numba_funcify_CheckAndRaise(op, node, **kwargs):
     error = op.exc_type
     msg = op.msg
@@ -114,4 +131,5 @@ def check_and_raise(x, *conditions):
                 raise error(msg)
         return x
 
-    return check_and_raise
+    cache_key = sha256(str((type(op), error, msg)).encode()).hexdigest()
+    return check_and_raise, cache_key