warp-lang 1.7.2__py3-none-win_amd64.whl → 1.8.0__py3-none-win_amd64.whl

warp/native/warp.h

@@ -19,6 +19,7 @@
 
 // defines all crt + builtin types
 #include "builtin.h"
+#include <cstdint>
 
 #define WP_CURRENT_STREAM ((void*)0xffffffffffffffff)
 
@@ -177,95 +178,56 @@ extern "C"
     WP_API void runlength_encode_int_host(uint64_t values, uint64_t run_values, uint64_t run_lengths, uint64_t run_count, int n);
     WP_API void runlength_encode_int_device(uint64_t values, uint64_t run_values, uint64_t run_lengths, uint64_t run_count, int n);
 
-    WP_API void bsr_matrix_from_triplets_float_host(
-        int rows_per_block,
-        int cols_per_block,
+    WP_API void bsr_matrix_from_triplets_host(
+        int block_size,
+        int scalar_size_in_bytes,
         int row_count,
-        int tpl_nnz,
-        int* tpl_rows,
-        int* tpl_columns,
-        void* tpl_values,
-        bool prune_numerical_zeros,
-        bool masked,
+        int col_count,
+        int tpl_nnz_upper_bound,
+        const int* tpl_nnz,
+        const int* tpl_rows,
+        const int* tpl_columns,
+        const void* tpl_values,
+        const uint64_t scalar_zero_mask,
+        bool masked_topology,
+        int* summed_block_offsets,
+        int* summed_block_indices,
         int* bsr_offsets,
         int* bsr_columns,
-        void* bsr_values,
         int* bsr_nnz,
         void* bsr_nnz_event);
-    WP_API void bsr_matrix_from_triplets_double_host(
-        int rows_per_block,
-        int cols_per_block,
+    WP_API void bsr_matrix_from_triplets_device(
+        int block_size,
+        int scalar_size_in_bytes,
         int row_count,
-        int tpl_nnz,
-        int* tpl_rows,
-        int* tpl_columns,
-        void* tpl_values,
-        bool prune_numerical_zeros,
-        bool masked,
+        int col_count,
+        int tpl_nnz_upper_bound,
+        const int* tpl_nnz,
+        const int* tpl_rows,
+        const int* tpl_columns,
+        const void* tpl_values,
+        const uint64_t scalar_zero_mask,
+        bool masked_topology,
+        int* summed_block_offsets,
+        int* summed_block_indices,
         int* bsr_offsets,
         int* bsr_columns,
-        void* bsr_values,
-        int* bsr_nnz,
-        void* bsr_nnz_event);
-    WP_API void bsr_matrix_from_triplets_float_device(
-        int rows_per_block,
-        int cols_per_block,
-        int row_count,
-        int tpl_nnz,
-        int* tpl_rows,
-        int* tpl_columns,
-        void* tpl_values,
-        bool prune_numerical_zeros,
-        bool masked,
-        int* bsr_offsets,
-        int* bsr_columns,
-        void* bsr_values,
-        int* bsr_nnz,
-        void* bsr_nnz_event);
-    WP_API void bsr_matrix_from_triplets_double_device(
-        int rows_per_block,
-        int cols_per_block,
-        int row_count,
-        int tpl_nnz,
-        int* tpl_rows,
-        int* tpl_columns,
-        void* tpl_values,
-        bool prune_numerical_zeros,
-        bool masked,
-        int* bsr_offsets,
-        int* bsr_columns,
-        void* bsr_values,
         int* bsr_nnz,
         void* bsr_nnz_event);
 
-    WP_API void bsr_transpose_float_host(int rows_per_block, int cols_per_block,
-        int row_count, int col_count, int nnz,
-        int* bsr_offsets, int* bsr_columns,
-        void* bsr_values,
-        int* transposed_bsr_offsets,
-        int* transposed_bsr_columns,
-        void* transposed_bsr_values);
-    WP_API void bsr_transpose_double_host(int rows_per_block, int cols_per_block,
+    WP_API void bsr_transpose_host(
         int row_count, int col_count, int nnz,
-        int* bsr_offsets, int* bsr_columns,
-        void* bsr_values,
+        const int* bsr_offsets, const int* bsr_columns,
         int* transposed_bsr_offsets,
         int* transposed_bsr_columns,
-        void* transposed_bsr_values);
-    WP_API void bsr_transpose_float_device(int rows_per_block, int cols_per_block,
+        int* src_block_indices
+        );
+    WP_API void bsr_transpose_device(
         int row_count, int col_count, int nnz,
-        int* bsr_offsets, int* bsr_columns,
-        void* bsr_values,
+        const int* bsr_offsets, const int* bsr_columns,
         int* transposed_bsr_offsets,
         int* transposed_bsr_columns,
-        void* transposed_bsr_values);
-    WP_API void bsr_transpose_double_device(int rows_per_block, int cols_per_block,
-        int row_count, int col_count, int nnz,
-        int* bsr_offsets, int* bsr_columns,
-        void* bsr_values,
-        int* transposed_bsr_offsets,
-        int* transposed_bsr_columns,
-        void* transposed_bsr_values);
+        int* src_block_indices);
 
 
     WP_API int cuda_driver_version();   // CUDA driver version
@@ -279,6 +241,7 @@ extern "C"
     WP_API void* cuda_device_get_primary_context(int ordinal);
     WP_API const char* cuda_device_get_name(int ordinal);
     WP_API int cuda_device_get_arch(int ordinal);
+    WP_API int cuda_device_get_sm_count(int ordinal);
     WP_API void cuda_device_get_uuid(int ordinal, char uuid[16]);
     WP_API int cuda_device_get_pci_domain_id(int ordinal);
     WP_API int cuda_device_get_pci_bus_id(int ordinal);
@@ -345,13 +308,23 @@ extern "C"
 
     WP_API bool cuda_graph_begin_capture(void* context, void* stream, int external);
     WP_API bool cuda_graph_end_capture(void* context, void* stream, void** graph_ret);
+    WP_API bool cuda_graph_create_exec(void* context, void* graph, void** graph_exec_ret);
     WP_API bool cuda_graph_launch(void* graph, void* stream);
     WP_API bool cuda_graph_destroy(void* context, void* graph);
+    WP_API bool cuda_graph_exec_destroy(void* context, void* graph_exec);
+    WP_API bool capture_debug_dot_print(void* graph, const char *path, uint32_t flags);
+
+    WP_API bool cuda_graph_insert_if_else(void* context, void* stream, int* condition, void** if_graph_ret, void** else_graph_ret);
+    WP_API bool cuda_graph_insert_while(void* context, void* stream, int* condition, void** body_graph_ret, uint64_t* handle_ret);
+    WP_API bool cuda_graph_set_condition(void* context, void* stream, int* condition, uint64_t handle);
+    WP_API bool cuda_graph_pause_capture(void* context, void* stream, void** graph_ret);
+    WP_API bool cuda_graph_resume_capture(void* context, void* stream, void* graph);
+    WP_API bool cuda_graph_insert_child_graph(void* context, void* stream, void* child_graph);
 
-    WP_API size_t cuda_compile_program(const char* cuda_src, const char* program_name, int arch, const char* include_dir, int num_cuda_include_dirs, const char** cuda_include_dirs, bool debug, bool verbose, bool verify_fp, bool fast_math, bool fuse_fp, bool lineinfo, const char* output_path, size_t num_ltoirs, char** ltoirs, size_t* ltoir_sizes, int* ltoir_input_types);
+    WP_API size_t cuda_compile_program(const char* cuda_src, const char* program_name, int arch, const char* include_dir, int num_cuda_include_dirs, const char** cuda_include_dirs, bool debug, bool verbose, bool verify_fp, bool fast_math, bool fuse_fp, bool lineinfo, bool compile_time_trace, const char* output_path, size_t num_ltoirs, char** ltoirs, size_t* ltoir_sizes, int* ltoir_input_types);
     WP_API bool cuda_compile_fft(const char* ltoir_output_path, const char* symbol_name, int num_include_dirs, const char** include_dirs, const char* mathdx_include_dir, int arch, int size, int elements_per_thread, int direction, int precision, int* shared_memory_size);
     WP_API bool cuda_compile_dot(const char* ltoir_output_path, const char* symbol_name, int num_include_dirs, const char** include_dirs, const char* mathdx_include_dir, int arch, int M, int N, int K, int precision_A, int precision_B, int precision_C, int type, int arrangement_A, int arrangement_B, int arrangement_C, int num_threads);
-    WP_API bool cuda_compile_solver(const char* fatbin_output_path, const char* ltoir_output_path, const char* symbol_name, int num_include_dirs, const char** include_dirs, const char* mathdx_include_dir, int arch, int M, int N, int function, int precision, int fill_mode, int num_threads);
+    WP_API bool cuda_compile_solver(const char* fatbin_output_path, const char* ltoir_output_path, const char* symbol_name, int num_include_dirs, const char** include_dirs, const char* mathdx_include_dir, int arch, int M, int N, int NRHS, int function, int side, int diag, int precision, int arrangement_A, int arrangement_B, int fill_mode, int num_threads);
 
     WP_API void* cuda_load_module(void* context, const char* ptx);
     WP_API void cuda_unload_module(void* context, void* module);

warp/optim/linear.py

@@ -39,7 +39,11 @@ class LinearOperator:
     .. code-block:: python
 
         def matvec(x: wp.array, y: wp.array, z: wp.array, alpha: Scalar, beta: Scalar):
-            '''Performs the operation z = alpha * x + beta * y'''
+            '''Perform a generalized matrix-vector product.
+
+            This function computes the operation z = alpha * (A @ x) + beta * y, where 'A'
+            is the linear operator represented by this class.
+            '''
             ...
 
     For performance reasons, by default the iterative linear solvers in this module will try to capture the calls

warp/paddle.py

@@ -16,7 +16,7 @@
 from __future__ import annotations
 
 import ctypes
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 import numpy
 
@@ -29,7 +29,7 @@ if TYPE_CHECKING:
 
 
 # return the warp device corresponding to a paddle device
-def device_from_paddle(paddle_device: Union[Place, CPUPlace, CUDAPinnedPlace, CUDAPlace, str]) -> warp.context.Device:
+def device_from_paddle(paddle_device: Place | CPUPlace | CUDAPinnedPlace | CUDAPlace | str) -> warp.context.Device:
     """Return the Warp device corresponding to a Paddle device.
 
     Args:
@@ -218,9 +218,9 @@ dtype_is_compatible.compatible_sets = None
 # wrap a paddle tensor to a wp array, data is not copied
 def from_paddle(
     t: paddle.Tensor,
-    dtype: Optional[paddle.dtype] = None,
-    requires_grad: Optional[bool] = None,
-    grad: Optional[paddle.Tensor] = None,
+    dtype: paddle.dtype | None = None,
+    requires_grad: bool | None = None,
+    grad: paddle.Tensor | None = None,
     return_ctype: bool = False,
 ) -> warp.array:
     """Convert a Paddle tensor to a Warp array without copying the data.
@@ -339,9 +339,8 @@ def from_paddle(
         return a
 
 
-def to_paddle(a: warp.array, requires_grad: bool = None) -> paddle.Tensor:
-    """
-    Convert a Warp array to a Paddle tensor without copying the data.
+def to_paddle(a: warp.array, requires_grad: bool | None = None) -> paddle.Tensor:
+    """Convert a Warp array to a Paddle tensor without copying the data.
 
     Args:
         a (warp.array): The Warp array to convert.

warp/render/render_opengl.py

@@ -691,7 +691,7 @@ class ShapeInstancer:
             cls.gl = gl
 
     def __new__(cls, *args, **kwargs):
-        instance = super(ShapeInstancer, cls).__new__(cls)
+        instance = super().__new__(cls)
         instance.instance_transform_gl_buffer = None
         instance.vao = None
         return instance
@@ -1858,7 +1858,7 @@ class OpenGLRenderer:
         self._scaling = scaling
         self.update_model_matrix()
 
-    def begin_frame(self, t: float = None):
+    def begin_frame(self, t: float | None = None):
         self._last_begin_frame_time = time.time()
         self.time = t or self.clock_time
 
@@ -2295,9 +2295,9 @@ Instances: {len(self._instances)}"""
         name: str,
         shape: int,
         body,
-        pos,
-        rot,
-        scale=(1.0, 1.0, 1.0),
+        pos: tuple,
+        rot: tuple,
+        scale: tuple = (1.0, 1.0, 1.0),
         color1=None,
         color2=None,
         custom_index: int = -1,
@@ -2720,10 +2720,11 @@ Instances: {len(self._instances)}"""
         length: float,
         color: tuple = (1.0, 1.0, 1.0),
         color2=None,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
         u_scaling=1.0,
         v_scaling=1.0,
+        visible: bool = True,
     ):
         """Add a plane for visualization
 
@@ -2811,9 +2812,10 @@ Instances: {len(self._instances)}"""
         pos: tuple,
         rot: tuple,
         radius: float,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
-        color=None,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
         """Add a sphere for visualization
 
@@ -2843,10 +2845,11 @@ Instances: {len(self._instances)}"""
         rot: tuple,
         radius: float,
         half_height: float,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
         up_axis: int = 1,
-        color: tuple = None,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
         """Add a capsule for visualization
 
@@ -2878,10 +2881,11 @@ Instances: {len(self._instances)}"""
         rot: tuple,
         radius: float,
         half_height: float,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
         up_axis: int = 1,
-        color: tuple = None,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
         """Add a cylinder for visualization
 
@@ -2913,10 +2917,11 @@ Instances: {len(self._instances)}"""
         rot: tuple,
         radius: float,
         half_height: float,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
         up_axis: int = 1,
-        color: tuple = None,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
         """Add a cone for visualization
 
@@ -2947,9 +2952,10 @@ Instances: {len(self._instances)}"""
         pos: tuple,
         rot: tuple,
         extents: tuple,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
-        color: tuple = None,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
         """Add a box for visualization
 
@@ -2982,9 +2988,10 @@ Instances: {len(self._instances)}"""
         rot=(0.0, 0.0, 0.0, 1.0),
         scale=(1.0, 1.0, 1.0),
         update_topology=False,
-        parent_body: str = None,
+        parent_body: str | None = None,
         is_template: bool = False,
         smooth_shading: bool = True,
+        visible: bool = True,
     ):
         """Add a mesh for visualization
 
@@ -3094,12 +3101,13 @@ Instances: {len(self._instances)}"""
         rot: tuple,
         base_radius: float,
         base_height: float,
-        cap_radius: float = None,
-        cap_height: float = None,
-        parent_body: str = None,
+        cap_radius: float | None = None,
+        cap_height: float | None = None,
+        parent_body: str | None = None,
         is_template: bool = False,
         up_axis: int = 1,
-        color: tuple[float, float, float] = None,
+        color: tuple[float, float, float] | None = None,
+        visible: bool = True,
     ):
         """Add a arrow for visualization
 
@@ -3127,10 +3135,16 @@ Instances: {len(self._instances)}"""
             self.add_shape_instance(name, shape, body, pos, rot, color1=color, color2=color)
         return shape
 
