Add Fine-Grained Circuit Partitioning#189
Open
linusschulte wants to merge 16 commits into
Open
Conversation
…ter gate to pz assignments; improved qasm parsing stability
…e_id-aware. Instead of the graph.sequence it now works with graph.gate_ids and graph.gate_info which allows reconstruction of gates upon compilation ouput and also enables reliable handling of gate-partitioning donwstream. Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <copilot@github.com>
… and /outside shuttler which can be used to predetermine the pz in which a gate is executed. If not set, the scheduling falls back to the standard gate-pz-assignment routines, reproducing the pre-existing behaviour.
…s an optimized gate_to_pz_assignment which can be fed into shuttle() to predetmerine which gate is executed in which pz, which in most cases should reduce shuttling overhead over the baseline local gate-pz-assignments decisions.
…fault parameters settings by setting the config flag 'use_fine_grained_gate_partition' to True
📝 WalkthroughSummary by CodeRabbitRelease Notes
WalkthroughAdds a QASM parser producing a ParsedCircuit with per-gate GateInfo and integer gate IDs, refactors graph/scheduler/shuttle codepaths to use gate IDs and GateRef, and implements a fine-grained tabu-based gate→processing-zone partitioner integrated into main runtime and tests. Changes
Sequence Diagram(s)sequenceDiagram
rect rgba(200,220,255,0.5)
participant User
end
rect rgba(200,255,200,0.5)
participant Main
participant Parser
participant Partitioner
participant Graph
participant Scheduler
end
User->>Main: invoke main(...) / run pipeline
Main->>Parser: create_initial_circuit(filename)
Parser-->>Main: ParsedCircuit(sequence:[gate_id...], gate_info)
Main->>Graph: graph.sequence = parsed.sequence\ngraph.gate_info = parsed.gate_info
alt fine-grained enabled
Main->>Partitioner: compute_fine_grained_gate_partition(sequence, gate_info, pz_names, distances)
Partitioner-->>Main: GatePartitionResult (gate_pz_assignment)
Main->>Graph: graph.gate_pz_assignment = result.gate_assignment
end
Main->>Scheduler: create_priority_queue(graph)
Scheduler->>Graph: graph.gate_qubits(gate_id) / preferred_pz_for_gate(gate_id)
Scheduler-->>Main: scheduling structures
Main->>Scheduler: run shuttle loop / execute steps
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/mqt/ionshuttler/multi_shuttler/outside/scheduling.py (1)
110-138: 🧹 Nitpick | 🔵 TrivialOptional: collapse identical
legacy/dynamicpolicy branches into one.Today both arms of
if policy == "legacy" / elif policy == "dynamic"(Lines 112–119 and 126–133) execute the same code paths, with a TODO-style comment noting that "dynamic" will diverge later. Until that scaffolding is filled in, you could reduce duplication and simplify by collapsing them, e.g.:if policy in {"legacy", "dynamic"}: chosen_pz = graph.map_to_pz[ion] else: msg = f"Unknown pz_assignment_policy: {policy}" raise ValueError(msg)This would also reduce
assign_gate_to_pzbelow the PLR0912 branch threshold reported by ruff. Feel free to defer until "dynamic" actually diverges.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/mqt/ionshuttler/multi_shuttler/outside/scheduling.py` around lines 110 - 138, Collapse the duplicate "legacy" and "dynamic" branches in assign_gate_to_pz so both policies follow the same path until the dynamic heuristic is implemented: for the single-qubit case use if policy in {"legacy","dynamic"} to select chosen_pz = graph.map_to_pz[ion] (otherwise raise the same ValueError), and for the two-qubit case use if gate_id in graph.locked_gates else handle ion0,ion1 with if policy in {"legacy","dynamic"} to call chosen_pz = pick_pz_for_2_q_gate(graph, ion0, ion1) (then set graph.locked_gates[gate_id] if gate_id present); keep the Unknown pz_assignment_policy ValueError behavior unchanged.src/mqt/ionshuttler/multi_shuttler/inside/scheduling.py (1)
144-186:⚠️ Potential issue | 🔴 CriticalCritical: empty-PZ fill is inside the gate loop and 1-qubit branch lacks
== ()check —next_gate_at_pzbecomes incorrect.Two coupled bugs in
create_priority_queue:
- The "fill empty pzs with
()" block (Lines 182–185) is indented inside thefor gate in sequence_to_useloop. As a result, after the first gate is processed, every other PZ already has the value()planted innext_gate_at_pz.- The 1-qubit branch (Line 155) only checks
if pz_name not in next_gate_at_pz, missing theor next_gate_at_pz[pz_gate] == ()guard that the 2-qubit branch correctly uses on Line 169.Concrete failing scenario (two PZs, gates
[cx q[1],q[2] → pz1, x q[0] → pz2]):
- After iter 1 (2q): the inner empty-fill plants
next_gate_at_pz["pz2"] = ().- In iter 2 (1q): Line 155 sees
"pz2" in next_gate_at_pzand skips the assignment, sonext_gate_at_pz["pz2"]stays()instead of becoming the 1q gate id.Additionally, when
sequence_to_useis empty, the for-body never runs, so the empty-fill never executes andnext_gate_at_pzis returned as{}, which causesKeyErrorwhen consumers index by PZ name.The analogous code in
outside/scheduling.pycorrectly has the fill block outside the loop and== ()on both arity branches.🐛 Proposed fix
for gate in sequence_to_use: qubits = graph.gate_qubits(gate) # 1-qubit gate if len(qubits) == 1: elem = qubits[0] pz_name = assign_gate_to_pz(graph, gate) # add first gate of pz to next_gate_at_pz (if not already there) - if pz_name not in next_gate_at_pz: + if pz_name not in next_gate_at_pz or next_gate_at_pz[pz_name] == (): next_gate_at_pz[pz_name] = gate # add ion to unique_sequence if elem not in unique_sequence: unique_sequence[elem] = pz_name if len(unique_sequence) == max_length: break # 2-qubit gate elif len(qubits) == 2: pz_for_2_q_gate = assign_gate_to_pz(graph, gate) # add first gate of pz to next_gate_at_pz (if not already there) if pz_for_2_q_gate not in next_gate_at_pz or next_gate_at_pz[pz_for_2_q_gate] == (): next_gate_at_pz[pz_for_2_q_gate] = gate # add ions to unique_sequence for elem in qubits: if elem not in unique_sequence: unique_sequence[elem] = pz_for_2_q_gate if len(unique_sequence) == max_length: break else: msg = "len gate 0 or > 2? - can only process 1 or 2-qubit gates" raise ValueError(msg) - # at the end fill all empty pzs with () - for pz in graph.pzs: - if pz.name not in next_gate_at_pz: - next_gate_at_pz[pz.name] = () + # at the end fill all empty pzs with () + for pz in graph.pzs: + if pz.name not in next_gate_at_pz: + next_gate_at_pz[pz.name] = () return unique_sequence, next_gate_at_pz🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/mqt/ionshuttler/multi_shuttler/inside/scheduling.py` around lines 144 - 186, In create_priority_queue, fix two bugs in how next_gate_at_pz is populated: move the "for pz in graph.pzs: if pz.name not in next_gate_at_pz: next_gate_at_pz[pz.name] = ()" block so it executes after the main "for gate in sequence_to_use" loop (ensuring it runs even when sequence_to_use is empty), and change the 1-qubit branch check from "if pz_name not in next_gate_at_pz" to "if pz_name not in next_gate_at_pz or next_gate_at_pz[pz_name] == ()" to match the 2-qubit branch; keep references to assign_gate_to_pz, unique_sequence, next_gate_at_pz, and graph.pzs when locating the code to modify.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.py`:
- Around line 1052-1099: The parameter current_cluster in _distance_delta is
unused; remove it from the function signature and from all call sites (e.g.
calls from _consider_supernode_moves) so the function uses the per-qubit current
= current_slice_assignments[qubit] logic that already exists; update the
_distance_delta definition to drop current_cluster and remove that argument
where _distance_delta(...) is invoked, and run a quick search to ensure no
remaining references expecting that parameter.
- Line 909: The return expression at the end using best_assignments,
assignments_by_slice, initial_cost, best_cost is ambiguous due to operator
precedence; change the return to explicitly parenthesize the conditional part so
it reads like return (best_assignments if best_assignments else
assignments_by_slice), initial_cost, best_cost to make the intent unambiguous
(referencing the variables best_assignments, assignments_by_slice, initial_cost,
best_cost in the existing function).
- Around line 530-531: The loop currently iterates required_edges.items() and
discards the weight; change it to iterate the dict keys directly so you only
unpack the (left,right) tuple (e.g., replace the items() iteration with
iterating required_edges itself) in the block where union_find.union(left,
right) is called (look for the loop around union_find.union in
gate_partitioning_tabu.py) to remove the unused _weight and satisfy the PERF102
suggestion.
In `@src/mqt/ionshuttler/multi_shuttler/outside/compilation.py`:
- Around line 113-134: build_dag_gate_id_lookup currently uses per-register
q._index which mismatches the globally-normalized GateInfo.qubits; change it to
compute a qubit->global_index map from the DAG's qregs (iterate dag.qregs
values, accumulate offset per register, map each qubit object to
offset+local_idx) and then replace tuple(q._index for q in node.qargs) with
tuple(qubit_to_global[q] for q in node.qargs) in build_dag_gate_id_lookup; apply
the same fix in find_best_gate where it likewise uses q._index so both functions
use the same qubit_to_global mapping and produce globally-normalized qubit
indices that match GateInfo.qubits.
In `@src/mqt/ionshuttler/multi_shuttler/outside/shuttle.py`:
- Around line 371-372: The assignment graph.locked_gates = locked_gates is
redundant inside the timestep loop; remove the repeated re-binding (the in-loop
graph.locked_gates = locked_gates) and keep the single initial binding where
locked_gates: dict[int, str] = {} and graph.locked_gates = locked_gates are set,
since assign_gate_to_pz and other code mutate the dict in place and never
reassign graph.locked_gates to a new dict.
In `@tests/test_multi_shuttler.py`:
- Around line 349-402: The create_priority_queue implementation fails to update
next_gate_at_pz for 1-qubit gates when a PZ slot was prefilled with the sentinel
() (the 2-qubit branch checks for "== ()" but the 1-qubit branch does not), so
update inside/scheduling.py:create_priority_queue to treat a prefilled () as
empty for 1-qubit gates and assign the gate id into next_gate_at_pz[pz] (mirror
the sentinel check used for 2-qubit gates), and add the suggested regression
test (sequence = [2q_for_pz1, 1q_for_pz2] with two PZs) to
tests/test_multi_shuttler.py to assert next_gate_at_pz["pz2"] equals the 1-qubit
gate id rather than () so this bug cannot regress.
---
Outside diff comments:
In `@src/mqt/ionshuttler/multi_shuttler/inside/scheduling.py`:
- Around line 144-186: In create_priority_queue, fix two bugs in how
next_gate_at_pz is populated: move the "for pz in graph.pzs: if pz.name not in
next_gate_at_pz: next_gate_at_pz[pz.name] = ()" block so it executes after the
main "for gate in sequence_to_use" loop (ensuring it runs even when
sequence_to_use is empty), and change the 1-qubit branch check from "if pz_name
not in next_gate_at_pz" to "if pz_name not in next_gate_at_pz or
next_gate_at_pz[pz_name] == ()" to match the 2-qubit branch; keep references to
assign_gate_to_pz, unique_sequence, next_gate_at_pz, and graph.pzs when locating
the code to modify.
In `@src/mqt/ionshuttler/multi_shuttler/outside/scheduling.py`:
- Around line 110-138: Collapse the duplicate "legacy" and "dynamic" branches in
assign_gate_to_pz so both policies follow the same path until the dynamic
heuristic is implemented: for the single-qubit case use if policy in
{"legacy","dynamic"} to select chosen_pz = graph.map_to_pz[ion] (otherwise raise
the same ValueError), and for the two-qubit case use if gate_id in
graph.locked_gates else handle ion0,ion1 with if policy in {"legacy","dynamic"}
to call chosen_pz = pick_pz_for_2_q_gate(graph, ion0, ion1) (then set
graph.locked_gates[gate_id] if gate_id present); keep the Unknown
pz_assignment_policy ValueError behavior unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 6fc25bc9-dd73-4b46-9d35-39d19403db94
📒 Files selected for processing (16)
src/mqt/ionshuttler/multi_shuttler/circuit_parsing.pysrc/mqt/ionshuttler/multi_shuttler/circuit_types.pysrc/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.pysrc/mqt/ionshuttler/multi_shuttler/inside/compilation.pysrc/mqt/ionshuttler/multi_shuttler/inside/graph.pysrc/mqt/ionshuttler/multi_shuttler/inside/scheduling.pysrc/mqt/ionshuttler/multi_shuttler/inside/shuttle.pysrc/mqt/ionshuttler/multi_shuttler/main.pysrc/mqt/ionshuttler/multi_shuttler/outside/compilation.pysrc/mqt/ionshuttler/multi_shuttler/outside/graph.pysrc/mqt/ionshuttler/multi_shuttler/outside/scheduling.pysrc/mqt/ionshuttler/multi_shuttler/outside/shuttle.pytests/test_gate_partitioning_tabu.pytests/test_main_mode_config.pytests/test_multi_shuttler.pytests/test_outside_scheduling.py
Comment on lines
+530
to
+531
| for (left, right), _weight in required_edges.items(): | ||
| union_find.union(left, right) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Optional: drop the unused weight when only the keys are needed.
Iterating required_edges.items() and discarding _weight is unnecessary; iterate the dict directly:
- for (left, right), _weight in required_edges.items():
- union_find.union(left, right)
+ for left, right in required_edges:
+ union_find.union(left, right)Also resolves the ruff PERF102 hint.
🧰 Tools
🪛 Ruff (0.15.12)
[warning] 530-530: When using only the keys of a dict use the keys() method
Replace .items() with .keys()
(PERF102)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.py` around lines
530 - 531, The loop currently iterates required_edges.items() and discards the
weight; change it to iterate the dict keys directly so you only unpack the
(left,right) tuple (e.g., replace the items() iteration with iterating
required_edges itself) in the block where union_find.union(left, right) is
called (look for the loop around union_find.union in gate_partitioning_tabu.py)
to remove the unused _weight and satisfy the PERF102 suggestion.
Comment on lines
+371
to
372
| locked_gates: dict[int, str] = {} | ||
| graph.locked_gates = locked_gates |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Redundant re-binding of graph.locked_gates.
Line 372 assigns graph.locked_gates = locked_gates, and Line 390 re-binds it to the same object every iteration of the timestep loop. Since neither the function nor assign_gate_to_pz ever rebinds graph.locked_gates to a new dict (they mutate in place), Line 390 is dead. Drop it to remove the false impression that re-syncing is necessary.
Also applies to: 389-391
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/mqt/ionshuttler/multi_shuttler/outside/shuttle.py` around lines 371 -
372, The assignment graph.locked_gates = locked_gates is redundant inside the
timestep loop; remove the repeated re-binding (the in-loop graph.locked_gates =
locked_gates) and keep the single initial binding where locked_gates: dict[int,
str] = {} and graph.locked_gates = locked_gates are set, since assign_gate_to_pz
and other code mutate the dict in place and never reassign graph.locked_gates to
a new dict.
Comment on lines
+349
to
+402
| def test_inside_priority_queue_accepts_gate_ids(self): | ||
| """The inside scheduler should accept gate-id sequences backed by metadata.""" | ||
| from unittest.mock import patch | ||
|
|
||
| from mqt.ionshuttler.multi_shuttler.circuit_types import GateInfo | ||
| from mqt.ionshuttler.multi_shuttler.inside.scheduling import create_priority_queue | ||
|
|
||
| gate_info = { | ||
| 0: GateInfo(qubits=(0,), qasm="x q[0];"), | ||
| 1: GateInfo(qubits=(1, 2), qasm="cx q[1],q[2];"), | ||
| } | ||
| graph = SimpleNamespace( | ||
| sequence=[0, 1], | ||
| gate_info=gate_info, | ||
| gate_qubits=lambda gate: gate_info[gate].qubits if isinstance(gate, int) else gate, | ||
| preferred_pz_for_gate=lambda _gate_id: None, | ||
| map_to_pz={0: "pz1", 1: "pz1", 2: "pz2"}, | ||
| locked_gates={}, | ||
| pzs=[SimpleNamespace(name="pz1"), SimpleNamespace(name="pz2")], | ||
| ) | ||
|
|
||
| with patch( | ||
| "mqt.ionshuttler.multi_shuttler.inside.scheduling.pick_pz_for_2_q_gate", | ||
| return_value="pz2", | ||
| ): | ||
| priority_queue, next_gate_at_pz = create_priority_queue(cast("Any", graph)) | ||
|
|
||
| assert priority_queue == {0: "pz1", 1: "pz2", 2: "pz2"} | ||
| assert next_gate_at_pz == {"pz1": 0, "pz2": 1} | ||
|
|
||
| def test_inside_priority_queue_prefers_explicit_gate_assignment(self): | ||
| """The inside scheduler should honor explicit gate-to-PZ overrides.""" | ||
| from mqt.ionshuttler.multi_shuttler.circuit_types import GateInfo | ||
| from mqt.ionshuttler.multi_shuttler.inside.scheduling import create_priority_queue | ||
|
|
||
| gate_info = { | ||
| 0: GateInfo(qubits=(0,), qasm="x q[0];"), | ||
| 1: GateInfo(qubits=(1, 2), qasm="cx q[1],q[2];"), | ||
| } | ||
| graph = SimpleNamespace( | ||
| sequence=[0, 1], | ||
| gate_info=gate_info, | ||
| gate_qubits=lambda gate: gate_info[gate].qubits if isinstance(gate, int) else gate, | ||
| gate_pz_assignment={0: "pz2", 1: "pz1"}, | ||
| preferred_pz_for_gate={0: "pz2", 1: "pz1"}.get, | ||
| map_to_pz={0: "pz1", 1: "pz1", 2: "pz2"}, | ||
| locked_gates={}, | ||
| pzs=[SimpleNamespace(name="pz1"), SimpleNamespace(name="pz2")], | ||
| ) | ||
|
|
||
| priority_queue, next_gate_at_pz = create_priority_queue(cast("Any", graph)) | ||
|
|
||
| assert priority_queue == {0: "pz2", 1: "pz1", 2: "pz1"} | ||
| assert next_gate_at_pz == {"pz1": 1, "pz2": 0} |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Test ordering masks a real next_gate_at_pz bug in inside/scheduling.py.
Both create_priority_queue tests put the 2-qubit gate after the 1-qubit gate (or use explicit gate_pz_assignment). Because the 1-qubit branch in inside/scheduling.py:create_priority_queue lacks the == () sentinel check that the 2-qubit branch uses, a 1-qubit gate that arrives after a different PZ has been pre-filled with () (see the inside-scheduling indentation issue I'll flag separately) silently fails to update next_gate_at_pz.
Consider adding a regression test where:
sequence = [<2q_gate_for_pz1>, <1q_gate_for_pz2>], with two PZs.- Assert
next_gate_at_pz["pz2"]equals the 1-qubit gate id (and is not()).
This would catch the bug noted in inside/scheduling.py Lines 147–186 today and guard against regressions later.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_multi_shuttler.py` around lines 349 - 402, The
create_priority_queue implementation fails to update next_gate_at_pz for 1-qubit
gates when a PZ slot was prefilled with the sentinel () (the 2-qubit branch
checks for "== ()" but the 1-qubit branch does not), so update
inside/scheduling.py:create_priority_queue to treat a prefilled () as empty for
1-qubit gates and assign the gate id into next_gate_at_pz[pz] (mirror the
sentinel check used for 2-qubit gates), and add the suggested regression test
(sequence = [2q_for_pz1, 1q_for_pz2] with two PZs) to
tests/test_multi_shuttler.py to assert next_gate_at_pz["pz2"] equals the 1-qubit
gate id rather than () so this bug cannot regress.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/mqt/ionshuttler/multi_shuttler/circuit_parsing.py`:
- Around line 97-101: The helper _load_quantum_circuit currently catches all
Exceptions when calling QuantumCircuit.from_qasm_str and falls back silently to
load_qasm3; change this to only catch the QASM2ParseError raised by
from_qasm_str (import QASM2ParseError from qiskit.qasm2 or the appropriate
qiskit module) so genuine runtime/MemoryError/etc. are not swallowed, and keep
the existing fallback to load_qasm3 when a QASM2ParseError occurs; reference
_load_quantum_circuit, QuantumCircuit.from_qasm_str, load_qasm3, and
QASM2ParseError when making the change.
In `@src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.py`:
- Around line 154-170: The warm-start chain breaks because seed_assignment is
only set when index == 0, so subsequent calls to _contract_slice always receive
slice 0's assignment as previous; to fix, after each contraction append, compute
and set seed_assignment = _build_qubit_assignment(contraction, num_qubits,
num_pzs) so the next iteration passes the immediately prior slice's assignment
into _contract_slice (use time_slices/index, _contract_slice,
_build_qubit_assignment, and seed_assignment to locate the change).
- Around line 698-870: Replace the current list-based tabu_list with a FIFO
deque and a parallel set to get O(1) eviction and membership: initialize
tabu_list as collections.deque(maxlen=tabu_list_length) and create a tabu_set
set() kept in sync; when applying a move append the tuple (slice_index,
supernode.id, previous_cluster) to the deque and add it to tabu_set, and when
the deque evicts an element remove it from tabu_set; update callers and
signatures of _consider_supernode_moves to accept tabu_set (instead of
tabu_list) and use it for membership checks (the hotspot check `move_key in
tabu_list` becomes `move_key in tabu_set`), and ensure anywhere the code relied
on list ordering uses the deque equivalently and that tabu_list_length is used
to construct the deque.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: fd2c1ded-9c10-450a-a16d-c8d53dd7795d
📒 Files selected for processing (5)
src/mqt/ionshuttler/multi_shuttler/circuit_parsing.pysrc/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.pysrc/mqt/ionshuttler/multi_shuttler/inside/scheduling.pysrc/mqt/ionshuttler/multi_shuttler/outside/compilation.pytests/test_multi_shuttler.py
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.py`:
- Around line 1340-1372: Hoist the computation of mean_load out of the
per-supernode loop in gate_partitioning_tabu (the block iterating "for supernode
in contraction.supernodes") because mean_load only depends on the slice-scoped
"loads"; compute mean_load once (and optionally precompute loads_len or a
boolean guard for empty loads) before the loop and then use that precomputed
value inside the loop when applying balance_penalty, leaving the rest of the
logic (overflow_penalty, distance sums, capacity_weight application, and writing
to scores[supernode.id]) unchanged.
- Around line 859-868: The balance delta is being recomputed after the move
using the already-mutated slice_loads, producing an incorrect value; modify
_consider_supernode_moves to return the pre-move balance delta (e.g.,
balance_delta alongside capacity/distance deltas) and extend best_move_state/its
return tuple to carry that value, then in the caller (where
capacity_cost/distance_cost are updated) use the returned balance_delta to
increment balance_cost instead of calling
_balance_delta(slice_loads[slice_index], ...); keep the existing _balance_delta
implementation but only call it inside _consider_supernode_moves where you still
have the pre-move load_vector to compute the correct delta.
- Around line 1238-1249: The loop in _build_slack_weights builds
next_active_per_slice using next_active_per_slice.insert(0, list(next_active)),
which makes it O(num_slices^2); change it to use
next_active_per_slice.append(list(next_active)) inside the for loop and then
call next_active_per_slice.reverse() once after the loop (or alternatively use
collections.deque and appendleft) so the same order is preserved but the
algorithm becomes linear; keep the same updates to next_active and the
list(next_active) copy to avoid aliasing, and ensure any downstream code
expecting the original order still works after the reverse.
In `@tests/test_multi_shuttler.py`:
- Around line 678-696: The test relies on assertions inside the side-effect
function _capture_graph but never asserts that run_shuttle_main was actually
invoked; update the test to patch run_shuttle_main using the context manager
form (e.g., with patch("mqt.ionshuttler.multi_shuttler.main.run_shuttle_main",
side_effect=_capture_graph) as mock_run:) and after calling main(config) assert
mock_run.assert_called_once() (or assert_called_once_with if you want to verify
arguments) so the test fails if run_shuttle_main/_capture_graph are not
executed. Ensure references to main, run_shuttle_main, and _capture_graph are
used exactly as in the diff.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 0ef53c1d-6ac7-4d61-9913-ee43c65d8688
📒 Files selected for processing (2)
src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.pytests/test_multi_shuttler.py
Comment on lines
+1238
to
+1249
| next_active_per_slice: list[list[int]] = [] | ||
| next_active = [num_slices] * num_qubits | ||
| for slice_index in range(num_slices - 1, -1, -1): | ||
| contraction = slice_contractions[slice_index] | ||
| active_qubits = set(contraction.required_unary) | ||
| for left, right in contraction.required_edges: | ||
| active_qubits.add(left) | ||
| active_qubits.add(right) | ||
| for qubit in active_qubits: | ||
| if 0 <= qubit < num_qubits: | ||
| next_active[qubit] = slice_index | ||
| next_active_per_slice.insert(0, list(next_active)) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
list.insert(0, ...) makes _build_slack_weights quadratic in the number of slices.
Each next_active_per_slice.insert(0, list(next_active)) shifts the whole list, turning the loop into O(num_slices²). For circuits decomposed into many slices this dominates _build_slack_weights. Build the list with append and reverse once at the end (or use collections.deque with appendleft).
♻️ Proposed change
- next_active_per_slice: list[list[int]] = []
+ next_active_per_slice: list[list[int]] = []
next_active = [num_slices] * num_qubits
for slice_index in range(num_slices - 1, -1, -1):
contraction = slice_contractions[slice_index]
active_qubits = set(contraction.required_unary)
for left, right in contraction.required_edges:
active_qubits.add(left)
active_qubits.add(right)
for qubit in active_qubits:
if 0 <= qubit < num_qubits:
next_active[qubit] = slice_index
- next_active_per_slice.insert(0, list(next_active))
+ next_active_per_slice.append(list(next_active))
+ next_active_per_slice.reverse()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| next_active_per_slice: list[list[int]] = [] | |
| next_active = [num_slices] * num_qubits | |
| for slice_index in range(num_slices - 1, -1, -1): | |
| contraction = slice_contractions[slice_index] | |
| active_qubits = set(contraction.required_unary) | |
| for left, right in contraction.required_edges: | |
| active_qubits.add(left) | |
| active_qubits.add(right) | |
| for qubit in active_qubits: | |
| if 0 <= qubit < num_qubits: | |
| next_active[qubit] = slice_index | |
| next_active_per_slice.insert(0, list(next_active)) | |
| next_active_per_slice: list[list[int]] = [] | |
| next_active = [num_slices] * num_qubits | |
| for slice_index in range(num_slices - 1, -1, -1): | |
| contraction = slice_contractions[slice_index] | |
| active_qubits = set(contraction.required_unary) | |
| for left, right in contraction.required_edges: | |
| active_qubits.add(left) | |
| active_qubits.add(right) | |
| for qubit in active_qubits: | |
| if 0 <= qubit < num_qubits: | |
| next_active[qubit] = slice_index | |
| next_active_per_slice.append(list(next_active)) | |
| next_active_per_slice.reverse() |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.py` around lines
1238 - 1249, The loop in _build_slack_weights builds next_active_per_slice using
next_active_per_slice.insert(0, list(next_active)), which makes it
O(num_slices^2); change it to use
next_active_per_slice.append(list(next_active)) inside the for loop and then
call next_active_per_slice.reverse() once after the loop (or alternatively use
collections.deque and appendleft) so the same order is preserved but the
algorithm becomes linear; keep the same updates to next_active and the
list(next_active) copy to avoid aliasing, and ensure any downstream code
expecting the original order still works after the reverse.
Comment on lines
+1340
to
+1372
| scores = [0.0] * len(contraction.supernodes) | ||
| current = qubit_assignments_by_slice[slice_index] | ||
| previous = qubit_assignments_by_slice[slice_index - 1] if slice_index > 0 else None | ||
| nxt = qubit_assignments_by_slice[slice_index + 1] if slice_index < len(qubit_assignments_by_slice) - 1 else None | ||
| counts = slice_counts[slice_index] | ||
| loads = slice_loads[slice_index] | ||
|
|
||
| for supernode in contraction.supernodes: | ||
| total = 0.0 | ||
| for qubit in supernode.qubits: | ||
| if qubit < 0 or qubit >= len(current): | ||
| continue | ||
| current_cluster = current[qubit] | ||
| if previous is not None: | ||
| total += _distance_between_clusters(previous[qubit], current_cluster, distance_matrix) | ||
| if nxt is not None: | ||
| total += _distance_between_clusters(current_cluster, nxt[qubit], distance_matrix) | ||
|
|
||
| cluster = current[supernode.qubits[0]] if supernode.qubits else -1 | ||
| overflow_penalty = 0.0 | ||
| if ( | ||
| capacity is not None | ||
| and 0 <= cluster < len(loads) | ||
| and 0 <= cluster < len(counts) | ||
| and counts[cluster] > capacity | ||
| ): | ||
| overflow_penalty = (counts[cluster] - capacity) * loads[cluster] | ||
| total += capacity_weight * overflow_penalty | ||
|
|
||
| mean_load = sum(loads) / len(loads) if loads else 0.0 | ||
| total += balance_penalty * max(0.0, loads[cluster] - mean_load) | ||
| scores[supernode.id] = total | ||
| return scores |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Hoist mean_load out of the supernode loop.
mean_load (Line 1369) only depends on loads, which is fixed for a given slice — but it's recomputed once per supernode in the inner loop at Lines 1347–1371. Compute it once before the loop. Same applies to anything else slice-scoped (e.g., the mean computation could be paired with a precomputed len(loads) guard).
♻️ Proposed change
counts = slice_counts[slice_index]
loads = slice_loads[slice_index]
+ mean_load = sum(loads) / len(loads) if loads else 0.0
@@
- cluster = current[supernode.qubits[0]] if supernode.qubits else -1
+ cluster = current[supernode.qubits[0]] if supernode.qubits else -1
@@
- mean_load = sum(loads) / len(loads) if loads else 0.0
total += balance_penalty * max(0.0, loads[cluster] - mean_load)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.py` around lines
1340 - 1372, Hoist the computation of mean_load out of the per-supernode loop in
gate_partitioning_tabu (the block iterating "for supernode in
contraction.supernodes") because mean_load only depends on the slice-scoped
"loads"; compute mean_load once (and optionally precompute loads_len or a
boolean guard for empty loads) before the loop and then use that precomputed
value inside the loop when applying balance_penalty, leaving the rest of the
logic (overflow_penalty, distance sums, capacity_weight application, and writing
to scores[supernode.id]) unchanged.
Comment on lines
+678
to
+696
| def test_main_threads_explicit_gate_assignment_to_graph(self, heuristic_config_1pz): | ||
| """main() should pass explicit gate assignments through to the runtime graph.""" | ||
| from unittest.mock import patch | ||
|
|
||
| from mqt.ionshuttler.multi_shuttler.main import main | ||
|
|
||
| config = dict(heuristic_config_1pz) | ||
| config["use_dag"] = False | ||
| config["gate_pz_assignment"] = {0: "pz1"} | ||
|
|
||
| def _capture_graph(graph, dag, use_cycle_or_paths, *, use_dag): | ||
| assert dag is None | ||
| assert use_cycle_or_paths == "cycles" | ||
| assert use_dag is False | ||
| assert graph.gate_pz_assignment == {0: "pz1"} | ||
| return 0 | ||
|
|
||
| with patch("mqt.ionshuttler.multi_shuttler.main.run_shuttle_main", side_effect=_capture_graph): | ||
| assert main(config) == 0 |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Captured-graph assertions only fire if run_shuttle_main is actually reached.
_capture_graph does the real validation here (and in the two tests below), but if main() short-circuits earlier — e.g., the new fine-grained validation rejecting gate_pz_assignment, a config default kicking in, or any SystemExit/early return — pytest will mark the test green even though run_shuttle_main was never called. Add compute_assignment.assert_called_once() style guarding for the side-effect callable too:
- with patch("mqt.ionshuttler.multi_shuttler.main.run_shuttle_main", side_effect=_capture_graph):
- assert main(config) == 0
+ with patch(
+ "mqt.ionshuttler.multi_shuttler.main.run_shuttle_main",
+ side_effect=_capture_graph,
+ ) as run_main:
+ assert main(config) == 0
+ run_main.assert_called_once()Otherwise the assertions inside _capture_graph are silently optional, which weakens the regression value of these new tests.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_multi_shuttler.py` around lines 678 - 696, The test relies on
assertions inside the side-effect function _capture_graph but never asserts that
run_shuttle_main was actually invoked; update the test to patch run_shuttle_main
using the context manager form (e.g., with
patch("mqt.ionshuttler.multi_shuttler.main.run_shuttle_main",
side_effect=_capture_graph) as mock_run:) and after calling main(config) assert
mock_run.assert_called_once() (or assert_called_once_with if you want to verify
arguments) so the test fails if run_shuttle_main/_capture_graph are not
executed. Ensure references to main, run_shuttle_main, and _capture_graph are
used exactly as in the diff.
added 2 commits
April 27, 2026 16:08
fix to tabu list management Addressing feedback: Minor speedup through better tabu list management
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.py (1)
1258-1269: 🧹 Nitpick | 🔵 Trivial
insert(0, ...)keeps_build_slack_weightsquadratic.This loop still does head insertion per iteration, which is O(n²) in number of slices. Appending then reversing preserves behavior in linear time.
♻️ Suggested linear-time variant
- next_active_per_slice: list[list[int]] = [] + next_active_per_slice: list[list[int]] = [] @@ - next_active_per_slice.insert(0, list(next_active)) + next_active_per_slice.append(list(next_active)) + next_active_per_slice.reverse()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.py` around lines 1258 - 1269, The current loop in _build_slack_weights uses next_active_per_slice.insert(0, list(next_active)) causing O(n^2) behavior; change to append each list with next_active_per_slice.append(list(next_active)) and after the for loop call next_active_per_slice.reverse() to preserve the original order, keeping all references to next_active_per_slice, next_active, slice_contractions, slice_index and num_qubits intact.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.py`:
- Around line 219-237: The validator _validate_gate_inputs currently allows
duplicate gate IDs in sequence which are later collapsed by a set in
_build_time_slices_relaxed; add an explicit duplicate check in
_validate_gate_inputs that detects repeated IDs in the sequence (e.g., by
comparing len(sequence) to len(set(sequence)) or by collecting ids with count>1)
and raise a ValueError listing the duplicated gate IDs so callers fail fast;
apply the same duplicate-ID rejection in the analogous validation location noted
in the review (the place that currently performs similar input checks for
_build_time_slices_relaxed) so duplicate sequences are consistently rejected.
In `@tests/test_gate_partitioning_tabu.py`:
- Around line 18-29: Add Google-style docstrings to the two helper functions
_distance_matrix and _sample_gate_info: for _distance_matrix document its
purpose (create a symmetric distance matrix with zeros on diagonal), list args
(num_pzs: int) and returns (list[list[float]]), and for _sample_gate_info
document its purpose (return sample mapping of gate IDs to GateInfo), list
returns (dict[int, GateInfo]) and describe content of the dict; keep them
concise and follow Google-style sections (Args, Returns) and reference the
GateInfo type used in _sample_gate_info.
- Line 99: The test uses zip(result.time_slices,
result.qubit_assignments_by_slice, strict=False) which allows length mismatches
to be silently ignored; change the loop to use strict zip (strict=True) so the
test fails if the two iterables differ in length—update the for loop iterating
over result.time_slices and result.qubit_assignments_by_slice to use strict=True
to enforce the invariant.
---
Duplicate comments:
In `@src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.py`:
- Around line 1258-1269: The current loop in _build_slack_weights uses
next_active_per_slice.insert(0, list(next_active)) causing O(n^2) behavior;
change to append each list with next_active_per_slice.append(list(next_active))
and after the for loop call next_active_per_slice.reverse() to preserve the
original order, keeping all references to next_active_per_slice, next_active,
slice_contractions, slice_index and num_qubits intact.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 59c093ee-c1dc-458c-8b50-5ae49b3440ce
📒 Files selected for processing (2)
src/mqt/ionshuttler/multi_shuttler/gate_partitioning_tabu.pytests/test_gate_partitioning_tabu.py
| capacity=2, | ||
| ) | ||
|
|
||
| for slice_gate_ids, qubit_assignment in zip(result.time_slices, result.qubit_assignments_by_slice, strict=False): |
There was a problem hiding this comment.
Use strict zip in this invariant test.
Using strict=False can silently pass when result.time_slices and result.qubit_assignments_by_slice diverge in length, which weakens this correctness check.
✅ Suggested test hardening
- for slice_gate_ids, qubit_assignment in zip(result.time_slices, result.qubit_assignments_by_slice, strict=False):
+ for slice_gate_ids, qubit_assignment in zip(result.time_slices, result.qubit_assignments_by_slice, strict=True):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| for slice_gate_ids, qubit_assignment in zip(result.time_slices, result.qubit_assignments_by_slice, strict=False): | |
| for slice_gate_ids, qubit_assignment in zip(result.time_slices, result.qubit_assignments_by_slice, strict=True): |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_gate_partitioning_tabu.py` at line 99, The test uses
zip(result.time_slices, result.qubit_assignments_by_slice, strict=False) which
allows length mismatches to be silently ignored; change the loop to use strict
zip (strict=True) so the test fails if the two iterables differ in length—update
the for loop iterating over result.time_slices and
result.qubit_assignments_by_slice to use strict=True to enforce the invariant.
Author
There was a problem hiding this comment.
Output consistency is ensured elsewhere, I dont think this change is required
There was a problem hiding this comment.
Only users with a collaborator, contributor, member, or owner role can interact with CodeRabbit.
Author
|
I have addressed all bot feedback that I deem relevant. PR should now be good for human review @danielschoenberger |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR implements Fine-Grained Partitioning (FGP) for the ionshuttler (accompanying paper submitted), an optional pre-compilation step that assigns gates to processing zones to optimize shuttling overhead (accounting for both ion transport between zones and "congestion" effects in overfull zones). Additionally, some scaffolding and minor maintainability improvements were introduced as part of the implementation.
Everything is designed to be backwards compatible, and FGP is toggled off by default, preserving the previous behaviour. Toggling it on (see example below) was shown to reduce shuttling overhead and compilation time in most instances.
Changes/Additions:
1. Circuit parsing was made more robust and maintainable, and moved to a centralised
circuit_parsing.py, s.t.inside/andoutside/compilation.pynow call on the same functions. Tests assert backwards compatibility and robustness. (commit 1)2. The internal representation of the circuits and gates in the
insideandoutsideshuttlers was changed toParsedCircuit, which identifies gates by an ID instead of a mere qubit-tuple. This is critical, since the previously ambigious representation via qubit-tuples is unfit for the gate-to-pz-assignment nature of FGP. Additionally, the class carries thegate_infometadata (including the qasm string) alongside the gate ID, which enables reconstruction of the gates after compilation (not strictly relevant to FGP, but a general maintainability improvement). This step is purely scaffolding and doesn't change the compilation flow itself. Tests assert backwards compatibility. (commits 2 & 3)3. Added optional
graph.gate_pz_assignment, which can be set to predetermine in what pzs the gates are executed. The output of FGP later feeds into this field, which then overwrites the "real-time" pz-decision-making process. While/whengate_pz_assignmentis unavailable, the compiler falls back to the existing mechanism. Tests assert intended handling ofgate_pz_assignement. (commit 4)4. Introduced new Fine-Grained Circuit partitioning algorithm based on Tabu-Search in
gate_partitioning_tabu.py,which decomposes a circuit into layers and assigns gates to pzs based on a shuttling-aware cost function. Output containsgate_pz_assignmentand metadata. Includes tests. (commits 5 & 6)5. Wired the new FGP algorithm into
outside/main(), such that it can be toggled on with a single config booluse_fine_grained_gate_partition = True(default False), that automatically uses adequate default parameters. (commits 7 & 8)Usage
The new feature can be tested by setting
use_fine_grained_gate_partition = Truein the following sample python call.AI Assistants were used in the creation of all submitted code (primarily GPT 5.4-Codex via the VS Code Codex plugin & Claude Opus & Sonnet 4.6 via GitHub Copilot). I confirm that I have personally reviewed and understood all AI-generated content, and accept full responsibility for it.
Checklist
If PR contains AI-assisted content:
Assisted-by: [Model Name] via [Tool Name]footer. -> Too late, but attributed above