[DRAFT] MemoryQubits via Std.ResourceEstimation.Load/Store

library/std/src/Std/ResourceEstimation.qs

@@ -211,6 +211,37 @@ function LeastFrequentlyUsed() : Int {
     return 1;
 }
 
+/// Signals to resource estimator that this qubit is stored to memory, i.e.
+/// transitions from "compute qubit" to "memory qubit".
+/// Each qubit is allocated as "compute" by default.
+/// MemoryQubitStore can be applied only to "compute" qubit. That is, qubit on which
+/// MemoryQubitLoad wasn't applied or if it was applied, it was followed by
+/// MemoryQubitLoad.
+/// While qubit is in "memory" state, no gates or measurements can be applied to it.
+function MemoryQubitStore(q: Qubit) : Unit {
+    body intrinsic;
+}
+
+/// Signals to resource estimator that this qubit is loaded from memory, i.e.
+/// transitions from "memory qubit" to "compute qubit".
+/// Can be applied only to "memory qubit", i.e. on which previously MemoryQubitStore
+/// was applied.
+function MemoryQubitLoad(q: Qubit) : Unit {
+    body intrinsic;
+}
+
+
+/// All subsequent `use` statements will allocate compute qubits.
+function AllocateComputeQubits() : Unit {
+    body intrinsic;
+}
+
+/// All subsequent `use` statements will allocate memory qubits.
+function AllocateMemoryQubits() : Unit {
+    body intrinsic;
+}
+
+
 export
     SingleVariant,
     BeginEstimateCaching,
@@ -228,4 +259,8 @@ export
     RepeatEstimates,
     EnableMemoryComputeArchitecture,
     LeastRecentlyUsed,
-    LeastFrequentlyUsed;
+    LeastFrequentlyUsed,
+    MemoryQubitLoad,
+    MemoryQubitStore,
+    AllocateComputeQubits,
+    AllocateMemoryQubits;

source/compiler/qsc_eval/src/backend.rs

@@ -1033,7 +1033,11 @@ impl Backend for SparseSim {
             | "AccountForEstimatesInternal"
             | "BeginRepeatEstimatesInternal"
             | "EndRepeatEstimatesInternal"
-            | "EnableMemoryComputeArchitecture" => Some(Ok(Value::unit())),
+            | "EnableMemoryComputeArchitecture"
+            | "MemoryQubitStore"
+            | "MemoryQubitLoad"
+            | "AllocateComputeQubits"
+            | "AllocateMemoryQubits" => 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,10 @@ impl<'a> PartialEvaluator<'a> {
             | "BeginRepeatEstimatesInternal"
             | "EndRepeatEstimatesInternal"
             | "EnableMemoryComputeArchitecture"
+            | "MemoryQubitStore"
+            | "MemoryQubitLoad"
+            | "AllocateComputeQubits"
+            | "AllocateMemoryQubits"
             | "ApplyIdleNoise"
             | "GlobalPhase"
             | "Message"

source/compiler/qsc_partial_eval/src/management.rs

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

source/resource_estimator/src/counts.rs

@@ -14,7 +14,13 @@ use rustc_hash::FxHashMap;
 use std::{array, cell::RefCell, f64::consts::PI, fmt::Debug, iter::Sum};
 
 use crate::{counts::memory_compute::CachingStrategy, system::LogicalResourceCounts};
-use memory_compute::MemoryComputeInfo;
+use memory_compute::{MemoryComputeInfo, QubitPool};
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+enum QubitType {
+    Compute,
+    Memory,
+}
 
 /// Resource counter implementation
 ///
@@ -53,6 +59,21 @@ pub struct LogicalCounter {
     /// Map to track any post-select measurements by their associated qubit.
     /// This value is used in a measurement, if present, before generating a random result.
     post_select_measurements: FxHashMap<usize, bool>,
+
+    /// Pool of typed compute qubits used by manual MemoryQubitStore/Load.
+    compute_qubit_pool: QubitPool,
+    /// Pool of typed memory qubits used by manual MemoryQubitStore/Load.
+    memory_qubit_pool: QubitPool,
+    /// Mapping from untyped allocated qubit id to its typed counterpart/state.
+    typed_qubit_map: FxHashMap<usize, (QubitType, usize)>,
+    /// Peak number of typed qubits alive at once in manual memory-qubit mode.
+    typed_qubit_peak: usize,
+    /// Type of next alloctaed qubit.
+    next_allocation_qubit_type: QubitType,
+
+    manual_memory_reads: usize,
+    manual_memory_writes: usize,
+    manual_memory_mode: bool,
 }
 
 impl Default for LogicalCounter {
@@ -73,6 +94,14 @@ impl Default for LogicalCounter {
             memory_compute: None,
             rnd: RefCell::new(StdRng::seed_from_u64(0)),
             post_select_measurements: FxHashMap::default(),
+            compute_qubit_pool: QubitPool::default(),
+            memory_qubit_pool: QubitPool::default(),
+            typed_qubit_map: FxHashMap::default(),
+            typed_qubit_peak: 0,
+            next_allocation_qubit_type: QubitType::Compute,
+            manual_memory_reads: 0,
+            manual_memory_writes: 0,
+            manual_memory_mode: false,
         }
     }
 }
@@ -87,12 +116,24 @@ impl LogicalCounter {
                     Some(memory_compute.read_from_memory_count() as u64),
                     Some(memory_compute.write_to_memory_count() as u64),
                 )
+            } else if self.manual_memory_mode {
+                (
+                    Some(self.compute_qubit_pool.max_in_use() as u64),
+                    Some(self.manual_memory_reads as u64),
+                    Some(self.manual_memory_writes as u64),
+                )
             } else {
                 (None, None, None)
             };
 
