Skip to content

Optimize CPU operand materialization with blocked pooled copies #55

Description

@exAClior

Context

Profiling the Complex64 TDVP contractions from #53 / #54 found that the remaining right-environment gap is not in faer's complex GEMM. It is in CPU operand materialization.

For abcd,bde->ace and abcde,cef->abdf, OMEinsum must pack the left operand. The current path:

  1. default-initializes fresh scratch storage;
  2. walks coordinates with a scalar permutation loop;
  3. clones the complete packed payload into another fresh Vec;
  4. drops both allocations after the contraction.

Tenferro performs one blocked permutation into reusable thread-local storage.

At chi=64, OMEinsum's pre-fix pack measured 137/424 us versus Tenferro's 16/31 us. Prepacking outside the public call removed 87%/97% of the total gaps. OMEinsum also requested two full payload allocations per call: 2.36 MiB for h1-right and 4.72 MiB for h2-right.

Scope

Optimize CPU operand materialization independently of the Complex64 GEMM dispatch work in #54:

  • use a blocked, cache-aware strided permutation;
  • write directly into the owned packed operand (no second full clone);
  • reuse initialized packing storage with a bounded per-thread pool;
  • preserve scalar, singleton, padded-stride, arbitrary-permutation, and zero-sized layouts;
  • keep the copy path sequential; GEMM threading policy is separate.

Acceptance criteria

  • Existing CPU contraction tests pass.
  • Add permutation-vs-scalar-reference coverage across representative layouts.
  • Add a zero-sized-layout regression test.
  • Add packing-buffer reuse coverage.
  • make check and cargo package --allow-dirty pass.
  • Warm right-layout packing no longer allocates a payload-sized buffer per call.
  • Rerun the preserved single-thread comparison with RAYON_NUM_THREADS=1.

Measured target

A local prototype using strided-perm, direct buffer ownership, and a bounded thread-local pool improved OMEinsum by 8-20% in the reverse-order all-size run. Post-fix OMEinsum/Tenferro ratios were 1.00-1.04x across all six right-layout cases, and warm allocation deltas versus prepacked input fell to 800/944 metadata bytes.

This issue deliberately excludes GEMM dispatch and benchmark-harness changes; those remain in #53 / #54.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions