feat: Add a with_loop method to Program (#1717)

* feat: Add a `apply_numshots_loop` method to `Program`

* update docs

* update snapshot

* Update docs/source/programs_and_gates.rst

Co-authored-by: Mark Skilbeck <[email protected]>

* remove unused section

* expand example to include execution

* refactor around suggested API

* doc cleanup

* formatting

* remove shot_count from memory_map

* add doc for num_iterations param

---------

Co-authored-by: Mark Skilbeck <[email protected]>

docs/source/advanced_usage.rst

@@ -317,6 +317,8 @@ register of qubits to build your program.
     H 12
     H 14
 
+.. _classical_control_flow: 
+
 **********************
 Classical control flow
 **********************
@@ -389,6 +391,13 @@ classical register.  There are several classical commands that can be used in th
 - ``MOVE`` which moves the value of a classical bit at one classical address into another
 - ``EXCHANGE`` which swaps the value of two classical bits
 
+.. note::
+
+   The approach documented here can be used to construct a "numshots" loop in pure Quil. See the
+   :py:meth:`~pyquil.quil.Program.with_loop` method and :ref:`build_a_fixed_count_loop` for more
+   information.
+
+
 If, then
 ========
 

docs/source/programs_and_gates.rst

@@ -213,6 +213,8 @@ use the compiler to :ref:`re-index <rewiring>` your qubits):
     for i, q in enumerate(qubits):
         p += MEASURE(q, ro[i])
 
+.. _specifying_trials:
+
 Specifying the number of trials
 ===============================
 
@@ -240,6 +242,63 @@ program should be executed 1000 times:
     The word “shot” comes from experimental physics where an experiment is
     performed many times, and each result is called a shot.
 
+.. _build_a_fixed_count_loop:
+
+Build a fixed-count loop with Quil
+----------------------------------
+
+Specifying trials with :py:meth:`~pyquil.quil.Program.wrap_in_numshots_loop` doesn't modify the Quil in your program in
+any way. Instead, the number of shots you specify is included in your job request and tells the executor how many times
+to run your program. However, with Quil's :ref:`classical_control_flow`, instructions it is possible to write a program
+that itself defines a loop over a number of shots. The :py:meth:`~pyquil.quil.Program.with_loop` method will help you
+do just that. It wraps the body of your program in a loop over a number of iterations you specify and returns the looped
+program.
+
+Let's see an example. We'll construct a classic bell state program and measure it 1000 times by wrapping the program in a
+Quil loop.
+
+.. testcode:: with_loop
+
+    from pyquil import Program, get_qc
+    from pyquil.quilatom import Label
+    from pyquil.gates import H, CNOT
+
+    # Setup the bell state program
+    p = Program(
+        H(0),
+        CNOT(0, 1),
+    )
+    ro = p.declare("ro", "BIT", 2)
+    p.measure(0, ro[0])
+    p.measure(1, ro[1])
+
+    # Declare a memory region to hold the number of shots
+    shot_count = p.declare("shot_count", "INTEGER")
+
+    # Wrap the program in a loop by specifying the number of iterations, a memory reference to
+    # hold the number of iterations, and two labels to mark the beginning and end of the loop.
+    looped_program = p.with_loop(1000, shot_count, Label("start-loop"), Label("end-loop"))
+    print(looped_program.out())
+
+    qc = get_qc("2q-qvm")
+    # Specify your desired shot count in the memory map.
+    results = qc.run(looped_program)
+
+.. testoutput:: with_loop
+
+    DECLARE ro BIT[2]
+    DECLARE shot_count INTEGER[1]
+    MOVE shot_count[0] 1000
+    LABEL @start-loop
+    H 0
+    CNOT 0 1
+    MEASURE 0 ro[0]
+    MEASURE 1 ro[1]
+    SUB shot_count[0] 1
+    JUMP-UNLESS @end-loop shot_count[0]
+    JUMP @start-loop
+    LABEL @end-loop
+
 
 .. _parametric_compilation:
 

pyquil/quil.py

@@ -320,6 +320,50 @@ def inst(self, *instructions: Union[InstructionDesignator, RSProgram]) -> "Progr
 
         return self
 
