[DRAFT] Allow manual and automatic memory/compute working together

library/src/lib.rs

@@ -61,6 +61,10 @@ pub const STD_LIB: &[(&str, &str)] = &[
         "qsharp-library-source:Std/Measurement.qs",
         include_str!("../std/src/Std/Measurement.qs"),
     ),
+    (
+        "qsharp-library-source:Std/Memory.qs",
+        include_str!("../std/src/Std/Memory.qs"),
+    ),
     (
         "qsharp-library-source:QIR/Intrinsic.qs",
         include_str!("../std/src/QIR/Intrinsic.qs"),

library/std/qsharp.json

@@ -13,6 +13,7 @@
     "src/Std/Logical.qs",
     "src/Std/Math.qs",
     "src/Std/Measurement.qs",
+    "src/Std/Memory.qs",
     "src/Std/Random.qs",
     "src/Std/Range.qs",
     "src/Std/ResourceEstimation.qs",

library/std/src/Std/Memory.qs

@@ -0,0 +1,17 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+
+/// Loads a qubit from "memory/cold" qubit to "compute/hot" qubit.
+/// Does nothing if qubit is already "hot".
+function MemoryQubitLoad(q : Qubit) : Unit {
+    body intrinsic;
+}
+
+/// Stores a qubit from "compute/hot" qubit to "memory/cold" qubit.
+/// Does nothing if qubit is already "cold".
+function MemoryQubitStore(q : Qubit) : Unit {
+    body intrinsic;
+}
+
+export MemoryQubitLoad, MemoryQubitStore;

source/compiler/qsc_eval/src/backend.rs

@@ -1033,7 +1033,9 @@ impl Backend for SparseSim {
             | "AccountForEstimatesInternal"
             | "BeginRepeatEstimatesInternal"
             | "EndRepeatEstimatesInternal"
-            | "EnableMemoryComputeArchitecture" => Some(Ok(Value::unit())),
+            | "EnableMemoryComputeArchitecture"
+            | "MemoryQubitLoad"
+            | "MemoryQubitStore" => Some(Ok(Value::unit())),
             "ConfigurePauliNoise" => {
                 let [xv, yv, zv] = &*arg.unwrap_tuple() else {
                     panic!("tuple arity for ConfigurePauliNoise intrinsic should be 3");

source/compiler/qsc_partial_eval/src/lib.rs

@@ -1798,6 +1798,8 @@ impl<'a> PartialEvaluator<'a> {
             | "BeginRepeatEstimatesInternal"
             | "EndRepeatEstimatesInternal"
             | "EnableMemoryComputeArchitecture"
+            | "MemoryQubitLoad"
+            | "MemoryQubitStore"
             | "ApplyIdleNoise"
             | "GlobalPhase"
             | "Message"

source/compiler/qsc_partial_eval/src/management.rs

@@ -149,6 +149,8 @@ impl Backend for QuantumIntrinsicsChecker {
             "BeginEstimateCaching" => Some(Ok(Value::Bool(true))),
             "EndEstimateCaching"
             | "EnableMemoryComputeArchitecture"
+            | "MemoryQubitStore"
+            | "MemoryQubitLoad"
             | "GlobalPhase"
             | "ConfigurePauliNoise"
             | "ConfigureQubitLoss"

source/resource_estimator/src/counts.rs

@@ -83,16 +83,20 @@ impl LogicalCounter {
         let (num_compute_qubits, read_from_memory_count, write_to_memory_count) =
             if let Some(memory_compute) = &self.memory_compute {
                 (
-                    Some(memory_compute.compute_size() as u64),
+                    Some(memory_compute.max_compute_qubits_count as u64),
                     Some(memory_compute.read_from_memory_count() as u64),
                     Some(memory_compute.write_to_memory_count() as u64),
                 )
             } else {
                 (None, None, None)
             };
+        let num_qubits = match &self.memory_compute {
+            Some(mc) => mc.max_compute_qubits_count + mc.max_memory_qubits_count,
+            None => self.next_free,
+        };
 
         LogicalResourceCounts {
-            num_qubits: self.next_free as _,
+            num_qubits: num_qubits as _,
             t_count: self.t_count as _,
             rotation_count: self.r_count as _,
             rotation_depth: self.layers.iter().filter(|layer| layer.r != 0).count() as _,
@@ -595,25 +599,36 @@ impl Backend for LogicalCounter {
         self.schedule_t(q);
     }
 
-    fn x(&mut self, _q: usize) {}
+    fn x(&mut self, q: usize) {
+        self.assert_compute_qubits([q]);
+    }
 
-    fn y(&mut self, _q: usize) {}
+    fn y(&mut self, q: usize) {
+        self.assert_compute_qubits([q]);
+    }
 
-    fn z(&mut self, _q: usize) {}
+    fn z(&mut self, q: usize) {
+        self.assert_compute_qubits([q]);
+    }
 
     fn qubit_allocate(&mut self) -> usize {
-        if let Some(index) = self.free_list.pop() {
+        let index = if let Some(index) = self.free_list.pop() {
             index
         } else {
             let index = self.next_free;
             self.next_free += 1;
             self.max_layer.push(self.allocation_barrier);
             index
-        }
+        };
+
+        index
     }
 
     fn qubit_release(&mut self, q: usize) -> bool {
         self.free_list.push(q);
+        if let Some(memory_compute) = &mut self.memory_compute {
+            memory_compute.release(q);
+        }
         true
     }
 
@@ -692,6 +707,18 @@ impl Backend for LogicalCounter {
                 self.enable_memory_compute(compute_capacity, strategy);
                 Some(Ok(Value::unit()))
             }
+            "MemoryQubitLoad" => {
+                let qid = arg.unwrap_qubit().deref().0;
+                self.assert_compute_qubits([qid]);
+                Some(Ok(Value::unit()))
+            }
+            "MemoryQubitStore" => {
+                if let Some(memory_compute) = &mut self.memory_compute {
+                    let qid = arg.unwrap_qubit().deref().0;
+                    memory_compute.move_to_memory(qid);
+                }
+                Some(Ok(Value::unit()))
+            }
             "GlobalPhase" | "ConfigurePauliNoise" | "ConfigureQubitLoss" | "ApplyIdleNoise" => {
                 Some(Ok(Value::unit()))
             }

source/resource_estimator/src/counts/memory_compute.rs

@@ -1,4 +1,5 @@
 use rustc_hash::{FxHashMap, FxHashSet};
+use std::cmp::max;
 use std::collections::VecDeque;
 use std::hash::Hash;
 
@@ -8,6 +9,7 @@ mod tests;
 pub enum CachingStrategy {
     LeastRecentlyUsed(LeastRecentlyUsedPriorityQueue<usize>),
     LeastFrequentlyUsed(LeastFrequentlyUsedPriorityQueue<usize>),
+    Manual,
 }
 
 impl CachingStrategy {
@@ -18,61 +20,105 @@ impl CachingStrategy {
     pub fn least_frequently_used(capacity: usize) -> Self {
         CachingStrategy::LeastFrequentlyUsed(LeastFrequentlyUsedPriorityQueue::new(capacity))
     }
+
+    pub fn manual() -> Self {
+        CachingStrategy::Manual
+    }
+
+    /// Inserts given qubits ids into the set of compute qubits.
+    /// Returns which compute qubits must be evicted to memory before given qubits can
+    /// be read or allocated.
+    fn insert_all(&mut self, qubit_ids: impl IntoIterator<Item = usize>) -> Vec<usize> {
+        match self {
+            CachingStrategy::LeastRecentlyUsed(lru) => lru.insert_all(qubit_ids),
+            CachingStrategy::LeastFrequentlyUsed(lfu) => lfu.insert_all(qubit_ids),
+            CachingStrategy::Manual => vec![],
+        }
+    }
 }
 
 pub struct MemoryComputeInfo {
-    /// LRU or LFU set with qubits currently in compute mode
-    compute_qubits: CachingStrategy,
-
-    /// Additional reads/writes not captured by the LRU or LFU set (e.g. when
-    /// manually counted for caching functions)
-    pub(crate) rfm_extra: usize,
-    pub(crate) wtm_extra: usize,
+    /// LRU or LFU set with qubits currently in compute mode.
+    strategy: CachingStrategy,
+    compute_qubits: FxHashSet<usize>,
+    memory_qubits: FxHashSet<usize>,
+    pub(crate) max_memory_qubits_count: usize,
+    pub(crate) max_compute_qubits_count: usize,
+    reads_count: usize,
+    writes_count: usize,
 }
 
 impl MemoryComputeInfo {
     pub fn new(strategy: CachingStrategy) -> Self {
         Self {
-            compute_qubits: strategy,
-            rfm_extra: 0,
-            wtm_extra: 0,
+            strategy: strategy,
+            compute_qubits: FxHashSet::default(),
+            memory_qubits: FxHashSet::default(),
+            max_memory_qubits_count: 0,
+            max_compute_qubits_count: 0,
+            reads_count: 0,
+            writes_count: 0,
         }
     }
 
-    pub fn assert_compute_qubits(&mut self, qubits: impl IntoIterator<Item = usize>) {
-        match &mut self.compute_qubits {
-            CachingStrategy::LeastRecentlyUsed(lru) => lru.insert_all(qubits),
-            CachingStrategy::LeastFrequentlyUsed(lfu) => lfu.insert_all(qubits),
+    /// Moves this qubit to set of compute qubits.
+    /// If it was a memory qubit, records that as "read" operation.
+    /// External callers must call `assert_compute_qubits` instead of this method (so
+    /// automatic eviction can happen).
+    fn move_to_compute(&mut self, qid: usize) {
+        if self.memory_qubits.contains(&qid) {
+            self.memory_qubits.remove(&qid);
+            self.reads_count += 1;
         }
+        self.compute_qubits.insert(qid);
+        self.max_compute_qubits_count =
+            max(self.max_compute_qubits_count, self.compute_qubits.len());
     }
 
-    pub fn compute_size(&self) -> usize {
-        match &self.compute_qubits {
-            CachingStrategy::LeastRecentlyUsed(lru) => lru.max_size(),
-            CachingStrategy::LeastFrequentlyUsed(lfu) => lfu.max_size(),
+    /// Moves this qubit to set of memory qubits.
+    /// If it was a compute qubit, records that as "write" operation.
+    pub fn move_to_memory(&mut self, qid: usize) {
+        if self.compute_qubits.contains(&qid) {
+            self.compute_qubits.remove(&qid);
+            self.writes_count += 1;
         }
+        self.memory_qubits.insert(qid);
+        self.max_memory_qubits_count = max(self.max_memory_qubits_count, self.memory_qubits.len());
     }
 
-    pub fn read_from_memory_count(&self) -> usize {
-        match &self.compute_qubits {
-            CachingStrategy::LeastRecentlyUsed(lru) => lru.inserted_new_count() + self.rfm_extra,
-            CachingStrategy::LeastFrequentlyUsed(lfu) => lfu.inserted_new_count() + self.rfm_extra,
+    /// Releases the qubit.
+    pub fn release(&mut self, qid: usize) {
+        self.compute_qubits.remove(&qid);
+        self.memory_qubits.remove(&qid);
+    }
+
+    /// Ensures that these qubits are compute qubits by doing reads if necessary.
+    /// If this requires evicting (writing) some qubits to memory, does that first.
+    pub fn assert_compute_qubits(&mut self, qubits: impl IntoIterator<Item = usize>) {
+        let qubits_copy: Vec<usize> = qubits.into_iter().collect();
+        let qubits_to_evict = self.strategy.insert_all(qubits_copy.clone());
+        for qid in qubits_to_evict {
+            self.move_to_memory(qid);
+        }
+        for qid in qubits_copy {
+            self.move_to_compute(qid);
         }
     }
 
+    pub fn read_from_memory_count(&self) -> usize {
+        self.reads_count
+    }
+
     pub fn write_to_memory_count(&self) -> usize {
-        match &self.compute_qubits {
-            CachingStrategy::LeastRecentlyUsed(lru) => lru.removed_count() + self.wtm_extra,
-            CachingStrategy::LeastFrequentlyUsed(lfu) => lfu.removed_count() + self.wtm_extra,
-        }
+        self.writes_count
     }
 
     pub fn increase_read_from_memory_count(&mut self, count: usize) {
-        self.rfm_extra += count;
+        self.reads_count += count;
     }
 
     pub fn increase_write_to_memory_count(&mut self, count: usize) {
-        self.wtm_extra += count;
+        self.writes_count += count;
     }
 }
 
@@ -125,9 +171,9 @@ impl<K: Eq + Hash + Clone> LeastRecentlyUsedPriorityQueue<K> {
     /// Insert multiple keys ensuring they are all present afterwards. If more
     /// unique new keys are provided than capacity, only the most recently
     /// processed up to `capacity` will remain.
-    pub fn insert_all<I: IntoIterator<Item = K>>(&mut self, keys: I) {
+    pub fn insert_all<I: IntoIterator<Item = K>>(&mut self, keys: I) -> Vec<K> {
         if self.capacity == 0 {
-            return;
+            return vec![];
         }
         // Collect unique keys from input preserving order of first occurrence.
         let mut seen_input = FxHashSet::default();
@@ -144,6 +190,7 @@ impl<K: Eq + Hash + Clone> LeastRecentlyUsedPriorityQueue<K> {
 
         // Process each key in order; we evict as we go and since new elements
         // are moved front they will be retained if we exceed capacity.
+        let mut removed_keys: Vec<K> = Vec::new();
         for k in ordered {
             if self.contains(&k) {
                 // Just update recency by moving element to front of deque
@@ -162,6 +209,7 @@ impl<K: Eq + Hash + Clone> LeastRecentlyUsedPriorityQueue<K> {
                 {
                     self.map.remove(&key);
                     self.removed += 1;
+                    removed_keys.push(key);
                 }
                 self.map.insert(k.clone());
                 self.nodes.push_front(k);
@@ -172,6 +220,8 @@ impl<K: Eq + Hash + Clone> LeastRecentlyUsedPriorityQueue<K> {
         if self.map.len() > self.max_size {
             self.max_size = self.map.len();
         }
+
+        return removed_keys;
     }
 }
 
@@ -222,9 +272,10 @@ impl<K: Eq + Hash + Clone> LeastFrequentlyUsedPriorityQueue<K> {
 
     /// Insert multiple keys ensuring they are all present afterwards. If unique
     /// keys exceed capacity, only a subset up to capacity will remain.
-    pub fn insert_all<I: IntoIterator<Item = K>>(&mut self, keys: I) {
+    /// Returns evicted keys.
+    pub fn insert_all<I: IntoIterator<Item = K>>(&mut self, keys: I) -> Vec<K> {
         if self.capacity == 0 {
-            return;
+            return vec![];
         }
         let mut seen = FxHashSet::default();
         let mut ordered: Vec<K> = Vec::new();
@@ -241,6 +292,7 @@ impl<K: Eq + Hash + Clone> LeastFrequentlyUsedPriorityQueue<K> {
         // Evict as needed to make space for new keys.  We need to evict before
         // adding the new elements, since frequency counters are low for new
         // elements and we risk to evict them before processing the whole input.
+        let mut removed_keys: Vec<K> = Vec::new();
         let new_missing = ordered
             .iter()
             .filter(|k| !self.map.contains_key(*k))
@@ -270,6 +322,7 @@ impl<K: Eq + Hash + Clone> LeastFrequentlyUsedPriorityQueue<K> {
                 }
                 if let Some(v) = victim {
                     self.remove_key_internal(&v);
+                    removed_keys.push(v);
                     needed -= 1;
                 } else {
                     break;
@@ -295,6 +348,8 @@ impl<K: Eq + Hash + Clone> LeastFrequentlyUsedPriorityQueue<K> {
         if self.map.len() > self.max_size {
             self.max_size = self.map.len();
         }
+
+        return removed_keys;
     }
 
     fn bump_bucket(&mut self, key: K, old_freq: u64, new_freq: u64) {

source/resource_estimator/src/counts/memory_compute/tests.rs

@@ -1,6 +1,7 @@
 use super::{LeastFrequentlyUsedPriorityQueue, LeastRecentlyUsedPriorityQueue};
 
 // ---------------- LRU tests -----------------
+
 #[test]
 fn lru_insert_all_all_existing() {
     let mut lru = LeastRecentlyUsedPriorityQueue::new(4);