cuequivariance-ops-cu12 0.4.0__py3-none-manylinux_2_39_aarch64.whl → 0.5.0__py3-none-manylinux_2_39_aarch64.whl

cuequivariance_ops/gated_gemm_triton.py

@@ -0,0 +1,340 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+import enum
+
+import triton
+import triton.language as tl
+
+
+class Precision(enum.Enum):
+    DEFAULT = 0
+    TF32 = 1
+    TF32x3 = 2
+    IEEE = 3
+
+
+@triton.jit
+def fused_sigmoid_gated_dual_gemm_forward_kernel(
+    x1_ptr,
+    x2_ptr,
+    w1_ptr,
+    w2_ptr,
+    mask_ptr,
+    o_ptr,
+    M,
+    N,
+    K,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    PRECISION: tl.constexpr,
+    APPLY_MASK: tl.constexpr,
+    TRANSPOSE_OUT: tl.constexpr,
+    TWO_INPUTS: tl.constexpr,
+):
+    # fully gated GEMM kernel with optional mask at the end
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+
+    start_m = pid_m * TILE_M
+    start_n = pid_n * TILE_N
+
+    offs_xm = start_m + tl.arange(0, TILE_M)
+    offs_wn = start_n + tl.arange(0, TILE_N)
+    offs_k = tl.arange(0, TILE_K)
+
+    x1_ptrs = x1_ptr + (offs_xm[:, None] * K + offs_k[None, :])
+    if TWO_INPUTS:
+        x2_ptrs = x2_ptr + (offs_xm[:, None] * K + offs_k[None, :])
+
+    w_tile_offs = offs_wn[None, :] * K + offs_k[:, None]
+
+    acc_1 = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
+    acc_2 = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
+
+    mask_m = offs_xm < M
+
+    if TWO_INPUTS:
+        for _ in range(0, tl.cdiv(K, TILE_K)):
+            x1 = tl.load(x1_ptrs, mask=mask_m[:, None], other=0.0).to(
+                w1_ptr.type.element_ty
+            )
+            w1_ptrs = w1_ptr + w_tile_offs
+            w1 = tl.load(w1_ptrs)
+            if PRECISION == 0:
+                acc_1 = tl.dot(x1, w1, acc_1)
+            elif PRECISION == 1:
+                acc_1 = tl.dot(x1, w1, acc_1, input_precision="tf32")
+            elif PRECISION == 2:
+                acc_1 = tl.dot(x1, w1, acc_1, input_precision="tf32x3")
+            elif PRECISION == 3:
+                acc_1 = tl.dot(x1, w1, acc_1, input_precision="ieee")
+            else:
+                tl.static_assert(
+                    False,
+                    "PRECISION must be 0 (default), 1 (tf32), 2 (tf32x3) or 3 (ieee)",
+                )
+
+            x1_ptrs += TILE_K
+            w1_ptr += TILE_K
+
+        for _ in range(0, tl.cdiv(K, TILE_K)):
+            x2 = tl.load(x2_ptrs, mask=mask_m[:, None], other=0.0).to(
+                w2_ptr.type.element_ty
+            )
+            w2_ptrs = w2_ptr + w_tile_offs
+            w2 = tl.load(w2_ptrs)
+            if PRECISION == 0:
+                acc_2 = tl.dot(x2, w2, acc_2)
+            elif PRECISION == 1:
+                acc_2 = tl.dot(x2, w2, acc_2, input_precision="tf32")
+            elif PRECISION == 2:
+                acc_2 = tl.dot(x2, w2, acc_2, input_precision="tf32x3")
+            elif PRECISION == 3:
+                acc_2 = tl.dot(x2, w2, acc_2, input_precision="ieee")
+            else:
+                tl.static_assert(
+                    False,
+                    "PRECISION must be 0 (default), 1 (tf32), 2 (tf32x3) or 3 (ieee)",
+                )
+
+            x2_ptrs += TILE_K
+            w2_ptr += TILE_K
+
+    else:
+        for _ in range(0, tl.cdiv(K, TILE_K)):
+            x = tl.load(x1_ptrs, mask=mask_m[:, None], other=0.0).to(
+                w1_ptr.type.element_ty
+            )
+
+            w1_ptrs = w1_ptr + w_tile_offs
+            w1 = tl.load(w1_ptrs)
+            if PRECISION == 0:
+                acc_1 = tl.dot(x, w1, acc_1)
+            elif PRECISION == 1:
+                acc_1 = tl.dot(x, w1, acc_1, input_precision="tf32")
+            elif PRECISION == 2:
+                acc_1 = tl.dot(x, w1, acc_1, input_precision="tf32x3")
+            elif PRECISION == 3:
+                acc_1 = tl.dot(x, w1, acc_1, input_precision="ieee")
+            else:
+                tl.static_assert(
+                    False,
+                    "PRECISION must be 0 (default), 1 (tf32), 2 (tf32x3) or 3 (ieee)",
+                )
+
+            w2_ptrs = w2_ptr + w_tile_offs
+            w2 = tl.load(w2_ptrs)
+            if PRECISION == 0:
+                acc_2 = tl.dot(x, w2, acc_2)
+            elif PRECISION == 1:
+                acc_2 = tl.dot(x, w2, acc_2, input_precision="tf32")
+            elif PRECISION == 2:
+                acc_2 = tl.dot(x, w2, acc_2, input_precision="tf32x3")
+            elif PRECISION == 3:
+                acc_2 = tl.dot(x, w2, acc_2, input_precision="ieee")
+            else:
+                tl.static_assert(
+                    False,
+                    "PRECISION must be 0 (default), 1 (tf32), 2 (tf32x3) or 3 (ieee)",
+                )
+
+            x1_ptrs += TILE_K
+            w1_ptr += TILE_K
+            w2_ptr += TILE_K
+
+    offs_om = pid_m * TILE_M + tl.arange(0, TILE_M)
+    offs_on = pid_n * TILE_N + tl.arange(0, TILE_N)
+
+    acc_1 = 1.0 / (1.0 + tl.exp(-acc_1))
+    acc_gated = acc_1 * acc_2
+
+    if APPLY_MASK:
+        mask = tl.load(mask_ptr + offs_om, mask=mask_m, other=0.0)
+        acc_gated = acc_gated * mask[:, None]
+
+    acc_gated = acc_gated.to(o_ptr.dtype.element_ty)
+
+    if TRANSPOSE_OUT:
+        o_ptrs = o_ptr + offs_on[None, :] * M + offs_om[:, None]
+    else:
+        o_ptrs = o_ptr + offs_om[:, None] * N + offs_on[None, :]
+
+    o_mask = offs_om[:, None] < M
+    tl.store(o_ptrs, acc_gated, mask=o_mask)
+
+
+@triton.jit
+def fused_sigmoid_gated_dual_gemm_backward_pregemm_kernel(
+    grad_xw1_ptr,
+    grad_xw2_ptr,
+    grad_mask_ptr,
+    grad_o_ptr,
+    x1_ptr,
+    x2_ptr,
+    w1_ptr,
+    w2_ptr,
+    mask_ptr,
+    M,
+    N,
+    K,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    PRECISION: tl.constexpr,
+    APPLY_MASK: tl.constexpr,
+    TRANSPOSE_OUT: tl.constexpr,
+    TWO_INPUTS: tl.constexpr,
+):
+    # fully gated GEMM kernel with optional mask at the end
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+
+    start_m = pid_m * TILE_M
+    start_n = pid_n * TILE_N
+
+    offs_xm = start_m + tl.arange(0, TILE_M)
+    offs_wn = start_n + tl.arange(0, TILE_N)
+    offs_k = tl.arange(0, TILE_K)
+
+    x1_ptrs = x1_ptr + (offs_xm[:, None] * K + offs_k[None, :])
+    if TWO_INPUTS:
+        x2_ptrs = x2_ptr + (offs_xm[:, None] * K + offs_k[None, :])
+    w_tile_offs = offs_wn[None, :] * K + offs_k[:, None]
+
+    acc_1 = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
+    acc_2 = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
+
+    mask_m = offs_xm < M
+
+    if TWO_INPUTS:
+        # recompute acc1 and acc2
+        for _ in range(0, tl.cdiv(K, TILE_K)):
+            x1 = tl.load(x1_ptrs, mask=mask_m[:, None], other=0.0).to(
+                w1_ptr.type.element_ty
+            )
+            w1_ptrs = w1_ptr + w_tile_offs
+            w1 = tl.load(w1_ptrs)
+
+            if PRECISION == 0:
+                acc_1 = tl.dot(x1, w1, acc_1)
+            elif PRECISION == 1:
+                acc_1 = tl.dot(x1, w1, acc_1, input_precision="tf32")
+            elif PRECISION == 2:
+                acc_1 = tl.dot(x1, w1, acc_1, input_precision="tf32x3")
+            elif PRECISION == 3:
+                acc_1 = tl.dot(x1, w1, acc_1, input_precision="ieee")
+            else:
+                tl.static_assert(
+                    False,
+                    "PRECISION must be 0 (default), 1 (tf32), 2 (tf32x3) or 3 (ieee)",
+                )
+
+            x1_ptrs += TILE_K
+            w1_ptr += TILE_K
+
+        for _ in range(0, tl.cdiv(K, TILE_K)):
+            x2 = tl.load(x2_ptrs, mask=mask_m[:, None], other=0.0).to(
+                w2_ptr.type.element_ty
+            )
+            w2_ptrs = w2_ptr + w_tile_offs
+            w2 = tl.load(w2_ptrs)
+
+            if PRECISION == 0:
+                acc_2 = tl.dot(x2, w2, acc_2)
+            elif PRECISION == 1:
+                acc_2 = tl.dot(x2, w2, acc_2, input_precision="tf32")
+            elif PRECISION == 2:
+                acc_2 = tl.dot(x2, w2, acc_2, input_precision="tf32x3")
+            elif PRECISION == 3:
+                acc_2 = tl.dot(x2, w2, acc_2, input_precision="ieee")
+            else:
+                tl.static_assert(
+                    False,
+                    "PRECISION must be 0 (default), 1 (tf32), 2 (tf32x3) or 3 (ieee)",
+                )
+
+            x2_ptrs += TILE_K
+            w2_ptr += TILE_K
+
+    else:
+        # recompute acc1 and acc2
+        for _ in range(0, tl.cdiv(K, TILE_K)):
+            x = tl.load(x1_ptrs, mask=mask_m[:, None], other=0.0).to(
+                w1_ptr.type.element_ty
+            )
+
+            w1_ptrs = w1_ptr + w_tile_offs
+            w1 = tl.load(w1_ptrs)
+            if PRECISION == 0:
+                acc_1 = tl.dot(x, w1, acc_1)
+            elif PRECISION == 1:
+                acc_1 = tl.dot(x, w1, acc_1, input_precision="tf32")
+            elif PRECISION == 2:
+                acc_1 = tl.dot(x, w1, acc_1, input_precision="tf32x3")
+            elif PRECISION == 3:
+                acc_1 = tl.dot(x, w1, acc_1, input_precision="ieee")
+            else:
+                tl.static_assert(
+                    False,
+                    "PRECISION must be 0 (default), 1 (tf32), 2 (tf32x3) or 3 (ieee)",
+                )
+
+            w2_ptrs = w2_ptr + w_tile_offs
+            w2 = tl.load(w2_ptrs)
+            if PRECISION == 0:
+                acc_2 = tl.dot(x, w2, acc_2)
+            elif PRECISION == 1:
+                acc_2 = tl.dot(x, w2, acc_2, input_precision="tf32")
+            elif PRECISION == 2:
+                acc_2 = tl.dot(x, w2, acc_2, input_precision="tf32x3")
+            elif PRECISION == 3:
+                acc_2 = tl.dot(x, w2, acc_2, input_precision="ieee")
+            else:
+                tl.static_assert(
+                    False,
+                    "PRECISION must be 0 (default), 1 (tf32), 2 (tf32x3) or 3 (ieee)",
+                )
+
+            x1_ptrs += TILE_K
+            w1_ptr += TILE_K
+            w2_ptr += TILE_K
+
+    offs_om = pid_m * TILE_M + tl.arange(0, TILE_M)
+    offs_on = pid_n * TILE_N + tl.arange(0, TILE_N)
+    if TRANSPOSE_OUT:
+        grad_o_ptrs = grad_o_ptr + offs_on[None, :] * M + offs_om[:, None]
+    else:
+        grad_o_ptrs = grad_o_ptr + offs_om[:, None] * N + offs_on[None, :]
+
+    grad_o = tl.load(grad_o_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32)
+
+    acc_sig = 1.0 / (1.0 + tl.exp(-acc_1))
+
+    if APPLY_MASK:
+        tmp = acc_sig * acc_2
+        grad_mask = grad_o * tmp
+        grad_mask = tl.sum(grad_mask, axis=1)
+        grad_mask_ptrs = grad_mask_ptr + pid_n * M + offs_om
+        tl.store(grad_mask_ptrs, grad_mask.to(grad_mask.type.element_ty), mask=mask_m)
+
+        mask = tl.load(mask_ptr + offs_om, mask=mask_m, other=0.0)
+        grad_o = grad_o * mask[:, None]
+
+    tmp = (1.0 - acc_sig) * acc_sig
+
+    grad_xw1 = grad_o * acc_2 * tmp
+    grad_xw2 = grad_o * acc_sig
+
+    grad_xw1_ptrs = grad_xw1_ptr + offs_om[:, None] * N + offs_on[None, :]
+    grad_xw2_ptrs = grad_xw2_ptr + offs_om[:, None] * N + offs_on[None, :]
+    tl.store(grad_xw1_ptrs, grad_xw1.to(grad_xw1.type.element_ty), mask=mask_m[:, None])
+    tl.store(grad_xw2_ptrs, grad_xw2.to(grad_xw2.type.element_ty), mask=mask_m[:, None])

