chore: Bump version to 0.0.41 and update module imports for improved functionality

Author: erfanzar

eformer/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "0.0.40"
+__version__ = "0.0.41"
 
 __all__ = (
     "aparser",

eformer/callib/__init__.py

@@ -12,42 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ._cjit import cjit, compile_function, lower_function
-from ._triton_call import (
-    cdiv,
-    get_triton_type,
-    next_power_of_2,
-    normalize_grid,
-    strides_from_shape,
-    triton_call,
-)
-
-try:
-    from ._suppress_triton import (
-        enable_all_triton_output,
-        silence_all_triton_output,
-    )
-except ImportError:
-
-    def silence_all_triton_output():
-        """Fallback function when real suppression isn't available."""
-        return False
-
-    def enable_all_triton_output():
-        """Fallback function when real suppression isn't available."""
-        pass
-
+from ._cjit import cjit, compile_function, load_cached_functions, lower_function
+from ._suppress_triton import disable_cpp_logs
+from ._triton_call import cdiv, get_triton_type, next_power_of_2, normalize_grid, strides_from_shape, triton_call
 
 __all__ = (
     "cdiv",
     "cjit",
     "compile_function",
-    "enable_all_triton_output",
+    "disable_cpp_logs",
     "get_triton_type",
+    "load_cached_functions",
     "lower_function",
     "next_power_of_2",
     "normalize_grid",
-    "silence_all_triton_output",
     "strides_from_shape",
     "triton_call",
 )

eformer/callib/_cjit.py

@@ -19,6 +19,7 @@
 import hashlib
 import os
 import pickle
+import re
 import typing as tp
 import warnings
 from pathlib import Path
@@ -88,6 +89,7 @@ def get_cache_dir() -> Path:
 COMPILE_FUNC_DIR = CACHE_DIR / "compiled_funcs"
 COMPILE_FUNC_DIR.mkdir(parents=True, exist_ok=True)
 COMPILED_FILE_NAME = "compiled.func"
+SIGNATURE_FILE_NAME = "compiled.signature"
 
 COMPILED_CACHE: dict[tuple, tp.Any] = {}
 
@@ -103,38 +105,75 @@ def is_jit_wrapped(fn: tp.Any) -> bool:
     )
 
 
-def cjit(
-    fn: tp.Callable[P, R],
-    static_argnums: tuple[int, ...] | None = None,
-    static_argnames: tuple[str, ...] | None = None,
-    verbose: bool = True,
-):
+def remove_memory_addresses(input_string: str) -> str:
+    """
+    Removes hexadecimal memory address patterns (e.g., 0x736a142445e0) from a string.
+
+    Args:
+        input_string: The string to process.
+
+    Returns:
+        The string with memory address patterns removed.
+    """
+    # Regex to find hexadecimal memory addresses:
+    # 0x     : matches the literal "0x"
+    # [0-9a-fA-F]+ : matches one or more hexadecimal characters (0-9, a-f, A-F)
+
+    return re.sub(r"0x[0-9a-fA-F]+", "", input_string)
+
+
+def cjit(fn: tp.Callable[P, R], verbose: bool = True):
     """
     A decorator that adds caching to a JAX JIT-compiled function.
     The input `fn` must already be a JIT-transformed function (e.g., from @jax.jit).
     """
     assert is_jit_wrapped(fn=fn), "function should be jit wrapped already"
 
+    static_argnums = fn._jit_info.static_argnums
+    static_argnames = fn._jit_info.static_argnames
+
+    if len(static_argnames) == 0:
+        static_argnames = None
+    if len(static_argnums) == 0:
+        static_argnums = None
+
+    state_signature = remove_memory_addresses(
+        str(
+            (
+                fn.__signature__._hash_basis(),
+                fn.__annotations__,
+                fn._fun.__annotations__,
+                fn._fun.__kwdefaults__,
+                fn._fun.__name__,
+                static_argnums,
+                static_argnames,
+            )
+        )
+    )
+
+    static_arg_indices = set(static_argnums) if static_argnums is not None else set()
+
     @functools.wraps(fn)
     def wrapped(*args, **kwargs):
-        static_arg_indices = set(static_argnums) if static_argnums is not None else set()
         dynamic_args = tuple(arg for i, arg in enumerate(args) if i not in static_arg_indices)
         dynamic_kwargs = kwargs.copy()
         if static_argnames is not None:
             for key in static_argnames:
                 dynamic_kwargs.pop(key, None)
         signature = get_signature_tree_util(dynamic_args, dynamic_kwargs)
-        cache_key = (fn, signature)
+        cache_key = (state_signature, signature)
         if cache_key in COMPILED_CACHE:
             compiled_func = COMPILED_CACHE[cache_key]
             return compiled_func(*dynamic_args, **dynamic_kwargs)
         lowered_func: Lowered = fn.lower(*args, **kwargs)
-        compiled_func = smart_compile(
+        compiled_func, signature = smart_compile(
             lowered_func=lowered_func,
             tag="cached-jit",
             verbose=verbose,
+            cache_key=cache_key,
         )
-        COMPILED_CACHE[cache_key] = compiled_func
+
+        COMPILED_CACHE[signature] = compiled_func
 
         return compiled_func(*dynamic_args, **dynamic_kwargs)
 
@@ -256,7 +295,8 @@ def smart_compile(
     lowered_func: Lowered,
     tag: str | None = None,
     verbose: bool = True,
-) -> Compiled:
+    cache_key: tuple[str, tuple] | None = None,
+) -> tuple[Compiled, tuple[str, tuple] | None]:
     """Compile a lowered JAX function with caching.
 
     Args:
@@ -271,16 +311,19 @@ def smart_compile(
     foldername = str(func_hash) if tag is None else f"{tag}-{func_hash}"
     func_dir = COMPILE_FUNC_DIR / foldername
     filepath = func_dir / COMPILED_FILE_NAME
+    signature_filepath = func_dir / SIGNATURE_FILE_NAME
     post_fix = f" (TAG : {tag})" if tag else ""
+    signature = cache_key
     if filepath.exists() and not RECOMPILE_FORCE:
         try:
             (serialized, in_tree, out_tree) = pickle.load(open(filepath, "rb"))
+            signature = pickle.load(open(signature_filepath, "rb"))
             compiled_func = deserialize_and_load(
                 serialized=serialized,
                 in_tree=in_tree,
                 out_tree=out_tree,
             )
-            return compiled_func
+            return compiled_func, signature
         except Exception as e:
             if verbose:
                 warnings.warn(
@@ -293,27 +336,29 @@ def smart_compile(
                 func_dir.mkdir(parents=True, exist_ok=True)
                 try:
                     pickle.dump((serialized, in_tree, out_tree), open(filepath, "wb"))
+                    pickle.dump(cache_key, open(signature_filepath, "wb"))
                 except Exception as e:
                     if verbose:
                         warnings.warn(
                             f"couldn't save compiled function due to {e}" + post_fix,
                             stacklevel=4,
                         )
-            return compiled_func
+            return compiled_func, signature
     else:
         compiled_func: Compiled = lowered_func.compile()
         if ECACHE_COMPILES:
             try:
                 serialized, in_tree, out_tree = serialize(compiled_func)
                 func_dir.mkdir(parents=True, exist_ok=True)
                 pickle.dump((serialized, in_tree, out_tree), open(filepath, "wb"))
+                pickle.dump(cache_key, open(signature_filepath, "wb"))
             except Exception as e:
                 if verbose:
                     warnings.warn(
                         f"couldn't save and serialize compiled function due to {e}" + post_fix,
                         stacklevel=4,
                     )
-        return compiled_func
+        return compiled_func, signature
 
 
 def save_compiled_fn(
@@ -341,10 +386,7 @@ def save_compiled_fn(
         warnings.warn(f"couldn't save compiled function due to {e}", stacklevel=4)
 
 
-def load_compiled_fn(
-    path: str | os.PathLike,
-    prefix: str | None = None,
-):
+def load_compiled_fn(path: str | os.PathLike, prefix: str | None = None):
     """Load a compiled function from disk.
 
     Args:
@@ -524,8 +566,20 @@ def compile_function(
     ).compile()
 
 
+def load_cached_functions() -> None:
+    files = [o for o in os.listdir(COMPILE_FUNC_DIR) if os.path.exists(os.path.join(COMPILE_FUNC_DIR, o))]
+    for file in files:
+        target_compiled_function = Path(COMPILE_FUNC_DIR) / file / COMPILED_FILE_NAME
+        target_signature = Path(COMPILE_FUNC_DIR) / file / SIGNATURE_FILE_NAME
+        (serialized, in_tree, out_tree) = pickle.loads(target_compiled_function.read_bytes())
+        signature = pickle.loads(target_signature.read_bytes())
+        compiled_function = deserialize_and_load(serialized=serialized, in_tree=in_tree, out_tree=out_tree)
+        COMPILED_CACHE[signature] = compiled_function
+
+
 if __name__ == "__main__":
     jnp = jax.numpy
+    load_cached_functions()
 
     @cjit
     @jax.jit

eformer/callib/_suppress_triton.py

@@ -44,32 +44,34 @@ def __init__(self):
 
     def start(self):
         """Start suppressing C-level stderr output."""
-        if not SUPPRESSION_AVAILABLE or self.suppressing:
-            return False
+        if os.environ.get("LET_TRITON_TALK", "true").lower() in ["1", "true", "on"]:
+            if not SUPPRESSION_AVAILABLE or self.suppressing:
+                return False
 
-        try:
-            self.old_stderr_fd = libc.dup(STDERR_FILENO)
-            if sys.platform == "win32":
-                self.null_fd = os.open("NUL", os.O_WRONLY)
-            else:
-                self.null_fd = os.open("/dev/null", os.O_WRONLY)
-            libc.dup2(self.null_fd, STDERR_FILENO)
-            self.suppressing = True
-            return True
-        except (OSError, AttributeError):
-            self.cleanup()
-            return False
+            try:
+                self.old_stderr_fd = libc.dup(STDERR_FILENO)
+                if sys.platform == "win32":
+                    self.null_fd = os.open("NUL", os.O_WRONLY)
+                else:
+                    self.null_fd = os.open("/dev/null", os.O_WRONLY)
+                libc.dup2(self.null_fd, STDERR_FILENO)
+                self.suppressing = True
+                return True
+            except (OSError, AttributeError):
+                self.cleanup()
+                return False
 
     def stop(self):
         """Stop suppressing C-level stderr output and restore original stderr."""
-        if not self.suppressing:
-            return
+        if os.environ.get("LET_TRITON_TALK", "true").lower() in ["1", "true", "on"]:
+            if not self.suppressing:
+                return
 
-        try:
-            if self.old_stderr_fd is not None:
-                libc.dup2(self.old_stderr_fd, STDERR_FILENO)
-        finally:
-            self.cleanup()
+            try:
+                if self.old_stderr_fd is not None:
+                    libc.dup2(self.old_stderr_fd, STDERR_FILENO)
+            finally:
+                self.cleanup()
 
     def cleanup(self):
         """Clean up resources."""
@@ -141,18 +143,21 @@ def enable_all_triton_output():
 
 
 @contextmanager
-def no_triton_logs():
+def disable_cpp_logs(verbose: bool = False):
     """Context manager to temporarily suppress Triton kernel autotuning logs."""
-    started = suppress_triton_logs()
-    try:
-        yield started
-    finally:
-        if started:
-            restore_triton_logs()
+    if verbose:
+        yield False
+    else:
+        started = suppress_triton_logs()
+        try:
+            yield started
+        finally:
+            if started:
+                restore_triton_logs()
 
 
 if __name__ == "__main__":
-    with no_triton_logs():
+    with disable_cpp_logs():
         print("Inside the context manager - C/C++ stderr is suppressed")
 
     suppress_triton_logs()