[REQ] Add adjoint for out-of-place tile Cholesky factorization #1316

Activity

1. nv-snidhan

   added

   feature requestRequest for something to be addedRequest for something to be added

   on Mar 23, 2026
2. shi-eric

   assigned

   nv-snidhan

   on Mar 23, 2026
3. shi-eric

   added the

   Enhancement

   issue type on Mar 23, 2026
4. shi-eric

   added this to the 1.13.0 milestone on Mar 23, 2026
5. shi-eric

   added

   autodiffTapes, adjoint generation, custom gradients, and gradient correctness in WarpTapes, adjoint generation, custom gradients, and gradient correctness in Warp

   tileTile programming modelTile programming model

   on Mar 23, 2026
6. pull

   added a commit that references this issue on Apr 15, 2026

   Add adjoint for out-of-place tile Cholesky factorization [NVIDIAGH-1316]

   ...

   d683d4d
7. shi-eric commented on Apr 19, 2026

   

   shi-eric

   on Apr 19, 2026

   Last edited by shi-eric

   Contributor

   More actions

   Implemented in commit d683d4d. wp.tile_cholesky() (out-of-place variant) is now differentiable and can be used under wp.Tape().

   @wp.kernel
   def cholesky_kernel(A: wp.array2d(dtype=float),
                       L: wp.array2d(dtype=float)):
       a = wp.tile_load(A, shape=(N, N))
       l = wp.tile_cholesky(a)
       wp.tile_store(L, l)
   
   A = wp.array(A_np, requires_grad=True)
   L = wp.zeros((N, N), dtype=float, requires_grad=True)
   
   with wp.Tape() as tape:
       wp.launch_tiled(cholesky_kernel, dim=[1], inputs=[A, L], block_dim=64)
   
   tape.backward(grads={L: wp.array(adj_L_np)})
   # A.grad now contains the symmetric-matrix Cholesky adjoint

   The adjoint implements the symmetric-matrix derivative of Cholesky (Murray 2016): it solves L^T \ Phi(L^T adj_L) via two triangular solves and accumulates gradients into the lower triangle of A. Three execution paths are wired up:

   • CPU / GPU-without-MathDx: single-thread scalar back-substitution.
   • GPU + MathDx with backward disabled: dummy LTO symbols (no-op).
   • GPU + MathDx: LTO-accelerated GEMM + TRSM with threaded element-wise steps using shared-memory scratch.

   Note that only the out-of-place variant got an adjoint. The in-place tile_cholesky_inplace form is still non-differentiable.
8. shi-eric

   closed this as completedon Apr 19, 2026