Comments

src/stim/util_top/reference_sample_tree.cc

@@ -22,6 +22,7 @@ void ReferenceSampleTree::flatten_and_simplify_into(std::vector<ReferenceSampleT
         return;
     }
 
+    // Flatten children.
     std::vector<ReferenceSampleTree> flattened;
     if (!prefix_bits.empty()) {
         flattened.push_back(ReferenceSampleTree{
@@ -35,49 +36,76 @@ void ReferenceSampleTree::flatten_and_simplify_into(std::vector<ReferenceSampleT
     }
 
     // Fuse children.
-    auto &f = flattened;
-    for (size_t k = 0; k < f.size(); k++) {
+    std::vector<ReferenceSampleTree> fused;
+    if (!flattened.empty()) {
+        fused.push_back(std::move(flattened[0]));
+    }
+    for (size_t k = 1; k < flattened.size(); k++) {
+        auto &dst = fused.back();
+        auto &src = flattened[k];
+
         // Combine children with identical contents by adding their rep counts.
-        while (k + 1 < f.size() && f[k].prefix_bits == f[k + 1].prefix_bits && f[k].suffix_children == f[k + 1].suffix_children) {
-            f[k].repetitions += f[k + 1].repetitions;
-            f.erase(f.begin() + k + 1);
-        }
+        if (dst.prefix_bits == src.prefix_bits && dst.suffix_children == src.suffix_children) {
+            dst.repetitions += src.repetitions;
 
         // Fuse children with unrepeated contents if they can be fused.
-        while (k + 1 < f.size() && f[k].repetitions == 1 && f[k].suffix_children.empty() && f[k + 1].repetitions == 1) {
-            f[k].suffix_children = std::move(f[k + 1].suffix_children);
-            f[k].prefix_bits.insert(f[k].prefix_bits.end(), f[k + 1].prefix_bits.begin(), f[k + 1].prefix_bits.end());
-            f.erase(f.begin() + k + 1);
+        } else if (src.repetitions == 1 && dst.repetitions == 1 && dst.suffix_children.empty()) {
+            dst.suffix_children = std::move(src.suffix_children);
+            dst.prefix_bits.insert(dst.prefix_bits.end(), src.prefix_bits.begin(), src.prefix_bits.end());
+
+        } else {
+            fused.push_back(std::move(src));
         }
     }
 
     if (repetitions == 1) {
         // Un-nest all the children.
-        for (auto &e : flattened) {
+        for (auto &e : fused) {
             out.push_back(e);
         }
-    } else if (flattened.size() == 1) {
+    } else if (fused.size() == 1) {
         // Merge with single child.
-        flattened[0].repetitions *= repetitions;
-        out.push_back(std::move(flattened[0]));
-    } else if (flattened.empty()) {
+        fused[0].repetitions *= repetitions;
+        out.push_back(std::move(fused[0]));
+    } else if (fused.empty()) {
         // Nothing to report.
-    } else if (flattened[0].suffix_children.empty() && flattened[0].repetitions == 1) {
+    } else if (fused[0].suffix_children.empty() && fused[0].repetitions == 1) {
         // Take payload from first child.
-        auto result = std::move(flattened[0]);
-        flattened.erase(flattened.begin());
+        ReferenceSampleTree result = std::move(fused[0]);
+        fused.erase(fused.begin());
         result.repetitions = repetitions;
-        result.suffix_children = std::move(flattened);
+        result.suffix_children = std::move(fused);
         out.push_back(std::move(result));
     } else {
         out.push_back(ReferenceSampleTree{
             .prefix_bits={},
-            .suffix_children=std::move(flattened),
+            .suffix_children=std::move(fused),
             .repetitions=repetitions,
         });
     }
 }
 
+/// Finds how far back feedback operations ever look, within the loop.
+uint64_t stim::max_feedback_lookback_in_loop(const Circuit &loop) {
+    uint64_t furthest_lookback = 0;
+    for (const auto &inst : loop.operations) {
+        if (inst.gate_type == GateType::REPEAT) {
+            furthest_lookback = std::max(furthest_lookback, max_feedback_lookback_in_loop(inst.repeat_block_body(loop)));
+        } else {
+            auto f = GATE_DATA[inst.gate_type].flags;
+            if ((f & GateFlags::GATE_CAN_TARGET_BITS) && (f & GateFlags::GATE_TARGETS_PAIRS)) {
+                // Feedback-capable operation. Check for any measurement record targets.
+                for (auto t : inst.targets) {
+                    if (t.is_measurement_record_target()) {
+                        furthest_lookback = std::max(furthest_lookback, (uint64_t)-t.rec_offset());
+                    }
+                }
+            }
+        }
+    }
+    return furthest_lookback;
+}
+
 void ReferenceSampleTree::try_factorize(size_t period_factor) {
     if (prefix_bits.size() != 0 || suffix_children.size() % period_factor != 0) {
         return;

src/stim/util_top/reference_sample_tree.h

@@ -7,11 +7,11 @@ namespace stim {
 
 /// A compressed tree representation of a reference sample.
 struct ReferenceSampleTree {
-    /// Bits to repeatedly output before outputting bits for the children.
+    /// Raw bits to output before bits from the children.
     std::vector<bool> prefix_bits;
     /// Compressed representations of additional bits to output after the prefix.
     std::vector<ReferenceSampleTree> suffix_children;
-    /// The number of times to repeatedly output the prefix and suffix bits.
+    /// The number of times to repeatedly output the prefix+suffix bits.
     size_t repetitions = 0;
 
     /// Initializes a reference sample tree containing a reference sample for the given circuit.
@@ -20,22 +20,30 @@ struct ReferenceSampleTree {
     /// Returns a tree with the same compressed contents, but a simpler tree structure.
     ReferenceSampleTree simplified() const;
 
-    /// Determines whether the tree contains any bits at all.
-    bool empty() const;
-
+    /// Checks if two trees are exactly the same, including structure (not just uncompressed contents).
     bool operator==(const ReferenceSampleTree &other) const;
+    /// Checks if two trees are not exactly the same, including structure (not just uncompressed contents).
     bool operator!=(const ReferenceSampleTree &other) const;
+    /// Returns a simple description of the tree's structure, like "5*('101'+6*('11'))".
     std::string str() const;
 
+    /// Determines whether the tree contains any bits at all.
+    bool empty() const;
     /// Computes the total size of the uncompressed bits represented by the tree.
     size_t size() const;
 
     /// Writes the contents of the tree into the given output vector.
     void decompress_into(std::vector<bool> &output) const;
 
+    /// Folds redundant children into the repetition count, if they repeat this many times.
+    ///
+    /// For example, if the tree's children are [A, B, C, A, B, C] and the tree has no
+    /// prefix, then `try_factorize(2)` will reduce the children to [A, B, C] and double
+    /// the repetition count.
     void try_factorize(size_t period_factor);
 
    private:
+    /// Helper method for `simplified`.
     void flatten_and_simplify_into(std::vector<ReferenceSampleTree> &out) const;
 };
 std::ostream &operator<<(std::ostream &out, const ReferenceSampleTree &v);
@@ -69,6 +77,8 @@ struct CompressedReferenceSampleHelper {
         bool allow_false_negative) const;
 };
 
+uint64_t max_feedback_lookback_in_loop(const Circuit &loop);
+
 }  // namespace stim
 
 #include "stim/util_top/reference_sample_tree.inl"

src/stim/util_top/reference_sample_tree.inl

@@ -40,29 +40,10 @@ ReferenceSampleTree CompressedReferenceSampleHelper<W>::do_loop_with_no_folding(
     return result;
 }
 
-static uint64_t max_feedback_lookback_in_loop(const Circuit &loop) {
-    uint64_t furthest_lookback = 0;
-    for (const auto &inst : loop.operations) {
-        if (inst.gate_type == GateType::REPEAT) {
-            furthest_lookback = std::max(furthest_lookback, max_feedback_lookback_in_loop(inst.repeat_block_body(loop)));
-        } else {
-            auto f = GATE_DATA[inst.gate_type].flags;
-            if ((f & GateFlags::GATE_CAN_TARGET_BITS) && (f & GateFlags::GATE_TARGETS_PAIRS)) {
-                // Feedback-capable operation. Check for any measurement record targets.
-                for (auto t : inst.targets) {
-                    if (t.is_measurement_record_target()) {
-                        furthest_lookback = std::max(furthest_lookback, (uint64_t)-t.rec_offset());
-                    }
-                }
-            }
-        }
-    }
-    return furthest_lookback;
-}
-
 template <size_t W>
 ReferenceSampleTree CompressedReferenceSampleHelper<W>::do_loop_with_tortoise_hare_folding(const Circuit &loop, uint64_t reps) {
     if (reps < 10) {
+        // Probably not worth the overhead of tortoise-and-hare. Just run it raw.
         return do_loop_with_no_folding(loop, reps);
     }
 
@@ -96,7 +77,7 @@ ReferenceSampleTree CompressedReferenceSampleHelper<W>::do_loop_with_tortoise_ha
         return result;
     }
 
-    // Advance until the remaining iterations are a multiple of the period.
+    // Run more loop iterations until the remaining iterations are a multiple of the found period.
     assert(result.suffix_children.size() == hare_steps);
     uint64_t period = hare_steps - tortoise_steps;
     size_t period_steps_left = (reps - hare_steps) / period;

src/stim/util_top/reference_sample_tree.perf.cc

@@ -15,8 +15,37 @@ BENCHMARK(reference_sample_tree_surface_code_d31_r1000000000) {
         auto result = ReferenceSampleTree::from_circuit_reference_sample(circuit);
         total += result.empty();
     })
-        .goal_millis(25)
-        .show_rate("Samples", circuit.count_measurements());
+        .goal_millis(25);
+    if (total) {
+        std::cerr << "data dependence";
+    }
+}
+
+BENCHMARK(reference_sample_tree_nested_circuit) {
+    Circuit circuit(R"CIRCUIT(
+        M 0
+        REPEAT 100000 {
+            REPEAT 100000 {
+                REPEAT 100000 {
+                    X 0
+                    M 0
+                }
+                X 0
+                M 0
+            }
+            X 0
+            M 0
+        }
+        X 0
+        M 0
+    )CIRCUIT");
+    simd_bits<MAX_BITWORD_WIDTH> ref(0);
+    auto total = 0;
+    benchmark_go([&]() {
+        auto result = ReferenceSampleTree::from_circuit_reference_sample(circuit);
+        total += result.empty();
+    })
+        .goal_micros(230);
     if (total) {
         std::cerr << "data dependence";
     }

src/stim/util_top/reference_sample_tree.test.cc

@@ -183,6 +183,46 @@ TEST(ReferenceSampleTree, feedback) {
     ASSERT_EQ(ref.str(), "1*('0010'+2*('0')+4*('1')+1*('01011001000111')+12*('101011001000111'))");
 }
 
+TEST(max_feedback_lookback_in_loop, simple) {
+    ASSERT_EQ(max_feedback_lookback_in_loop(Circuit()), 0);
+
+    ASSERT_EQ(max_feedback_lookback_in_loop(Circuit(R"CIRCUIT(
+        REPEAT 100 {
+            REPEAT 100 {
+                M 0
+                X 0
+                M 0
+            }
+            REPEAT 200 {
+                M 0
+                DETECTOR rec[-1]
+            }
+            X 1
+            CX 1 0
+        }
+    )CIRCUIT")), 0);
+
+    ASSERT_EQ(max_feedback_lookback_in_loop(Circuit(R"CIRCUIT(
+        CX rec[-1] 0
+    )CIRCUIT")), 1);
+
+    ASSERT_EQ(max_feedback_lookback_in_loop(Circuit(R"CIRCUIT(
+        CZ 0 rec[-2]
+    )CIRCUIT")), 2);
+
+    ASSERT_EQ(max_feedback_lookback_in_loop(Circuit(R"CIRCUIT(
+        CZ 0 rec[-2]
+        CY 0 rec[-3]
+    )CIRCUIT")), 3);
+
+    ASSERT_EQ(max_feedback_lookback_in_loop(Circuit(R"CIRCUIT(
+        CZ 0 rec[-2]
+        REPEAT 100 {
+            CX rec[-5] 0
+        }
+    )CIRCUIT")), 5);
+}
+
 TEST(ReferenceSampleTree, nested_loops) {
     Circuit circuit(R"CIRCUIT(
         REPEAT 100 {