Skip to content

perf(linalg): route block-sparse per-block transpose through the backend - #457

Merged
ultimatile merged 2 commits into
mainfrom
perf/311-route-block-transpose-hptt
Jul 7, 2026
Merged

perf(linalg): route block-sparse per-block transpose through the backend#457
ultimatile merged 2 commits into
mainfrom
perf/311-route-block-transpose-hptt

Conversation

@ultimatile

@ultimatile ultimatile commented Jul 7, 2026

Copy link
Copy Markdown
Owner

Summary

Block-sparse per-block transpose was a hand-written naive per-element kernel (transpose_block_data) that bypassed the backend's transpose, so HPTT was never used per block even when the build enabled it. This routes all four per-block transpose sites through ComputeBackend::transpose, which selects HPTT or the native naive kernel at compile time. The benchmark in #456 showed HPTT beats the naive kernel at every representative block shape, so the routing is a measured win.

Closes #311

Changes

  • block_sparse_contract.rs: route the contract_to_scalar and contract_to_tensor (lhs / rhs) transposes through a new transpose_block helper over ComputeBackend::transpose; memoize each rhs block's transpose by block index so a block shared across several lhs blocks is transposed at most once per contraction; remove transpose_block_data and its compute_strides helper.
  • block_sparse_permute.rs: transpose each block straight into its freshly-zeroed output block, dropping an intermediate buffer and copy.
  • block_sparse_contract/tests/mod.rs: compute_strides moves here as a test-only helper — its only remaining consumers are tests.
  • block_sparse_permute/tests.rs: add a ground-truth permute test asserting the exact element mapping, closing a gap where the site was only round-trip-tested.
  • benches/block_sparse_ops.rs: add an end-to-end permute_bsp benchmark — the existing contraction groups all take the GEMM trans-flag fast path and never exercise the physical transpose, so nothing tracked the changed path end to end; also update the bench_block_transpose doc comment, which now measures the production per-block transpose path.

Impact

contract_to_scalar and contract_to_tensor are file-private; the only signature change (contract_to_scalar gains backend, returns Result) is matched at its sole in-file caller. No public API changes; net removal of two pub(crate) internals.

Test plan

  • Numerical equivalence is guarded by the existing contract physical-transpose tests and the new ground-truth permute test, which pin the routed transpose against hand-computed values.
  • The per-block transpose stays sequential (ExecPolicy::Sequential): the block_transpose benchmark's parallel variant is slower at these block sizes, where Rayon dispatch overhead dominates.
  • cargo test -p ariadnetor-linalg and cargo test -p ariadnetor-linalg --features hptt: 292 pass each, so the ground-truth tests certify the HPTT kernel as well as the naive one.
  • Full workspace tests green; cargo make clippy (all targets, deny warnings) clean.

Benchmark

End-to-end permute_bsp (rank-3, cyclic perm, physical transpose per block), native naive kernel vs HPTT (--features hptt), single-machine criterion medians — the relative ratio is the load-bearing figure:

case naive (hptt off) HPTT (hptt on) change
q4 d64 85 us 19 us -78%
q4 d128 338 us 67 us -81%
q8 d64 179 us 42 us -77%
q8 d128 727 us 153 us -81%

The default build keeps the native naive kernel, so the routing is performance-neutral there; the roughly 4-5x speedup is realized when building with --features hptt.

Notes

The rhs transpose cache trades memory for recomputation: at most one transposed copy per distinct rhs block lives for the contraction's duration.

@coderabbitai ignore

Route the four per-block transpose sites (permute, contract-to-scalar,
and the contract-to-tensor lhs/rhs branches) through
ComputeBackend::transpose instead of the hand-written naive
transpose_block_data, so builds with HPTT use it per block. The
benchmark showed HPTT beats the naive kernel at every representative
block shape.

Memoize each rhs block's physical transpose by block index so a block
shared across multiple lhs blocks is transposed at most once per
contraction, and transpose permute blocks straight into the output
block to drop an intermediate buffer and copy. The per-block transpose
stays sequential: Rayon dispatch loses at these block sizes.

Remove transpose_block_data and its compute_strides helper (the native
backend already provides the naive path as its non-HPTT fallback);
compute_strides moves into the contract test module as a test-only
helper. Add a ground-truth block-sparse permute test, since the site
was previously only round-trip-tested.

Closes #311

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR improves block-sparse linalg performance by routing per-block transposes through ComputeBackend::transpose, enabling HPTT (when built with --features hptt) and removing the previous hand-written per-element transpose path that bypassed the backend. It also reduces redundant work in block-sparse contraction by caching RHS block transposes, and strengthens permutation correctness testing with a ground-truth mapping assertion.

Changes:

  • Route block-sparse per-block transpose in both contraction and permutation through ComputeBackend::transpose (HPTT-enabled), removing the in-tree naive transpose_block_data.
  • Cache each RHS block’s physical transpose during contraction so a shared RHS block is transposed at most once per contraction.
  • Add a rank-3 ground-truth permute test that asserts exact element mapping; update the transpose benchmark doc comment to reflect the new production path.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.

Show a summary per file
File Description
crates/ariadnetor-linalg/src/block_sparse_permute/tests.rs Adds a ground-truth rank-3 permutation test with explicit element mapping checks.
crates/ariadnetor-linalg/src/block_sparse_permute.rs Uses ComputeBackend::transpose directly into the output block to avoid an intermediate buffer/copy.
crates/ariadnetor-linalg/src/block_sparse_contract/tests/mod.rs Moves stride computation into tests as a test-only helper for index math assertions.
crates/ariadnetor-linalg/src/block_sparse_contract.rs Routes per-block transpose through the backend and adds RHS transpose memoization; removes the in-tree transpose kernel.
crates/ariadnetor-linalg/benches/block_sparse_ops.rs Updates benchmark documentation to state it measures the production per-block transpose path.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

The existing contraction benchmarks all take the GEMM trans-flag fast
path and never exercise the physical per-block transpose, so the win
from routing that path through the backend was untracked end to end.
`permute_bsp` permutes a rank-3 block-sparse tensor by a cyclic perm,
which physically transposes every block, giving a benchmark that moves
with the routed transpose (and with HPTT under `--features hptt`).

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

@ultimatile
ultimatile merged commit 5b6336e into main Jul 7, 2026
2 checks passed
@ultimatile
ultimatile deleted the perf/311-route-block-transpose-hptt branch July 8, 2026 07:28
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.

perf: block-sparse per-block transpose bypasses HPTT; needs a benchmark

2 participants