warp-lang 1.6.2__py3-none-macosx_10_13_universal2.whl → 1.7.0__py3-none-macosx_10_13_universal2.whl

warp/__init__.py

@@ -84,7 +84,12 @@ from warp.context import Stream, get_stream, set_stream, wait_stream, synchroniz
 from warp.context import Event, record_event, wait_event, synchronize_event, get_event_elapsed_time
 from warp.context import RegisteredGLBuffer
 from warp.context import is_mempool_supported, is_mempool_enabled, set_mempool_enabled
-from warp.context import set_mempool_release_threshold, get_mempool_release_threshold
+from warp.context import (
+    set_mempool_release_threshold,
+    get_mempool_release_threshold,
+    get_mempool_used_mem_current,
+    get_mempool_used_mem_high,
+)
 from warp.context import is_mempool_access_supported, is_mempool_access_enabled, set_mempool_access_enabled
 from warp.context import is_peer_access_supported, is_peer_access_enabled, set_peer_access_enabled
 
@@ -120,6 +125,7 @@ from warp.paddle import device_from_paddle, device_to_paddle
 from warp.paddle import stream_from_paddle
 
 from warp.build import clear_kernel_cache
+from warp.build import clear_lto_cache
 
 from warp.constants import *
 

warp/build.py

@@ -14,10 +14,16 @@
 # limitations under the License.
 
 import ctypes
+import errno
+import hashlib
+import json
 import os
+import time
+from pathlib import Path
 
 import warp.config
 from warp.thirdparty import appdirs
+from warp.types import *
 
 # From nvJitLink.h
 nvJitLink_input_type = {"cubin": 1, "ptx": 2, "ltoir": 3, "fatbin": 4, "object": 5, "library": 6}
@@ -131,6 +137,7 @@ def clear_kernel_cache() -> None:
 
     Only directories beginning with ``wp_`` will be deleted.
     This function only clears the cache for the current Warp version.
