Skip to content

[WIP] Better BTA without allocation - #117

Open
dance858 wants to merge 3 commits into
mainfrom
btda-fold-diag
Open

[WIP] Better BTA without allocation#117
dance858 wants to merge 3 commits into
mainfrom
btda-fold-diag

Conversation

@dance858

@dance858 dance858 commented Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Claude:

Removes the last "must remove this allocation. Very important" TODOs: BTDA_pd_pd_fill_values and BTDA_pd_spd_fill_values allocated and freed a diag(d) @ A intermediate on every Hessian evaluation, and the blockwise spd BTDA variants multiplied that by n_blocks.

Each BTA kernel is now a shared core with an optional d (NULL = identity) folded into the existing row-gather into pre-sized kernel_dwork — the pattern BTDA_pd_csc/BTDA_csc_pd already used — so the BTDA entry points are thin wrappers and no fill path allocates. The blockwise BTDA_spd_pd/BTDA_spd_spd inherit the fix automatically.

Tests:

  • Direct BTDA_pd_pd correctness on both the matching-row_perm (fast) and gather paths vs the diag-decomposition oracle — the fast path with d is new code.
  • A no-transient-allocation regression test (pd_pd both paths, pd_spd, blockwise spd_pd) using g_peak_bytes/g_allocated_bytes; it fails against the previous kernels.
  • End-to-end multiply Hessian with dense operators (new_left_matmul_dense) — the first wsum_hess-level coverage of multiply's pd/pd BTDA dispatch, verified by numerical differentiation.

Verified: 448 tests in Release, UBSan clean, guard-malloc + leaks --atExit clean, PROFILE_ONLY builds, clang-format clean on all edited regions (the one pre-existing violation in the dead _via_transpose_dead function is untouched).

🤖 Generated with Claude Code

BTDA_pd_pd_fill_values and BTDA_pd_spd_fill_values allocated, filled,
and freed a diag(d) @ A intermediate on every Hessian evaluation (the
"must remove this allocation. Very important" TODOs), and the blockwise
BTDA_spd_pd/spd_spd variants compounded that to n_blocks temps per
fill. Restructure each BTA kernel into a shared core taking an optional
d (NULL = identity) that scales A's rows by d[global_row] while they
are being gathered into the pre-sized kernel_dwork scratch — the same
pattern BTDA_pd_csc/BTDA_csc_pd already use — so the BTDA entry points
become thin wrappers and no fill allocates. The pd_pd fast path
(identical row_perms) scales into A->kernel_dwork, which
BTA_pd_pd_alloc already sizes to exactly m0*n0 on that path.

New tests: direct BTDA_pd_pd correctness on both the matching-row_perm
and gather paths against the diag-decomposition oracle; a
no-transient-allocation regression test over pd_pd (both paths),
pd_spd, and blockwise spd_pd using the tracked allocator's peak-bytes
counter (fails on the previous kernels); and an end-to-end multiply
Hessian test with dense operators — the first wsum_hess coverage of
multiply's pd/pd dispatch.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HRnTPZsCXK6x2PSGUzanmg
@dance858 dance858 changed the title Fold diag(d) into the BTA gathers, delete the per-eval DA temps [WIP] Better BTA without allocation Sep 1, 2026
dance858 and others added 2 commits August 31, 2026 20:19
Replace the hand-written scale-copy loops with an unconditional
memcpy gather followed by a per-row cblas_dscal pass when d is given,
matching the existing DA_pd_fill_values / ATDA_pd_fill_values idiom.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HRnTPZsCXK6x2PSGUzanmg
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HRnTPZsCXK6x2PSGUzanmg
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant