Performance improvements for Kamino's direct solvers (#3610)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Milad Rakhsha <mrakhsha@nvidia.com>

Author: 3 people

CHANGELOG.md

@@ -23,7 +23,7 @@
 - Add viewer layer system to overlay multiple solvers/models in supported rendering viewers; call `ViewerBase.activate(layer_id)` to route subsequent `set_model` / `log_state` / `log_*` calls into a named layer, `ViewerBase.set_layer_visible()` to toggle layers independently, and `ViewerBase.set_layer_transform()` to position layers side-by-side. See `example_basic_multi_solver_overlay.py`
 - Add `Heightfield.create_from_mesh()` and `newton.utils.rasterize_mesh_to_heightfield()` to build a heightfield collider by ray-casting a `wp.Mesh`, replacing a large static terrain mesh with an equivalent heightfield.
 - Add `ViewerBase.camera_speed` to configure keyboard translation speed for interactive viewers. (#3439)
-- Add opt-in DVI forward dynamics to `SolverKamino` through `SolverKamino.Config(dynamics_solver="dvi")`, with sparse and dense execution, DVI-specific diagnostics, and warm-starting. PADMM remains the default.
+- Add opt-in DVI forward dynamics to `SolverKamino` through `SolverKamino.Config(dynamics_solver="dvi")`, with sparse and dense execution, DVI-specific convergence diagnostics, warm-starting, bounded contact-recovery controls, and RCM-reordered bilateral factorization with reusable ordering and panel-parallel numeric factorization for large systems. PADMM remains the default.
 - Add SDF contact support for convex-hull shapes with mesh-attached SDFs and opt-in SDF contact generation for box shapes.
 - Add opt-in filtering of static-static, static-kinematic, and kinematic-kinematic contacts during broad-phase collision detection. Set `CollisionPipeline(include_static_kinematic_pairs=False)` to enable filtering; the default preserves existing contact generation. `Model.shape_contact_pairs` remains an unfiltered superset for direct consumers such as `SolverKamino` and hydroelastic SDF setup.
 - Add opt-in `body_frame_origin="com"` to `ModelBuilder.add_rod()` and `ModelBuilder.add_rod_graph()` for COM-centered cable capsule body frames.
@@ -79,6 +79,8 @@
 
 ### Fixed
 
+- Complete Kamino RCM traversal for large and disconnected systems and reuse the resulting permutation by default; set `reuse_permutation=False` to recompute it for changing matrix topology.
+- Fix panel-parallel RCM-blocked LLT factorization hanging when a matrix ends in a partial tile.
 - Fix USD capsule, cylinder, and cone visual and site scaling to follow the authored primitive axis.
 - Fix USD plane visual width and length to scale along the axes defined by the `UsdGeomPlane` schema, and orient X- and Y-axis plane visuals along the authored axis.
 - Validate `ArticulationView` mask shapes and devices before launching selection kernels. (#3448)

docs/solvers/kamino.rst

@@ -56,3 +56,18 @@ Set ``sparse_jacobian=False`` for fully dense DVI, or set
 solver. With
 ``collect_solver_info=True``, DVI stores terminal residual status that should
 not be interpreted as PADMM ADMM residuals.
+
+For large bilateral systems, opt into RCM-reordered factorization explicitly:
+
+.. code-block:: python
+
+   config.dvi.bilateral_solver_type = "LLTBRCM"
+   config.dvi.bilateral_solver_kwargs = {
+       "block_size": 32,
+       "reuse_permutation": True,
+       "parallel_factorization": True,
+   }
+
+The cached permutation remains mathematically valid when matrix values or
+sparsity change and is recomputed automatically if the active dimension
+changes. Keep the default ``"LLTB"`` solver for small systems.

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

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers
 # SPDX-License-Identifier: Apache-2.0
 
-"""Availability checks and native fallbacks for optional Warp tile builtins."""
+"""Availability checks and fallback for optional Warp tile builtins."""
 
 import inspect
 import os
@@ -26,157 +26,39 @@ def _has_warp_builtin(name: str) -> bool:
     return name in builtin_functions
 
 
-def _has_native_tile_arg_support() -> bool:
+def _has_native_left_transpose_update_support() -> bool:
     if not _tile_transpose_update_enabled():
         return False
 
     try:
         from warp._src import codegen  # noqa: PLC0415
 
-        source = inspect.getsource(codegen.codegen_snippet)
-    except Exception:
-        return False
-
-    return "template_params" in source and "is_tile(arg.type)" in source and "& {arg.emit()" in source
-
-
-def _has_native_tile_access_helpers() -> bool:
-    if not _tile_transpose_update_enabled():
-        return False
-
-    try:
+        codegen_source = inspect.getsource(codegen.codegen_snippet)
         tile_header = Path(wp.__file__).resolve().parent / "native" / "tile.h"
-        source = tile_header.read_text(encoding="utf-8")
+        tile_source = tile_header.read_text(encoding="utf-8")
     except Exception:
         return False
 
     return (
-        "tile_read(const tile_register_t" in source
-        and "tile_read(const tile_shared_t" in source
-        and "tile_add(tile_register_t" in source
-        and "tile_add(tile_shared_t" in source
+        "template_params" in codegen_source
+        and "is_tile(arg.type)" in codegen_source
+        and "& {arg.emit()" in codegen_source
+        and "tile_add(tile_register_t" in tile_source
+        and "tile_add(tile_shared_t" in tile_source
     )
 
 
-def _has_native_tile_update_support() -> bool:
-    return _has_native_tile_arg_support() and _has_native_tile_access_helpers()
-
-
-def _copy_dense_2d_snippet(tile_name: str, layout_name: str, values_name: str, cols_name: str, storage: str) -> str:
-    if storage == "register":
-        return (
-            f"{tile_name}.apply([&](int reg, auto c) {{ "
-            f"{values_name}[c[0] * {cols_name} + c[1]] = {tile_name}.data[reg]; }});"
-        )
-    if storage == "generic":
-        return f"""for (int linear = WP_TILE_THREAD_IDX; linear < {layout_name}::Size; linear += WP_TILE_BLOCK_DIM) {{
-    auto c = {layout_name}::coord_from_linear(linear);
-    int reg = linear / WP_TILE_BLOCK_DIM;
-    {values_name}[c[0] * {cols_name} + c[1]] = tile_read({tile_name}, reg, linear);
-}}"""
-    raise ValueError(f"Unsupported tile storage specialization: {storage!r}")
-
-
-def _update_output_snippet(
-    layout_name: str,
-    output_name: str,
-    rows_name: str,
-    cols_name: str,
-    k_name: str,
-    left_values_name: str,
-    right_values_name: str,
-    left_transposed: bool,
-    storage: str,
-) -> str:
-    if left_transposed:
-        product = f"{left_values_name}[k * {rows_name} + c[0]] * {right_values_name}[k * {cols_name} + c[1]]"
-    else:
-        product = f"{left_values_name}[c[0] * {k_name} + k] * {right_values_name}[c[1] * {k_name} + k]"
-
-    if storage == "shared":
-        write = f"""const T value = a * sum;
-    if constexpr ({layout_name}::Unique)
-        {output_name}.data(linear) += value;
-    else
-        wp::atomic_add(&{output_name}.data(linear), value);"""
-    elif storage == "register":
-        return f"""const T a = static_cast<T>(alpha);
-{output_name}.apply([&](int reg, auto c) {{
-    T sum = T{{}};
-    WP_PRAGMA_UNROLL
-    for (int k = 0; k < {k_name}; ++k) {{
-        sum += {product};
-    }}
-    {output_name}.data[reg] += a * sum;
-}});
-WP_TILE_SYNC();"""
-    elif storage == "generic":
-        write = f"""int reg = linear / WP_TILE_BLOCK_DIM;
-    tile_add({output_name}, reg, linear, a * sum);"""
-    else:
-        raise ValueError(f"Unsupported tile storage specialization: {storage!r}")
-
-    return f"""const T a = static_cast<T>(alpha);
-for (int linear = WP_TILE_THREAD_IDX; linear < {layout_name}::Size; linear += WP_TILE_BLOCK_DIM) {{
-    auto c = {layout_name}::coord_from_linear(linear);
-    T sum = T{{}};
-    WP_PRAGMA_UNROLL
-    for (int k = 0; k < {k_name}; ++k) {{
-        sum += {product};
-    }}
-    {write}
-}}
-WP_TILE_SYNC();"""
-
-
-def _make_tile_matmul_transpose_update_snippet(out_storage: str, input_storage: str) -> str:
-    copy_left = _copy_dense_2d_snippet("left", "LeftLayout", "left_values", "K", input_storage)
-    copy_right = _copy_dense_2d_snippet("right", "RightLayout", "right_values", "K", input_storage)
-    update_out = _update_output_snippet(
-        "OutLayout", "out", "Rows", "Cols", "K", "left_values", "right_values", False, out_storage
-    )
-    return f"""using OutTile = tile_out;
-using LeftTile = tile_left;
-using RightTile = tile_right;
-using T = typename OutTile::Type;
-using OutLayout = typename OutTile::Layout;
-using LeftLayout = typename LeftTile::Layout;
-using RightLayout = typename RightTile::Layout;
-
-static_assert(OutLayout::Shape::N == 2, "out must be 2D");
-static_assert(LeftLayout::Shape::N == 2, "left must be 2D");
-static_assert(RightLayout::Shape::N == 2, "right must be 2D");
-static_assert(LeftLayout::Shape::dim(0) == OutLayout::Shape::dim(0), "left rows must match out rows");
-static_assert(RightLayout::Shape::dim(0) == OutLayout::Shape::dim(1), "right rows must match out cols");
-static_assert(LeftLayout::Shape::dim(1) == RightLayout::Shape::dim(1), "left/right cols must match");
-
-constexpr int Rows = OutLayout::Shape::dim(0);
-constexpr int Cols = OutLayout::Shape::dim(1);
-constexpr int K = LeftLayout::Shape::dim(1);
-
-#if defined(__CUDA_ARCH__)
-__shared__ T left_values[Rows * K];
-__shared__ T right_values[Cols * K];
-#else
-T left_values[Rows * K];
-T right_values[Cols * K];
-#endif
-
-{copy_left}
-{copy_right}
-WP_TILE_SYNC();
-
-{update_out}
-"""
+HAS_TILE_MATMUL_TRANSPOSE_UPDATE = _has_warp_builtin("tile_matmul_transpose_update")
+HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE = _has_warp_builtin("tile_matmul_left_transpose_update")
+HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE = (
+    not HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE and _has_native_left_transpose_update_support()
+)
 
 
-def _make_tile_matmul_left_transpose_update_snippet(out_storage: str, input_storage: str) -> str:
-    copy_left = _copy_dense_2d_snippet("left", "LeftLayout", "left_values", "Rows", input_storage)
-    copy_right = _copy_dense_2d_snippet("right", "RightLayout", "right_values", "Cols", input_storage)
-    update_out = _update_output_snippet(
-        "OutLayout", "out", "Rows", "Cols", "K", "left_values", "right_values", True, out_storage
-    )
-    return f"""using OutTile = tile_out;
+@cache
+def make_tile_matmul_left_transpose_update_func(block_size: int):
+    """Create ``out += alpha * transpose(left) @ right`` for the LLT solve."""
+    snippet = """using OutTile = tile_out;
 using LeftTile = tile_left;
 using RightTile = tile_right;
 using T = typename OutTile::Type;
@@ -203,46 +85,22 @@ def _make_tile_matmul_left_transpose_update_snippet(out_storage: str, input_stor
 T right_values[K * Cols];
 #endif
 
-{copy_left}
-{copy_right}
+left.apply([&](int reg, auto c) { left_values[c[0] * Rows + c[1]] = left.data[reg]; });
+right.apply([&](int reg, auto c) { right_values[c[0] * Cols + c[1]] = right.data[reg]; });
 WP_TILE_SYNC();
 
-{update_out}
-"""
-
-
-HAS_TILE_MATMUL_TRANSPOSE_UPDATE = _has_warp_builtin("tile_matmul_transpose_update")
-HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE = _has_warp_builtin("tile_matmul_left_transpose_update")
-HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE = not HAS_TILE_MATMUL_TRANSPOSE_UPDATE and _has_native_tile_update_support()
-HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE = (
-    not HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE and _has_native_tile_update_support()
-)
-
-
-@cache
-def make_tile_matmul_transpose_update_func(
-    block_size: int, out_storage: str = "shared", input_storage: str = "register"
-):
-    """Create ``out += alpha * left @ transpose(right)`` as a native tile function."""
-    snippet = _make_tile_matmul_transpose_update_snippet(out_storage, input_storage)
-
-    @wp.func_native(snippet)
-    def tile_matmul_transpose_update(
-        out: wp.tile[float, block_size, block_size],
-        left: wp.tile[float, block_size, block_size],
-        right: wp.tile[float, block_size, block_size],
-        alpha: float,
-    ): ...
-
-    return tile_matmul_transpose_update
-
-
-@cache
-def make_tile_matmul_left_transpose_update_func(
-    block_size: int, out_storage: str = "generic", input_storage: str = "register"
-):
-    """Create ``out += alpha * transpose(left) @ right`` as a native tile function."""
-    snippet = _make_tile_matmul_left_transpose_update_snippet(out_storage, input_storage)
+const T a = static_cast<T>(alpha);
+for (int linear = WP_TILE_THREAD_IDX; linear < OutLayout::Size; linear += WP_TILE_BLOCK_DIM) {
+    auto c = OutLayout::coord_from_linear(linear);
+    T sum = T{};
+    WP_PRAGMA_UNROLL
+    for (int k = 0; k < K; ++k) {
+        sum += left_values[k * Rows + c[0]] * right_values[k * Cols + c[1]];
+    }
+    int reg = linear / WP_TILE_BLOCK_DIM;
+    tile_add(out, reg, linear, a * sum);
+}
+WP_TILE_SYNC();"""
 
     @wp.func_native(snippet)
     def tile_matmul_left_transpose_update(

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

@@ -10,11 +10,9 @@
 
 from ._tile_builtins import (
     HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE,
-    HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE,
     HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE,
     HAS_TILE_MATMUL_TRANSPOSE_UPDATE,
     make_tile_matmul_left_transpose_update_func,
-    make_tile_matmul_transpose_update_func,
 )
 
 ###
@@ -166,10 +164,6 @@ def llt_blocked_factorize_kernel(
                     L_block = wp.tile_load(L_i, shape=(block_size, block_size), offset=(k, j))
                     if wp.static(HAS_TILE_MATMUL_TRANSPOSE_UPDATE):
                         wp.tile_matmul_transpose_update(A_kk_tile, L_block, L_block, alpha=-1.0)
-                    elif wp.static(HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE):
-                        wp.static(make_tile_matmul_transpose_update_func(block_size, "shared", "register"))(
-                            A_kk_tile, L_block, L_block, -1.0
-                        )
                     else:
                         L_block_T = wp.tile_transpose(L_block)
                         wp.tile_matmul(L_block, L_block_T, A_kk_tile, alpha=-1.0)
@@ -206,10 +200,6 @@ def llt_blocked_factorize_kernel(
                         L_2_tile = wp.tile_load(L_i, shape=(block_size, block_size), offset=(k, j))
                         if wp.static(HAS_TILE_MATMUL_TRANSPOSE_UPDATE):
                             wp.tile_matmul_transpose_update(A_ik_tile, L_tile, L_2_tile, alpha=-1.0)
-                        elif wp.static(HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE):
-                            wp.static(make_tile_matmul_transpose_update_func(block_size, "shared", "register"))(
-                                A_ik_tile, L_tile, L_2_tile, -1.0
-                            )
                         else:
                             L_T_tile = wp.tile_transpose(L_2_tile)
                             wp.tile_matmul(L_tile, L_T_tile, A_ik_tile, alpha=-1.0)
@@ -306,7 +296,7 @@ def llt_blocked_solve_kernel(
                     if wp.static(HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE):
                         wp.tile_matmul_left_transpose_update(rhs_tile, L_tile, x_tile, alpha=-1.0)
                     elif wp.static(HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE):
-                        wp.static(make_tile_matmul_left_transpose_update_func(block_size, "generic", "register"))(
+                        wp.static(make_tile_matmul_left_transpose_update_func(block_size))(
                             rhs_tile, L_tile, x_tile, -1.0
                         )
                     else:
@@ -395,7 +385,7 @@ def llt_blocked_solve_inplace_kernel(
                     if wp.static(HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE):
                         wp.tile_matmul_left_transpose_update(rhs_tile, L_tile, x_tile, alpha=-1.0)
                     elif wp.static(HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE):
-                        wp.static(make_tile_matmul_left_transpose_update_func(block_size, "generic", "register"))(
+                        wp.static(make_tile_matmul_left_transpose_update_func(block_size))(
                             rhs_tile, L_tile, x_tile, -1.0
                         )
                     else: