Cache keys for numba Op dispatches

Author: ricardoV94

pytensor/link/numba/dispatch/blockwise.py

@@ -1,23 +1,26 @@
-import sys
+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.cache import compile_numba_function_src
 from pytensor.link.numba.dispatch import basic as numba_basic
-from pytensor.link.numba.dispatch.basic import numba_funcify
+from pytensor.link.numba.dispatch.basic import (
+    numba_funcify_and_cache_key,
+    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
@@ -30,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,
@@ -58,36 +61,56 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
     src += ")"
 
     to_tuple = numba_basic.numba_njit(
-        compile_function_src(
+        compile_numba_function_src(
             src,
             "to_tuple",
             global_env={"to_fixed_tuple": to_fixed_tuple},
-        ),
-        # cache=True leads to a numba.cloudpickle dump failure in Python 3.10
-        # May be fine in Python 3.11, but I didn't test. It was fine in 3.12
-        cache=sys.version_info >= (3, 12),
-    )
-
-    def blockwise_wrapper(*inputs_and_core_shapes):
-        inputs, core_shapes = inputs_and_core_shapes[:nin], inputs_and_core_shapes[nin:]
-        tuple_core_shapes = to_tuple(core_shapes)
-        return _vectorized(
-            core_op_fn,
-            input_bc_patterns,
-            output_bc_patterns,
-            output_dtypes,
-            inplace_pattern,
-            (),  # constant_inputs
-            inputs,
-            tuple_core_shapes,
-            None,  # size
         )
+    )
 
     def blockwise(*inputs_and_core_shapes):
-        raise NotImplementedError("Non-jitted BlockwiseWithCoreShape not implemented")
+        raise NotImplementedError(
+            "Numba implementation of Blockwise cannot be evaluated in Python (non-JIT) mode."
+        )
 
     @overload(blockwise, jit_options=_jit_options)
     def ov_blockwise(*inputs_and_core_shapes):
-        return blockwise_wrapper
+        def impl(*inputs_and_core_shapes):
+            inputs, core_shapes = (
+                inputs_and_core_shapes[:nin],
+                inputs_and_core_shapes[nin:],
+            )
+            tuple_core_shapes = to_tuple(core_shapes)
+            return _vectorized(
+                core_op_fn,
+                input_bc_patterns,
+                output_bc_patterns,
+                output_dtypes,
+                inplace_pattern,
+                (),  # constant_inputs
+                inputs,
+                tuple_core_shapes,
+                None,  # size
+            )
+
+        return impl
 
-    return blockwise
+    if core_op_key is None:
+        # We were told the core op cannot be cached
+        blockwise_key = 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,15 @@
 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_njit,
+    numba_funcify_and_cache_key,
+    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,10 +33,27 @@ def numba_funcify_OpFromGraph(op, node=None, **kwargs):
         accept_inplace=True,
     )
     NUMBA.optimizer(fgraph)
-    return numba_funcify(op.fgraph, squeeze_output=True, **kwargs)
+    fgraph_fn, fgraph_cache_key = numba_funcify_and_cache_key(
+        op.fgraph, squeeze_output=True, **kwargs
+    )
+
+    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 fgraph_fn, ofg_cache_key
 
 
-@numba_funcify.register(TypeCastingOp)
+@register_funcify_default_op_cache_key(TypeCastingOp)
 def numba_funcify_type_casting(op, **kwargs):
     @numba_basic.numba_njit
     def identity(x):
@@ -42,7 +62,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):
 
@@ -59,7 +79,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
 
@@ -88,7 +108,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
@@ -100,4 +120,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