Add optional prune_classical_qubits setting for circuit generation (#2802)

This change adds a parameter to Q# and OpenQASM circuit generation
called `prune_classical_qubits` that will generate a simplified circuit
diagram where qubits that are unused or purely classical (never enter
superposition) are removed. This can be handy for places where the
program comes from generated logic that may use more qubits than
strictly needed. This is purely for visualizing the simplified circuit
and does not perform any kind of execution optimization.

So a program like this that has nine qubits but doesn't really use all
of them:
<img width="376" height="297" alt="image"
src="https://github.com/user-attachments/assets/1ebf9bbd-8485-47d9-ad8b-26f392e2d789"
/>

will by default show the full circuit as it does today:
<img width="1081" height="650" alt="image"
src="https://github.com/user-attachments/assets/bbbe4238-d8ac-4a75-9ff2-00cfc30225ca"
/>

but with `prune_classical_qubits = True` will remove the purely
classical q1, q2, q4, q5, and q8:
<img width="663" height="422" alt="image"
src="https://github.com/user-attachments/assets/7713fdea-8ade-417d-bf75-15dde3534c56"
/>

Author: swernli

source/compiler/qsc/src/interpret/circuit_tests.rs

@@ -213,6 +213,7 @@ fn grouping_nested_callables() {
             max_operations: usize::MAX,
             source_locations: false,
             group_by_scope: true,
+            prune_classical_qubits: false,
         },
     )
     .expect("circuit generation should succeed");
@@ -1072,6 +1073,7 @@ fn operation_declared_in_eval() {
                 max_operations: usize::MAX,
                 source_locations: false,
                 group_by_scope: true,
+                prune_classical_qubits: false,
             },
         )
         .expect("circuit generation should succeed");

source/compiler/qsc_circuit/src/builder.rs

@@ -42,6 +42,8 @@ pub struct CircuitTracer {
     circuit_builder: OperationListBuilder,
     next_result_id: usize,
     user_package_ids: Vec<PackageId>,
+    superposition_qubits: FxHashSet<QubitWire>,
+    classical_one_qubits: FxHashSet<QubitWire>,
 }
 
 impl Tracer for CircuitTracer {
@@ -69,6 +71,13 @@ impl Tracer for CircuitTracer {
     ) {
         let called_at = map_stack_frames_to_locations(stack);
         let display_args: Vec<String> = theta.map(|p| format!("{p:.4}")).into_iter().collect();
+        let controls = if self.config.prune_classical_qubits {
+            // Any controls that are known to be classically one can be removed, so this
+            // will return the updated controls list.
+            &self.update_qubit_status(name, targets, controls)
+        } else {
+            controls
+        };
         self.circuit_builder.gate(
             self.wire_map_builder.current(),
             name,
@@ -91,6 +100,8 @@ impl Tracer for CircuitTracer {
         };
         self.wire_map_builder.link_result_to_qubit(q, r);
         if name == "MResetZ" {
+            self.classical_one_qubits
+                .remove(&self.wire_map_builder.wire_map.qubit_wire(q));
             self.circuit_builder
                 .mresetz(self.wire_map_builder.current(), q, r, called_at);
         } else {
@@ -101,6 +112,8 @@ impl Tracer for CircuitTracer {
 
     fn reset(&mut self, stack: &[Frame], q: usize) {
         let called_at = map_stack_frames_to_locations(stack);
+        self.classical_one_qubits
+            .remove(&self.wire_map_builder.wire_map.qubit_wire(q));
         self.circuit_builder
             .reset(self.wire_map_builder.current(), q, called_at);
     }
@@ -152,6 +165,8 @@ impl CircuitTracer {
             ),
             next_result_id: 0,
             user_package_ids: user_package_ids.to_vec(),
+            superposition_qubits: FxHashSet::default(),
+            classical_one_qubits: FxHashSet::default(),
         }
     }
 
@@ -191,6 +206,8 @@ impl CircuitTracer {
             ),
             next_result_id: 0,
             user_package_ids: user_package_ids.to_vec(),
+            superposition_qubits: FxHashSet::default(),
+            classical_one_qubits: FxHashSet::default(),
         }
     }
 
@@ -220,16 +237,57 @@ impl CircuitTracer {
         operations: &[OperationOrGroup],
         source_lookup: &impl SourceLookup,
     ) -> Circuit {
+        let mut operations: Vec<OperationOrGroup> = operations.to_vec();
+        let mut qubits = self.wire_map_builder.wire_map.to_qubits();
+        // We need to pass the original number of qubits, before any trimming, to finish the circuit below.
+        let num_qubits = qubits.len();
+
+        if self.config.prune_classical_qubits {
+            // Remove qubits that are always classical.
+            qubits.retain(|q| self.superposition_qubits.contains(&q.id.into()));
+
+            // Remove operations that don't use any non-classical qubits.
+            operations.retain_mut(|op| self.should_keep_operation_mut(op));
+        }
+
         let operations = operations
             .iter()
             .map(|o| o.clone().into_operation(source_lookup))
             .collect();
 
-        finish_circuit(
-            self.wire_map_builder.wire_map.to_qubits(),
-            operations,
-            source_lookup,
-        )
+        finish_circuit(qubits, operations, num_qubits, source_lookup)
+    }
+
+    fn should_keep_operation_mut(&self, op: &mut OperationOrGroup) -> bool {
+        if matches!(op.kind, OperationOrGroupKind::Single) {
+            // This is a normal gate operation, so only keep it if all the qubits are non-classical.
+            op.all_qubits()
+                .iter()
+                .all(|q| self.superposition_qubits.contains(q))
+        } else {
+            // This is a grouped operation, so process the children recursively.
+            let mut used_qubits = FxHashSet::default();
+            op.children_mut()
+                .expect("operation should be a group with children")
+                .retain_mut(|child_op| {
+                    // Prune out child ops that don't use any non-classical qubits.
+                    // This has the side effect of updating each child op's target qubits.
+                    if self.should_keep_operation_mut(child_op) {
+                        for q in child_op.all_qubits() {
+                            used_qubits.insert(q);
+                        }
+                        true
+                    } else {
+                        false
+                    }
+                });
+            // Update the targets of this grouped operation to only include qubits actually used by child operations.
+            op.op
+                .targets_mut()
+                .retain(|q| used_qubits.contains(&q.qubit.into()));
+            // Only keep this grouped operation if any of its targets were kept.
+            !op.op.targets_mut().is_empty()
+        }
     }
 
     /// Splits the qubit arguments from classical arguments so that the qubits
@@ -311,6 +369,113 @@ impl CircuitTracer {
         }
         first_user_code_location(&self.user_package_ids, stack)
     }
+
+    fn mark_qubit_in_superposition(&mut self, wire: QubitWire) {
+        assert!(
+            self.config.prune_classical_qubits,
+            "should only be called when pruning is enabled"
+        );
+        self.superposition_qubits.insert(wire);
+        self.classical_one_qubits.remove(&wire);
+    }
+
+    fn flip_classical_qubit(&mut self, wire: QubitWire) {
+        assert!(
+            self.config.prune_classical_qubits,
+            "should only be called when pruning is enabled"
+        );
+        if self.classical_one_qubits.contains(&wire) {
+            self.classical_one_qubits.remove(&wire);
+        } else {
+            self.classical_one_qubits.insert(wire);
+        }
+    }
+
+    fn update_qubit_status(
+        &mut self,
+        name: &str,
+        targets: &[usize],
+        controls: &[usize],
+    ) -> Vec<usize> {
+        match name {
+            "H" | "Rx" | "Ry" | "SX" | "Rxx" | "Ryy" => {
+                // These gates create superpositions, so mark the qubits as non-trimmable
+                for &q in targets {
+                    let mapped_q = self.wire_map_builder.wire_map.qubit_wire(q);
+                    self.mark_qubit_in_superposition(mapped_q);
+                }
+            }
+            "X" | "Y" => {
+                let mapped_target = self.wire_map_builder.wire_map.qubit_wire(targets[0]);
+                let controls: Vec<usize> = controls
+                    .iter()
+                    .filter(|c| !self.classical_one_qubits.contains(&(**c).into()))
+                    .copied()
+                    .collect();
+                if !self.superposition_qubits.contains(&mapped_target) {
+                    // The target is not yet marked as non-trimmable, so check the controls.
+                    let superposition_controls_count = controls
+                        .iter()
+                        .filter(|c| self.superposition_qubits.contains(&(**c).into()))
+                        .count();
+
+                    if controls.is_empty() {
+                        // If all controls are classical 1 or there are no controls, the target is flipped
+                        self.flip_classical_qubit(mapped_target);
+                    } else if superposition_controls_count == controls.len() {
+                        // If all controls are in superposition, the target is also in superposition
+                        self.mark_qubit_in_superposition(mapped_target);
+                    }
+                }
+                return controls;
+            }
+            "Z" => {
+                // Only clean up the classical 1 qubits from the controls list. No need to update the target,
+                // since Z does not introduce superpositions.
+                return controls
+                    .iter()
+                    .filter(|c| !self.classical_one_qubits.contains(&(**c).into()))
+                    .copied()
+                    .collect();
+            }
+            "SWAP" => {
+                // If either qubit is non-trimmable, both become non-trimmable
+                let q0_mapped = self.wire_map_builder.wire_map.qubit_wire(targets[0]);
+                let q1_mapped = self.wire_map_builder.wire_map.qubit_wire(targets[1]);
+                if self.superposition_qubits.contains(&q0_mapped)
+                    || self.superposition_qubits.contains(&q1_mapped)
+                {
+                    self.mark_qubit_in_superposition(q0_mapped);
+                    self.mark_qubit_in_superposition(q1_mapped);
+                } else {
+                    match (
+                        self.classical_one_qubits.contains(&q0_mapped),
+                        self.classical_one_qubits.contains(&q1_mapped),
+                    ) {
+                        (true, false) | (false, true) => {
+                            self.flip_classical_qubit(q0_mapped);
+                            self.flip_classical_qubit(q1_mapped);
+                        }
+                        _ => {
+                            // Nothing to do if both are classical 0 or both are in superposition
+                        }
+                    }
+                }
+            }
+            "S" | "T" | "Rz" | "Rzz" => {
+                // These gates don't create superpositions on their own, so do nothing
+            }
+            _ => {
+                // For any other gate, conservatively mark all target qubits as non-trimmable
+                for &q in targets.iter().chain(controls.iter()) {
+                    let mapped_q = self.wire_map_builder.wire_map.qubit_wire(q);
+                    self.mark_qubit_in_superposition(mapped_q);
+                }
+            }
+        }
+        // Return the normal controls list if no changes were made.
+        controls.to_vec()
+    }
 }
 
 fn first_user_code_location(
@@ -332,6 +497,7 @@ fn first_user_code_location(
 fn finish_circuit(
     mut qubits: Vec<Qubit>,
     mut operations: Vec<Operation>,
+    num_qubits: usize,
     source_location_lookup: &impl SourceLookup,
 ) -> Circuit {
     resolve_locations(&mut operations, source_location_lookup);
@@ -342,7 +508,7 @@ fn finish_circuit(
         }
     }
 
-    let component_grid = operation_list_to_grid(operations, qubits.len());
+    let component_grid = operation_list_to_grid(operations, num_qubits);
     Circuit {
         qubits,
         component_grid,
@@ -458,6 +624,8 @@ pub struct TracerConfig {
     pub source_locations: bool,
     /// Group operations according to call graph in the circuit diagram
     pub group_by_scope: bool,
+    /// Prune purely classical or unused qubits
+    pub prune_classical_qubits: bool,
 }
 
 impl TracerConfig {
@@ -476,6 +644,7 @@ impl Default for TracerConfig {
             max_operations: Self::DEFAULT_MAX_OPERATIONS,
             source_locations: true,
             group_by_scope: true,
+            prune_classical_qubits: false,
         }
     }
 }

source/compiler/qsc_circuit/src/builder/tests.rs

@@ -1,7 +1,10 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+#![allow(clippy::unicode_not_nfc)]
+
 mod group_scopes;
+mod prune_classical_qubits;
 
 use std::vec;
 
@@ -128,6 +131,7 @@ fn exceed_max_operations() {
             max_operations: 2,
             source_locations: false,
             group_by_scope: false,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
     );
@@ -156,6 +160,7 @@ fn source_locations_enabled() {
             max_operations: 10,
             source_locations: true,
             group_by_scope: false,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
     );
@@ -193,6 +198,7 @@ fn source_locations_disabled() {
             max_operations: 10,
             source_locations: false,
             group_by_scope: false,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
     );
@@ -224,6 +230,7 @@ fn source_locations_multiple_user_frames() {
             max_operations: 10,
             source_locations: true,
             group_by_scope: false,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
     );
@@ -262,6 +269,7 @@ fn source_locations_library_frames_excluded() {
             max_operations: 10,
             source_locations: true,
             group_by_scope: false,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
     );
@@ -294,6 +302,7 @@ fn source_locations_only_library_frames() {
             max_operations: 10,
             source_locations: true,
             group_by_scope: false,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
     );
@@ -326,6 +335,7 @@ fn source_locations_enabled_no_stack() {
             max_operations: 10,
             source_locations: true,
             group_by_scope: false,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
     );
@@ -351,6 +361,7 @@ fn qubit_source_locations_via_stack() {
             max_operations: 10,
             source_locations: true,
             group_by_scope: false,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
     );
@@ -375,6 +386,7 @@ fn qubit_labels_for_preallocated_qubits() {
             max_operations: 10,
             source_locations: true,
             group_by_scope: false,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
         Some((
@@ -414,6 +426,7 @@ fn measurement_target_propagated_to_group() {
             max_operations: usize::MAX,
             source_locations: false,
             group_by_scope: true,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
     );

source/compiler/qsc_circuit/src/builder/tests/group_scopes.rs

@@ -12,6 +12,7 @@ fn check_groups(c: &FakeCompilation, instructions: &[(Vec<Frame>, &str)], expect
             max_operations: usize::MAX,
             source_locations: false,
             group_by_scope: true,
+            prune_classical_qubits: false,
         },
         &FakeCompilation::user_package_ids(),
     );