cuequivariance_ops/tuning_decorator.py

@@ -0,0 +1,328 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+import inspect
+
+# import logging  # Added logging import
+import os
+from enum import Enum
+from typing import Any, Callable
+
+import torch
+from tqdm import tqdm
+
+from .cache_manager import CacheManager
+
+# Configure logging
+# logger = logging.getLogger(__name__)
+
+
+class BenchmarkMode(Enum):
+    FLUSH_CACHE = 0
+    FLUSH_CACHE_PEAK_PROXY = 1
+    ROT_BUFFER = 2
+    ROT_BUFFER_PEAK_PROXY = 3
+
+
+def run_bench(
+    f, input_dict, warmup_iter=250, run_iter=250, bench_mode=BenchmarkMode.ROT_BUFFER
+):
+    initial_output = f(**input_dict)
+
+    if bench_mode in (BenchmarkMode.ROT_BUFFER, BenchmarkMode.ROT_BUFFER_PEAK_PROXY):
+        len_rot = 4
+        inputs_rot = [None] * len_rot
+        for r in range(len_rot):
+            r_inputs = []
+            for key, value in input_dict.items():
+                if isinstance(value, torch.Tensor):
+                    if bench_mode == BenchmarkMode.ROT_BUFFER_PEAK_PROXY:
+                        r_inputs.append(
+                            (
+                                key,
+                                torch.ones_like(
+                                    value, requires_grad=value.requires_grad
+                                ),
+                            )
+                        )
+                    else:
+                        r_inputs.append(
+                            (
+                                key,
+                                torch.randn_like(
+                                    value, requires_grad=value.requires_grad
+                                ),
+                            )
+                        )
+                else:
+                    r_inputs.append((key, value))
+            r_inputs = dict(r_inputs)
+            inputs_rot[r] = r_inputs
+
+        for it in range(warmup_iter):
+            _ = f(**inputs_rot[it % len_rot])
+
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for it in range(run_iter):
+            _ = f(**inputs_rot[it % len_rot])
+        end.record()
+        torch.cuda.synchronize()
+        elapsed = start.elapsed_time(end)
+
+    elif bench_mode in (
+        BenchmarkMode.FLUSH_CACHE,
+        BenchmarkMode.FLUSH_CACHE_PEAK_PROXY,
+    ):
+        cache_filler = torch.empty(1024 * 1024 * 256, dtype=torch.int8, device="cuda")
+
+        if bench_mode == BenchmarkMode.FLUSH_CACHE_PEAK_PROXY:
+            _inputs = {}
+            for key, value in input_dict.items():
+                if isinstance(value, torch.Tensor):
+                    _inputs.append(
+                        (key, torch.ones_like(value, requires_grad=value.requires_grad))
+                    )
+                else:
+                    _inputs.append((key, value))
+            input_dict = _inputs
+
+        for _ in range(warmup_iter):
+            cache_filler.zero_()
+            _ = f(**input_dict)
+
+        starts = [torch.cuda.Event(enable_timing=True) for _ in range(run_iter)]
+        ends = [torch.cuda.Event(enable_timing=True) for _ in range(run_iter)]
+        for i in range(run_iter):
+            cache_filler.zero_()
+            starts[i].record()
+            _ = f(**input_dict)
+            ends[i].record()
+        torch.cuda.synchronize()
+        elapsed = sum(s.elapsed_time(e) for s, e in zip(starts, ends))
+
+    return elapsed / run_iter, initial_output
+
+
+def input_to_key_default(**args) -> str:
+    key_parts = []
+    for arg in args:
+        if isinstance(arg, torch.Tensor):
+            key_parts.append(f"{list(arg.shape)}_{arg.dtype}")
+        elif isinstance(arg, bool):
+            key_parts.append("True" if arg else "False")
+        elif isinstance(arg, str):
+            key_parts.append(arg)
+        else:
+            key_parts.append(str(arg.__class__.__name__))
+
+    return "_".join(key_parts)
+
+
+def combine_all_kwargs(
+    fn: Callable,
+    args: tuple,
+    kwargs: dict[str, Any],
+) -> dict[str, Any]:
+    # Get the function signature
+    sig = inspect.signature(fn)
+    params = sig.parameters
+    param_names = list(params.keys())
+
+    # Create dictionary of default values
+    defaults = {
+        name: param.default
+        for name, param in params.items()
+        if param.default is not inspect.Parameter.empty
+    }
+    # Create dictionary mapping positional args to parameter names
+    args_as_kwargs = {
+        param_names[i]: args[i] for i in range(min(len(args), len(param_names)))
+    }
+    # Create combined dictionary of all parameters
+    all_kwargs = defaults.copy()  # Start with defaults
+    all_kwargs.update(args_as_kwargs)  # Override with positional args
+    all_kwargs.update(kwargs)  # Override with explicit kwargs
+
+    return all_kwargs
+
+
+def autotune_aot(
+    input_generator: Callable,
+    input_to_key: Callable | None,
+    input_configs: list[dict[str, Any]],
+    tunable_configs: list[dict[str, Any]],
+    prune_configs_fn: Callable[
+        [list[dict[str, Any]], dict[str, Any]], list[dict[str, Any]]
+    ]
+    | None,
+    bench_mode=BenchmarkMode.ROT_BUFFER,
+    warmup_iter=25,
+    run_iter=100,
+) -> None:
+    def decorator(fn: Callable) -> Callable:
+        def wrapper(*args, **kwargs):
+            all_kwargs = combine_all_kwargs(fn, args, kwargs)
+            nonlocal input_to_key
+            nonlocal input_configs
+            # Early exit if AOT is disabled
+            if os.environ.get("CUEQ_DISABLE_AOT_TUNING") == "1":
+                if input_to_key is None:
+                    input_to_key = input_to_key_default
+
+                function_key = fn.__name__
+                cache_manager = CacheManager()
+                input_key = input_to_key(**all_kwargs)
+
+                try:
+                    best_cached_config = cache_manager.get(function_key, input_key)
+                    if best_cached_config is not None:
+                        all_kwargs.update(best_cached_config["config"])
+                        return fn(**all_kwargs)
+                    elif os.environ.get("CUEQ_DEFAULT_CONFIG") != "1":
+                        input_configs = [None]
+                        # Continue to rest of function for tuning
+                        pass
+                except (PermissionError, OSError):
+                    pass
+
+                if os.environ.get("CUEQ_DEFAULT_CONFIG") == "1":
+                    return fn(**all_kwargs)
+
+            if input_to_key is None:
+                input_to_key = input_to_key_default
+
+            # Check if the function is already cached
+            function_key = fn.__name__
+            cache_manager = CacheManager()
+            input_key = input_to_key(**all_kwargs)
+
+            try:
+                best_cached_config = cache_manager.get(function_key, input_key)
+            except (PermissionError, OSError):
+                # If there's a permission error, fall back to JIT compilation
+                # logger.warning(f"Permission error accessing cache: {e}. Falling back to JIT compilation.")
+                return fn(**all_kwargs)
+
+            if best_cached_config is not None:
+                # If cached, return the cached result
+                all_kwargs.update(best_cached_config["config"])
+                result = fn(**all_kwargs)
+                return result
+            else:
+                # start autotuning process
+                input_configs = input_configs + [None]
+
+                try:
+                    # Initialize the progress bar
+                    progress_bar = tqdm(
+                        input_configs, desc="Autotuning Progress", unit="config"
+                    )
+
+                    for input_config in progress_bar:
+                        # generate input based on the config
+                        input_data = (
+                            input_generator(**input_config)
+                            if input_config is not None
+                            else all_kwargs
+                        )
+
+                        # Make a copy of all_kwargs to avoid modifying the original
+                        current_kwargs = all_kwargs.copy()
+                        current_kwargs.update(input_data)
+                        current_input_key = input_to_key(**current_kwargs)
+
+                        try:
+                            best_cached_config = cache_manager.get(
+                                function_key, current_input_key
+                            )
+                        except (PermissionError, OSError):
+                            # If there's a permission error, skip this config and continue
+                            # logger.warning(f"Permission error accessing cache for config: {e}. Skipping.")
+                            continue
+
+                        if best_cached_config is not None:
+                            continue
+
+                        # prune the tunable configs based on the all_kwargs
+                        pruned_tunable_configs = (
+                            prune_configs_fn(tunable_configs, **all_kwargs)
+                            if prune_configs_fn is not None
+                            else tunable_configs
+                        )
+
+                        best_config = None
+                        best_time = float("inf")
+                        working_config = []
+                        for tunable in pruned_tunable_configs:
+                            try:
+                                current_kwargs.update(tunable)
+                                fn(**current_kwargs)
+                                torch.cuda.synchronize()
+                                working_config.append(tunable)
+                            except Exception:
+                                pass
+
+                        if not working_config:
+                            continue
+
+                        for tunable in working_config:
+                            current_kwargs.update(tunable)
+                            elapse, _ = run_bench(
+                                fn,
+                                current_kwargs,
+                                warmup_iter=warmup_iter,
+                                run_iter=run_iter,
+                                bench_mode=bench_mode,
+                            )
+                            if elapse < best_time:
+                                best_time = elapse
+                                best_config = tunable
+
+                        try:
+                            cache_manager.set(
+                                function_key,
+                                current_input_key,
+                                {"config": best_config, "time": best_time},
+                            )
+                        except (PermissionError, OSError):
+                            # If there's a permission error saving the cache, continue without saving
+                            # logger.warning(f"Permission error saving cache: {e}. Continuing without saving.")
+                            pass
+
+                    try:
+                        # Save the cache to a file
+                        cache_manager.save_cache(function_key)
+                    except (PermissionError, OSError):
+                        # If there's a permission error saving the final cache, continue without saving
+                        # logger.warning(f"Permission error saving final cache: {e}. Continuing without saving.")
+                        pass
+
+                    # After tuning, try to get the best config
+                    try:
+                        best_cached_config = cache_manager.get(function_key, input_key)
+                        if best_cached_config is not None:
+                            all_kwargs.update(best_cached_config["config"])
+                    except (PermissionError, OSError):
+                        # If there's a permission error getting the final config, continue without it
+                        # logger.warning(f"Permission error getting final config: {e}. Continuing without config.")
+                        pass
+
+                except Exception:
+                    # If any other error occurs during tuning, fall back to JIT compilation
+                    # logger.warning(f"Error during tuning: {e}. Falling back to JIT compilation.")
+                    pass
+
+                return fn(**all_kwargs)
+
+        return wrapper
+
+    return decorator

{cuequivariance_ops_cu12-0.4.0.dist-info → cuequivariance_ops_cu12-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: cuequivariance-ops-cu12
-Version: 0.4.0
+Version: 0.5.0
 Summary: cuequivariance-ops - GPU Accelerated Extensions for Equivariant Primitives
 Author: NVIDIA Corporation
 License: # Software License Agreement
@@ -177,6 +177,9 @@ Classifier: Programming Language :: Python
 Project-URL: Homepage, https://github.com/nvidia/cuEquivariance
 Project-URL: Documentation, https://github.com/nvidia/cuEquivariance
 Requires-Python: >=3.10
+Requires-Dist: nvidia-cublas-cu12>=12.5.0
+Requires-Dist: tqdm
+Requires-Dist: pynvml
 Provides-Extra: test
 Requires-Dist: numpy; extra == "test"
 Requires-Dist: pytest; extra == "test"