+    LTO artifacts are not affected.
     """
 
     warp.context.init()
@@ -145,3 +152,406 @@ def clear_kernel_cache() -> None:
         if os.path.isdir(item_path) and item.startswith("wp_"):
             # Remove the directory and its contents
             shutil.rmtree(item_path, ignore_errors=True)
+
+
+def clear_lto_cache() -> None:
+    """Clear the LTO cache directory of previously generated LTO code.
+
+    The LTO cache is stored within a subdirectory of the kernel cache directory.
+    This function only clears the cache for the current Warp version.
+    """
+
+    warp.context.init()
+
+    import shutil
+
+    is_intialized = warp.context.runtime is not None
+    assert is_intialized, "The kernel cache directory is not configured; wp.init() has not been called yet or failed."
+
+    lto_path = os.path.join(warp.config.kernel_cache_dir, "lto")
+    if os.path.isdir(lto_path):
+        # Remove the lto directory and its contents
+        shutil.rmtree(lto_path, ignore_errors=True)
+
+
+def safe_rename(src, dst, attempts=5, delay=0.1):
+    for i in range(attempts):
+        try:
+            os.rename(src, dst)
+            return
+        except FileExistsError:
+            return
+        except OSError as e:
+            if e.errno == errno.ENOTEMPTY:
+                # if directory exists we assume another process
+                # got there first, in which case we will copy
+                # our output to the directory manually in second step
+                return
+            else:
+                # otherwise assume directory creation failed e.g.: access denied
+                # on Windows we see occasional failures to rename directories due to
+                # some process holding a lock on a file to be moved to workaround
+                # this we make multiple attempts to rename with some delay
+                if i < attempts - 1:
+                    time.sleep(delay)
+                else:
+                    print(
+                        f"Could not update Warp cache with compiled binaries, trying to rename {src} to {dst}, error {e}"
+                    )
+                    raise e
+
+
+def hash_symbol(symbol):
+    ch = hashlib.sha256()
+    ch.update(symbol.encode("utf-8"))
+    return ch.hexdigest()
+
+
+def get_lto_cache_dir():
+    lto_dir = os.path.join(warp.config.kernel_cache_dir, "lto")
+    return lto_dir
+
+
+def get_cached_lto(path):
+    if os.path.exists(path):
+        with open(path, "rb") as f:
+            lto_code_data = f.read()
+        return lto_code_data
+    else:
+        return None
+
+
+def get_cached_lto_meta(path, symbol):
+    if os.path.exists(path):
+        with open(path, "r") as f:
+            keys = json.load(f)
+        value = keys[symbol]
+        return value
+    else:
+        return None
+
+
+def build_lto_dot(M, N, K, adtype, bdtype, cdtype, alayout, blayout, clayout, arch, num_threads, builder):
+    # TODO: MathDx doesn't yet have heuristics for Blackwell
+    arch = min(arch, 90)
+
+    # Maps Python/Warp types to C++ types and enums
+    def cublasdx_type_map(dtype):
+        if dtype == float16:
+            return ("wp::float16", 3, 0)
+        if dtype == float32:
+            return ("wp::float32", 5, 0)
+        if dtype == float64:
+            return ("wp::float64", 6, 0)
+        if dtype == vec2h:
+            return ("wp::vec2h", 3, 1)
+        if dtype == vec2f:
+            return ("wp::vec2f", 5, 1)
+        if dtype == vec2d:
+            return ("wp::vec2d", 6, 1)
+        raise TypeError("Unsupported input type in tile_matmul")
+
+    def cublasdx_arrangement_map(layout):
+        if layout == "colmajor":
+            return 0  # CUBLASDX_ARRANGEMENT_COL_MAJOR
+        if layout == "rowmajor":
+            return 1  # CUBLASDX_ARRANGEMENT_ROW_MAJOR
+        raise ValueError("Unsupported layout in tile_matmul")
+
+    (a_dtype, a_prec, a_type) = cublasdx_type_map(adtype)
+    (b_dtype, b_prec, b_type) = cublasdx_type_map(bdtype)
+    (c_dtype, c_prec, c_type) = cublasdx_type_map(cdtype)
+    a_arrangement = cublasdx_arrangement_map(alayout)
+    b_arrangement = cublasdx_arrangement_map(blayout)
+    c_arrangement = cublasdx_arrangement_map(clayout)
+
+    if a_type != b_type or a_type != c_type:
+        raise TypeError("time_matmul(A, B, C) requires all inputs to be real or complex")
+
+    element_type = a_type
+
+    lto_symbol = f"dot_{M}_{N}_{K}_{arch}_{num_threads}_{a_arrangement}_{b_arrangement}_{c_arrangement}_{a_prec}_{b_prec}_{c_prec}_{element_type}"
+
+    # early out if LTO for this symbol is already cached in current module
+    if lto_symbol in builder.ltoirs:
+        return lto_symbol, builder.ltoirs[lto_symbol]
+
+    # hash symbol and determine output path
+    h = hash_symbol(lto_symbol)
+
+    lto_dir = get_lto_cache_dir()
+    lto_name = f"{h[:7]}.lto"
+    lto_path = os.path.join(lto_dir, lto_name)
+
+    # early out if LTO for this symbol is already built but not cached in current module
+    lto_code_data = get_cached_lto(lto_path)
+
+    if lto_code_data is not None:
+        builder.ltoirs[lto_symbol] = lto_code_data
+        builder.ltoirs_decl[lto_symbol] = (
+            f"void {lto_symbol}({c_dtype}, {a_dtype}*, {b_dtype}*, {c_dtype}, {c_dtype}*);"
+        )
+
+        return lto_symbol, lto_code_data
+
+    # create a temporary (process unique) dir for build outputs before moving to the binary dir
+    build_dir = f"{lto_dir}_p{os.getpid()}"
+
+    # dir may exist from previous attempts / runs / archs
+    Path(build_dir).mkdir(parents=True, exist_ok=True)
+
+    # temporary path to compile to in build_dir
+    temp_lto_path = os.path.join(build_dir, lto_name)
+
+    # compile LTO
+    result = warp.context.runtime.core.cuda_compile_dot(
+        temp_lto_path.encode("utf-8"),
+        lto_symbol.encode("utf-8"),
+        0,
+        None,
+        None,
+        arch,
+        M,
+        N,
+        K,
+        a_prec,
+        b_prec,
+        c_prec,
+        element_type,
+        a_arrangement,
+        b_arrangement,
+        c_arrangement,
+        num_threads,
+    )
+
+    if not result:
+        if Path(temp_lto_path).exists():
+            Path(temp_lto_path).unlink()
+        raise RuntimeError("Failed to compile tile_matmul")
+    else:
+        with open(temp_lto_path, "rb") as f:
+            lto_code_data = f.read()
+
+    builder.ltoirs[lto_symbol] = lto_code_data
+    builder.ltoirs_decl[lto_symbol] = f"void {lto_symbol}({c_dtype}, {a_dtype}*, {b_dtype}*, {c_dtype}, {c_dtype}*);"
+
+    # try to move process outputs to cache
+    safe_rename(build_dir, lto_dir)
+
+    if os.path.exists(lto_dir):
+        if not os.path.exists(lto_path):
+            # copy output file to the destination lto dir
+            try:
+                os.rename(temp_lto_path, lto_path)
+            except (OSError, FileExistsError):
+                # another process likely updated the lto dir first
+                pass
+
+    if build_dir:
+        import shutil
+
+        # clean up build_dir used for this process
+        shutil.rmtree(build_dir, ignore_errors=True)
+
+    return lto_symbol, lto_code_data
+
+
+def build_lto_solver(M, N, solver, solver_enum, fill_mode, arch, precision_enum, num_threads, parameter_list, builder):
+    # TODO: MathDx doesn't yet have heuristics for Blackwell
+    arch = min(arch, 90)
+
+    lto_symbol = f"{solver}_{M}_{N}_{arch}_{precision_enum}"
+    ltoir_decl = f"void {lto_symbol}{parameter_list};"
+
+    # early out if LTO for this symbol is already cached in current module
+    if lto_symbol in builder.ltoirs:
+        return lto_symbol, builder.ltoirs[lto_symbol]
+
+    # hash symbol and determine output path
+    h = hash_symbol(lto_symbol)
+
+    lto_dir = get_lto_cache_dir()
+    lto_name = f"{h[:7]}.lto"
+    lto_path = os.path.join(lto_dir, lto_name)
+
+    # we also cache a universal fatbin binary for this symbol
+    universal_fatbin_name = f"{h[:7]}_fatbin.lto"
+    universal_fatbin_path = os.path.join(lto_dir, universal_fatbin_name)
+
+    lto_code_data = get_cached_lto(lto_path)
+    universal_fatbin_code_data = get_cached_lto(universal_fatbin_path)
+
+    # early out if LTO for this symbol is already built but not cached in current module
+    if lto_code_data is not None and universal_fatbin_code_data is not None:
+        builder.ltoirs[lto_symbol] = lto_code_data
+        builder.ltoirs_decl[lto_symbol] = ltoir_decl
+        builder.fatbins[lto_symbol] = universal_fatbin_code_data
+
+        return lto_symbol, lto_code_data
+
+    # create a temporary (process unique) dir for build outputs before moving to the binary dir
+    build_dir = f"{lto_dir}_p{os.getpid()}"
+
+    # dir may exist from previous attempts / runs / archs
+    Path(build_dir).mkdir(parents=True, exist_ok=True)
+
+    # temporary paths to compile to in build_dir
+    temp_lto_path = os.path.join(build_dir, lto_name)
+    temp_universal_fatbin_path = os.path.join(build_dir, universal_fatbin_name)
+
+    # compile LTO
+    result = warp.context.runtime.core.cuda_compile_solver(
+        temp_universal_fatbin_path.encode("utf-8"),
+        temp_lto_path.encode("utf-8"),
+        lto_symbol.encode("utf-8"),
+        0,
+        None,
+        None,
+        arch,
+        M,
+        N,
+        solver_enum,
+        precision_enum,
+        fill_mode,
+        num_threads,
+    )
+
+    if not result:
+        for path in [temp_universal_fatbin_path, temp_lto_path]:
+            if Path(path).exists():
+                Path(path).unlink()
+        raise RuntimeError("Failed to compile tile_cholesky")
+
+    else:
+        with open(temp_lto_path, "rb") as f:
+            lto_code_data = f.read()
+        with open(temp_universal_fatbin_path, "rb") as f:
+            universal_fatbin_code_data = f.read()
+
+    builder.ltoirs[lto_symbol] = lto_code_data
+    builder.ltoirs_decl[lto_symbol] = ltoir_decl
+    builder.fatbins[lto_symbol] = universal_fatbin_code_data
+
+    # try to move process outputs to lto cache
+    safe_rename(build_dir, lto_dir)
+
+    if os.path.exists(lto_dir):
+        for p in [(lto_path, temp_lto_path), (universal_fatbin_path, temp_universal_fatbin_path)]:
+            path, temp_path = p
+            if not os.path.exists(path):
+                # copy output file to the destination lto dir
+                try:
+                    os.rename(temp_path, path)
+                except (OSError, FileExistsError):
+                    # another process likely updated the lto dir first
+                    pass
+
+    if build_dir:
+        import shutil
+
+        # clean up build_dir used for this process
+        shutil.rmtree(build_dir, ignore_errors=True)
+
+    return lto_symbol, lto_code_data
+
+
+def build_lto_fft(arch, size, ept, direction, dir, precision, builder):
+    # TODO: MathDx doesn't yet have heuristics for Blackwell
+    arch = min(arch, 90)
+
+    lto_symbol = f"fft_{size}_{ept}_{arch}_{direction}_{precision}"
+
+    # early out if LTO for this symbol is already cached in current module
+    if lto_symbol in builder.ltoirs:
+        return lto_symbol, builder.ltoirs[lto_symbol], builder.shared_memory_bytes[lto_symbol]
+
+    # hash symbol and determine output path
+    h = hash_symbol(lto_symbol)
+
+    lto_dir = get_lto_cache_dir()
+    lto_name = f"{h[:7]}.lto"
+    lto_path = os.path.join(lto_dir, lto_name)
+
+    # we also cache shared memory requirements for this kernel in a .meta file
+    meta_name = f"{h[:7]}.meta"
+    meta_path = os.path.join(lto_dir, meta_name)
+
+    # early out if LTO for this symbol is already built but not cached in current module
+    lto_code_data = get_cached_lto(lto_path)
+    shared_memory_bytes = get_cached_lto_meta(meta_path, lto_symbol)
+
+    if lto_code_data is not None and shared_memory_bytes is not None:
+        builder.ltoirs[lto_symbol] = lto_code_data
+        builder.shared_memory_bytes[lto_symbol] = shared_memory_bytes
+
+        return lto_symbol, lto_code_data, shared_memory_bytes
+
+    # create a temporary (process unique) dir for build outputs before moving to the binary dir
+    build_dir = f"{lto_dir}_p{os.getpid()}"
+
+    # dir may exist from previous attempts / runs / archs
+    Path(build_dir).mkdir(parents=True, exist_ok=True)
+
+    # temporary paths to compile to in build_dir
+    temp_lto_path = os.path.join(build_dir, lto_name)
+    temp_meta_path = os.path.join(build_dir, meta_name)
+
+    # compile LTO
+    shared_memory_size = ctypes.c_int(0)
+
+    result = warp.context.runtime.core.cuda_compile_fft(
+        temp_lto_path.encode("utf-8"),
+        lto_symbol.encode("utf-8"),
+        0,
+        None,
+        None,
+        arch,
+        size,
+        ept,
+        dir,
+        precision,
+        ctypes.byref(shared_memory_size),
+    )
+
+    shared_memory_bytes = Tile.round_up(shared_memory_size.value)
+
+    if not result:
+        if Path(temp_lto_path).exists():
+            Path(temp_lto_path).unlink()
+        raise RuntimeError("Failed to compile tile_fft")
+
+    else:
+        with open(temp_lto_path, "rb") as f:
+            lto_code_data = f.read()
+
+        # output meta file with shared memory requirements for this lto_symbol
+        meta = {}
+        meta[lto_symbol] = shared_memory_bytes
+
+        with open(temp_meta_path, "w") as meta_file:
+            json.dump(meta, meta_file)
+
+    builder.ltoirs[lto_symbol] = lto_code_data
+    builder.shared_memory_bytes[lto_symbol] = shared_memory_bytes
+
+    # try to move process outputs to cache
+    safe_rename(build_dir, lto_dir)
+
+    if os.path.exists(lto_dir):
+        for p in [(lto_path, temp_lto_path), (meta_path, temp_meta_path)]:
+            path, temp_path = p
+            if not os.path.exists(path):
+                # copy output file to the destination lto dir
+                try:
+                    os.rename(temp_path, path)
+                except (OSError, FileExistsError):
+                    # another process likely updated the lto dir first
+                    pass
+
+    if build_dir:
+        import shutil
+
+        # clean up build_dir used for this process
+        shutil.rmtree(build_dir, ignore_errors=True)
+
+    return lto_symbol, lto_code_data, shared_memory_bytes

warp/build_dll.py

@@ -147,14 +147,6 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
     cuda_home = args.cuda_path
     cuda_cmd = None
 
-    if args.quick:
-        cutlass_includes = ""
-        cutlass_enabled = "WP_ENABLE_CUTLASS=0"
-    else:
-        cutlass_home = "warp/native/cutlass"
-        cutlass_includes = f'-I"{cutlass_home}/include" -I"{cutlass_home}/tools/util/include"'
-        cutlass_enabled = "WP_ENABLE_CUTLASS=1"
-
     if args.quick or cu_path is None:
         cuda_compat_enabled = "WP_ENABLE_CUDA_COMPATIBILITY=0"
     else:
@@ -270,7 +262,7 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
             iter_dbg = "_ITERATOR_DEBUG_LEVEL=2"
             debug = "_DEBUG"
 
-        cpp_flags = f'/nologo /std:c++17 /GR- {runtime} /D "{debug}" /D "{cuda_enabled}" /D "{cutlass_enabled}" /D "{mathdx_enabled}" /D "{cuda_compat_enabled}" /D "{iter_dbg}" /I"{native_dir}" {includes} '
+        cpp_flags = f'/nologo /std:c++17 /GR- {runtime} /D "{debug}" /D "{cuda_enabled}" /D "{mathdx_enabled}" /D "{cuda_compat_enabled}" /D "{iter_dbg}" /I"{native_dir}" {includes} '
 
         if args.mode == "debug":
             cpp_flags += "/Zi /Od /D WP_ENABLE_DEBUG=1"
@@ -299,10 +291,10 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
             cu_out = cu_path + ".o"
 
             if mode == "debug":
-                cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 --compiler-options=/MT,/Zi,/Od -g -G -O0 -DNDEBUG -D_ITERATOR_DEBUG_LEVEL=0 -I"{native_dir}" -line-info {" ".join(nvcc_opts)} -DWP_ENABLE_CUDA=1 -D{cutlass_enabled} {cutlass_includes} -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
+                cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 --compiler-options=/MT,/Zi,/Od -g -G -O0 -DNDEBUG -D_ITERATOR_DEBUG_LEVEL=0 -I"{native_dir}" -line-info {" ".join(nvcc_opts)} -DWP_ENABLE_CUDA=1 -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
 
             elif mode == "release":
-                cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -O3 {" ".join(nvcc_opts)} -I"{native_dir}" -DNDEBUG -DWP_ENABLE_CUDA=1 -D{cutlass_enabled} {cutlass_includes} -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
+                cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -O3 {" ".join(nvcc_opts)} -I"{native_dir}" -DNDEBUG -DWP_ENABLE_CUDA=1 -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
 
             with ScopedTimer("build_cuda", active=args.verbose):
                 run_cmd(cuda_cmd)
@@ -329,7 +321,7 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
         else:
             version = "-fabi-version=13"  # GCC 8.2+
 
-        cpp_flags = f'{version} --std=c++17 -fno-rtti -D{cuda_enabled} -D{cutlass_enabled} -D{mathdx_enabled} -D{cuda_compat_enabled} -fPIC -fvisibility=hidden -D_GLIBCXX_USE_CXX11_ABI=0 -I"{native_dir}" {includes} '
+        cpp_flags = f'{version} --std=c++17 -fno-rtti -D{cuda_enabled} -D{mathdx_enabled} -D{cuda_compat_enabled} -fPIC -fvisibility=hidden -D_GLIBCXX_USE_CXX11_ABI=0 -I"{native_dir}" {includes} '
 
         if mode == "debug":
             cpp_flags += "-O0 -g -D_DEBUG -DWP_ENABLE_DEBUG=1 -fkeep-inline-functions"
@@ -357,10 +349,10 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
             cu_out = cu_path + ".o"
 
             if mode == "debug":
-                cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -g -G -O0 --compiler-options -fPIC,-fvisibility=hidden -D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 -line-info {" ".join(nvcc_opts)} -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{cutlass_enabled} {cutlass_includes} -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
+                cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -g -G -O0 --compiler-options -fPIC,-fvisibility=hidden -D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 -line-info {" ".join(nvcc_opts)} -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
 
             elif mode == "release":
-                cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -O3 --compiler-options -fPIC,-fvisibility=hidden {" ".join(nvcc_opts)} -DNDEBUG -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{cutlass_enabled} {cutlass_includes} -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
+                cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -O3 --compiler-options -fPIC,-fvisibility=hidden {" ".join(nvcc_opts)} -DNDEBUG -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
 
             with ScopedTimer("build_cuda", active=args.verbose):
                 run_cmd(cuda_cmd)