Use new Tile Cholesky API

NVIDIA · Feb 4, 2025 · cff4549 · cff4549
1 parent 6f1767b
commit cff4549
Showing 1 changed file with 4 additions and 11 deletions.
diff --git a/warp/sim/integrator_featherstone.py b/warp/sim/integrator_featherstone.py
@@ -1177,14 +1177,7 @@ def eval_dense_gemm_tile(
             P_body = wp.tile_matmul(M_body, J_body)
 
             # assign to the P slice
-            wp.tile_assign(
-                P,
-                P_body,
-                offset=(
-                    i * 6,
-                    0,
-                ),
-            )
+            wp.tile_assign(P, P_body, offset=(i * 6, 0))
 
         # compute H = J^T*P
         H = wp.tile_matmul(wp.tile_transpose(J), P)
@@ -1205,9 +1198,9 @@ def eval_tiled_dense_cholesky_batched(
 
         a = wp.tile_load(A[articulation], shape=(num_dofs, num_dofs), storage="shared")
         r = wp.tile_load(R[articulation], shape=num_dofs, storage="shared")
-        wp.tile_diag_add(a, r)
-        wp.tile_cholesky(a)
-        wp.tile_store(L[articulation], wp.tile_transpose(a))
+        a_r = wp.tile_diag_add(a, r)
+        l = wp.tile_cholesky(a_r)
+        wp.tile_store(L[articulation], wp.tile_transpose(l))
 
     return eval_tiled_dense_cholesky_batched