warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/__init__.py

@@ -10,8 +10,9 @@
 
 from warp.types import array, array1d, array2d, array3d, array4d, constant
 from warp.types import indexedarray, indexedarray1d, indexedarray2d, indexedarray3d, indexedarray4d
+from warp.fabric import fabricarray, fabricarrayarray, indexedfabricarray, indexedfabricarrayarray
 
-from warp.types import int8, uint8, int16, uint16, int32, uint32, int64, uint64, float16, float32, float64
+from warp.types import bool, int8, uint8, int16, uint16, int32, uint32, int64, uint64, float16, float32, float64
 from warp.types import vec2, vec2b, vec2ub, vec2s, vec2us, vec2i, vec2ui, vec2l, vec2ul, vec2h, vec2f, vec2d
 from warp.types import vec3, vec3b, vec3ub, vec3s, vec3us, vec3i, vec3ui, vec3l, vec3ul, vec3h, vec3f, vec3d
 from warp.types import vec4, vec4b, vec4ub, vec4s, vec4us, vec4i, vec4ui, vec4l, vec4ul, vec4h, vec4f, vec4d
@@ -25,7 +26,9 @@ from warp.types import spatial_matrix, spatial_matrixh, spatial_matrixf, spatial
 
 # geometry types
 from warp.types import Bvh, Mesh, HashGrid, Volume, MarchingCubes
-from warp.types import bvh_query_t, mesh_query_aabb_t, hash_grid_query_t
+from warp.types import bvh_query_t, hash_grid_query_t, mesh_query_aabb_t, mesh_query_point_t, mesh_query_ray_t
+
+
 
 # device-wide gemms
 from warp.types import matmul, adj_matmul, batched_matmul, adj_batched_matmul, from_ptr
@@ -34,7 +37,7 @@ from warp.types import matmul, adj_matmul, batched_matmul, adj_batched_matmul, f
 from warp.types import vector as vec
 from warp.types import matrix as mat
 
-from warp.context import init, func, kernel, struct, overload
+from warp.context import init, func, func_grad, func_replay, func_native, kernel, struct, overload
 from warp.context import is_cpu_available, is_cuda_available, is_device_available
 from warp.context import get_devices, get_preferred_device
 from warp.context import get_cuda_devices, get_cuda_device_count, get_cuda_device, map_cuda_device, unmap_cuda_device
@@ -56,7 +59,6 @@ from warp.context import (
 )
 from warp.context import set_module_options, get_module_options, get_module
 from warp.context import capture_begin, capture_end, capture_launch
-from warp.context import print_builtins, export_builtins, export_stubs
 from warp.context import Kernel, Function, Launch
 from warp.context import Stream, get_stream, set_stream, synchronize_stream
 from warp.context import Event, record_event, wait_event, wait_stream
@@ -78,3 +80,7 @@ from warp.dlpack import from_dlpack, to_dlpack
 from warp.constants import *
 
 from . import builtins
+
+import warp.config
+
+__version__ = warp.config.version

warp/__init__.pyi

@@ -0,0 +1 @@
+from .stubs import *

warp/build.py

@@ -26,8 +26,8 @@ def build_cuda(cu_path, arch, output_path, config="release", verify_fp=False, fa
             err = warp.context.runtime.core.cuda_compile_program(
                 src, arch, inc_path, config == "debug", warp.config.verbose, verify_fp, fast_math, output_path
             )
-            if err:
-                raise Exception("CUDA build failed")
+            if err != 0:
+                raise Exception(f"CUDA kernel build failed with error code {err}")
 
 
 # load PTX or CUBIN as a CUDA runtime module (input type determined by input_path extension)
@@ -45,7 +45,9 @@ 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")
 
-        warp.context.runtime.llvm.compile_cpp(src, cpp_path, inc_path, obj_path, mode == "debug")
+        err = warp.context.runtime.llvm.compile_cpp(src, cpp_path, inc_path, obj_path, mode == "debug")
+        if err != 0:
+            raise Exception(f"CPU kernel build failed with error code {err}")
 
 
 kernel_bin_dir = None

warp/build_dll.py

@@ -8,11 +8,24 @@
 import sys
 import os
 import subprocess
+import platform
 
 import warp.config
 from warp.utils import ScopedTimer
 
 
+# returns a canonical machine architecture string
+# - "x86_64" for x86-64, aka. AMD64, aka. x64
+# - "aarch64" for AArch64, aka. ARM64
+def machine_architecture() -> str:
+    machine = platform.machine()
+    if machine == "x86_64" or machine == "AMD64":
+        return "x86_64"
+    if machine == "aarch64" or machine == "arm64":
+        return "aarch64"
+    raise RuntimeError(f"Unrecognized machine architecture {machine}")
+
+
 def run_cmd(cmd, capture=False):
     if warp.config.verbose:
         print(cmd)
@@ -182,12 +195,15 @@ def build_dll_for_arch(dll_path, cpp_paths, cu_path, libs, mode, arch, verify_fp
                 "-gencode=arch=compute_75,code=sm_75",  # Turing
                 "-gencode=arch=compute_80,code=sm_80",  # Ampere
                 "-gencode=arch=compute_86,code=sm_86",
-                # SASS for supported mobile architectures (e.g. Tegra/Jetson)
-                # "-gencode=arch=compute_53,code=sm_53",
-                # "-gencode=arch=compute_62,code=sm_62",
-                # "-gencode=arch=compute_72,code=sm_72",
-                # "-gencode=arch=compute_87,code=sm_87",
             ]
+            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
+                ]
 
             # support for Ada and Hopper is available with CUDA Toolkit 11.8+
             if ctk_version >= (11, 8):
@@ -354,11 +370,15 @@ def build_dll(dll_path, cpp_paths, cu_path, libs=[], mode="release", verify_fp=F
     if sys.platform == "darwin":
         # create a universal binary by combining x86-64 and AArch64 builds
         build_dll_for_arch(dll_path + "-x86_64", cpp_paths, cu_path, libs, mode, "x86_64", verify_fp, fast_math, quick)
-        build_dll_for_arch(dll_path + "-arm64", cpp_paths, cu_path, libs, mode, "arm64", verify_fp, fast_math, quick)
+        build_dll_for_arch(
+            dll_path + "-aarch64", cpp_paths, cu_path, libs, mode, "aarch64", verify_fp, fast_math, quick
+        )
 
-        run_cmd(f"lipo -create -output {dll_path} {dll_path}-x86_64 {dll_path}-arm64")
+        run_cmd(f"lipo -create -output {dll_path} {dll_path}-x86_64 {dll_path}-aarch64")
         os.remove(f"{dll_path}-x86_64")
-        os.remove(f"{dll_path}-arm64")
+        os.remove(f"{dll_path}-aarch64")
 
     else:
-        build_dll_for_arch(dll_path, cpp_paths, cu_path, libs, mode, "x86_64", verify_fp, fast_math, quick)
+        build_dll_for_arch(
+            dll_path, cpp_paths, cu_path, libs, mode, machine_architecture(), verify_fp, fast_math, quick
+        )