torchzero 0.1.8__py3-none-any.whl → 0.3.2__py3-none-any.whl

torchzero/utils/compile.py

@@ -1,39 +1,177 @@
-"""Experimental and I need to test this on Windows."""
-import warnings
-import functools
+import time
+
 import torch
+import torch.utils.benchmark
+
+class _OptionalCompiler:
+    """this holds .enable attribute, set to True to enable compiling library wise"""
+    def __init__(self):
+        self.enable = False
+
+    def enable_compilation(
+        self,
+        x,
+        fullgraph: bool = False,
+        dynamic: bool | None = None,
+        backend="inductor",
+        mode: str | None = "max-autotune-no-cudagraphs",
+        options: dict[str, str | int | bool] | None = None,
+        disable: bool = False,
+    ):
+        """compiles if self.compile is True otherwise returns uncompiled `x`"""
+        return _MaybeCompiledFunc(x, self, fullgraph=fullgraph, dynamic=dynamic, backend=backend, mode=mode, options=options, disable=disable)
+
+class _MaybeCompiledFunc:
+    def __init__(self, func, compiler: _OptionalCompiler, **kwargs):
+        self.func = func
+        self.kwargs = kwargs
+        self.compiled = False
+        self.compiler = compiler
+
+    def __call__(self, *args, **kwargs):
+        if self.compiler.enable and not self.compiled:
+            self.func = torch.compile(self.func, **self.kwargs)
+            self.compiled = True
+        return self.func(*args, **kwargs)
+
+
+_optional_compiler = _OptionalCompiler()
+"""this holds .enable attribute, set to True to enable compiling for a few functions that benefit from it."""
+
+def set_compilation(enable: bool):
+    """`enable` is False by default. When True, certain functions will be compiled, which may not work on some systems like Windows, but it usually improves performance."""
+    _optional_compiler.enable = enable
+
+def enable_compilation(fn): return _optional_compiler.enable_compilation(fn)
+
+def benchmark_compile_cuda(fn, n: int, **kwargs):
+    # warmup
+    for _ in range(n):
+        fn(**kwargs)
+
+    compiled = torch.compile(fn, mode = 'max-autotune-no-cudagraphs')
+
+    # compiled warmup
+    for _ in range(n):
+        if _ == 0:
+            start = time.perf_counter()
+            compiled(**kwargs)
+            print(f'Compiling took {time.perf_counter() - start} s.')
+        else:
+            compiled(**kwargs)
+
+    # UNCOMPILED
+    starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    torch.cuda.synchronize()
+    starter.record() # type:ignore
+    start = time.perf_counter()
+
+    for _ in range(n):
+        fn(**kwargs)
+
+    ender.record() # type:ignore
+    torch.cuda.synchronize()
+    sec = 1e-3 * starter.elapsed_time(ender)
+
+    print(f'Uncompiled took {sec} CUDA s, {time.perf_counter() - start} perf_counter s.')
+
+    # COMPILED
+    starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    torch.cuda.synchronize()
+    starter.record() # type:ignore
+    start = time.perf_counter()
+
+    for _ in range(n):
+        compiled(**kwargs)
+
+    ender.record() # type:ignore
+    torch.cuda.synchronize()
+    sec = 1e-3 * starter.elapsed_time(ender)
+
+    print(f'Compiled took {sec} CUDA s, {time.perf_counter() - start} perf_counter s.')
+
+    # UNCOMPILED
+    starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    torch.cuda.synchronize()
+    starter.record() # type:ignore
+    start = time.perf_counter()
+
+    for _ in range(n):
+        fn(**kwargs)
+
+    ender.record() # type:ignore
+    torch.cuda.synchronize()
+    sec = 1e-3 * starter.elapsed_time(ender)
+
+    print(f'Uncompiled took {sec} CUDA s, {time.perf_counter() - start} perf_counter s.')
+
+    # COMPILED
+    starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    torch.cuda.synchronize()
+    starter.record() # type:ignore
+    start = time.perf_counter()
+
+    for _ in range(n):
+        compiled(**kwargs)
+
+    ender.record() # type:ignore
+    torch.cuda.synchronize()
+    sec = 1e-3 * starter.elapsed_time(ender)
+
+    print(f'Compiled took {sec} CUDA s, {time.perf_counter() - start} perf_counter s.')
+
+def benchmark_compile_cpu(fn, n: int, **kwargs):
+    # warmup
+    for _ in range(n):
+        fn(**kwargs)
+
+    compiled = torch.compile(fn, mode = 'max-autotune-no-cudagraphs')
+
+    # compiled warmup
+    for _ in range(n):
+        if _ == 0:
+            start = time.perf_counter()
+            compiled(**kwargs)
+            print(f'Compiling took {time.perf_counter() - start} s.')
+        else:
+            compiled(**kwargs)
+
+    # UNCOMPILED
+    start = time.perf_counter()
+
+    for _ in range(n):
+        fn(**kwargs)
+
+    sec = time.perf_counter() - start
+
+    print(f'Uncompiled took {sec} s., {sec/n} per call')
+
+    # COMPILED
+    start = time.perf_counter()
+
+    for _ in range(n):
+        compiled(**kwargs)
+
+    sec = time.perf_counter() - start
+
+    print(f'Compiled took {sec} s., {sec/n} per call')
+
+    # UNCOMPILED
+    start = time.perf_counter()
+
+    for _ in range(n):
+        fn(**kwargs)
+
+    sec = time.perf_counter() - start
+
+    print(f'Uncompiled took {sec} s., {sec/n} per call')
+
+    # COMPILED
+    start = time.perf_counter()
+
+    for _ in range(n):
+        compiled(**kwargs)
+
+    sec = time.perf_counter() - start
 
-ENABLE_COMPILING = True
-
-def _try_compiling(warn=False):
-    def add(x,y): return x + y
-    compled_add = torch.compile(add)
-    try:
-        res = compled_add(torch.tensor(1.), torch.tensor(2.))
-    except Exception as e:
-        if warn: warnings.warn(f'Compiling failed so no further functions will be compiled:\n{e}')
-        return False
-    if res == 3: return True
-    return False
-
-class _Compiler:
-    def __init__(self, warn=False):
-        self.can_compile = None
-        self.warn = warn
-
-    def maybe_compile(self, fn, **kwargs):
-        if self.can_compile is None: self.can_compile = _try_compiling(self.warn)
-        if self.can_compile: return torch.compile(fn, **kwargs)
-        return fn
-
-_COMPILER = _Compiler(False)
-
-@functools.wraps(torch.compile)
-def maybe_compile(*args, **kwargs):
-    """Compiles a function if possible. Same usage as `torch.compile`.
-
-    On first try this will attempt to compile a simple test function. If that fails, all subsequent functions will not be compiled.
-    I need to actually test this on windows.
-    """
-    if ENABLE_COMPILING: return _COMPILER.maybe_compile(*args, **kwargs)
-    return args[0]
+    print(f'Compiled took {sec} s., {sec/n} per call')