+    def with_loop(
+        self,
+        num_iterations: int,
+        iteration_count_reference: MemoryReference,
+        start_label: Union[Label, LabelPlaceholder],
+        end_label: Union[Label, LabelPlaceholder],
+    ) -> "Program":
+        """
+        Return a copy of the ``Program`` wrapped in a Quil loop that will execute ``num_iterations`` times.
+
+        This loop is implemented with Quil and should not be confused with the ``num_shots`` property set by
+        :py:meth:`~pyquil.quil.Program.wrap_in_numshots_loop`. The latter is metadata on the program that
+        can tell an executor how many times to run the program. In comparison, this method adds Quil
+        instructions to your program to specify a loop in the Quil program itself.
+
+        The loop is constructed by wrapping the body of the program in classical Quil instructions.
+        The given ``iteration_count_reference`` must refer to an INTEGER memory region. The value at the
+        reference given will be set to ``num_iterations`` and decremented in the loop. The loop will
+        terminate when the reference reaches 0. For this reason your program should not itself
+        modify the value at the reference unless you intend to modify the remaining number of
+        iterations (i.e. to break the loop).
+
+        The given ``start_label`` and ``end_label`` will be used as the entry and exit points for
+        the loop, respectively. You should provide unique ``JumpTarget``\\s that won't be used elsewhere
+        in the program.
+
+        If ``num_iterations`` is 1, then a copy of the program is returned without any changes. Raises a
+        ``TypeError`` if ``num_iterations`` is negative.
+
+        :param num_iterations: The number of times to execute the loop.
+        :param loop_count_reference: The memory reference to use as the loop counter.
+        :param start_label: The ``JumpTarget`` to use at the start of the loop.
+        :param end_label: The ``JumpTarget`` to use to at the end of the loop.
+        """
+        looped_program = Program(
+            self._program.wrap_in_loop(
+                iteration_count_reference._to_rs_memory_reference(),
+                start_label.target,
+                end_label.target,
+                num_iterations,
+            )
+        )
+        return looped_program
+
     @_invalidates_cached_properties
     def resolve_placeholders(self) -> None:
         """

test/unit/__snapshots__/test_program.ambr

@@ -0,0 +1,24 @@
+# name: test_with_loop
+  '''
+  DECLARE ro BIT[1]
+  DECLARE shot_count INTEGER[1]
+  DEFFRAME 0 "rx":
+  	HARDWARE-OBJECT: "hardware"
+  DEFWAVEFORM custom:
+  	1, 2
+  DEFCAL I 0:
+  	DELAY 0 1
+  MOVE shot_count[0] 100
+  LABEL @start-loop
+  MEASURE q ro[0]
+  JUMP-UNLESS @end-reset ro[0]
+  X q
+  LABEL @end-reset
+  I 0
+  SUB shot_count[0] 1
+  JUMP-UNLESS @end-loop shot_count[0]
+  JUMP @start-loop
+  LABEL @end-loop
+
+  '''
+# ---

test/unit/test_program.py

@@ -1,7 +1,10 @@
 import numpy as np
+from syrupy.assertion import SnapshotAssertion
 
 from pyquil import Program
-from pyquil.quil import AbstractInstruction, Declare, Measurement, MemoryReference
+from pyquil.quilatom import Label
+from pyquil.quilbase import MemoryReference
+from pyquil.quil import AbstractInstruction, Declare, Measurement
 from pyquil.experiment._program import (
     measure_qubits,
     parameterized_single_qubit_measurement_basis,
@@ -87,6 +90,32 @@ def test_adding_does_not_mutate():
     assert p1.calibrations != p_all.calibrations
 
 
+def test_with_loop(snapshot: SnapshotAssertion):
+    p = Program(
+        """DECLARE ro BIT
+DECLARE shot_count INTEGER
+MEASURE q ro
+JUMP-UNLESS @end-reset ro
+X q
+LABEL @end-reset
+
+DEFCAL I 0:
+    DELAY 0 1.0
+DEFFRAME 0 \"rx\":
+    HARDWARE-OBJECT: \"hardware\"
+DEFWAVEFORM custom:
+    1,2
+I 0
+"""
+    )
+    p_copy = p.copy()
+
+    looped = p.with_loop(100, MemoryReference("shot_count"), Label("start-loop"), Label("end-loop"))
+
+    assert p_copy == p
+    assert looped.out() == snapshot
+
+
 def test_filter_program():
     program = Program(Declare("ro", "BIT", 1), MEASURE(0, MemoryReference("ro", 1)), H(0))