-    def render_ref(self, name: str, path: str, pos: tuple, rot: tuple, scale: tuple, color: tuple = None):
-        """
-        Create a reference (instance) with the given name to the given path.
-        """
+    def render_ref(
+        self,
+        name: str,
+        path: str,
+        pos: tuple,
+        rot: tuple,
+        scale: tuple,
+        color: tuple[float, float, float] | None = None,
+    ):
+        """Create a reference (instance) with the given name to the given path."""
 
         if path in self._instances:
             _, body, shape, _, original_scale, color1, color2 = self._instances[path]
@@ -3141,7 +3155,7 @@ Instances: {len(self._instances)}"""
 
         raise Exception("Cannot create reference to path: " + path)
 
-    def render_points(self, name: str, points, radius, colors=None):
+    def render_points(self, name: str, points, radius, colors=None, as_spheres: bool = True, visible: bool = True):
         """Add a set of points
 
         Args:
@@ -3199,7 +3213,7 @@ Instances: {len(self._instances)}"""
         if name not in self._shape_instancers:
             instancer = ShapeInstancer(self._shape_shader, self._device)
             vertices, indices = self._create_capsule_mesh(radius, 0.5)
-            if color is None or isinstance(color, list) and len(color) > 0 and isinstance(color[0], list):
+            if color is None or (isinstance(color, list) and len(color) > 0 and isinstance(color[0], list)):
                 color = tab10_color_map(len(self._shape_geo_hash))
             instancer.register_shape(vertices, indices, color, color)
             instancer.allocate_instances(np.zeros((len(lines), 3)))
@@ -3220,7 +3234,15 @@ Instances: {len(self._instances)}"""
                 device=self._device,
             )
 
-    def render_line_list(self, name: str, vertices, indices, color: tuple = None, radius: float = 0.01):
+    def render_line_list(
+        self,
+        name: str,
+        vertices,
+        indices,
+        color: tuple[float, float, float] | None = None,
+        radius: float = 0.01,
+        visible: bool = True,
+    ):
         """Add a line list as a set of capsules
 
         Args:
@@ -3235,7 +3257,14 @@ Instances: {len(self._instances)}"""
         lines = np.array(lines)
         self._render_lines(name, lines, color, radius)
 
-    def render_line_strip(self, name: str, vertices, color: tuple = None, radius: float = 0.01):
+    def render_line_strip(
+        self,
+        name: str,
+        vertices,
+        color: tuple[float, float, float] | None = None,
+        radius: float = 0.01,
+        visible: bool = True,
+    ):
         """Add a line strip as a set of capsules
 
         Args: