perf: eliminate TQ transform allocations

[LHT129]

Summary: remove per-vector transform metadata allocation and reuse caller-owned scratch buffers. Query scratch is private to each computer; PCA metadata code layout remains unchanged.\n\nValidation: targeted TransformQuantizer and transformer tests pass after rebuilding against main; changed transformer implementations pass clang-tidy-15.\n\nRefs #2702

[vsag-bot]

/label status/waiting-for-review
/waiting-on reviewer
/request-review @jiaweizone
/request-review @wxyucs
/request-review @inabao

[mergify]

Merge Protections

🟢 All 2 merge protections satisfied — ready to merge.

Show 2 satisfied protections

🟢 Require kind label

• label~=^kind/

🟢 Require version label

• label~=^version/

[Copilot]

Pull request overview

This PR optimizes the TransformQuantizer hot path by removing per-vector transformer metadata allocations and enabling reuse of caller-owned scratch buffers during transform execution, while keeping PCA metadata byte layout stable.

Changes:

• Replace VectorTransformer::Transform’s metadata object return with an in-place uint8_t* meta output buffer.
• Rework TransformQuantizer to execute the transformer chain using ping-pong scratch buffers and return a pointer to the transformed vector data.
• Add per-Computer scratch buffers for query-time transforms and update affected transformer implementations/tests.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/quantization/transform_quantization/transform_quantizer.h	Reworks chain transform execution to use caller-provided scratch and avoid per-vector meta allocations.
src/quantization/computer.h	Adds per-computer scratch buffers to reuse during query preprocessing.
src/impl/transform/vector_transformer.h	Updates transformer interface to write metadata into caller-provided storage.
src/impl/transform/random_orthogonal_transformer.h	Updates ROM transformer signature to match new interface.
src/impl/transform/random_orthogonal_transformer.cpp	Drops per-call ROM meta allocation; adapts to new Transform signature.
src/impl/transform/pca_transformer.h	Updates PCA transformer signature; keeps meta size/align contract.
src/impl/transform/pca_transformer.cpp	Writes PCA meta bytes directly into provided buffer (preserving layout/value).
src/impl/transform/mrle_transformer.h	Updates MRLE transformer signature; removes unused meta allocation.
src/impl/transform/mrle_transformer_test.cpp	Updates MRLE tests to match void Transform API.
src/impl/transform/fht_kac_rotate_transformer.h	Updates FHT transformer signature to match new interface.
src/impl/transform/fht_kac_rotate_transformer.cpp	Drops per-call FHT meta allocation; adapts to new Transform signature.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[LHT129]

Thanks for this PR. The change to eliminate per-vector make_shared allocations in the transform hot path is a solid performance improvement. The ping-pong scratch buffer pattern in ExecuteChainTransform is clean and correctly handles the input/output pointer swapping.

I have reviewed the full diff and the existing inline comments from Copilot and the author already cover the notable points:

1. TransformBaseVector still allocates secondary_scratch per call — this path is used by BottomQuantizerAccessor::PrepareBottomInput, so the allocation overhead remains there. Consider plumbing a caller-owned buffer.

2. PCATransformer::Transform still allocates centralized_vec per call — this undermines the goal of eliminating hot-path allocations for PCA-based configurations. A pre-allocated member buffer or an additional scratch parameter would help.

3. meta_offsets null-safety in ExecuteChainTransform — when codes != nullptr but meta_offsets is null, the pointer arithmetic codes + meta_offsets[i] would dereference null. In practice the current callers always pass valid offsets, but a defensive check or assertion would make the contract clearer.

4. VectorTransformer::Transform empty default — making this pure-virtual would catch missing overrides at compile time rather than silently producing uninitialized output.

5. Missing <cstring> include in mrle_transformer.h — currently relies on transitive includes.

6. PCA residual_norm hardcoded to 0.0F — the comment explains this preserves the existing serialized layout and the field is not consumed by RecoveryDistance. A TODO or tracking issue reference would help ensure this placeholder is not forgotten if RecoveryDistance starts consuming it.

