Merge pull request #195 from BQSKit/targeted_foreach

pass down foreach

bqskit/passes/control/foreach.py

@@ -38,6 +38,24 @@ class ForEachBlockPass(BasePass):
     pass_down_key_prefix = 'ForEachBlockPass_pass_down_'
     """If a key exists in the pass data with this prefix, pass it to blocks."""
 
+    pass_down_block_specific_key_prefix = (
+        'ForEachBlockPass_specific_pass_down_'
+    )
+    """
+    Data specific to the processing of individual blocks in a partitioned
+    circuit can be injected into the `PassData` in `run` by using this prefix.
+
+    The expected type of the associated value is `dict[int, Any]`, where
+    integer (sub-)keys correspond to block numbers in a partitioned quantum
+    circuit.
+
+    Pseudocode example for seed circuits:
+        seeds = {block_id: [seed_circuit_a, seed_circuit_b, ...], ...}
+        key = self.pass_down_block_specific_key_prefix + 'seed_circuits'
+        seed_updater = UpdateDataPass(key, seeds)
+        workflow = Workflow([..., seed_updater, ForEachBlockPass(...), ...])
+    """
+
     def __init__(
         self,
         loop_body: WorkflowLike,
@@ -197,6 +215,10 @@ async def run(self, circuit: Circuit, data: PassData) -> None:
             for key in data:
                 if key.startswith(self.pass_down_key_prefix):
                     block_data[key] = data[key]
+                elif key.startswith(
+                    self.pass_down_block_specific_key_prefix,
+                ) and i in data[key]:
+                    block_data[key] = data[key][i]
             block_data.seed = data.seed
 
             subcircuits.append(subcircuit)

tests/passes/control/test_foreachblock.py

@@ -5,15 +5,19 @@
 from bqskit.compiler.basepass import BasePass
 from bqskit.compiler.compiler import Compiler
 from bqskit.compiler.passdata import PassData
+from bqskit.compiler.workflow import Workflow
 from bqskit.ir.circuit import Circuit
 from bqskit.ir.gates import CircuitGate
+from bqskit.ir.gates import CNOTGate
 from bqskit.ir.gates import HGate
 from bqskit.ir.gates import XGate
 from bqskit.ir.gates import YGate
 from bqskit.ir.gates import ZGate
 from bqskit.ir.operation import Operation
 from bqskit.passes import UnfoldPass
 from bqskit.passes.control.foreach import ForEachBlockPass
+from bqskit.passes.partitioning.quick import QuickPartitioner
+from bqskit.passes.util.update import UpdateDataPass
 
 
 @pytest.fixture
@@ -47,6 +51,18 @@ async def run(self, circuit: Circuit, data: PassData) -> None:
             circuit.append_gate(HGate(), 1)
 
 
+class TestForEachPassDownSpecificSeeds(BasePass):
+    async def run(self, circuit: Circuit, data: PassData) -> None:
+        key = ForEachBlockPass.pass_down_block_specific_key_prefix
+        key += 'test_info'
+        assert key in data
+        assert data[key] == 'a' or data[key] == 'b'
+        if data[key] == 'a':
+            data['test_response'] = 'a'
+        else:
+            data['test_response'] = 'b'
+
+
 def never_replace(c: Circuit, o: Operation) -> bool:
     return False
 
@@ -130,3 +146,37 @@ def test_no_hang_on_empty_collection(compiler: Compiler) -> None:
     circuit.append_gate(XGate(), 0)
     feb_pass = ForEachBlockPass(RemoveXGatePass(), collection_filter=empty_coll)
     compiler.compile(circuit, feb_pass)
+
+
+def test_pass_down_seeds(compiler: Compiler) -> None:
+    circuit = Circuit(4)
+    circuit.append_gate(CNOTGate(), (0, 2))
+    circuit.append_gate(CNOTGate(), (1, 2))
+    circuit.append_gate(CNOTGate(), (1, 3))
+    circuit.append_gate(CNOTGate(), (2, 3))
+
+    # Manually set seed for blocks 0 and 1
+    input_info = {0: 'a', 1: 'b'}
+
+    partitioner = QuickPartitioner()
+    check_specific = TestForEachPassDownSpecificSeeds()
+    foreach = ForEachBlockPass(check_specific)
+
+    key = foreach.pass_down_block_specific_key_prefix + 'test_info'
+    updater = UpdateDataPass(key, input_info)
+
+    # For checking specific pass down data exists
+    workflow = Workflow([partitioner, updater, foreach])
+
+    # Check usability of block specific keys
+    compiled, data = compiler.compile(
+        circuit, workflow, request_data=True,
+    )
+
+    block0_data = data['ForEachBlockPass_data'][0][0]
+    block1_data = data['ForEachBlockPass_data'][0][1]
+    response_key = 'test_response'
+    assert response_key in block0_data
+    assert response_key in block1_data
+    assert block0_data[response_key] == 'a'
+    assert block1_data[response_key] == 'b'