Add full Expr support to StochasticSwap (#10506) (#10510)

This was mostly already there (and why I had missed problems when
testing before), the missing piece was just for `switch` statements.
This commit also adds some final tests of the pass.

(cherry picked from commit 7af335e)

Co-authored-by: Jake Lishman <jake.lishman@ibm.com>

qiskit/transpiler/passes/routing/stochastic_swap.py

@@ -18,6 +18,7 @@
 import numpy as np
 
 from qiskit.converters import dag_to_circuit, circuit_to_dag
+from qiskit.circuit.classical import expr, types
 from qiskit.circuit.quantumregister import QuantumRegister
 from qiskit.transpiler.basepasses import TransformationPass
 from qiskit.transpiler.exceptions import TranspilerError
@@ -474,7 +475,16 @@ def _controlflow_exhaustive_acyclic(operation: ControlFlowOp):
         return len(operation.blocks) == 2
     if isinstance(operation, SwitchCaseOp):
         cases = operation.cases()
-        max_matches = 2 if isinstance(operation.target, Clbit) else 1 << len(operation.target)
+        if isinstance(operation.target, expr.Expr):
+            type_ = operation.target.type
+            if type_.kind is types.Bool:
+                max_matches = 2
+            elif type_.kind is types.Uint:
+                max_matches = 1 << type_.width
+            else:
+                raise RuntimeError(f"unhandled target type: '{type_}'")
+        else:
+            max_matches = 2 if isinstance(operation.target, Clbit) else 1 << len(operation.target)
         return CASE_DEFAULT in cases or len(cases) == max_matches
     return False
 

test/python/transpiler/test_stochastic_swap.py

@@ -29,6 +29,7 @@
 from qiskit.providers.fake_provider import FakeMumbai, FakeMumbaiV2
 from qiskit.compiler.transpiler import transpile
 from qiskit.circuit import ControlFlowOp, Clbit, CASE_DEFAULT
+from qiskit.circuit.classical import expr
 
 
 @ddt
@@ -882,6 +883,44 @@ def test_pre_intra_post_if_else(self):
         expected.measure(qreg, creg[[2, 4, 0, 3, 1]])
         self.assertEqual(dag_to_circuit(cdag), expected)
 
+    def test_if_expr(self):
+        """Test simple if conditional with an `Expr` condition."""
+        coupling = CouplingMap.from_line(4)
+
+        body = QuantumCircuit(4)
+        body.cx(0, 1)
+        body.cx(0, 2)
+        body.cx(0, 3)
+        qc = QuantumCircuit(4, 2)
+        qc.if_test(expr.logic_and(qc.clbits[0], qc.clbits[1]), body, [0, 1, 2, 3], [])
+
+        dag = circuit_to_dag(qc)
+        cdag = StochasticSwap(coupling, seed=58).run(dag)
+        check_map_pass = CheckMap(coupling)
+        check_map_pass.run(cdag)
+        self.assertTrue(check_map_pass.property_set["is_swap_mapped"])
+
+    def test_if_else_expr(self):
+        """Test simple if/else conditional with an `Expr` condition."""
+        coupling = CouplingMap.from_line(4)
+
+        true = QuantumCircuit(4)
+        true.cx(0, 1)
+        true.cx(0, 2)
+        true.cx(0, 3)
+        false = QuantumCircuit(4)
+        false.cx(3, 0)
+        false.cx(3, 1)
+        false.cx(3, 2)
+        qc = QuantumCircuit(4, 2)
+        qc.if_else(expr.logic_and(qc.clbits[0], qc.clbits[1]), true, false, [0, 1, 2, 3], [])
+
+        dag = circuit_to_dag(qc)
+        cdag = StochasticSwap(coupling, seed=58).run(dag)
+        check_map_pass = CheckMap(coupling)
+        check_map_pass.run(cdag)
+        self.assertTrue(check_map_pass.property_set["is_swap_mapped"])
+
     def test_no_layout_change(self):
         """test controlflow with no layout change needed"""
         num_qubits = 5
@@ -996,6 +1035,23 @@ def test_while_loop(self):
         expected.measure(qreg, creg)
         self.assertEqual(dag_to_circuit(cdag), expected)
 
+    def test_while_loop_expr(self):
+        """Test simple while loop with an `Expr` condition."""
+        coupling = CouplingMap.from_line(4)
+
+        body = QuantumCircuit(4)
+        body.cx(0, 1)
+        body.cx(0, 2)
+        body.cx(0, 3)
+        qc = QuantumCircuit(4, 2)
+        qc.while_loop(expr.logic_and(qc.clbits[0], qc.clbits[1]), body, [0, 1, 2, 3], [])
+
+        dag = circuit_to_dag(qc)
+        cdag = StochasticSwap(coupling, seed=58).run(dag)
+        check_map_pass = CheckMap(coupling)
+        check_map_pass.run(cdag)
+        self.assertTrue(check_map_pass.property_set["is_swap_mapped"])
+
     def test_switch_single_case(self):
         """Test routing of 'switch' with just a single case."""
         qreg = QuantumRegister(5, "q")
@@ -1110,6 +1166,89 @@ def test_switch_exhaustive(self, labels):
 
         self.assertEqual(canonicalize_control_flow(test), canonicalize_control_flow(expected))
 
+    def test_switch_nonexhaustive_expr(self):
+        """Test routing of 'switch' with an `Expr` target and several but nonexhaustive cases."""
+        qreg = QuantumRegister(5, "q")
+        creg = ClassicalRegister(3, "c")
+
+        qc = QuantumCircuit(qreg, creg)
+        case0 = QuantumCircuit(qreg, creg[:])
+        case0.cx(0, 1)
+        case0.cx(1, 2)
+        case0.cx(2, 0)
+        case1 = QuantumCircuit(qreg, creg[:])
+        case1.cx(1, 2)
+        case1.cx(2, 3)
+        case1.cx(3, 1)
+        case2 = QuantumCircuit(qreg, creg[:])
+        case2.cx(2, 3)
+        case2.cx(3, 4)
+        case2.cx(4, 2)
+        qc.switch(expr.bit_or(creg, 5), [(0, case0), ((1, 2), case1), (3, case2)], qreg, creg)
+
+        coupling = CouplingMap.from_line(len(qreg))
+        pass_ = StochasticSwap(coupling, seed=58)
+        test = pass_(qc)
+
+        check = CheckMap(coupling)
+        check(test)
+        self.assertTrue(check.property_set["is_swap_mapped"])
+
+        expected = QuantumCircuit(qreg, creg)
+        case0 = QuantumCircuit(qreg, creg[:])
+        case0.cx(0, 1)
+        case0.cx(1, 2)
+        case0.swap(0, 1)
+        case0.cx(2, 1)
+        case0.swap(0, 1)
+        case1 = QuantumCircuit(qreg, creg[:])
+        case1.cx(1, 2)
+        case1.cx(2, 3)
+        case1.swap(1, 2)
+        case1.cx(3, 2)
+        case1.swap(1, 2)
+        case2 = QuantumCircuit(qreg, creg[:])
+        case2.cx(2, 3)
+        case2.cx(3, 4)
+        case2.swap(3, 4)
+        case2.cx(3, 2)
+        case2.swap(3, 4)
+        expected.switch(expr.bit_or(creg, 5), [(0, case0), ((1, 2), case1), (3, case2)], qreg, creg)
+
+        self.assertEqual(canonicalize_control_flow(test), canonicalize_control_flow(expected))
+
+    @data((0, 1, 2, 3), (CASE_DEFAULT,))
+    def test_switch_exhaustive_expr(self, labels):
+        """Test routing of 'switch' with exhaustive cases on an `Expr` target; we should not require
+        restoring the layout afterwards."""
+        qreg = QuantumRegister(5, "q")
+        creg = ClassicalRegister(2, "c")
+
+        qc = QuantumCircuit(qreg, creg)
+        case0 = QuantumCircuit(qreg[[0, 1, 2]], creg[:])
+        case0.cx(0, 1)
+        case0.cx(1, 2)
+        case0.cx(2, 0)
+        qc.switch(expr.bit_or(creg, 3), [(labels, case0)], qreg[[0, 1, 2]], creg)
+
+        coupling = CouplingMap.from_line(len(qreg))
+        pass_ = StochasticSwap(coupling, seed=58)
+        test = pass_(qc)
+
+        check = CheckMap(coupling)
+        check(test)
+        self.assertTrue(check.property_set["is_swap_mapped"])
+
+        expected = QuantumCircuit(qreg, creg)
+        case0 = QuantumCircuit(qreg[[0, 1, 2]], creg[:])
+        case0.cx(0, 1)
+        case0.cx(1, 2)
+        case0.swap(0, 1)
+        case0.cx(2, 1)
+        expected.switch(expr.bit_or(creg, 3), [(labels, case0)], qreg[[0, 1, 2]], creg)
+
+        self.assertEqual(canonicalize_control_flow(test), canonicalize_control_flow(expected))
+
     def test_nested_inner_cnot(self):
         """test swap in nested if else controlflow construct; swap in inner"""
         seed = 1