No blocking issues found. The core change — replacing TransformerMetaPtr returns with caller-owned uint8_t* scratch — is correct and well-executed.

[Copilot]

Pull request overview

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

Suppressed comments (4)

src/quantization/transform_quantization/transform_quantizer.h:260

• Prefer std::memcpy (from ) over unqualified memcpy to avoid relying on implementation-defined global declarations.

        memcpy(transformed_data.data() + i * transformed_dim,
               transformed,
               transformed_dim * sizeof(float));

src/quantization/transform_quantization/transform_quantizer.h:314

• TransformBaseVector still creates a new secondary_scratch vector per call. For multi-stage transform chains, that will reintroduce per-call heap allocations (capacity can't be reused across calls), which partially undermines the goal of eliminating hot-path allocations. Consider extending TransformBaseVector (or adding an overload) to accept a caller-owned secondary scratch buffer (like ExecuteChainTransform does), and update the call sites that already pass a scratch vector to provide both buffers.

    Vector<float> secondary_scratch(this->allocator_);
    const auto* transformed =
        ExecuteChainTransform(input, nullptr, nullptr, output, secondary_scratch);

src/impl/transform/mrle_transformer.h:43

• Use std::memcpy instead of unqualified memcpy to ensure the function is declared (and to match the rest of the codebase which uses std::memcpy in other transformers).

              float* transformed_vec,
              uint8_t* /*meta*/ = nullptr) const override {
        memcpy(transformed_vec, original_vec, this->output_dim_ * sizeof(float));
        if constexpr (metric == MetricType::METRIC_TYPE_COSINE) {

src/quantization/transform_quantization/transform_quantizer.h:21

• This header uses memcpy (see TrainImpl) but does not include . Relying on transitive includes can fail on stricter toolchains; include here so the header is self-contained.

This issue also appears on line 258 of the same file.

#include <cassert>
#include <sstream>
#include <string>

[LHT129]

I have reviewed this PR. The changes look well-structured and correctly eliminate per-vector shared_ptr allocations in the TransformQuantizer hot path by:

1. Interface change: VectorTransformer::Transform() now takes a raw uint8_t* meta pointer instead of returning a shared_ptr, and is pure virtual — forcing all derived classes to explicitly implement it.
2. Scratch buffer reuse: Computer<TransformQuantizer> now owns primary_scratch_ and secondary_scratch_ members, which are reused across ProcessQuery calls.
3. Ping-pong buffering: ExecuteChainTransform uses caller-provided scratch buffers with ping-pong swapping between transform chain steps, eliminating internal Vector<float> allocations.
4. Batch path optimization: EncodeBatchImpl allocates scratch once outside the loop instead of per-vector, and now correctly checks EncodeOne return values (a silent bug fix from the old code which ignored failures).

The previously raised comments (Copilot and LHT129) have been addressed in the updated commit 8ce20dd: <cstring> include added to mrle_transformer.h, meta_offsets null guard with assert added, and Transform is now pure virtual.

No blocking issues found. The TransformBaseVector secondary scratch allocation noted by Copilot is acknowledged as follow-up work requiring broader API changes to the BottomQuantizerAccessor interface.

[LHT129]

Overall this is a well-structured change that eliminates per-vector heap allocations in the transform-quantizer hot paths by reusing caller-owned scratch buffers. The approach of passing scratch Vector<float>& references through ExecuteChainTransform and storing them in the Computer specialization is clean and avoids the previous std::make_shared allocations on every transform call.

Key design decisions that look correct:

• The ping-pong double-buffering in ExecuteChainTransform correctly handles transform chains of any length without extra allocations.
• The Computer<TransformQuantizer> specialization stores both scratch buffers as members, so the query hot path (ProcessQueryImpl) has zero per-query allocations.
• EncodeBatchImpl now reuses scratch buffers across the entire batch instead of allocating per-vector inside EncodeOneImpl.
• The Transform signature change from TransformerMetaPtr return to void with an uint8_t* meta out-parameter is consistent across all transformer implementations.
• PCA metadata write is guarded by a null check, preserving backward compatibility for callers that do not need metadata.

The existing review comments have covered the notable points (PCA centralized_vec internal allocation, TransformBaseVector secondary scratch, base-class pure-virtual consideration), and the author has addressed or acknowledged each one. No blocking issues remain.

[LHT129]

Overall this is a clean and well-structured change. The core idea — replacing per-vector std::make_shared<Meta> allocations with caller-owned uint8_t* scratch buffers and reusing ping-pong scratch across the transform chain — is sound and directly addresses the hot-path allocation problem.

Summary of review findings:

1. The VectorTransformer::Transform signature change to pure virtual (= 0) is the correct design choice; every derived transformer must implement it, and the empty default was a latent footgun. (Addressed in the latest commit.)

2. The ExecuteChainTransform refactor correctly uses caller-provided primary_scratch/secondary_scratch for ping-pong buffering. The assert + if guard on meta_offsets is reasonable defense-in-depth. (Addressed in the latest commit.)

3. The Computer<TransformQuantizer> specialization now owns primary_scratch_ and secondary_scratch_ as members, which are reused across ProcessQueryImpl calls — this covers the query hot path well.

4. EncodeBatchImpl now reuses scratch buffers across the batch loop and properly propagates EncodeOne failures, which is a nice correctness improvement over the previous code.

5. The TransformBaseVector path still allocates a fresh secondary_scratch per call. This is a known limitation acknowledged by the author as follow-up work requiring a broader API change to BottomQuantizerAccessor.

The PR description mentions validation against TransformQuantizer and transformer tests, and clang-tidy-15 passing. No blocking issues remain.

[LHT129]

Review Summary

This PR eliminates per-vector TransformerMeta heap allocations by changing VectorTransformer::Transform() from returning shared_ptr<TransformerMeta> to accepting a caller-owned uint8_t* scratch buffer, and refactoring ExecuteChainTransform to use ping-pong scratch buffers.

Changes reviewed

• 11 files, +115/-63 lines
• Core changes in transform_quantizer.h (ExecuteChainTransform, EncodeOneImpl, EncodeBatchImpl, ProcessQueryImpl)
• Interface change in vector_transformer.h (Transform now pure virtual with uint8_t* meta parameter)
• All transformer implementations updated (PCA, FHT, ROM, MRLE)
• Computer<TransformQuantizer> now owns primary_scratch_ and secondary_scratch_ members

Issues addressed in commit 8ce20ddb

• VectorTransformer::Transform is now pure virtual (= 0) — prevents silent no-op if a derived class forgets to override
• Added <cstring> include to mrle_transformer.h — fixes reliance on transitive includes
• Added meta_offsets null guard with assert + runtime check in ExecuteChainTransform — prevents null dereference

Remaining notes (non-blocking)

• TransformBaseVector still allocates a fresh secondary_scratch per call; author has acknowledged this as follow-up work
• EncodeBatchImpl duplicates transform+encode logic from EncodeOneImpl; author has self-noted this for potential future refactoring
• PCA residual_norm is hardcoded to 0.0F to preserve serialized code layout; author has noted this is intentional

Verdict

The core optimization is sound: scratch buffers are now caller-owned and reused across calls in the hot path (ProcessQueryImpl, EncodeOneImpl, EncodeBatchImpl). All previously raised concerns have been addressed or acknowledged. No blocking issues found.

[LHT129]

I have reviewed this PR and the changes look solid overall. The core idea — eliminating per-vector TransformerMetaPtr heap allocations by passing caller-owned scratch buffers through ExecuteChainTransform — is well executed.

Summary of what was reviewed:

1. vector_transformer.h: Transform signature changed from returning TransformerMetaPtr to void with an optional uint8_t* meta out-parameter. Making it pure virtual (= 0) is the right call — it prevents silent no-ops from derived classes that forget to override.

2. transform_quantizer.h: The main logic change. ExecuteChainTransform now takes two caller-owned Vector<float>& scratch buffers and ping-pongs between them, avoiding the per-iteration Vector<float> next_data(...) allocation. The meta pointer is now only computed when codes != nullptr AND the transformer actually has metadata (GetMetaSize() > 0), which is correct. The assert(meta_offsets != nullptr) guard in debug builds is a reasonable safety net.

3. computer.h: primary_scratch_ and secondary_scratch_ are added as Computer members, which is the right place — the Computer object is per-query and reused across SetQuery calls, so the scratch buffers are amortized over the query lifetime.

4. Transformer implementations (FHT, PCA, ROM, MRLE): All updated consistently to the new Transform signature. PCA correctly handles the meta != nullptr case by writing the residual norm placeholder.

5. EncodeBatchImpl: Now reuses scratch buffers across the batch loop instead of calling EncodeOneImpl (which would allocate fresh scratch each time). The trade-off between code deduplication and allocation avoidance has been noted in existing review comments.

The existing review thread has already covered the main discussion points (scratch buffer plumbing, TransformBaseVector secondary scratch, code deduplication in EncodeBatchImpl), and the author has addressed the actionable items. No blocking issues remain from my review.

[LHT129]

Review Summary

This PR eliminates per-vector TransformerMetaPtr allocations by changing the Transform interface to accept a caller-owned uint8_t* meta scratch buffer and reusing ping-pong scratch buffers (primary_scratch / secondary_scratch) across the transform chain. The approach is sound and well-scoped.

What was reviewed

• Interface change (vector_transformer.h): Transform now returns void and is pure-virtual (= 0), which is the correct design choice.
• Transformer implementations (FHT, ROM, MRLE, PCA): All correctly updated to the new signature. PCA preserves backward-compatible metadata layout.
• TransformQuantizer (transform_quantizer.h): ExecuteChainTransform now accepts caller-owned scratch buffers, eliminating the per-call Vector<float> allocation. The ping-pong buffer strategy between primary_scratch and secondary_scratch is correct.
• Computer (computer.h): Scratch buffers are now members of the TQ Computer specialization, reused across ProcessQuery calls — this is the key hot-path win.
• Tests (mrle_transformer_test.cpp): Updated to match the new void-return signature.

Existing review comments

Several inline comments from Copilot and the author (LHT129) have already been addressed or discussed:

• VectorTransformer::Transform pure-virtual → resolved (= 0)
• Missing <cstring> include in mrle_transformer.h → resolved
• PCA residual_norm hardcoded to 0.0F → acknowledged with rationale
• TransformBaseVector secondary scratch allocation → author explained concurrency-safe design choice
• EncodeBatchImpl logic duplication with EncodeOneImpl → noted for future refactoring

No blocking issues remain. The changes are consistent with the stated goal and the existing codebase style.

[LHT129]

Thanks for this PR. The approach of eliminating per-vector allocations by reusing caller-owned scratch buffers is sound and well-executed.

I reviewed the full diff across all 11 changed files. The key changes look correct:

• VectorTransformer::Transform is now pure virtual (= 0), which is the right design choice — no silent no-ops.
• ExecuteChainTransform correctly implements ping-pong buffering between primary_scratch and secondary_scratch, and the meta_offsets null-safety with GetMetaSize() > 0 guard is appropriate.
• Computer<TransformQuantizer> now owns primary_scratch_ and secondary_scratch_ as members, which is the right place for query-path scratch reuse.
• PCA Transform preserves the existing serialized code layout by writing residual_norm = 0.0F when meta != nullptr, with a clear comment explaining the current state.

One observation (non-blocking): EncodeBatchImpl now inlines the transform+encode logic that was previously delegated to EncodeOneImpl. This is a reasonable trade-off for scratch buffer reuse, but it means future changes to the encode path need to be kept in sync across both methods. A follow-up could extract a shared helper if this becomes a maintenance concern.

[LHT129]

[suggestion] PCATransformer::Transform still allocates a Vector<float> centralized_vec on every call, which undermines the goal of eliminating per-vector allocations in the hot path. This is called from ExecuteChainTransform for every query and every base vector during encoding.

The centralized_vec buffer is sized to input_dim_ and used as scratch for the centralized data before the BLAS gemv. Consider accepting a caller-owned scratch buffer (or reusing one of the ping-pong buffers already plumbed through ExecuteChainTransform) to avoid this allocation.

// Current code still allocates:
vsag::Vector<float> centralized_vec(allocator_);
centralized_vec.resize(input_dim_, 0.0F);

[LHT129]

[suggestion] EncodeOneImpl creates fresh primary_scratch and secondary_scratch vectors on every call, while ProcessQueryImpl and EncodeBatchImpl reuse caller-owned scratch buffers. If EncodeOneImpl is called in a hot path (e.g., incremental index insertion), this per-call allocation pair defeats the purpose of eliminating hot-path allocations.

Consider either:

1. Accepting scratch buffers as parameters (similar to how ProcessQueryImpl uses computer.primary_scratch_), or
2. Adding a note that EncodeOneImpl is not performance-critical and the allocation is acceptable.

Vector<float> primary_scratch(this->allocator_);
Vector<float> secondary_scratch(this->allocator_);

[LHT129]

Overall this is a well-structured performance improvement that eliminates per-vector shared_ptr allocations in the transform chain by switching to caller-owned scratch buffers. The ping-pong buffer approach in ExecuteChainTransform is clean and avoids the per-iteration Vector<float> allocation that existed before.

Key observations:

• The VectorTransformer::Transform signature change to void with uint8_t* meta is a breaking API change, but all in-tree transformers have been updated consistently.
• The Computer<TransformQuantizer> now owns scratch buffers as members, which is the right place for query-path reuse.
• EncodeBatchImpl inlining the loop (instead of calling EncodeOneImpl) avoids the extra scratch allocation pair per element — good tradeoff even with the code duplication.
• The meta_offsets null-safety guard in ExecuteChainTransform is reasonable for internal callers.

Two suggestions were left as inline comments regarding remaining allocation opportunities in PCATransformer::Transform and EncodeOneImpl.

[LHT129]

[suggestion] EncodeBatchImpl now inlines the transform+encode logic (creating scratch buffers, calling ExecuteChainTransform, then quantizer_->EncodeOne) instead of delegating to EncodeOneImpl. This duplicates the logic between single-encode and batch-encode paths. If the encoding logic changes in the future (e.g., additional pre/post processing), both paths need to be updated independently, which is a maintenance risk.

Consider having EncodeBatchImpl call EncodeOneImpl in the loop and check the return value, so there is a single source of truth for the encode path. The scratch allocation concern in EncodeOneImpl (noted in a separate comment) can be addressed by accepting scratch buffers as parameters or using thread-local storage.

[LHT129]

Review Summary

This PR eliminates per-vector TransformerMeta heap allocations by changing the Transform() interface from returning shared_ptr<Meta> to accepting a caller-owned uint8_t* scratch buffer, and by introducing ping-pong scratch buffers in ExecuteChainTransform and the Computer<TransformQuantizer> specialization.

What was reviewed

• Interface change: VectorTransformer::Transform is now pure virtual (= 0) with a uint8_t* meta = nullptr parameter — all derived transformers updated consistently.
• Scratch buffer management: ExecuteChainTransform uses caller-provided primary_scratch/secondary_scratch with correct ping-pong swapping between transform steps.
• Query path: ProcessQueryImpl reuses computer.primary_scratch_/computer.secondary_scratch_ — no per-query allocations.
• Encode paths: EncodeOneImpl and EncodeBatchImpl both create local scratch buffers and pass them to ExecuteChainTransform.
• PCA metadata: PCATransformer::Transform writes a residual_norm = 0.0F placeholder into the caller-provided meta buffer, preserving the existing serialized code layout.

Issues already addressed (commit 8ce20ddb)

• VectorTransformer::Transform made pure virtual (was empty default impl)
• <cstring> include added to mrle_transformer.h
• meta_offsets null-pointer guard with assert + runtime check added in ExecuteChainTransform

Remaining discussion points (acknowledged by author)

• TransformBaseVector creates a fresh secondary_scratch per call — author notes this is a generic API constraint (used by BottomQuantizerAccessor) and out of scope for this PR.
• PCATransformer::Transform still allocates centralized_vec internally — author notes a member scratch buffer would introduce data races under concurrent use.
• EncodeBatchImpl inlines transform+encode instead of calling EncodeOneImpl — author confirms this is intentional to retain batch-owned scratch buffers.

Verdict

No blocking issues found. The core design (caller-owned scratch buffers, ping-pong in ExecuteChainTransform, Computer-owned scratch for the query path) is sound and directly eliminates the targeted hot-path allocations. The remaining allocation sites have been discussed and are either constrained by API design or deferred to follow-up work.

[Copilot]

Pull request overview

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

Suppressed comments (1)

src/quantization/transform_quantization/transform_quantizer.h:260

• The linked issue (#2702) calls out capturing before/after allocation + latency benchmarks (e.g., 100k x 768 with the rom,pca,fp32 chain) before merging. This PR changes the hot Train/Encode/Query transform path, but the PR description doesn’t include any benchmark numbers or allocation evidence yet; without that it’s hard to validate the intended perf win and catch regressions (e.g., accidental reallocations in the scratch ping-pong logic). Please add the requested benchmark results (or a lightweight allocation/throughput report) to the PR description before merge.

    // 2. execute transform on original data
    const uint64_t transformed_dim = this->GetTransformedDim();
    Vector<float> transformed_data(transformed_dim * count, 0, this->allocator_);
    Vector<float> primary_scratch(this->allocator_);
    Vector<float> secondary_scratch(this->allocator_);
    for (uint64_t i = 0; i < count; ++i) {
        const auto* transformed = ExecuteChainTransform(
            data + i * this->dim_, nullptr, nullptr, primary_scratch, secondary_scratch);
        memcpy(transformed_data.data() + i * transformed_dim,
               transformed,
               transformed_dim * sizeof(float));

[LHT129]

[suggestion] PCATransformer::Transform still allocates a Vector<float> centralized_vec on every call (line 70). Since this PR aims to eliminate per-vector transform allocations, this remaining allocation inside the PCA transform hot path partially defeats the purpose.

The centralized_vec buffer has a fixed size of input_dim_. Consider accepting a caller-owned scratch buffer for this purpose, or pre-allocating it as a member (with appropriate synchronization if needed for thread safety).

// Current code still allocates:
vsag::Vector<float> centralized_vec(allocator_);
centralized_vec.resize(input_dim_, 0.0F);

[LHT129]

[note] The assert(meta_offsets != nullptr) on line 318 will abort in debug builds if codes != nullptr but meta_offsets is null. However, the subsequent if (meta_offsets != nullptr && ...) guard on line 319 handles this gracefully in release builds. This mixed assert+guard pattern is fine defensively, but if the invariant is truly that codes != nullptr implies meta_offsets != nullptr, consider making the assert unconditional (remove the if guard) to catch violations early. If null meta_offsets is a valid scenario, the assert should be removed to avoid debug-only crashes.

[LHT129]

[note] EncodeBatchImpl now propagates EncodeOne failures (returns false), whereas the original implementation always returned true regardless of individual encode results. This is a correctness improvement, but callers that previously assumed EncodeBatchImpl never fails may now encounter error paths that were previously silently ignored. Ensure all callers handle the false return appropriately.

// New behavior: returns false on first encode failure
if (not quantizer_->EncodeOne(transformed, output_codes)) {
    return false;
}