trnsparse 0.4.0 — fused Schwarz-screened SpMM
Closes #19.
Added
- `screened_spmm(A, diag_integrals, B, threshold)` — fused Schwarz-screened dense matmul. One NKI kernel fuses the full pipeline (outer-product pair bound → threshold → mask-apply → `nc_matmul`) into a single dispatch. Saves ~30–50% end-to-end vs the unfused `density_screen + from_dense + spmm` flow on Fock-build-sized inputs.
- `_screened_spmm_kernel` (`@nki.jit`) — stationary-A-tile-reuse GEMM extended with a per-tile pair-bound mask built from the 1-D Schwarz-bound vector.
- `_ScreenedSpMMFunction` — `torch.autograd.Function` wrapper. Third differentiable NKI kernel in the trnsci suite (after v0.2.0 CSR SpMM and v0.3.0 BSR SpMM). `torch.autograd.gradcheck` passes at `atol=1e-4` on hardware.
Validation
| Surface | Tests | Result |
|---|---|---|
| CPU suite | 4 `TestScreenedSpmm` | ✅ |
| Simulator (ubuntu-latest CI + trn1) | 2 `TestScreenedSpmmSimulator` | ✅ |
| Hardware (trn1.2xlarge) | 7 `TestNkiScreenedSpmmParity` + `TestNkiScreenedSpmmDifferentiability` | ✅ |
Total: 49 CPU + 23 hardware tests green; no regression across the suite.
Also closed this session
- #24 — fused CG NKI kernel not buildable under NKI 2.24/0.3.0 (no `break`, no iteration-carried scalar state across `affine_range`, no nested kernels). Per-iteration `_cg_step_kernel` reframe evaluated and found to save only 5–20% — closed honestly. If upstream NKI gains persistent-SBUF-across-calls or in-kernel control flow, the whole-loop CG kernel can be reopened.
Known limits
- `screened_spmm` currently restricted to square A (M == K) with 1-D `diag_integrals`. The common Fock-build case. Rectangular / asymmetric-bounds extension is a follow-up if asked for.