Merge branch 'sparse_opt' into 'main'

warp.sparse: asynchronous functions, block shape conversions, natural operators

See merge request omniverse/warp!575

CHANGELOG.md

@@ -16,6 +16,13 @@
   anticipate when a class instance is created but not instantiated before garbage collection.
 - Add code-completion support for `wp.config` variables.
 - Remove usage of a static task (thread) index for CPU kernels to address multithreading concerns.
+- The `mask` argument to `wp.sim.eval_fk` now accepts both integer and bool arrays
+- Support for NumPy >= 2.0
+- Fix hashing of replay functions and snippets
+- New `warp.sparse` features:
+  - Sparse matrix allocations (from `bsr_from_triplets`, `bsr_axpy`, etc) can now be captured in CUDA graphs; exact number of non-zeros can be optionally requested asynchronously.
+  - `bsr_assign` now supports changing block shape (including CSR/BSR conversions)
+  - Add Python operator overloads for common sparse matrix operations, e.g `A += 0.5 * B`, `y = x @ C`
 
 ## [1.2.1] - 2024-06-14
 

docs/modules/sparse.rst

@@ -14,6 +14,9 @@ Sparse Matrices
 
 Currently `warp.sparse` supports Block Sparse Row (BSR) matrices, the BSR format can also be used to represent Compressed Sparse Row (CSR) matrices as a special case with a 1x1 block size.
 
+Overloaded Python mathematical operators are supported for sparse matrix addition (`+`), subtraction (`-`), multiplication by a scalar (`*`) and matrix-matrix or matrix-vector multiplication (`@`),
+including in-place variants where possible.
+
 .. automodule:: warp.sparse
     :members:
 

warp/context.py

@@ -2794,39 +2794,38 @@ def __init__(self):
             self.core.volume_get_blind_data_info.restype = ctypes.c_char_p
 
             bsr_matrix_from_triplets_argtypes = [
-                ctypes.c_int,
-                ctypes.c_int,
-                ctypes.c_int,
-                ctypes.c_int,
-                ctypes.c_uint64,
-                ctypes.c_uint64,
-                ctypes.c_uint64,
-                ctypes.c_uint64,
-                ctypes.c_uint64,
-                ctypes.c_uint64,
+                ctypes.c_int,  # rows_per_bock
+                ctypes.c_int,  # cols_per_blocks
+                ctypes.c_int,  # row_count
+                ctypes.c_int,  # tpl_nnz
+                ctypes.POINTER(ctypes.c_int),  # tpl_rows
+                ctypes.POINTER(ctypes.c_int),  # tpl_cols
+                ctypes.c_void_p,  # tpl_values
+                ctypes.c_bool,  # prune_numerical_zeros
+                ctypes.POINTER(ctypes.c_int),  # bsr_offsets
+                ctypes.POINTER(ctypes.c_int),  # bsr_columns
+                ctypes.c_void_p,  # bsr_values
+                ctypes.POINTER(ctypes.c_int),  # bsr_nnz
+                ctypes.c_void_p,  # bsr_nnz_event
             ]
+
             self.core.bsr_matrix_from_triplets_float_host.argtypes = bsr_matrix_from_triplets_argtypes
             self.core.bsr_matrix_from_triplets_double_host.argtypes = bsr_matrix_from_triplets_argtypes
             self.core.bsr_matrix_from_triplets_float_device.argtypes = bsr_matrix_from_triplets_argtypes
             self.core.bsr_matrix_from_triplets_double_device.argtypes = bsr_matrix_from_triplets_argtypes
 
-            self.core.bsr_matrix_from_triplets_float_host.restype = ctypes.c_int
-            self.core.bsr_matrix_from_triplets_double_host.restype = ctypes.c_int
-            self.core.bsr_matrix_from_triplets_float_device.restype = ctypes.c_int
-            self.core.bsr_matrix_from_triplets_double_device.restype = ctypes.c_int
-
             bsr_transpose_argtypes = [
-                ctypes.c_int,
-                ctypes.c_int,
-                ctypes.c_int,
-                ctypes.c_int,
-                ctypes.c_int,
-                ctypes.c_uint64,
-                ctypes.c_uint64,
-                ctypes.c_uint64,
-                ctypes.c_uint64,
-                ctypes.c_uint64,
-                ctypes.c_uint64,
+                ctypes.c_int,  # rows_per_bock
+                ctypes.c_int,  # cols_per_blocks
+                ctypes.c_int,  # row_count
+                ctypes.c_int,  # col count
+                ctypes.c_int,  # nnz
+                ctypes.POINTER(ctypes.c_int),  # transposed_bsr_offsets
+                ctypes.POINTER(ctypes.c_int),  # transposed_bsr_columns
+                ctypes.c_void_p,  # bsr_values
+                ctypes.POINTER(ctypes.c_int),  # transposed_bsr_offsets
+                ctypes.POINTER(ctypes.c_int),  # transposed_bsr_columns
+                ctypes.c_void_p,  # transposed_bsr_values
             ]
             self.core.bsr_transpose_float_host.argtypes = bsr_transpose_argtypes
             self.core.bsr_transpose_double_host.argtypes = bsr_transpose_argtypes