Cache stuff hazardly

Author: ricardoV94

pytensor/link/numba/dispatch/elemwise.py

@@ -264,6 +264,12 @@ def axis_apply_fn(x):
 
 @numba_funcify.register(Elemwise)
 def numba_funcify_Elemwise(op, node, **kwargs):
+    # op = getattr(np, str(op.scalar_op).lower())
+    # @numba_njit
+    # def elemwise_is_numpy(x):
+    #     return op(x)
+    # return elemwise_is_numpy
+
     scalar_inputs = [get_scalar_type(dtype=input.dtype)() for input in node.inputs]
     scalar_node = op.scalar_op.make_node(*scalar_inputs)
 
@@ -276,7 +282,7 @@ def numba_funcify_Elemwise(op, node, **kwargs):
 
     nin = len(node.inputs)
     nout = len(node.outputs)
-    core_op_fn = store_core_outputs(scalar_op_fn, nin=nin, nout=nout)
+    core_op_fn = store_core_outputs(scalar_op_fn, op.scalar_op, nin=nin, nout=nout)
 
     input_bc_patterns = tuple(inp.type.broadcastable for inp in node.inputs)
     output_bc_patterns = tuple(out.type.broadcastable for out in node.outputs)

pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -3,6 +3,7 @@
 import base64
 import pickle
 from collections.abc import Callable, Sequence
+from hashlib import sha256
 from textwrap import indent
 from typing import Any, cast
 
@@ -15,15 +16,19 @@
 from numba.core.types.misc import NoneType
 from numba.np import arrayobj
 
+from pytensor.graph.op import HasInnerGraph
 from pytensor.link.numba.dispatch import basic as numba_basic
-from pytensor.link.utils import compile_function_src
+from pytensor.link.numba.super_utils import compile_function_src2
+from pytensor.scalar import ScalarOp
 
 
 def encode_literals(literals: Sequence) -> str:
     return base64.encodebytes(pickle.dumps(literals)).decode()
 
 
-def store_core_outputs(core_op_fn: Callable, nin: int, nout: int) -> Callable:
+def store_core_outputs(
+    core_op_fn: Callable, core_op: ScalarOp, nin: int, nout: int
+) -> Callable:
     """Create a Numba function that wraps a core function and stores its vectorized outputs.
 
     @njit
@@ -52,9 +57,14 @@ def store_core_outputs({inp_signature}, {out_signature}):
 {indent(store_outputs, " " * 4)}
 """
     global_env = {"core_op_fn": core_op_fn}
-    func = compile_function_src(
-        func_src, "store_core_outputs", {**globals(), **global_env}
-    )
+    # func = compile_function_src(
+    #     func_src, "store_core_outputs", {**globals(), **global_env},
+    # )
+    if isinstance(core_op, HasInnerGraph):
+        key = sha256(core_op.c_code_template.encode()).hexdigest()
+    else:
+        key = str(core_op)
+    func = compile_function_src2(key, func_src, "store_core_outputs", global_env)
     return cast(Callable, numba_basic.numba_njit(func))
 
 

pytensor/link/numba/super_utils.py

@@ -0,0 +1,138 @@
+import importlib
+import os
+import sys
+import tempfile
+from collections.abc import Callable
+from typing import Any
+
+import numba
+import numba.core.caching
+from numba.core.caching import CacheImpl
+
+
+class PyTensorLoader(importlib.abc.SourceLoader):
+    def __init__(self):
+        # Key is "pytensor_generated_" + hash of pytensor graph
+        self._module_sources = {}
+        self._module_globals = {}
+        self._module_locals = {}
+
+    def get_source(self, fullname):
+        if fullname not in self._module_sources:
+            raise ImportError()
+        return self._module_sources[fullname]
+
+    def get_data(self, path):
+        if path not in self._module_sources:
+            raise ImportError()
+        return self._module_sources[path].encode("utf-8")
+
+    def get_filename(self, path):
+        if path not in self._module_sources:
+            raise ImportError()
+        return path
+
+    def add_module(self, name, src, global_env, local_env):
+        self._module_sources[name] = src
+        self._module_globals[name] = global_env
+        self._module_locals[name] = local_env
+
+    def exec_module(self, module):
+        name = module.__name__
+        variables = module.__dict__
+        variables.update(self._module_globals[name])
+        variables.update(self._module_locals[name])
+        code = compile(self._module_sources[name], name, "exec")
+        exec(code, variables)
+
+    def create_module(self, spec):
+        return None
+
+
+pytensor_loader = PyTensorLoader()
+
+
+def load_module(key, src, global_env, local_env):
+    pytensor_loader.add_module(key, src, global_env, local_env)
+    spec = importlib.util.spec_from_loader(key, pytensor_loader)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    sys.modules[key] = module
+    return module
+
+
+class NumbaPyTensorCacheLocator(numba.core.caching._CacheLocator):
+    def __init__(self, py_func, py_file):
+        # print(f"New locator {py_func=}, {py_file=}")
+        self._py_func = py_func
+        self._py_file = py_file
+        self._hash = py_file
+        # src_hash = hash(pytensor_loader._module_sources[self._py_file])
+        # self._hash = hash((src_hash, py_file, pytensor.__version__))
+
+    def ensure_cache_path(self):
+        path = self.get_cache_path()
+        os.makedirs(path, exist_ok=True)
+        # Ensure the directory is writable by trying to write a temporary file
+        tempfile.TemporaryFile(dir=path).close()
+
+    def get_cache_path(self):
+        """
+        Return the directory the function is cached in.
+        """
+        return "~/.cache/pytensor"
+
+    def get_source_stamp(self):
+        """
+        Get a timestamp representing the source code's freshness.
+        Can return any picklable Python object.
+        """
+
+        return self._hash
+
+    def get_disambiguator(self):
+        """
+        Get a string disambiguator for this locator's function.
+        It should allow disambiguating different but similarly-named functions.
+        """
+        return None
+
+    @classmethod
+    def from_function(cls, py_func, py_file):
+        """
+        Create a locator instance for the given function located in the
+        given file.
+        """
+        if py_func.__module__ in pytensor_loader._module_sources:
+            return cls(py_func, py_file)
+
+
+CacheImpl._locator_classes.append(NumbaPyTensorCacheLocator)
+
+
+def compile_function_src2(
+    key: str,
+    src: str,
+    function_name: str,
+    global_env: dict[Any, Any] | None = None,
+    local_env: dict[Any, Any] | None = None,
+) -> Callable:
+    # with NamedTemporaryFile(delete=False) as f:
+    #     filename = f.name
+    #     f.write(src.encode())
+
+    if global_env is None:
+        global_env = {}
+
+    if local_env is None:
+        local_env = {}
+
+    # mod_code = compile(src, filename, mode="exec")
+    # exec(mod_code, global_env, local_env)
+    # print(key, src)
+    module = load_module(key, src, global_env, local_env)
+    res = getattr(module, function_name)
+
+    # res = cast(Callable, res)
+    # res.__source__ = src  # type: ignore
+    return res