Skip to content

Optimize loop hafnian: share powertrace computation between f_loop calls#409

Open
seba2390 wants to merge 4 commits intoXanaduAI:masterfrom
seba2390:optimize-loop-hafnian-batch-gamma
Open

Optimize loop hafnian: share powertrace computation between f_loop calls#409
seba2390 wants to merge 4 commits intoXanaduAI:masterfrom
seba2390:optimize-loop-hafnian-batch-gamma

Conversation

@seba2390
Copy link
Copy Markdown

Context:
In _calc_loop_hafnian_batch_gamma_even and _calc_loop_hafnian_batch_gamma_odd, the functions f_loop() and f_loop_odd() are called inside a loop over displacement vectors (k in range(n_D)). Both functions internally compute powertrace(AX), however, the matrix AX depends only on the edge configuration (j), and not on the displacement vector (D[k]). The original code redundantly computes the same powertrace 2 × n_D times per j iteration.

Description of the Change:

  1. Modified f_loop and f_loop_odd in _hafnian.py to accept powtrace_arr as a parameter instead of computing it internally.
  2. Updated callers in loop_hafnian_batch.py and loop_hafnian_batch_gamma.py to compute powertrace once per j iteration and pass it to both functions.

Benefits:

  • 2-6x speedup on loop_hafnian_batch_gamma for typical batch sizes (10-50 displacement vectors)
  • Speedup scales with batch size — minimal overhead for batch_size=1, significant gains for larger batches
  • Benchmark results on Apple M3 Pro and M1 Pro available at: https://github.com/seba2390/thewalrus_opt_benchmark

Possible Drawbacks:

  • Changes the internal API of f_loop and f_loop_odd (now require powtrace_arr parameter). These are internal functions not exposed in the public API.

Related GitHub Issues:

None

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant