Skip to content

Memoize compile_decomposition_rules per operator variant - #3174

Open
Qottmann wants to merge 3 commits into
mainfrom
qottmann/cache-compile-decomposition-rules
Open

Qottmann wants to merge 3 commits into
mainfrom
qottmann/cache-compile-decomposition-rules

Conversation

@Qottmann

@Qottmann Qottmann commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

Context:
Closes #3173.

Description of the Change:
Under qml.qjit(capture=True) graph decomposition, the compiler's on-demand rule loader rebuilds the full reachable-rule closure for every operator it queries. As a result, compile_decomposition_rules recompiles the same operator variant many times over — on the RZ phase-gradient demo it was invoked ~5870 times for only ~179 unique variants, making rule compilation the dominant cost of decomposition-heavy programs.

compile_decomposition_rules is a pure function of its arguments, so this PR memoizes its result in a module-level cache (_COMPILE_DECOMP_CACHE) keyed by the fully-resolved operator instance and modifier variant (op_name, op_id, dynamic_shape, wire_lens, static_data, extra_data, is_custom_op, wrap_adjoint, wrap_control, n_ctrl). Each unique variant is now compiled once.

Correctness: the cached inlined_module is only ever read downstream (walked and cloned), never mutated in place, so sharing the same module across callers is safe.

Benefits:
Eliminates redundant recompilation of identical decomposition-rule variants, substantially reducing compile time for decomposition-heavy programs under capture.

Possible Drawbacks:
The cache persists for the process lifetime. Values are keyed by a stable, hashable representation of the fully-resolved operator variant, so distinct variants remain distinct.

Related GitHub Issues:
Closes #3173.

Made with Cursor

Qottmann and others added 3 commits September 2, 2026 11:42
The on-demand rule loader rebuilds the full reachable-rule closure for every
operator the compiler queries, recompiling identical operator variants many
times over. Cache the compiled module keyed by the fully-resolved operator
instance and modifier variant; the returned module is only walked/cloned
downstream, never mutated, so sharing it across callers is safe.

Fixes #3173

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
The compiled module returned by compile_decomposition_rules depends on the
decomposition rules currently registered for the operator, not just the
operator variant. Keying the memoization cache on the operator variant alone
returned a stale module when the registered rule set changed between programs
(e.g. across qml.decomposition.local_decomps blocks), breaking
test_no_distribution_rule_for_non_invertible_body.

Fingerprint the registered rule set into the cache key and pin the rule
objects in the cache value so their ids cannot be recycled. Also apply black
formatting to the cache-key helper.

Co-authored-by: Cursor <cursoragent@cursor.com>
@codecov

codecov Bot commented Sep 2, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 89.47368% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.64%. Comparing base (d408f1f) to head (94bba5d).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
...tend/catalyst/decomposition/decomposition_rules.py 89.47% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #3174      +/-   ##
==========================================
- Coverage   95.65%   95.64%   -0.01%     
==========================================
  Files         177      177              
  Lines       20763    20782      +19     
  Branches     2080     2083       +3     
==========================================
+ Hits        19860    19877      +17     
- Misses        718      720       +2     
  Partials      185      185              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

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.

compile_decomposition_rules recompiles identical operator variants many times under capture

1 participant