Repository navigation
[REQ] Add upper Cholesky factorization support to tile operations #1318
Description
Summary
Add a fill_mode parameter to wp.tile_cholesky(), wp.tile_cholesky_inplace(), wp.tile_cholesky_solve(), and wp.tile_cholesky_solve_inplace() to support upper-triangle Cholesky factorization and solve alongside the existing lower-triangle default.
Motivation
The current lower-triangle Cholesky reads columns sequentially, which causes shared memory bank conflicts at power-of-2 tile sizes. Upper-triangle Cholesky reads rows instead, eliminating these conflicts and improving performance:
- 1.23xâ1.70x speedup in microbenchmarks at power-of-2 tile sizes (RTX PRO 6000 Blackwell)
- 2.7â3.1% wall-clock improvement in MuJoCo
three_humanoidssimulation (nworld=8192)
Proposed API
# Lower Cholesky (unchanged default)
L = wp.tile_cholesky(A)
x = wp.tile_cholesky_solve(L, y)
# Upper Cholesky (new)
U = wp.tile_cholesky(A, fill_mode="upper")
x = wp.tile_cholesky_solve(U, y, fill_mode="upper")
fill_mode is a compile-time constant string ("lower" or "upper"), following the same
pattern as wp.tile_sum(a, axis=...).
Activity
shi-eric commentedon Mar 31, 2026
Added in 0cda7e1
I believe @daedalus5 has some follow-up items regarding variable naming and deprecations.
Add fill_mode parameter to tile Cholesky operations [NVIDIAGH-1318]
Merge branch 'adenzler/upper-triangle-ops' into 'main'