perf(linalg): make batched einsum contraction allocation-free per slice

[ultimatile]

Summary

The batched einsum contraction ran a full contract_dense per batch slice — two operand copies, two memory-order reorders, and a notation re-parse plus plan rebuild around a small GEMM. That fixed overhead dominates when the batch and per-slice dimensions are small, as in the successive randomized compression (SRC) sweep's environment recursion, where these tiny batched contractions were the largest single cost of an MPO–MPS apply. This rewrites the per-slice loop as one bare backend GEMM into a preallocated output.

Closes #475

Changes

• crates/ariadnetor-linalg/src/einsum.rs — rewrite batched_contract: permute and reorder both operands once so the batch axes are the slowest-varying in the backend's preferred order (batch-first for row-major, batch-last for column-major), making every per-slice (m, k) / (k, n) block a contiguous range; loop one backend.gemm per slice into a preallocated buffer, with (m, n, k) and the execution policy computed once; validate operand ranks and shared batch/contracted extents up front (returning InvalidArgument); short-circuit degenerate (zero-extent / empty-batch) contractions; extract group_dims / group_extent; drop the now-unused build_batch_free_notation.
• crates/ariadnetor-linalg/src/einsum/tests.rs (new) — a differential test asserting a row-major-preferred and the native column-major backend agree, exercising the row-major layout branch the native backend never selects.
• crates/ariadnetor-linalg/tests/einsum.rs — add multi-batch asymmetric interleaved-output, mismatched-extent, rank-exceeds-arity, and zero-extent cases.

Impact

Any einsum whose pairwise decomposition carries a batch index routes through batched_contract (the MPS environment recursion, the top-crate einsum wrappers, host_ops). The GEMM stays in preferred_order (batch on the slowest-varying axis), so the ComputeBackend contract — compute kernels run in preferred_order, only transpose is layout-parametric — is not widened. Output is unchanged up to floating-point reassociation.

Test plan

• Existing batched einsum tests (bik,bkj->bij and its output-order variants, bi,bi->b, aij,ajk,ak->ai, bkli,bjlk->bij) pass unchanged; they pin the numerical result (integer cases bit-exact, multi-contracted to 1e-10).
• Added: multi-batch asymmetric (abik,abkj->baji), row-major / column-major backend agreement, mismatched contracted extent -> InvalidArgument, rank exceeds arity -> InvalidArgument, zero contracted extent -> zeros.
• The SRC apply and incremental-QR suites pass unchanged (tolerance-based).
• cargo bench -p ariadnetor-algorithms --bench mpo_mps_apply, before/after on the same machine: adaptive/n8_chi8 -18%, adaptive/n12_chi16 -21%, fixed arms -14%.
• Full workspace cargo make test green; cargo make clippy clean.

Plan-vs-actual delta

• Scope narrowed to one of the issue's six items. The issue's Scope lists six per-round copy sites. Profiling at several problem sizes put items 1–5 each under ~1% of an apply and non-load-bearing at scale, while item 6 (batched_contract) was 20–26%. Items 1–5 are deferred as measured noise; only item 6 landed. A comment on Eliminate the per-round copies in the SRC sweep #475 records the per-site breakdown and the sizes measured.
• Two adjacent dispatch bugs fixed. Review surfaced two pre-existing defects in einsum_pair, next to the rewritten path: dim_of ran before the rank check (panicking on a rank-below-arity operand instead of returning InvalidArgument), and the Hadamard path was gated on the extents being 1 rather than the index sets being empty (dropping unit free/contracted axes). Both are fixed here with regression tests, since they sit in the code this PR rewrites.
• The issue's Acceptance is met: adaptive/n8_chi8 no longer regresses against the pre-Phase 2: incremental QR and batched environment sketching for SRC #473 baseline (-18%) and adaptive/n12_chi16 improves further (-21%); the fixed arms improved too (-14%).

[Copilot]

Pull request overview

This PR optimizes the linalg einsum dispatcher’s batched 2-tensor contraction path by eliminating per-batch-slice operand copies/reorders and per-slice notation parsing, replacing it with a single pre-permute/reorder followed by one backend GEMM per slice into a preallocated output buffer. This targets small per-slice GEMMs common in SRC sweep environment recursion.

Changes:

• Reworked batched_contract to precompute permutations/layout once, then run ComputeBackend::gemm per batch slice into a preallocated buffer (plus upfront validation and degenerate short-circuits).
• Added an in-crate unit test to exercise the row-major-preferred backend branch and compare against the native column-major backend.
• Added integration tests covering multi-batch asymmetric batched contractions and additional batched error/degenerate cases.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
crates/ariadnetor-linalg/src/einsum.rs	Rewrites batched 2-tensor einsum contraction to an allocation-free-per-slice GEMM loop with upfront validation and new helpers.
crates/ariadnetor-linalg/src/einsum/tests.rs	Adds unit coverage for row-major vs column-major backend agreement on the batched contraction path.
crates/ariadnetor-linalg/tests/einsum.rs	Adds integration tests for multi-batch asymmetric output order, invalid-argument cases, and zero-extent contracted axis behavior.

---

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

[Copilot]

Pull request overview

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