+        let num_qubits = if self.manual_memory_mode {
+            self.compute_qubit_pool.max_in_use() + self.memory_qubit_pool.max_in_use()
+        } else {
+            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 _,
@@ -461,8 +502,17 @@ impl LogicalCounter {
     }
 
     fn assert_compute_qubits(&mut self, qubits: impl IntoIterator<Item = usize>) {
+        let qubits: Vec<_> = qubits.into_iter().collect();
         if let Some(memory_compute) = &mut self.memory_compute {
-            memory_compute.assert_compute_qubits(qubits);
+            memory_compute.assert_compute_qubits(qubits.iter().copied());
+        }
+
+        for qubit in qubits {
+            if let Some((qubit_type, _)) = self.typed_qubit_map.get(&qubit) {
+                if matches!(qubit_type, QubitType::Memory) {
+                    panic!("cannot apply compute operation to a memory qubit");
+                }
+            }
         }
     }
 
@@ -602,17 +652,33 @@ impl Backend for LogicalCounter {
     fn z(&mut self, _q: usize) {}
 
     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
-        }
+        };
+
+        let typed_index = match self.next_allocation_qubit_type {
+            QubitType::Compute => self.compute_qubit_pool.allocate(),
+            QubitType::Memory => self.memory_qubit_pool.allocate(),
+        };
+        self.typed_qubit_map
+            .insert(index, (self.next_allocation_qubit_type, typed_index));
+        self.typed_qubit_peak = self.typed_qubit_peak.max(self.typed_qubit_map.len());
+
+        index
     }
 
     fn qubit_release(&mut self, q: usize) -> bool {
+        if let Some((qubit_type, index)) = self.typed_qubit_map.remove(&q) {
+            match qubit_type {
+                QubitType::Compute => self.compute_qubit_pool.release(index),
+                QubitType::Memory => self.memory_qubit_pool.release(index),
+            }
+        }
         self.free_list.push(q);
         true
     }
@@ -630,6 +696,15 @@ impl Backend for LogicalCounter {
         if let Some(val) = q1_post_select {
             self.post_select_measurements.insert(q0, val);
         }
+
+        let q0_typed = self.typed_qubit_map.remove(&q0);
+        let q1_typed = self.typed_qubit_map.remove(&q1);
+        if let Some(typed_info) = q0_typed {
+            self.typed_qubit_map.insert(q1, typed_info);
+        }
+        if let Some(typed_info) = q1_typed {
+            self.typed_qubit_map.insert(q0, typed_info);
+        }
     }
 
     fn capture_quantum_state(&mut self) -> (Vec<(BigUint, Complex<f64>)>, usize) {
@@ -706,6 +781,56 @@ impl Backend for LogicalCounter {
 
                 Some(Ok(Value::unit()))
             }
+            "MemoryQubitLoad" => {
+                self.manual_memory_mode = true;
+                let q = arg.unwrap_qubit().deref().0;
+                let Some((qubit_type, index)) = self.typed_qubit_map.get(&q).copied() else {
+                    return Some(Err(
+                        "MemoryQubitLoad can only be applied to a memory qubit".to_string()
+                    ));
+                };
+                if !matches!(qubit_type, QubitType::Memory) {
+                    return Some(Err(
+                        "MemoryQubitLoad can only be applied to a memory qubit".to_string()
+                    ));
+                }
+                self.memory_qubit_pool.release(index);
+                let compute_index = self.compute_qubit_pool.allocate();
+                self.typed_qubit_map
+                    .insert(q, (QubitType::Compute, compute_index));
+                self.manual_memory_reads += 1;
+                Some(Ok(Value::unit()))
+            }
+            "MemoryQubitStore" => {
+                self.manual_memory_mode = true;
+                let q = arg.unwrap_qubit().deref().0;
+                let Some((qubit_type, index)) = self.typed_qubit_map.get(&q).copied() else {
+                    return Some(Err(
+                        "MemoryQubitStore can only be applied to a compute qubit".to_string(),
+                    ));
+                };
+                if !matches!(qubit_type, QubitType::Compute) {
+                    return Some(Err(
+                        "MemoryQubitStore can only be applied to a compute qubit".to_string(),
+                    ));
+                }
+                self.compute_qubit_pool.release(index);
+                let memory_index = self.memory_qubit_pool.allocate();
+                self.typed_qubit_map
+                    .insert(q, (QubitType::Memory, memory_index));
+                self.manual_memory_writes += 1;
+                Some(Ok(Value::unit()))
+            }
+            "AllocateComputeQubits" => {
+                self.manual_memory_mode = true;
+                self.next_allocation_qubit_type = QubitType::Compute;
+                Some(Ok(Value::unit()))
+            }
+            "AllocateMemoryQubits" => {
+                self.manual_memory_mode = true;
+                self.next_allocation_qubit_type = QubitType::Memory;
+                Some(Ok(Value::unit()))
+            }
             _ => None,
         }
     }

source/resource_estimator/src/counts/memory_compute.rs

@@ -5,6 +5,42 @@ use std::hash::Hash;
 #[cfg(test)]
 mod tests;
 
+#[derive(Default)]
+pub struct QubitPool {
+    free_list: Vec<usize>,
+    next_id: usize,
+    in_use: usize,
+    max_in_use: usize,
+}
+
+#[derive(Default)]
+pub struct QubitPool {
+    free_list: Vec<usize>,
+    next_id: usize,
+}
+
+impl QubitPool {
+    pub fn allocate(&mut self) -> usize {
+        let index = if let Some(index) = self.free_list.pop() {
+            index
+        } else {
+            let index = self.next_id;
+            self.next_id += 1;
+            index
+        };
+
+        index
+    }
+
+    pub fn release(&mut self, index: usize) {
+        self.free_list.push(index);
+    }
+
+    pub fn max_in_use(&self) -> usize {
+        self.next_id
+    }
+}
+
 pub enum CachingStrategy {
     LeastRecentlyUsed(LeastRecentlyUsedPriorityQueue<usize>),
     LeastFrequentlyUsed(LeastFrequentlyUsedPriorityQueue<usize>),