warp-lang 1.8.1__py3-none-macosx_10_13_universal2.whl → 1.9.0__py3-none-macosx_10_13_universal2.whl

warp/build.py

@@ -51,7 +51,7 @@ def build_cuda(
         output_path = output_path.encode("utf-8")
 
         if warp.config.llvm_cuda:
-            warp.context.runtime.llvm.compile_cuda(src, cu_path_bytes, inc_path, output_path, False)
+            warp.context.runtime.llvm.wp_compile_cuda(src, cu_path_bytes, inc_path, output_path, False)
 
         else:
             if ltoirs is None:
@@ -67,7 +67,7 @@ def build_cuda(
                 fatbins
             )
             arr_link_input_types = (ctypes.c_int * num_link)(*link_input_types)
-            err = warp.context.runtime.core.cuda_compile_program(
+            err = warp.context.runtime.core.wp_cuda_compile_program(
                 src,
                 program_name_bytes,
                 arch,
@@ -96,7 +96,7 @@ def load_cuda(input_path, device):
     if not device.is_cuda:
         raise RuntimeError("Not a CUDA device")
 
-    return warp.context.runtime.core.cuda_load_module(device.context, input_path.encode("utf-8"))
+    return warp.context.runtime.core.wp_cuda_load_module(device.context, input_path.encode("utf-8"))
 
 
 def build_cpu(obj_path, cpp_path, mode="release", verify_fp=False, fast_math=False, fuse_fp=True):
@@ -106,7 +106,7 @@ def build_cpu(obj_path, cpp_path, mode="release", verify_fp=False, fast_math=Fal
         inc_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "native").encode("utf-8")
         obj_path = obj_path.encode("utf-8")
 
-        err = warp.context.runtime.llvm.compile_cpp(
+        err = warp.context.runtime.llvm.wp_compile_cpp(
             src, cpp_path, inc_path, obj_path, mode == "debug", verify_fp, fuse_fp
         )
         if err != 0:
@@ -129,6 +129,15 @@ def init_kernel_cache(path=None):
     else:
         cache_root_dir = appdirs.user_cache_dir(appname="warp", appauthor="NVIDIA", version=warp.config.version)
 
+        if os.name == "nt" and os.path.isabs(cache_root_dir) and not cache_root_dir.startswith("\\\\?\\"):
+            # Add Windows long-path prefix, accounting for UNC shares.
+            if cache_root_dir.startswith("\\\\"):
+                # UNC path  \\server\share\…  →  \\?\UNC\server\share\…
+                cache_root_dir = "\\\\?\\UNC\\" + cache_root_dir.lstrip("\\")
+            else:
+                # Drive-letter path  C:\…  →  \\?\C:\…
+                cache_root_dir = "\\\\?\\" + cache_root_dir
+
     warp.config.kernel_cache_dir = cache_root_dir
 
     os.makedirs(warp.config.kernel_cache_dir, exist_ok=True)
@@ -246,7 +255,12 @@ def _build_lto_base(lto_symbol, compile_func, builder, extra_files=None):
             the cached file data.
 
     Returns:
-        Tuple containing lto_code_data followed by any extra data from extra_files
+        Tuple where the first element is a success flag (``bool``). The second
+        element is the LTO code as bytes (or ``None`` on failure).
+        If ``extra_files`` is provided, additional elements follow in the same
+        order as the keys in ``extra_files``:
+          - ``".meta"``: int (shared memory bytes).
+          - ``"_fatbin.lto"``: bytes (universal fatbin).
     """
     if extra_files is None:
         extra_files = {}
@@ -283,9 +297,9 @@ def _build_lto_base(lto_symbol, compile_func, builder, extra_files=None):
 
         if all_files_cached:
             if not extra_files:
-                return (lto_code_data,)
+                return (True, lto_code_data)
             else:
-                return (lto_code_data, *[extra_files[ext] for ext in extra_files.keys()])
+                return (True, lto_code_data, *[extra_files[ext] for ext in extra_files.keys()])
 
     # Create process-dependent temporary build directory
     build_dir = f"{lto_dir}_p{os.getpid()}"
@@ -303,21 +317,24 @@ def _build_lto_base(lto_symbol, compile_func, builder, extra_files=None):
         for path in temp_file_paths.values():
             if Path(path).exists():
                 Path(path).unlink()
-        raise RuntimeError(f"Failed to compile {lto_symbol}")
-
-    # Move outputs to cache
-    safe_rename(build_dir, lto_dir)
-
-    # If build_dir couldn't be moved by a rename, move the outputs one-by-one to lto_dir
-    if os.path.exists(lto_dir):
-        for ext, path in file_paths.items():
-            if not os.path.exists(path):
-                try:
-                    # copy output file to the destination lto dir
-                    os.rename(temp_file_paths[ext], path)
-                except (OSError, FileExistsError):
-                    # another process likely updated the lto dir first
-                    pass
+
+        outputs[".lto"] = None
+        for ext in extra_files.keys():
+            outputs[ext] = None
+    else:
+        # Move outputs to cache
+        safe_rename(build_dir, lto_dir)
+
+        # If build_dir couldn't be moved by a rename, move the outputs one-by-one to lto_dir
+        if os.path.exists(lto_dir):
+            for ext, path in file_paths.items():
+                if not os.path.exists(path):
+                    try:
+                        # copy output file to the destination lto dir
+                        os.rename(temp_file_paths[ext], path)
+                    except (OSError, FileExistsError):
+                        # another process likely updated the lto dir first
+                        pass
 
     # Clean up the temporary build directory
     if build_dir:
@@ -326,9 +343,9 @@ def _build_lto_base(lto_symbol, compile_func, builder, extra_files=None):
         shutil.rmtree(build_dir, ignore_errors=True)
 
     if not extra_files:
-        return (outputs[".lto"],)
+        return (result, outputs[".lto"])
     else:
-        return (outputs[".lto"], *[outputs[ext] for ext in extra_files.keys()])
+        return (result, outputs[".lto"], *[outputs[ext] for ext in extra_files.keys()])
 
 
 def build_lto_dot(M, N, K, adtype, bdtype, cdtype, alayout, blayout, clayout, arch, num_threads, builder):
@@ -372,7 +389,7 @@ def build_lto_dot(M, N, K, adtype, bdtype, cdtype, alayout, blayout, clayout, ar
     lto_symbol = f"dot_{M}_{N}_{K}_{arch}_{num_threads}_{a_arrangement}_{b_arrangement}_{c_arrangement}_{a_prec}_{b_prec}_{c_prec}_{element_type}"
 
     def compile_lto_dot(temp_paths):
-        result = warp.context.runtime.core.cuda_compile_dot(
+        result = warp.context.runtime.core.wp_cuda_compile_dot(
             temp_paths[".lto"].encode("utf-8"),
             lto_symbol.encode("utf-8"),
             0,
@@ -402,7 +419,13 @@ def build_lto_dot(M, N, K, adtype, bdtype, cdtype, alayout, blayout, clayout, ar
     if lto_symbol in builder.ltoirs:
         lto_code_data = builder.ltoirs[lto_symbol]
     else:
-        (lto_code_data,) = _build_lto_base(lto_symbol, compile_lto_dot, builder, {})
+        (result, lto_code_data) = _build_lto_base(lto_symbol, compile_lto_dot, builder, {})
+
+        if not result:
+            raise RuntimeError(
+                f"Failed to compile LTO '{lto_symbol}'. "
+                "Set the environment variable LIBMATHDX_LOG_LEVEL=5 and rerun for more details."
+            )
 
         # Update builder
         builder.ltoirs[lto_symbol] = lto_code_data
@@ -429,6 +452,7 @@ def build_lto_solver(
     num_threads,
     parameter_list,
     builder,
+    smem_estimate_bytes=None,
 ):
     arch = 120 if arch > 121 else arch
 
@@ -446,7 +470,7 @@ def build_lto_solver(
 
     def compile_lto_solver(temp_paths):
         # compile LTO
-        result = warp.context.runtime.core.cuda_compile_solver(
+        result = warp.context.runtime.core.wp_cuda_compile_solver(
             temp_paths["_fatbin.lto"].encode("utf-8"),
             temp_paths[".lto"].encode("utf-8"),
             lto_symbol.encode("utf-8"),
@@ -479,10 +503,43 @@ def build_lto_solver(
     if lto_symbol in builder.ltoirs:
         lto_code_data = builder.ltoirs[lto_symbol]
     else:
-        lto_code_data, universal_fatbin_code_data = _build_lto_base(
+        (result, lto_code_data, universal_fatbin_code_data) = _build_lto_base(
             lto_symbol, compile_lto_solver, builder, {"_fatbin.lto": get_cached_lto}
         )
 
+        if not result:
+            hint = ""
+            if smem_estimate_bytes:
+                max_smem_bytes = 232448
+                max_smem_is_estimate = True
+                for d in warp.get_cuda_devices():
+                    if d.arch == arch:
+                        # We can directly query the max shared memory for this device
+                        queried_bytes = warp.context.runtime.core.wp_cuda_get_max_shared_memory(d.context)
+                        if queried_bytes > 0:
+                            max_smem_bytes = queried_bytes
+                            max_smem_is_estimate = False
+                            break
+                if smem_estimate_bytes > max_smem_bytes:
+                    source = "estimated limit" if max_smem_is_estimate else "device-reported limit"
+                    hint = (
+                        f"Estimated shared memory requirement is {smem_estimate_bytes}B, "
+                        f"but the {source} is {max_smem_bytes}B. "
+                        "The tile size(s) may be too large for this device."
+                    )
+
+            if warp.context.runtime.toolkit_version < (12, 6):
+                raise RuntimeError(
+                    "cuSolverDx requires CUDA Toolkit 12.6.3 or later. This version of Warp was built against CUDA Toolkit "
+                    f"{warp.context.runtime.toolkit_version[0]}.{warp.context.runtime.toolkit_version[1]}. "
+                    "Upgrade your CUDA Toolkit and rebuild Warp, or install a Warp wheel built with CUDA >= 12.6.3."
+                )
+            else:
+                raise RuntimeError(
+                    f"Failed to compile LTO '{lto_symbol}'. {hint}"
+                    " Set the environment variable LIBMATHDX_LOG_LEVEL=5 and rerun for more details."
+                )
+
         # Update builder
         builder.ltoirs[lto_symbol] = lto_code_data
         builder.ltoirs_decl[lto_symbol] = f"void {lto_symbol}{parameter_list};"
@@ -499,7 +556,7 @@ def build_lto_fft(arch, size, ept, direction, dir, precision, builder):
     def compile_lto_fft(temp_paths):
         shared_memory_size = ctypes.c_int(0)
 
-        result = warp.context.runtime.core.cuda_compile_fft(
+        result = warp.context.runtime.core.wp_cuda_compile_fft(
             temp_paths[".lto"].encode("utf-8"),
             lto_symbol.encode("utf-8"),
             0,
@@ -535,10 +592,16 @@ def build_lto_fft(arch, size, ept, direction, dir, precision, builder):
         lto_code_data = builder.ltoirs[lto_symbol]
         shared_memory_bytes = builder.shared_memory_bytes[lto_symbol]
     else:
-        lto_code_data, shared_memory_bytes = _build_lto_base(
+        (result, lto_code_data, shared_memory_bytes) = _build_lto_base(
             lto_symbol, compile_lto_fft, builder, {".meta": lambda path: get_cached_lto_meta(path, lto_symbol)}
         )
 
+        if not result:
+            raise RuntimeError(
+                f"Failed to compile LTO '{lto_symbol}'."
+                "Set the environment variable LIBMATHDX_LOG_LEVEL=5 and rerun for more details."
+            )
+
         # Update builder
         builder.ltoirs[lto_symbol] = lto_code_data
         builder.shared_memory_bytes[lto_symbol] = shared_memory_bytes

warp/build_dll.py

@@ -13,16 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import os
 import platform
 import subprocess
 import sys
-from typing import List, Optional
 
 from warp.utils import ScopedTimer
 
 verbose_cmd = True  # print command lines before executing them
 
+MIN_CTK_VERSION = (12, 0)
+
 
 def machine_architecture() -> str:
     """Return a canonical machine architecture string.
@@ -120,7 +123,7 @@ def find_host_compiler():
         return run_cmd("which g++").decode()
 
 
-def get_cuda_toolkit_version(cuda_home):
+def get_cuda_toolkit_version(cuda_home) -> tuple[int, int]:
     try:
         # the toolkit version can be obtained by running "nvcc --version"
         nvcc_path = os.path.join(cuda_home, "bin", "nvcc")
@@ -128,14 +131,16 @@ def get_cuda_toolkit_version(cuda_home):
         # search for release substring (e.g., "release 11.5")
         import re
 
-        m = re.search(r"(?<=release )\d+\.\d+", nvcc_version_output)
+        m = re.search(r"release (\d+)\.(\d+)", nvcc_version_output)
         if m is not None:
-            return tuple(int(x) for x in m.group(0).split("."))
+            major, minor = map(int, m.groups())
+            return (major, minor)
         else:
             raise Exception("Failed to parse NVCC output")
 
     except Exception as e:
-        print(f"Failed to determine CUDA Toolkit version: {e}")
+        print(f"Warning: Failed to determine CUDA Toolkit version: {e}")
+        return MIN_CTK_VERSION
 
 
 def quote(path):
@@ -169,7 +174,7 @@ def add_llvm_bin_to_path(args):
     return True
 
 
-def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, arch, libs: Optional[List[str]] = None, mode=None):
+def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, arch, libs: list[str] | None = None, mode=None):
     mode = args.mode if (mode is None) else mode
     cuda_home = args.cuda_path
     cuda_cmd = None
@@ -197,17 +202,12 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, arch, libs: Optional[
 
     if cu_path:
         # check CUDA Toolkit version
-        min_ctk_version = (11, 5)
-        ctk_version = get_cuda_toolkit_version(cuda_home) or min_ctk_version
-        if ctk_version < min_ctk_version:
+        ctk_version = get_cuda_toolkit_version(cuda_home)
+        if ctk_version < MIN_CTK_VERSION:
             raise Exception(
-                f"CUDA Toolkit version {min_ctk_version[0]}.{min_ctk_version[1]}+ is required (found {ctk_version[0]}.{ctk_version[1]} in {cuda_home})"
+                f"CUDA Toolkit version {MIN_CTK_VERSION[0]}.{MIN_CTK_VERSION[1]}+ is required (found {ctk_version[0]}.{ctk_version[1]} in {cuda_home})"
             )
 
-        if ctk_version[0] < 12 and args.libmathdx_path:
-            print("MathDx support requires at least CUDA 12, skipping")
-            args.libmathdx_path = None
-
         # NVCC gencode options
         gencode_opts = []
 
@@ -216,91 +216,71 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, arch, libs: Optional[
 
         if args.quick:
             # minimum supported architectures (PTX)
-            gencode_opts += ["-gencode=arch=compute_52,code=compute_52", "-gencode=arch=compute_75,code=compute_75"]
-            clang_arch_flags += ["--cuda-gpu-arch=sm_52", "--cuda-gpu-arch=sm_75"]
+            if ctk_version >= (13, 0):
+                gencode_opts += ["-gencode=arch=compute_75,code=compute_75"]
+                clang_arch_flags += ["--cuda-gpu-arch=sm_75"]
+            else:
+                gencode_opts += ["-gencode=arch=compute_52,code=compute_52", "-gencode=arch=compute_75,code=compute_75"]
+                clang_arch_flags += ["--cuda-gpu-arch=sm_52", "--cuda-gpu-arch=sm_75"]
         else:
             # generate code for all supported architectures
             gencode_opts += [
                 # SASS for supported desktop/datacenter architectures
-                "-gencode=arch=compute_52,code=sm_52",  # Maxwell
-                "-gencode=arch=compute_60,code=sm_60",  # Pascal
-                "-gencode=arch=compute_61,code=sm_61",
-                "-gencode=arch=compute_70,code=sm_70",  # Volta
                 "-gencode=arch=compute_75,code=sm_75",  # Turing
                 "-gencode=arch=compute_75,code=compute_75",  # Turing (PTX)
                 "-gencode=arch=compute_80,code=sm_80",  # Ampere
                 "-gencode=arch=compute_86,code=sm_86",
+                "-gencode=arch=compute_89,code=sm_89",  # Ada
+                "-gencode=arch=compute_90,code=sm_90",  # Hopper
             ]
 
-            # TODO: Get this working with sm_52, sm_60, sm_61
             clang_arch_flags += [
                 # SASS for supported desktop/datacenter architectures
-                "--cuda-gpu-arch=sm_52",
-                "--cuda-gpu-arch=sm_60",
-                "--cuda-gpu-arch=sm_61",
-                "--cuda-gpu-arch=sm_70",  # Volta
                 "--cuda-gpu-arch=sm_75",  # Turing
                 "--cuda-gpu-arch=sm_80",  # Ampere
                 "--cuda-gpu-arch=sm_86",
+                "--cuda-gpu-arch=sm_89",  # Ada
+                "--cuda-gpu-arch=sm_90",  # Hopper
             ]
 
             if arch == "aarch64" and sys.platform == "linux":
-                gencode_opts += [
-                    # SASS for supported mobile architectures (e.g. Tegra/Jetson)
-                    "-gencode=arch=compute_53,code=sm_53",  # X1
-                    "-gencode=arch=compute_62,code=sm_62",  # X2
-                    "-gencode=arch=compute_72,code=sm_72",  # Xavier
-                    "-gencode=arch=compute_87,code=sm_87",  # Orin
-                ]
-
-                clang_arch_flags += [
-                    # SASS for supported mobile architectures
-                    "--cuda-gpu-arch=sm_53",  # X1
-                    "--cuda-gpu-arch=sm_62",  # X2
-                    "--cuda-gpu-arch=sm_72",  # Xavier
-                    "--cuda-gpu-arch=sm_87",  # Orin
-                ]
-
-                if ctk_version >= (12, 8):
-                    gencode_opts += ["-gencode=arch=compute_101,code=sm_101"]  # Thor (CUDA 12 numbering)
-                    clang_arch_flags += ["--cuda-gpu-arch=sm_101"]
+                # SASS for supported mobile architectures (e.g. Tegra/Jetson)
+                gencode_opts += ["-gencode=arch=compute_87,code=sm_87"]  # Orin
+                clang_arch_flags += ["--cuda-gpu-arch=sm_87"]
+
+                if ctk_version >= (13, 0):
+                    gencode_opts += ["-gencode=arch=compute_110,code=sm_110"]  # Thor
+                    clang_arch_flags += ["--cuda-gpu-arch=sm_110"]
+                else:
+                    gencode_opts += [
+                        "-gencode=arch=compute_53,code=sm_53",  # X1
+                        "-gencode=arch=compute_62,code=sm_62",  # X2
+                        "-gencode=arch=compute_72,code=sm_72",  # Xavier
+                    ]
+                    clang_arch_flags += [
+                        "--cuda-gpu-arch=sm_53",
+                        "--cuda-gpu-arch=sm_62",
+                        "--cuda-gpu-arch=sm_72",
+                    ]
+
+                    if ctk_version >= (12, 8):
+                        gencode_opts += ["-gencode=arch=compute_101,code=sm_101"]  # Thor (CUDA 12 numbering)
+                        clang_arch_flags += ["--cuda-gpu-arch=sm_101"]
 
             if ctk_version >= (12, 8):
                 # Support for Blackwell is available with CUDA Toolkit 12.8+
                 gencode_opts += [
-                    "-gencode=arch=compute_89,code=sm_89",  # Ada
-                    "-gencode=arch=compute_90,code=sm_90",  # Hopper
                     "-gencode=arch=compute_100,code=sm_100",  # Blackwell
                     "-gencode=arch=compute_120,code=sm_120",  # Blackwell
                     "-gencode=arch=compute_120,code=compute_120",  # PTX for future hardware
                 ]
 
                 clang_arch_flags += [
-                    "--cuda-gpu-arch=sm_89",  # Ada
-                    "--cuda-gpu-arch=sm_90",  # Hopper
                     "--cuda-gpu-arch=sm_100",  # Blackwell
                     "--cuda-gpu-arch=sm_120",  # Blackwell
                 ]
-            elif ctk_version >= (11, 8):
-                # Support for Ada and Hopper is available with CUDA Toolkit 11.8+
-                gencode_opts += [
-                    "-gencode=arch=compute_89,code=sm_89",  # Ada
-                    "-gencode=arch=compute_90,code=sm_90",  # Hopper
-                    "-gencode=arch=compute_90,code=compute_90",  # PTX for future hardware
-                ]
-
-                clang_arch_flags += [
-                    "--cuda-gpu-arch=sm_89",  # Ada
-                    "--cuda-gpu-arch=sm_90",  # Hopper
-                ]
             else:
-                gencode_opts += [
-                    "-gencode=arch=compute_86,code=compute_86",  # PTX for future hardware
-                ]
-
-                clang_arch_flags += [
-                    "--cuda-gpu-arch=sm_86",  # PTX for future hardware
-                ]
+                gencode_opts += ["-gencode=arch=compute_90,code=compute_90"]  # PTX for future hardware
 
         nvcc_opts = [
             *gencode_opts,