Add multi-RHS Kamino FK velocity solves (#3916)

Signed-off-by: hanxudong159 <42986190+han-xudong@users.noreply.github.com>
Co-authored-by: Guirec-Maloisel <25688871+Guirec-Maloisel@users.noreply.github.com>

Author: han-xudong

newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked_semi_sparse.py

@@ -76,15 +76,16 @@ def compute_inverse_ordering(ordering):
 
 
 @wp.kernel
-def reorder_rows_kernel(
+def reorder_rows_cols_kernel(
     src: wp.array3d[float],
     dst: wp.array3d[float],
     ordering: wp.array2d[int],
     n_rows_arr: wp.array[int],
     n_cols_arr: wp.array[int],
     batch_mask: wp.array[wp.bool],
 ):
-    batch_id, i, j = wp.tid()  # 2D launch: (n_rows, n_cols)
+    """Apply a permutation to each matrix along its rows and columns."""
+    batch_id, i, j = wp.tid()
     n_rows = n_rows_arr[batch_id]
     n_cols = n_cols_arr[batch_id]
     if i < n_rows and j < n_cols and batch_mask[batch_id]:
@@ -94,19 +95,19 @@ def reorder_rows_kernel(
 
 
 @wp.kernel
-def reorder_rows_kernel_col_vector(
+def reorder_rows_kernel(
     src: wp.array3d[float],
     dst: wp.array3d[float],
     ordering: wp.array2d[int],
     n_rows_arr: wp.array[int],
     batch_mask: wp.array[wp.bool],
 ):
-    batch_id, i = wp.tid()
+    """Apply a row permutation to each matrix of RHS column vectors."""
+    batch_id, i, rhs_id = wp.tid()
     n_rows = n_rows_arr[batch_id]
     if i < n_rows and batch_mask[batch_id]:
         src_row = ordering[batch_id, i]
-        # For column vectors (2d arrays with shape (n, 1)), just copy columns directly
-        dst[batch_id, i, 0] = src[batch_id, src_row, 0]
+        dst[batch_id, i, rhs_id] = src[batch_id, src_row, rhs_id]
 
 
 def to_binary_matrix(M):
@@ -320,7 +321,7 @@ def blocked_cholesky_solve_kernel(
         Uses forward/backward substitution with block size optimization.
         """
 
-        batch_id, _tid_block = wp.tid()
+        batch_id, rhs_id, _tid_block = wp.tid()
 
         if not batch_mask[batch_id]:
             return
@@ -342,8 +343,7 @@ def blocked_cholesky_solve_kernel(
             if L_tile_pattern[tile_i, tile_i] == 0:
                 continue
 
-            i_end = i + block_size
-            rhs_tile = wp.tile_load(b, shape=(block_size, 1), offset=(i, 0))
+            rhs_tile = wp.tile_load(b, shape=(block_size, 1), offset=(i, rhs_id))
             # Hoist the diagonal load above the j loop so the shared-memory fetch can
             # overlap with the gemm pipeline below.
             L_diag = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i))
@@ -354,11 +354,11 @@ def blocked_cholesky_solve_kernel(
                     if L_tile_pattern[tile_i, tile_j] == 0:
                         continue
                     L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j))
-                    y_block = wp.tile_load(y, shape=(block_size, 1), offset=(j, 0))
+                    y_block = wp.tile_load(y, shape=(block_size, 1), offset=(j, rhs_id))
                     wp.tile_matmul(L_block, y_block, rhs_tile, alpha=-1.0)
                     # equivalent to rhs_tile -= L_block * y_block, without intermediate allocation
             wp.tile_lower_solve_inplace(L_diag, rhs_tile)  # In-place solve to avoid extra allocation
-            wp.tile_store(y, rhs_tile, offset=(i, 0))
+            wp.tile_store(y, rhs_tile, offset=(i, rhs_id))
 
         # Backward substitution: solve L^T x = y
         for i in range(n - block_size, -1, -block_size):
@@ -369,7 +369,7 @@ def blocked_cholesky_solve_kernel(
 
             i_start = i
             i_end = i_start + block_size
-            rhs_tile = wp.tile_load(y, shape=(block_size, 1), offset=(i_start, 0))
+            rhs_tile = wp.tile_load(y, shape=(block_size, 1), offset=(i_start, rhs_id))
             # Hoist the diagonal load above the j loop (see forward sub above).
             L_diag = wp.tile_load(L, shape=(block_size, block_size), offset=(i_start, i_start))
             if i_end < n:
@@ -379,15 +379,15 @@ def blocked_cholesky_solve_kernel(
                     if L_tile_pattern[tile_j, tile_i] == 0:
                         continue
                     L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i_start))
-                    x_tile = wp.tile_load(x, shape=(block_size, 1), offset=(j, 0))
+                    x_tile = wp.tile_load(x, shape=(block_size, 1), offset=(j, rhs_id))
                     if wp.static(HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE):
                         wp.tile_matmul_left_transpose_update(rhs_tile, L_tile, x_tile, alpha=-1.0)
                     else:
                         L_T_tile = wp.tile_transpose(L_tile)
                         wp.tile_matmul(L_T_tile, x_tile, rhs_tile, alpha=-1.0)
                     # equivalent to rhs_tile -= L_T_tile * x_tile, without intermediate allocation
             wp.tile_upper_solve_inplace(wp.tile_transpose(L_diag), rhs_tile)  # In-place solve to avoid extra allocation
-            wp.tile_store(x, rhs_tile, offset=(i_start, 0))
+            wp.tile_store(x, rhs_tile, offset=(i_start, rhs_id))
 
     return blocked_cholesky_solve_kernel
 
@@ -399,10 +399,33 @@ class SemiSparseBlockCholeskySolverBatched:
     to skip unnecessary computation. Handles multiple systems in parallel with fixed matrix size.
     """
 
-    def __init__(self, num_batches: int, max_num_equations: int, block_size=16, device="cuda", enable_reordering=True):
+    def __init__(
+        self,
+        num_batches: int,
+        max_num_equations: int,
+        block_size: int = 16,
+        device: wp.DeviceLike = "cuda",
+        enable_reordering: bool = True,
+        rhs_sizes: list[int] | None = None,
+    ):
+        """
+        Initialize the solver given batch and matrix dimensions.
+
+        Args:
+            num_batches: Batch dimension, i.e. number of systems to solve in parallel.
+            max_num_equations: Maximal matrix size (i.e. n for a n x n matrix), among all systems.
+            block_size: Block size for the blocked Cholesky factorization.
+            device: Device on which to allocate the solver's data.
+            enable_reordering: Whether to enable Cuthill-McKee ordering as part of the factorization.
+            rhs_sizes: List of right-hand-side sizes for which to enable a solve (i.e. k for a n x k RHS).
+                Defaults to [1] (i.e. single-column rhs) if not provided.
+                More rhs sizes can be requested after construction with :meth:`request_rhs_size()`.
+                Note that only uniform RHS sizes (across the batch dimension) are supported currently.
+        """
         self.num_batches = num_batches
         self.max_num_equations = max_num_equations
         self.device = device
+        self.enable_reordering = enable_reordering
 
         self.num_threads_per_block_factorize = 128
         self.num_threads_per_block_solve = 64
@@ -412,34 +435,46 @@ def __init__(self, num_batches: int, max_num_equations: int, block_size=16, devi
         self.cholesky_kernel = create_blocked_cholesky_kernel(block_size)
         self.solve_kernel = create_blocked_cholesky_solve_kernel(block_size)
 
-        # Allocate workspace arrays for factorization and solve
         # Compute padded size rounded up to next multiple of block size
         self.padded_num_equations = (
             (self.max_num_equations + self.block_size - 1) // self.block_size
         ) * self.block_size
 
+        # Allocate workspace arrays for factorization and solve
         self.A_swizzled = wp.zeros(
             shape=(num_batches, self.padded_num_equations, self.padded_num_equations), dtype=float, device=self.device
         )
         self.L = wp.zeros(
             shape=(num_batches, self.padded_num_equations, self.padded_num_equations), dtype=float, device=self.device
         )
-        self.y = wp.zeros(
-            shape=(num_batches, self.padded_num_equations, 1), dtype=float, device=self.device
-        )  # temp memory
-        self.result_swizzled = wp.zeros(
-            shape=(num_batches, self.padded_num_equations, 1), dtype=float, device=self.device
-        )  # temp memory
-        self.rhs_swizzled = wp.zeros(
-            shape=(num_batches, self.padded_num_equations, 1), dtype=float, device=self.device
-        )  # temp memory
-
         self.num_tiles = (self.padded_num_equations + self.block_size - 1) // self.block_size
         self.L_tile_pattern = wp.zeros(
             shape=(num_batches, self.num_tiles, self.num_tiles), dtype=int, device=self.device
         )
 
-        self.enable_reordering = enable_reordering
+        # Allocate temporary buffers for requested RHS sizes
+        self.temp_buffers: dict[int, tuple[wp.array3d[float], wp.array3d[float], wp.array3d[float]]] = {}
+        """ Temporary buffers: (rhs_swizzled, result_swizzled, y) for each rhs size """
+        rhs_sizes = [1] if rhs_sizes is None else rhs_sizes
+        for rhs_size in rhs_sizes:
+            self.request_rhs_size(rhs_size)
+
+    def request_rhs_size(self, rhs_size: int) -> None:
+        """
+        Preallocate the necessary internal buffers for a solve with specified RHS size.
+
+        Must be called before the first call to :meth:`solve()` with this specific RHS size, if
+        that size was not provided upon construction as part of ``rhs_sizes``.
+        """
+        if rhs_size < 1:
+            raise ValueError("rhs_size must be positive")
+        if rhs_size not in self.temp_buffers:
+            shape = (self.num_batches, self.padded_num_equations, rhs_size)
+            self.temp_buffers[rhs_size] = (
+                wp.zeros(shape=shape, dtype=float, device=self.device),
+                wp.zeros(shape=shape, dtype=float, device=self.device),
+                wp.zeros(shape=shape, dtype=float, device=self.device),
+            )
 
     def capture_sparsity_pattern(
         self,
@@ -539,7 +574,7 @@ def factorize(
         # Reorder A and store in self.A_reordered
         if self.enable_reordering:
             wp.launch(
-                reorder_rows_kernel,
+                reorder_rows_cols_kernel,
                 dim=[self.num_batches, self.max_num_equations, self.max_num_equations],
                 inputs=[
                     A,
@@ -566,43 +601,77 @@ def solve(
         rhs: wp.array3d[float],
         result: wp.array3d[float],
         batch_mask: wp.array[wp.bool],
-    ):
+    ) -> None:
         """
         Solves A x = b given the Cholesky factor L (A = L L^T) using
         blocked forward and backward substitution.
 
         Args:
-            rhs: Input right-hand-side matrices of shape (batch_size, n, p).
-            result: Output solution matrices of shape (batch_size, n, p).
+            rhs: Input right-hand-side matrices of shape (batch_size, n, rhs_size).
+                Assumes rhs_size was requested beforehand, upon construction or with :meth:`request_rhs_size`.
+            result: Output solution matrices of shape (batch_size, n, rhs_size).
             batch_mask: Boolean flag for each matrix in the batch, indicating whether to process it (False = skip)
+
+        Raises:
+            ValueError: If shapes are incompatible or ``rhs_size`` was not requested.
         """
+        if rhs.shape != result.shape:
+            raise ValueError("rhs and result must have identical shapes")
+        if rhs.ndim != 3 or rhs.shape[0] != self.num_batches:
+            raise ValueError("rhs must have shape (batch_size, n, rhs_size)")
+        if rhs.shape[1] < self.max_num_equations:
+            raise ValueError(f"rhs must have at least {self.max_num_equations} rows")
+
+        # Retrieve temporary buffers
+        rhs_size = rhs.shape[2]
+        buffers = self.temp_buffers.get(rhs_size)
+        if buffers is None:
+            raise ValueError(
+                f"RHS size {rhs_size} was not preallocated; call request_rhs_size({rhs_size}) before solve"
+            )
+        rhs_swizzled, result_swizzled, y = buffers
 
-        R = result
+        # Apply reordering
+        solve_rhs = rhs
+        solve_result = result
         if self.enable_reordering:
-            R = self.result_swizzled
             wp.launch(
-                reorder_rows_kernel_col_vector,
-                dim=[self.num_batches, self.max_num_equations],
-                inputs=[rhs, self.rhs_swizzled, self.ordering, self.num_active_equations, batch_mask],
+                reorder_rows_kernel,
+                dim=[self.num_batches, self.max_num_equations, rhs_size],
+                inputs=[rhs, rhs_swizzled, self.ordering, self.num_active_equations, batch_mask],
                 device=self.device,
             )
-
-            rhs = self.rhs_swizzled
+            solve_rhs = rhs_swizzled
+            solve_result = result_swizzled
 
         # Then solve the system using blocked_cholesky_solve kernel
         wp.launch_tiled(
             self.solve_kernel,
-            dim=self.num_batches,
-            inputs=[self.L, self.L_tile_pattern, rhs, R, self.y, self.num_active_equations, batch_mask],
+            dim=(self.num_batches, rhs_size),
+            inputs=[
+                self.L,
+                self.L_tile_pattern,
+                solve_rhs,
+                solve_result,
+                y,
+                self.num_active_equations,
+                batch_mask,
+            ],
             block_dim=self.num_threads_per_block_solve,
             device=self.device,
         )
 
+        # Undo reordering
         if self.enable_reordering:
-            # Undo reordering
             wp.launch(
-                reorder_rows_kernel_col_vector,
-                dim=[self.num_batches, self.max_num_equations],
-                inputs=[R, result, self.inverse_ordering, self.num_active_equations, batch_mask],
+                reorder_rows_kernel,
+                dim=[self.num_batches, self.max_num_equations, rhs_size],
+                inputs=[
+                    solve_result,
+                    result,
+                    self.inverse_ordering,
+                    self.num_active_equations,
+                    batch_mask,
+                ],
                 device=self.device,
             )