feat: update handling qubit resources in repetitions

docs/concepts/compilation.md

@@ -111,8 +111,12 @@ In the example we discussed previously, this allows us to skip the definition of
 
 #### Step 2.6: Repetitions
 
-In case a routine is repeated (i.e. has a non-empty `repetition` field), its local resource definitions get updated according 
-to the repetition rules; the repetition specification itself gets updated using the parameter map.
+In case a routine is repeated (i.e. has a non-empty `repetition` field), its local resource definitions get updated according to the repetition rules; the repetition specification itself gets updated using the parameter map.
+
+Qubit-type resources are handled differently from additive or multiplicative resources during this step:
+
+- additive and multiplicative resources are aggregated across iterations according to the repetition sequence
+- qubit resources describe how many qubits are needed at the peak moment of one run, so they are not added up across sequential iterations. In the current implementation, qubit resources also can't depend on the iterator symbol.
 
 #### Step 2.7: Resource compilation
 
@@ -146,4 +150,4 @@ Finally, derived resources (provided through `derived_resources` field) are calc
 ### Step 3: Postprocessing
 
 After compilation is done, there might be certain operations that the user might want to perform on a compiled routine, e.g.: aggregating resources.
-Currently, there are no postprocessing steps by default.
+Currently, there are no postprocessing steps by default.

src/bartiq/compilation/_compile.py

@@ -267,6 +267,16 @@ def _process_repeated_resources(
 
     new_resources = {}
     for resource in child_resources.values():
+        iterator_symbol = getattr(repetition.sequence, "iterator_symbol", None)
+        if (
+            resource.type == ResourceType.qubits
+            and iterator_symbol is not None
+            and backend.serialize(iterator_symbol) in backend.free_symbols(resource.value)
+        ):
+            raise BartiqCompilationError(
+                f'Cannot process qubit resource "{resource.name}" in repetitive structure when it depends '
+                f'on iterator symbol "{backend.serialize(iterator_symbol)}".'
+            )
         replacement_value = backend.as_expression(f"{children[0].name}.{resource.name}")
         first_pass_value_stub = backend.as_expression(f"{children[0].name}.__fp__{resource.name}")
         use_default_first_pass_value = resource.name not in children[0].first_pass_resources
@@ -276,10 +286,13 @@ def _process_repeated_resources(
         elif resource.type == ResourceType.multiplicative:
             first_pass_value = 1 if use_default_first_pass_value else first_pass_value_stub
             new_value = repetition.sequence_prod(replacement_value / first_pass_value, backend) * first_pass_value
-        elif resource.type == ResourceType.qubits and repetition.sequence.type == "constant":
+        elif resource.type == ResourceType.qubits:
             # NOTE: Actually this could also be `new_value = resource.value`.
-            # The reason it's not, is that in such case local_ancillae are counted twice
-            # in calculate_highwater.
+            # The reason it's not, is that in such case qubit resources such as local_ancillae
+            # are counted twice in calculate_highwater.
+            #
+            # Qubit resources track peak concurrent usage, so sequential repetition should not
+            # aggregate them regardless of repetition sequence type.
             continue
         elif ast.literal_eval(os.environ.get(REPETITION_ALLOW_ARBITRARY_RESOURCES_ENV, "False")):
             new_value = replacement_value

tests/compilation/test_repetitions.py

@@ -21,6 +21,7 @@
 
 from bartiq import compile_routine, evaluate
 from bartiq.compilation import CompilationFlags
+from bartiq.compilation.derived_resources import calculate_highwater
 from bartiq.errors import BartiqCompilationError
 
 
@@ -44,6 +45,36 @@ def _routine_with_repetition(repetition_dict: dict) -> RoutineV1:
     )
 
 
+def _routine_with_qubit_repetition(repetition_dict: dict) -> RoutineV1:
+    return SchemaV1(
+        program={
+            "name": "root",
+            "input_params": ["N"],
+            "ports": [
+                {"name": "in_0", "direction": "input", "size": "N"},
+                {"name": "out_0", "direction": "output", "size": "N"},
+            ],
+            "children": [
+                {
+                    "name": "child",
+                    "ports": [
+                        {"name": "thru_0", "direction": "through", "size": "N"},
+                    ],
+                    "resources": [
+                        {"name": "local_ancillae", "type": "qubits", "value": 2},
+                    ],
+                }
+            ],
+            "connections": [
+                "in_0 -> child.thru_0",
+                "child.thru_0 -> out_0",
+            ],
+            "repetition": repetition_dict,
+        },
+        version="v1",
+    )
+
+
 def _constant_sequence_sum(unit_cost, count, multiplier):
     sum = 0
     for _ in range(count):
@@ -394,3 +425,63 @@ def test_evaluate_repetitions_with_custom_function():
     functions_map = {"my_fun": lambda x: x**2}
     evaluated_routine = evaluate(compiled_routine, assignments, functions_map=functions_map).routine
     assert evaluated_routine.repetition.count == 100
+
+@pytest.mark.parametrize(
+    "repetition_dict",
+    [
+        {"count": 5, "sequence": {"type": "constant", "multiplier": 1}},
+        {"count": 5, "sequence": {"type": "arithmetic", "initial_term": 1, "difference": 1}},
+        {"count": 5, "sequence": {"type": "geometric", "ratio": 2}},
+        {"count": 5, "sequence": {"type": "closed_form", "sum": "N", "prod": "1", "num_terms_symbol": "N"}},
+        {"count": 5, "sequence": {"type": "custom", "term_expression": "1", "iterator_symbol": "i"}},
+    ],
+)
+def test_qubit_resources_are_supported_when_independent_of_iterator(repetition_dict):
+    routine = _routine_with_qubit_repetition(repetition_dict)
+    derived_resources = [{"name": "qubit_highwater", "type": "qubits", "calculate": calculate_highwater}]
+
+    compiled_routine = compile_routine(routine, derived_resources=derived_resources).routine
+    evaluated_routine = evaluate(compiled_routine, {"N": 4}).routine
+
+    assert "qubit_highwater" in evaluated_routine.resources
+    assert evaluated_routine.resources["qubit_highwater"].value == 6
+
+
+def test_qubit_resources_depending_on_iterator_raise_for_custom_repetition():
+    routine = SchemaV1(
+        program={
+            "name": "root",
+            "input_params": ["N"],
+            "ports": [
+                {"name": "in_0", "direction": "input", "size": "N"},
+                {"name": "out_0", "direction": "output", "size": "N"},
+            ],
+            "children": [
+                {
+                    "name": "child",
+                    "ports": [
+                        {"name": "thru_0", "direction": "through", "size": "N"},
+                    ],
+                    "resources": [
+                        {"name": "local_ancillae", "type": "qubits", "value": "i + 1"},
+                    ],
+                }
+            ],
+            "connections": [
+                "in_0 -> child.thru_0",
+                "child.thru_0 -> out_0",
+            ],
+            "repetition": {
+                "count": 4,
+                "sequence": {"type": "custom", "term_expression": "1", "iterator_symbol": "i"},
+            },
+        },
+        version="v1",
+    )
+    derived_resources = [{"name": "qubit_highwater", "type": "qubits", "calculate": calculate_highwater}]
+
+    with pytest.raises(
+        BartiqCompilationError,
+        match='Cannot process qubit resource "local_ancillae" in repetitive structure when it depends on iterator symbol "i"',
+    ):
+        compile_routine(routine, derived_resources=derived_resources)