Skip to content

[Enhance]: MatrixRotator::rotate() strided column access — loop interchange gives ~12–15× single-thread speedup #533

Description

@cluster2600

Summary

MatrixRotator::rotate() computes the (1×dim) · (dim×dim) rotation as a
matrix-vector product with the reduction on the outer loop and the matrix
indexed by column
(matrix_[i * dim + j], i inner → stride dim). For a
row-major matrix this walks each column with a dim-element stride, which
thrashes cache/TLB for large dim and prevents the inner loop from
auto-vectorizing.

Interchanging the loops (input index i outer, output index j inner, with
out[] as the accumulator) makes the matrix read row-contiguous
(matrix_[i*dim + j] steps by 1). The summation over i stays in the same
order, so the result is numerically identical up to FMA/vectorization rounding.

Measured ~12–15× single-thread speedup, bit-equivalent.

Location

src/core/quantizer/rotator/matrix_rotator.cc, MatrixRotator::rotate().

// current — matrix read by column (stride = dim)
for (size_t j = 0; j < dim; ++j) {
  float sum = 0.0f;
  for (size_t i = 0; i < dim; ++i) {
    sum += in[i] * matrix_[i * dim + j];   // stride-dim access
  }
  out[j] = sum;
}

Note: the sibling unrotate() already uses the contiguous pattern
(matrix_[j * dim + i], i inner steps by 1), so only rotate() is affected.

Proposed fix (loop interchange)

for (size_t j = 0; j < dim; ++j) out[j] = 0.0f;
for (size_t i = 0; i < dim; ++i) {
  const float xi = in[i];
  const float *row = &matrix_[i * dim];   // contiguous row
  for (size_t j = 0; j < dim; ++j) {
    out[j] += xi * row[j];
  }
}

Pure scalar, no SIMD intrinsics and no OpenMP — portable across every target
(x86 / arm64 / RISC-V) and consistent with keeping ENABLE_OPENMP off. It also
lets the compiler auto-vectorize the now-unit-stride inner loop.

Benchmark

Standalone microbench of the two loop orders, clang++ -O3 -std=c++17,
single thread, Apple Silicon, best-of-5, random orthonormal-ish matrix.
Per-call time for one rotate():

dim current (strided) interchanged speedup
128 12.3 µs 1.0 µs 11.9×
256 60.5 µs 5.1 µs 11.8×
512 290.8 µs 20.5 µs 14.2×
768 691.3 µs 46.3 µs 14.9×

max_abs_diff vs the current order ≤ 1e-6 at every size (FMA/vectorization
rounding only; identical summation order).

Why it matters

rotate() is on the query path when matrix rotation is enabled (the random-
rotation rotator added in #483). At dim=768 the current order costs ~0.69 ms
per rotation; the interchange brings that under 0.05 ms with no accuracy change.

Environment

  • Discovered on macOS / Apple Silicon, clang++ -O3. The access-pattern issue
    is architecture-independent (column-stride vs unit-stride on a row-major
    array); the magnitude scales with dim.

Notes

Surfaced via a polyhedral auto-scheduling pass
(cluster_compilot, an
implementation of Agentic Auto-Scheduling, arXiv:2511.00592): loop
interchange is proven legal (no carried dependence) before the bit-equivalence
check. Happy to open a PR with the change plus a small benchmark if this looks
useful.

Metadata

Metadata

Labels

No labels
No labels

Type

No type

Fields

No fields configured for issues without a type.

Projects

Status
In progress

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions