warp-lang 1.4.2__py3-none-win_amd64.whl → 1.5.0__py3-none-win_amd64.whl

warp/__init__.py

@@ -26,6 +26,9 @@ from warp.types import transform, transformh, transformf, transformd
 from warp.types import spatial_vector, spatial_vectorh, spatial_vectorf, spatial_vectord
 from warp.types import spatial_matrix, spatial_matrixh, spatial_matrixf, spatial_matrixd
 
+# annotation types
+from warp.types import Int, Float, Scalar
+
 # geometry types
 from warp.types import Bvh, Mesh, HashGrid, Volume, MarchingCubes
 from warp.types import BvhQuery, HashGridQuery, MeshQueryAABB, MeshQueryPoint, MeshQueryRay
@@ -58,6 +61,7 @@ from warp.context import (
     copy,
     from_numpy,
     launch,
+    launch_tiled,
     synchronize,
     force_load,
     load_module,

warp/autograd.py

@@ -34,6 +34,7 @@ def gradcheck(
     input_output_mask: List[Tuple[Union[str, int], Union[str, int]]] = None,
     device: wp.context.Devicelike = None,
     max_blocks=0,
+    block_dim=256,
     max_inputs_per_var=-1,
     max_outputs_per_var=-1,
     plot_relative_error=False,
@@ -44,7 +45,8 @@ def gradcheck(
     Checks whether the autodiff gradient of a Warp kernel matches finite differences.
     Fails if the relative or absolute errors between the autodiff and finite difference gradients exceed the specified tolerance, or if the autodiff gradients contain NaN values.
 
-    The kernel function and its adjoint version are launched with the given inputs and outputs, as well as the provided ``dim`` and ``max_blocks`` arguments (see :func:`warp.launch` for more details).
+    The kernel function and its adjoint version are launched with the given inputs and outputs, as well as the provided
+    ``dim``, ``max_blocks``, and ``block_dim`` arguments (see :func:`warp.launch` for more details).
 
     Note:
         This function only supports Warp kernels whose input arguments precede the output arguments.
@@ -65,6 +67,7 @@ def gradcheck(
         input_output_mask: List of tuples specifying the input-output pairs to compute the Jacobian for. Inputs and outputs can be identified either by their integer indices of where they appear in the kernel input/output arguments, or by the respective argument names as strings. If None, computes the Jacobian for all input-output pairs.
         device: The device to launch on (optional)
         max_blocks: The maximum number of CUDA thread blocks to use.
+        block_dim: The number of threads per block.
         max_inputs_per_var: Maximum number of input dimensions over which to evaluate the Jacobians for the input-output pairs. Evaluates all input dimensions if value <= 0.
         max_outputs_per_var: Maximum number of output dimensions over which to evaluate the Jacobians for the input-output pairs. Evaluates all output dimensions if value <= 0.
         plot_relative_error: If True, visualizes the relative error of the Jacobians in a plot (requires ``matplotlib``).
@@ -85,6 +88,7 @@ def gradcheck(
         input_output_mask=input_output_mask,
         device=device,
         max_blocks=max_blocks,
+        block_dim=block_dim,
         max_inputs_per_var=max_inputs_per_var,
         eps=eps,
         plot_jacobians=False,
@@ -98,6 +102,7 @@ def gradcheck(
         input_output_mask=input_output_mask,
         device=device,
         max_blocks=max_blocks,
+        block_dim=block_dim,
         max_outputs_per_var=max_outputs_per_var,
         plot_jacobians=False,
     )
@@ -237,7 +242,6 @@ def gradcheck_tape(
         input_output_masks: Dictionary of input-output masks for each kernel in the tape, mapping from kernel keys to input-output masks. Inputs and outputs can be identified either by their integer indices of where they appear in the kernel input/output arguments, or by the respective argument names as strings. If None, computes the Jacobian for all input-output pairs.
         blacklist_kernels: List of kernel keys to exclude from the gradient check.
         whitelist_kernels: List of kernel keys to include in the gradient check. If not empty or None, only kernels in this list are checked.
-        max_blocks: The maximum number of CUDA thread blocks to use.
         max_inputs_per_var: Maximum number of input dimensions over which to evaluate the Jacobians for the input-output pairs. Evaluates all input dimensions if value <= 0.
         max_outputs_per_var: Maximum number of output dimensions over which to evaluate the Jacobians for the input-output pairs. Evaluates all output dimensions if value <= 0.
         plot_relative_error: If True, visualizes the relative error of the Jacobians in a plot (requires ``matplotlib``).
@@ -262,7 +266,7 @@ def gradcheck_tape(
     for launch in tape.launches:
         if not isinstance(launch[0], wp.Kernel):
             continue
-        kernel, dim, max_blocks, inputs, outputs, device = launch[:6]
+        kernel, dim, max_blocks, inputs, outputs, device, block_dim = launch[:7]
         if len(whitelist_kernels) > 0 and kernel.key not in whitelist_kernels:
             continue
         if kernel.key in blacklist_kernels:
@@ -280,6 +284,7 @@ def gradcheck_tape(
             input_output_mask=input_output_mask,
             device=device,
             max_blocks=max_blocks,
+            block_dim=block_dim,
             max_inputs_per_var=max_inputs_per_var,
             max_outputs_per_var=max_outputs_per_var,
             plot_relative_error=plot_relative_error,
@@ -611,13 +616,15 @@ def jacobian(
     input_output_mask: List[Tuple[Union[str, int], Union[str, int]]] = None,
     device: wp.context.Devicelike = None,
     max_blocks=0,
+    block_dim=256,
     max_outputs_per_var=-1,
     plot_jacobians=False,
 ) -> Dict[Tuple[int, int], wp.array]:
     """
     Computes the Jacobians of a Warp kernel launch for the provided selection of differentiable inputs to differentiable outputs.
 
-    The kernel adjoint function is launched with the given inputs and outputs, as well as the provided ``dim`` and ``max_blocks`` arguments (see :func:`warp.launch` for more details).
+    The kernel adjoint function is launched with the given inputs and outputs, as well as the provided ``dim``,
+    ``max_blocks``, and ``block_dim`` arguments (see :func:`warp.launch` for more details).
 
     Note:
         This function only supports Warp kernels whose input arguments precede the output arguments.
@@ -634,6 +641,7 @@ def jacobian(
         input_output_mask: List of tuples specifying the input-output pairs to compute the Jacobian for. Inputs and outputs can be identified either by their integer indices of where they appear in the kernel input/output arguments, or by the respective argument names as strings. If None, computes the Jacobian for all input-output pairs.
         device: The device to launch on (optional)
         max_blocks: The maximum number of CUDA thread blocks to use.
+        block_dim: The number of threads per block.
         max_outputs_per_var: Maximum number of output dimensions over which to evaluate the Jacobians for the input-output pairs. Evaluates all output dimensions if value <= 0.
         plot_jacobians: If True, visualizes the computed Jacobians in a plot (requires ``matplotlib``).
 
@@ -661,7 +669,15 @@ def jacobian(
         device = infer_device(inputs + outputs)
 
     tape = wp.Tape()
-    tape.record_launch(kernel=kernel, dim=dim, max_blocks=max_blocks, inputs=inputs, outputs=outputs, device=device)
+    tape.record_launch(
+        kernel=kernel,
+        dim=dim,
+        inputs=inputs,
+        outputs=outputs,
+        device=device,
+        max_blocks=max_blocks,
+        block_dim=block_dim,
+    )
 
     jacobians = {}
 
@@ -709,6 +725,7 @@ def jacobian_fd(
     input_output_mask: List[Tuple[Union[str, int], Union[str, int]]] = None,
     device: wp.context.Devicelike = None,
     max_blocks=0,
+    block_dim=256,
     max_inputs_per_var=-1,
     eps=1e-4,
     plot_jacobians=False,
@@ -717,7 +734,8 @@ def jacobian_fd(
     Computes the finite-difference Jacobian of a Warp kernel launch for the provided selection of differentiable inputs to differentiable outputs.
     The method uses a central difference scheme to approximate the Jacobian.
 
-    The kernel is launched multiple times in forward-only mode with the given inputs and outputs, as well as the provided ``dim`` and ``max_blocks`` arguments (see :func:`warp.launch` for more details).
+    The kernel is launched multiple times in forward-only mode with the given inputs and outputs, as well as the
+    provided ``dim``, ``max_blocks``, and ``block_dim`` arguments (see :func:`warp.launch` for more details).
 
     Note:
         This function only supports Warp kernels whose input arguments precede the output arguments.
@@ -734,6 +752,7 @@ def jacobian_fd(
         input_output_mask: List of tuples specifying the input-output pairs to compute the Jacobian for. Inputs and outputs can be identified either by their integer indices of where they appear in the kernel input/output arguments, or by the respective argument names as strings. If None, computes the Jacobian for all input-output pairs.
         device: The device to launch on (optional)
         max_blocks: The maximum number of CUDA thread blocks to use.
+        block_dim: The number of threads per block.
         max_inputs_per_var: Maximum number of input dimensions over which to evaluate the Jacobians for the input-output pairs. Evaluates all input dimensions if value <= 0.
         eps: The finite-difference step size.
         plot_jacobians: If True, visualizes the computed Jacobians in a plot (requires ``matplotlib``).
@@ -793,10 +812,26 @@ def jacobian_fd(
             input_num = min(input_num, max_inputs_per_var)
         for i in range(input_num):
             set_element(flat_input, i, -eps, relative=True)
-            wp.launch(kernel, dim=dim, max_blocks=max_blocks, inputs=inputs, outputs=left_outputs, device=device)
+            wp.launch(
+                kernel,
+                dim=dim,
+                inputs=inputs,
+                outputs=left_outputs,
+                device=device,
+                max_blocks=max_blocks,
+                block_dim=block_dim,
+            )
 
             set_element(flat_input, i, 2 * eps, relative=True)
-            wp.launch(kernel, dim=dim, max_blocks=max_blocks, inputs=inputs, outputs=right_outputs, device=device)
+            wp.launch(
+                kernel,
+                dim=dim,
+                inputs=inputs,
+                outputs=right_outputs,
+                device=device,
+                max_blocks=max_blocks,
+                block_dim=block_dim,
+            )
 
             set_element(flat_input, i, -eps, relative=True)
 

warp/build.py

@@ -5,6 +5,7 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import ctypes
 import os
 
 import warp.config
@@ -12,7 +13,7 @@ from warp.thirdparty import appdirs
 
 
 # builds cuda source to PTX or CUBIN using NVRTC (output type determined by output_path extension)
-def build_cuda(cu_path, arch, output_path, config="release", verify_fp=False, fast_math=False):
+def build_cuda(cu_path, arch, output_path, config="release", verify_fp=False, fast_math=False, ltoirs=None):
     with open(cu_path, "rb") as src_file:
         src = src_file.read()
         cu_path = cu_path.encode("utf-8")
@@ -23,8 +24,26 @@ def build_cuda(cu_path, arch, output_path, config="release", verify_fp=False, fa
             warp.context.runtime.llvm.compile_cuda(src, cu_path, inc_path, output_path, False)
 
         else:
+            if ltoirs is None:
+                ltoirs = []
+
+            num_ltoirs = len(ltoirs)
+            arr_lroirs = (ctypes.c_char_p * num_ltoirs)(*ltoirs)
+            arr_lroir_sizes = (ctypes.c_size_t * num_ltoirs)(*[len(l) for l in ltoirs])
             err = warp.context.runtime.core.cuda_compile_program(
-                src, arch, inc_path, config == "debug", warp.config.verbose, verify_fp, fast_math, output_path
+                src,
+                arch,
+                inc_path,
+                0,
+                None,
+                config == "debug",
+                warp.config.verbose,
+                verify_fp,
+                fast_math,
+                output_path,
+                num_ltoirs,
+                arr_lroirs,
+                arr_lroir_sizes,
             )
             if err != 0:
                 raise Exception(f"CUDA kernel build failed with error code {err}")

warp/build_dll.py

@@ -172,6 +172,10 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
                 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
+
         gencode_opts = []
 
         if args.quick:
@@ -223,6 +227,13 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
     # is the library being built with CUDA enabled?
     cuda_enabled = "WP_ENABLE_CUDA=1" if (cu_path is not None) else "WP_ENABLE_CUDA=0"
 
+    if args.libmathdx_path:
+        libmathdx_includes = f' -I"{args.libmathdx_path}/include"'
+        mathdx_enabled = "WP_ENABLE_MATHDX=1"
+    else:
+        libmathdx_includes = ""
+        mathdx_enabled = "WP_ENABLE_MATHDX=0"
+
     if os.name == "nt":
         if args.host_compiler:
             host_linker = os.path.join(os.path.dirname(args.host_compiler), "link.exe")
@@ -244,7 +255,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 "{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 "{cutlass_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"
@@ -273,10 +284,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" --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} -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{cutlass_enabled} {cutlass_includes} -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
 
             elif mode == "release":
-                cuda_cmd = f'"{cuda_home}/bin/nvcc" -O3 {" ".join(nvcc_opts)} -I"{native_dir}" -DNDEBUG -DWP_ENABLE_CUDA=1 -D{cutlass_enabled} {cutlass_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{cutlass_enabled} {cutlass_includes} -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
 
             with ScopedTimer("build_cuda", active=args.verbose):
                 run_cmd(cuda_cmd)
@@ -285,6 +296,9 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
                     f'cudart_static.lib nvrtc_static.lib nvrtc-builtins_static.lib nvptxcompiler_static.lib ws2_32.lib user32.lib /LIBPATH:"{cuda_home}/lib/x64"'
                 )
 
+                if args.libmathdx_path:
+                    linkopts.append(f'nvJitLink_static.lib /LIBPATH:"{args.libmathdx_path}/lib" mathdx_static.lib')
+
         with ScopedTimer("link", active=args.verbose):
             link_cmd = f'"{host_linker}" {" ".join(linkopts + libs)} /out:"{dll_path}"'
             run_cmd(link_cmd)
@@ -300,7 +314,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{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{cutlass_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"
@@ -328,10 +342,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} -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{cutlass_enabled} {cutlass_includes} -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} -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{cutlass_enabled} {cutlass_includes} -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
 
             with ScopedTimer("build_cuda", active=args.verbose):
                 run_cmd(cuda_cmd)
@@ -341,6 +355,9 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
                     f'-L"{cuda_home}/lib64" -lcudart_static -lnvrtc_static -lnvrtc-builtins_static -lnvptxcompiler_static -lpthread -ldl -lrt'
                 )
 
+                if args.libmathdx_path:
+                    ld_inputs.append(f"-lnvJitLink_static -L{args.libmathdx_path}/lib -lmathdx_static")
+
         if sys.platform == "darwin":
             opt_no_undefined = "-Wl,-undefined,error"
             opt_exclude_libs = ""