warp-lang 1.3.2__py3-none-manylinux2014_aarch64.whl → 1.4.0__py3-none-manylinux2014_aarch64.whl

warp/__init__.py

@@ -99,11 +99,17 @@ from warp.jax import device_from_jax, device_to_jax
 
 from warp.dlpack import from_dlpack, to_dlpack
 
+from warp.paddle import from_paddle, to_paddle
+from warp.paddle import dtype_from_paddle, dtype_to_paddle
+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.constants import *
 
 from . import builtins
+from warp.builtins import static
 
 import warp.config as config
 

warp/autograd.py

@@ -17,7 +17,7 @@ __all__ = [
     "jacobian_fd",
     "gradcheck",
     "gradcheck_tape",
-    "plot_kernel_jacobians",
+    "jacobian_plot",
 ]
 
 
@@ -182,7 +182,7 @@ def gradcheck(
         else:
             print(FontColors.OKGREEN + f"Gradient check for kernel {function.key} passed" + FontColors.ENDC)
     if plot_relative_error:
-        plot_kernel_jacobians(
+        jacobian_plot(
             relative_error_jacs,
             function,
             inputs,
@@ -190,7 +190,7 @@ def gradcheck(
             title=f"{function.key} kernel Jacobian relative error",
         )
     if plot_absolute_error:
-        plot_kernel_jacobians(
+        jacobian_plot(
             absolute_error_jacs,
             function,
             inputs,
@@ -307,7 +307,7 @@ def infer_device(xs: list):
     return wp.get_preferred_device()
 
 
-def plot_kernel_jacobians(
+def jacobian_plot(
     jacobians: Dict[Tuple[int, int], wp.array],
     kernel: wp.Kernel,
     inputs: Sequence,
@@ -515,6 +515,59 @@ def plot_kernel_jacobians(
     return fig
 
 
+def plot_kernel_jacobians(
+    jacobians: Dict[Tuple[int, int], wp.array],
+    kernel: wp.Kernel,
+    inputs: Sequence,
+    outputs: Sequence,
+    show_plot=True,
+    show_colorbar=True,
+    scale_colors_per_submatrix=False,
+    title: str = None,
+    colormap: str = "coolwarm",
+    log_scale=False,
+):
+    """
+    Visualizes the Jacobians computed by :func:`jacobian` or :func:`jacobian_fd` in a combined image plot.
+    Requires the ``matplotlib`` package to be installed.
+
+    Note:
+        This function is deprecated and will be removed in a future Warp version. Please call :func:`jacobian_plot` instead.
+
+    Args:
+        jacobians: A dictionary of Jacobians, where the keys are tuples of input and output indices, and the values are the Jacobian matrices.
+        kernel: The Warp kernel function, decorated with the ``@wp.kernel`` decorator.
+        inputs: List of input variables.
+        outputs: List of output variables.
+        show_plot: If True, displays the plot via ``plt.show()``.
+        show_colorbar: If True, displays a colorbar next to the plot (or a colorbar next to every submatrix if ).
+        scale_colors_per_submatrix: If True, considers the minimum and maximum of each Jacobian submatrix separately for color scaling. Otherwise, uses the global minimum and maximum of all Jacobians.
+        title: The title of the plot (optional).
+        colormap: The colormap to use for the plot.
+        log_scale: If True, uses a logarithmic scale for the matrix values shown in the image plot.
+
+    Returns:
+        The created Matplotlib figure.
+    """
+    wp.utils.warn(
+        "The function `plot_kernel_jacobians` is deprecated and will be removed in a future Warp version. Please call `jacobian_plot` instead.",
+        DeprecationWarning,
+        stacklevel=3,
+    )
+    return jacobian_plot(
+        jacobians,
+        kernel,
+        inputs,
+        outputs,
+        show_plot=show_plot,
+        show_colorbar=show_colorbar,
+        scale_colors_per_submatrix=scale_colors_per_submatrix,
+        title=title,
+        colormap=colormap,
+        log_scale=log_scale,
+    )
+
+
 def scalarize_array_1d(arr):
     # convert array to 1D array with scalar dtype
     if arr.dtype in wp.types.scalar_types:
@@ -638,7 +691,7 @@ def jacobian(
         jacobians[input_i, output_i] = jacobian
 
     if plot_jacobians:
-        plot_kernel_jacobians(
+        jacobian_plot(
             jacobians,
             kernel,
             inputs,
@@ -753,7 +806,7 @@ def jacobian_fd(
         jacobians[input_i, output_i] = jacobian
 
     if plot_jacobians:
-        plot_kernel_jacobians(
+        jacobian_plot(
             jacobians,
             kernel,
             inputs,

warp/build_dll.py

@@ -32,13 +32,12 @@ def run_cmd(cmd):
         print(cmd)
 
     try:
-        return subprocess.check_output(cmd, shell=True)
+        return subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
     except subprocess.CalledProcessError as e:
-        if e.stdout:
-            print(e.stdout.decode())
-        if e.stderr:
-            print(e.stderr.decode())
-        raise (e)
+        print("Command failed with exit code:", e.returncode)
+        print("Command output was:")
+        print(e.output.decode())
+        raise e
 
 
 # cut-down version of vcvars64.bat that allows using
@@ -157,7 +156,6 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
 
     warp_home_path = pathlib.Path(__file__).parent
     warp_home = warp_home_path.resolve()
-    nanovdb_home = warp_home_path.parent / "_build/host-deps/nanovdb/include"
 
     # output stale, rebuild
     if args.verbose:
@@ -246,7 +244,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}" /I"{nanovdb_home}" {includes} '
+        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} '
 
         if args.mode == "debug":
             cpp_flags += "/Zi /Od /D WP_ENABLE_DEBUG=1"
@@ -275,10 +273,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}" -I"{nanovdb_home}" -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" --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}"'
 
             elif mode == "release":
-                cuda_cmd = f'"{cuda_home}/bin/nvcc" -O3 {" ".join(nvcc_opts)} -I"{native_dir}" -I"{nanovdb_home}" -DNDEBUG -DWP_ENABLE_CUDA=1 -D{cutlass_enabled} {cutlass_includes} -o "{cu_out}" -c "{cu_path}"'
+                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}"'
 
             with ScopedTimer("build_cuda", active=args.verbose):
                 run_cmd(cuda_cmd)

warp/builtins.py

@@ -1533,7 +1533,7 @@ add_builtin(
 
 def spatial_vector_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     if arg_types is None:
-        return spatial_vector(dtype=Float)
+        return vector(length=6, dtype=Float)
 
     dtype = arg_values.get("dtype", None)
 
@@ -2758,6 +2758,7 @@ add_builtin(
     "rand_init",
     input_types={"seed": int},
     value_type=uint32,
+    export=False,
     group="Random",
     doc="Initialize a new random number generator given a user-defined seed. Returns a 32-bit integer representing the RNG state.",
 )
@@ -2766,6 +2767,7 @@ add_builtin(
     "rand_init",
     input_types={"seed": int, "offset": int},
     value_type=uint32,
+    export=False,
     group="Random",
     doc="""Initialize a new random number generator given a user-defined seed and an offset.
 
@@ -2777,6 +2779,7 @@ add_builtin(
     "randi",
     input_types={"state": uint32},
     value_type=int,
+    export=False,
     group="Random",
     doc="Return a random integer in the range [0, 2^32).",
 )
@@ -2784,6 +2787,7 @@ add_builtin(
     "randi",
     input_types={"state": uint32, "low": int, "high": int},
     value_type=int,
+    export=False,
     group="Random",
     doc="Return a random integer between [low, high).",
 )
@@ -2791,6 +2795,7 @@ add_builtin(
     "randf",
     input_types={"state": uint32},
     value_type=float,
+    export=False,
     group="Random",
     doc="Return a random float between [0.0, 1.0).",
 )
@@ -2798,17 +2803,24 @@ add_builtin(
     "randf",
     input_types={"state": uint32, "low": float, "high": float},
     value_type=float,
+    export=False,
     group="Random",
     doc="Return a random float between [low, high).",
 )
 add_builtin(
-    "randn", input_types={"state": uint32}, value_type=float, group="Random", doc="Sample a normal distribution."
+    "randn",
+    input_types={"state": uint32},
+    value_type=float,
+    export=False,
+    group="Random",
+    doc="Sample a normal distribution.",
 )
 
 add_builtin(
     "sample_cdf",
     input_types={"state": uint32, "cdf": array(dtype=float)},
     value_type=int,
+    export=False,
     group="Random",
     doc="Inverse-transform sample a cumulative distribution function.",
 )
@@ -2816,6 +2828,7 @@ add_builtin(
     "sample_triangle",
     input_types={"state": uint32},
     value_type=vec2,
+    export=False,
     group="Random",
     doc="Uniformly sample a triangle. Returns sample barycentric coordinates.",
 )
@@ -2823,6 +2836,7 @@ add_builtin(
     "sample_unit_ring",
     input_types={"state": uint32},
     value_type=vec2,
+    export=False,
     group="Random",
     doc="Uniformly sample a ring in the xy plane.",
 )
@@ -2830,6 +2844,7 @@ add_builtin(
     "sample_unit_disk",
     input_types={"state": uint32},
     value_type=vec2,
+    export=False,
     group="Random",
     doc="Uniformly sample a disk in the xy plane.",
 )
@@ -2837,6 +2852,7 @@ add_builtin(
     "sample_unit_sphere_surface",
     input_types={"state": uint32},
     value_type=vec3,
+    export=False,
     group="Random",
     doc="Uniformly sample a unit sphere surface.",
 )
@@ -2844,6 +2860,7 @@ add_builtin(
     "sample_unit_sphere",
     input_types={"state": uint32},
     value_type=vec3,
+    export=False,
     group="Random",
     doc="Uniformly sample a unit sphere.",
 )
@@ -2851,6 +2868,7 @@ add_builtin(
     "sample_unit_hemisphere_surface",
     input_types={"state": uint32},
     value_type=vec3,
+    export=False,
     group="Random",
     doc="Uniformly sample a unit hemisphere surface.",
 )
@@ -2858,6 +2876,7 @@ add_builtin(
     "sample_unit_hemisphere",
     input_types={"state": uint32},
     value_type=vec3,
+    export=False,
     group="Random",
     doc="Uniformly sample a unit hemisphere.",
 )
@@ -2865,6 +2884,7 @@ add_builtin(
     "sample_unit_square",
     input_types={"state": uint32},
     value_type=vec2,
+    export=False,
     group="Random",
     doc="Uniformly sample a unit square.",
 )
@@ -2872,6 +2892,7 @@ add_builtin(
     "sample_unit_cube",
     input_types={"state": uint32},
     value_type=vec3,
+    export=False,
     group="Random",
     doc="Uniformly sample a unit cube.",
 )
@@ -2880,6 +2901,7 @@ add_builtin(
     "poisson",
     input_types={"state": uint32, "lam": float},
     value_type=uint32,
+    export=False,
     group="Random",
     doc="""Generate a random sample from a Poisson distribution.
 
@@ -3639,6 +3661,30 @@ add_builtin(
 )
 
 
+def vector_assign_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
+    vec_type = arg_types["a"]
+    return vec_type
+
+
+# implements vector[index] = value
+add_builtin(
+    "assign",
+    input_types={"a": vector(length=Any, dtype=Scalar), "i": int, "value": Scalar},
+    value_func=vector_assign_value_func,
+    hidden=True,
+    group="Utility",
+)
+
+# implements quaternion[index] = value
+add_builtin(
+    "assign",
+    input_types={"a": quaternion(dtype=Scalar), "i": int, "value": Scalar},
+    value_func=vector_assign_value_func,
+    hidden=True,
+    group="Utility",
+)
+
+
 def matrix_index_row_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     mat_type = arg_types["a"]
     row_type = mat_type._wp_row_type_
@@ -3646,7 +3692,7 @@ def matrix_index_row_value_func(arg_types: Mapping[str, type], arg_values: Mappi
     return Reference(row_type)
 
 
-# implements matrix[i] = row
+# implements &matrix[i] = row
 add_builtin(
     "index",
     input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int},
@@ -3664,7 +3710,7 @@ def matrix_index_value_func(arg_types: Mapping[str, type], arg_values: Mapping[s
     return Reference(value_type)
 
 
-# implements matrix[i,j] = scalar
+# implements &matrix[i,j] = scalar
 add_builtin(
     "index",
     input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "j": int},
@@ -3674,6 +3720,41 @@ add_builtin(
     skip_replay=True,
 )
 
+
+def matrix_assign_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
+    mat_type = arg_types["a"]
+    return mat_type
+
+
+def matrix_vector_sametype(arg_types: Mapping[str, Any]):
+    mat_size = arg_types["a"]._shape_[0]
+    vec_size = arg_types["value"]._length_
+    mat_type = arg_types["a"]._type_
+    vec_type = arg_types["value"]._type_
+    return mat_size == vec_size and mat_type == vec_type
+
+
+# implements matrix[i,j] = scalar
+add_builtin(
+    "assign",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "j": int, "value": Scalar},
+    value_func=matrix_assign_value_func,
+    hidden=True,
+    group="Utility",
+)
+
+
+# implements matrix[i] = vector
+add_builtin(
+    "assign",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "value": vector(length=Any, dtype=Scalar)},
+    constraint=matrix_vector_sametype,
+    value_func=matrix_assign_value_func,
+    hidden=True,
+    group="Utility",
+)
+
+
 for t in scalar_types + vector_types + (bool,):
     if "vec" in t.__name__ or "mat" in t.__name__:
         continue
@@ -4107,6 +4188,14 @@ add_builtin(
     doc="Modulo operation using truncated division.",
     group="Operators",
 )
+add_builtin(
+    "mod",
+    input_types={"a": vector(length=Any, dtype=Scalar), "b": vector(length=Any, dtype=Scalar)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(Scalar),
+    doc="Modulo operation using truncated division.",
+    group="Operators",
+)
 
 add_builtin(
     "div",
@@ -4218,3 +4307,36 @@ for t in int_types:
 
 
 add_builtin("unot", input_types={"a": array(dtype=Any)}, value_type=builtins.bool, doc="", group="Operators")
+
+# ---------------------------------
+# Code Generation
+
+add_builtin(
+    "static",
+    input_types={"expr": Any},
+    value_type=Any,
+    doc="""Evaluates a static Python expression and replaces it with its result.
+
+    See the `codegen.html#static-expressions <section on code generation>`_ for more details.
+
+    Note:
+        The inner expression must only reference variables that are available from the current scope where the Warp kernel or function containing the expression is defined,
+        which includes constant variables and variables captured in the current closure in which the function or kernel is implemented.
+        The return type of the expression must be either a Warp function, a string, or a type that is supported inside Warp kernels and functions
+        (excluding Warp arrays since they cannot be created in a Warp kernel at the moment).""",
+    group="Code Generation",
+)
+
+
+def static(expr):
+    """
+    Evaluates a static expression and replaces the expression with its result.
+
+    Args:
+        expr: A Python expression to evaluate. Must return a non-null value which must be either a Warp function, a string, or a type that is supported inside Warp kernels and functions (excluding Warp arrays since they cannot be created in a Warp kernel at the moment).
+
+    Note:
+        The inner expression must only reference variables that are available from the current scope where the Warp kernel or function containing the expression is defined,
+        which includes constant variables and variables captured in the current closure in which the function or kernel is implemented.
+    """
+    return expr