Skip to content

Add gate-batched cache-local runner across SIMD backends#1092

Draft
xxie24 wants to merge 1 commit into
quantumlib:mainfrom
xxie24:blocked-v2
Draft

Add gate-batched cache-local runner across SIMD backends#1092
xxie24 wants to merge 1 commit into
quantumlib:mainfrom
xxie24:blocked-v2

Conversation

@xxie24

@xxie24 xxie24 commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

build_gate_batch.sh

Summary

This PR adds an experimental gate-batched, cache-local state-vector runner that works across NEON, AVX2, and AVX-512 backends.

Algorithm

The state vector is divided into contiguous blocks of 2^L amplitudes. Each block represents the state space of L physical qubits while the remaining n-L qubits have a fixed bit assignment.

While circuit gates remain:

  1. Select a causally executable gate batch whose logical qubits can be placed within the L block-local physical positions.
  2. Remap the required logical qubits into those positions. All disjoint qubit-position swaps are applied together in one state-vector pass.
  3. Fuse the selected gates once to produce the executable gates for the batch.
  4. Execute the batch using the original proposal's three-level loop:
while (gates_remain) {
  const auto plan = PlanNextGateBatch();
  ApplySwapsToState(plan);
  const auto executable_gates = FuseBatchGates(plan);

  #pragma omp parallel for schedule(dynamic, 1)
  for (int64_t block_index = 0; block_index < num_state_blocks; ++block_index) {
    auto block = GetStateBlock(block_index);

    for (const auto& gate : executable_gates) {
      sequential_simulator.ApplyGate(gate, block);
    }
  }
}

The outer block loop is parallelized with OpenMP. A worker thread handles one contiguous L-qubit block at a time and applies all executable gates to that block using a sequential SIMD simulator. After finishing the block, the thread receives another block.

The innermost local-index loop remains inside the existing simulator gate kernel. This preserves the architecture-specific NEON, SSE, AVX2, or AVX-512 implementation.

Frequently used logical qubits are initially placed in a fixed block-local region. SIMD lane positions remain pinned because state remapping moves whole SIMD chunks rather than individual lanes.

Performance

Circuit: circuit_q30, 1617 raw gates. Times are in seconds; lower is better.

Platform SIMD backend CPU allocation Parameters Base runner Gate-batch runner Gate passes Swap passes Speedup
Apple M2 Max NEON 12 threads L=19, f=3 ~56.00 46.18 41.72 4.46 ~1.21×
Google C4 / Intel SSE4.2 12 physical cores L=19, f=3 60.12 59.54 51.84 7.69 1.01×
Google C4 / Intel AVX-512 12 physical cores L=19, f=3 52.84 27.70 19.90 7.79 1.91×
Google C4D / AMD Turin AVX-512 8 physical cores, SMT disabled L=19, f=3 78.74 34.40 22.43 11.97 2.29×
Google C4A / Axion NEON 16 physical cores L=19, f=3 31.77 30.88 27.83 3.05 1.03×

Times are in seconds and lower is better. Speedup is the base-runner time divided by the gate-batch-runner time. Results should be compared within each platform rather than across different machines.

The SSE and AVX-512 C4 results were collected on the same c4-standard-24 VM using 12 threads pinned one per physical core. Gate batching provides almost no improvement with SSE, but improves AVX-512 performance by 1.91×. Swap time remains approximately unchanged; the primary difference is gate execution, where AVX-512 reduces the gate-pass time from 51.84 seconds to 19.90 seconds.

This indicates that the base runner is largely memory-bandwidth-bound. After gate batching improves cache locality, SIMD compute throughput becomes more important, allowing AVX-512 to benefit substantially from its wider vectors.

Validation

(build script for all platforms is attached)

  • Compared full state vectors against the existing qsim runner.
  • Tested multiple block and fusion sizes.
  • Tested NEON, AVX2, and AVX-512 builds.
  • Verified generated NEON and AVX-512 SIMD instructions.
  • Observed numerical differences around 1e-9 in the regression tests.

Introduce an experimental runner that executes circuits as gate batches
over cache-sized state blocks. Each batch selects causally executable
raw gates, maps their logical qubits into block-local physical positions,
fuses the mapped gates, and executes them across all state blocks in
parallel.

Bound planning cost by comparing greedy, resident, and upcoming-gate-
seeded candidates. Score candidates by circuit progress while accounting
for the number of required state-vector swap passes.

Apply each set of disjoint qubit-position swaps in one state-vector pass.
Keep architecture-specific SIMD lane positions pinned so remapping moves
whole chunks on NEON, SSE, AVX2, and AVX-512 layouts.

Add QubitMappedState to retain the logical-to-physical qubit layout
across execution. The mapped-state API returns without restoring
canonical qubit order, while the raw-state compatibility API restores
identity before returning.

Add a command-line runner for exercising the gate-batch implementation
and register the runner and supporting headers with the Bazel build.
Register a Makefile target that builds both the gate-batch runner and
the existing base runner.

Signed-off-by: Xin Xie <xinpin@google.com>
@github-actions github-actions Bot added the size: XL lines changed >1000 label Jul 23, 2026

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Code Review

This pull request introduces a gate-batched cache-local simulation runner (QSimGateBatchRunner) along with qubit remapping and layout utilities (QubitLayout, QubitMappedState, and ApplyBitPairSwaps) to optimize simulation performance. It also adds a benchmark application (qsim_gate_batch) and updates the build configurations. Feedback was provided to relax an assertion in run_qsim_gate_batch.h that would otherwise cause failures for small circuits where the number of qubits is less than the SIMD chunk size.

Comment thread lib/run_qsim_gate_batch.h
layout_(layout),
gate_batch_planner_(partition_.num_state_qubits,
partition_.block_qubits, chunk_qubits_) {
assert(partition_.block_qubits >= chunk_qubits_);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

high

The assertion assert(partition_.block_qubits >= chunk_qubits_); will fail for small circuits where the number of qubits is less than the SIMD chunk size (e.g., a 2-qubit circuit on an AVX-512 backend where chunk_qubits_ is 4). Since cache blocking is not required for circuits that fit entirely within a single SIMD chunk, we should relax this assertion to allow block_qubits to be smaller than chunk_qubits_ when the total number of state qubits is also smaller than chunk_qubits_.

Suggested change
assert(partition_.block_qubits >= chunk_qubits_);
assert(partition_.block_qubits >= std::min(chunk_qubits_, partition_.num_state_qubits));

@xxie24
xxie24 marked this pull request as draft July 23, 2026 21:12
@xxie24

xxie24 commented Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

@sergeisakov (I am OOO this week) this is pretty much as your proposal #1 step with minor change. This version's remapping is using the SIMD by auto vectorization (also in-place permutation, so no extra buffer), so I verified it actually support all the SIMD platform with minimal efforts.

I noticed X86 AVX512 has a very large improvement as you mentioned early (I was struggling to observe it on ARM NEON in the past weeks), so I pushed this PR as a draft to share this early result.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size: XL lines changed >1000

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant