Cache more Ops

Author: ricardoV94

pytensor/link/numba/dispatch/basic.py

@@ -3,6 +3,7 @@
 import warnings
 from copy import copy
 from functools import singledispatch
+from hashlib import sha256
 from textwrap import dedent
 
 import numba
@@ -14,7 +15,8 @@
 from numba import types
 from numba.core.errors import NumbaWarning, TypingError
 from numba.cpython.unsafe.tuple import tuple_setitem  # noqa: F401
-from numba.extending import box, overload, register_jitable as _register_jitable
+from numba.extending import box, overload
+from numba.extending import register_jitable as _register_jitable
 
 from pytensor import In, config
 from pytensor.compile import NUMBA
@@ -25,11 +27,9 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.type import Type
 from pytensor.ifelse import IfElse
+from pytensor.link.numba.cache import compile_and_cache_numba_function_src
 from pytensor.link.numba.dispatch.sparse import CSCMatrixType, CSRMatrixType
-from pytensor.link.utils import (
-    compile_function_src,
-    fgraph_to_python,
-)
+from pytensor.link.utils import fgraph_to_python
 from pytensor.scalar.basic import ScalarType
 from pytensor.sparse import SparseTensorType
 from pytensor.tensor.basic import Nonzero
@@ -558,7 +558,13 @@ def specify_shape(x, {create_arg_string(shape_input_names)}):
         """
     )
 
-    specify_shape = compile_function_src(func, "specify_shape", globals())
+    specify_shape = compile_and_cache_numba_function_src(
+        func,
+        "specify_shape",
+        globals(),
+        # use the sha256 of func_conditions in the key to avoid long keys
+        key=f"SpecifyShape_{sha256(';'.join(func_conditions).encode()).hexdigest()}",
+    )
     return numba_njit(specify_shape)
 
 

pytensor/link/numba/dispatch/elemwise.py

@@ -1,4 +1,5 @@
 from functools import singledispatch
+from hashlib import sha256
 from textwrap import dedent, indent
 
 import numba
@@ -276,7 +277,7 @@ def numba_funcify_Elemwise(op, node, **kwargs):
     )
 
     # TODO: Proper key
-    key = "_".join(
+    core_op_key = "_".join(
         map(
             str,
             (
@@ -287,7 +288,10 @@ def numba_funcify_Elemwise(op, node, **kwargs):
             ),
         )
     )
-    core_op_fn = store_core_outputs(scalar_op_fn, nin=nin, nout=nout, core_op_key=key)
+    core_op_key = sha256(core_op_key.encode()).hexdigest()
+    core_op_fn = store_core_outputs(
+        scalar_op_fn, nin=nin, nout=nout, core_op_key=core_op_key
+    )
 
     input_bc_patterns = tuple(inp.type.broadcastable for inp in node.inputs)
     output_bc_patterns = tuple(out.type.broadcastable for out in node.outputs)
@@ -348,11 +352,24 @@ def elemwise(*inputs):
     def ov_elemwise(*inputs):
         return elemwise_wrapper
 
+    # TODO: Also input dtypes in key
+    elemwise_key = "_".join(
+        map(
+            str,
+            (
+                "Elemwise",
+                core_op_key,
+                input_bc_patterns,
+                inplace_pattern,
+            ),
+        )
+    )
+    elemwise_key = sha256(elemwise_key.encode()).hexdigest()
     f = compile_and_cache_numba_function_src(
         "def f(*inputs): return elemwise(*inputs)",
         "f",
         {**globals(), **{"elemwise": elemwise}},
-        key=f"Elemwise_{key}",
+        key=elemwise_key,
     )
 
     return numba_njit(f)

pytensor/link/numba/dispatch/subtensor.py

@@ -1,9 +1,10 @@
 import numpy as np
 
 from pytensor.graph import Type
+from pytensor.link.numba.cache import compile_and_cache_numba_function_src
 from pytensor.link.numba.dispatch import numba_funcify
 from pytensor.link.numba.dispatch.basic import generate_fallback_impl, numba_njit
-from pytensor.link.utils import compile_function_src, unique_name_generator
+from pytensor.link.utils import unique_name_generator
 from pytensor.tensor import TensorType
 from pytensor.tensor.rewriting.subtensor import is_full_slice
 from pytensor.tensor.subtensor import (
@@ -95,10 +96,11 @@ def {function_name}({", ".join(input_names)}):
     return np.asarray(z)
     """
 
-    func = compile_function_src(
+    func = compile_and_cache_numba_function_src(
         subtensor_def_src,
         function_name=function_name,
         global_env=globals() | {"np": np},
+        key=f"{function_name}({', '.join(str(i.type) for i in node.inputs)})",
     )
     return numba_njit(func, boundscheck=True)
 

pytensor/link/numba/dispatch/tensor_basic.py

@@ -2,9 +2,10 @@
 
 import numpy as np
 
+from pytensor.link.numba.cache import compile_and_cache_numba_function_src
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import create_tuple_string, numba_funcify
-from pytensor.link.utils import compile_function_src, unique_name_generator
+from pytensor.link.utils import unique_name_generator
 from pytensor.tensor.basic import (
     Alloc,
     AllocEmpty,
@@ -49,8 +50,11 @@ def allocempty({", ".join(shape_var_names)}):
     return np.empty(scalar_shape, dtype)
     """
 
-    alloc_fn = compile_function_src(
-        alloc_def_src, "allocempty", {**globals(), **global_env}
+    alloc_fn = compile_and_cache_numba_function_src(
+        alloc_def_src,
+        "allocempty",
+        {**globals(), **global_env},
+        key=f"AllocEmpty({op.dtype})",
     )
 
     return numba_basic.numba_njit(alloc_fn)
@@ -93,7 +97,12 @@ def alloc(val, {", ".join(shape_var_names)}):
     res[...] = val
     return res
     """
-    alloc_fn = compile_function_src(alloc_def_src, "alloc", {**globals(), **global_env})
+    alloc_fn = compile_and_cache_numba_function_src(
+        alloc_def_src,
+        "alloc",
+        {**globals(), **global_env},
+        key="Alloc",
+    )
 
     return numba_basic.numba_njit(alloc_fn)
 
@@ -212,8 +221,11 @@ def makevector({", ".join(input_names)}):
     return np.array({create_list_string(input_names)}, dtype=dtype)
     """
 
-    makevector_fn = compile_function_src(
-        makevector_def_src, "makevector", {**globals(), **global_env}
+    makevector_fn = compile_and_cache_numba_function_src(
+        makevector_def_src,
+        "makevector",
+        {**globals(), **global_env},
+        key=f"MakeVector({op.dtype})",
     )
 
     return numba_basic.numba_njit(makevector_fn)