diff --git a/library/std/src/Std/ResourceEstimation.qs b/library/std/src/Std/ResourceEstimation.qs index 7d58ad9408..4d2ea65aea 100644 --- a/library/std/src/Std/ResourceEstimation.qs +++ b/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; diff --git a/source/compiler/qsc_eval/src/backend.rs b/source/compiler/qsc_eval/src/backend.rs index 36f4a42536..452e051525 100644 --- a/source/compiler/qsc_eval/src/backend.rs +++ b/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"); diff --git a/source/compiler/qsc_partial_eval/src/lib.rs b/source/compiler/qsc_partial_eval/src/lib.rs index 0994287cc9..d059294bdb 100644 --- a/source/compiler/qsc_partial_eval/src/lib.rs +++ b/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" diff --git a/source/compiler/qsc_partial_eval/src/management.rs b/source/compiler/qsc_partial_eval/src/management.rs index 46f3355e65..8315e74f0a 100644 --- a/source/compiler/qsc_partial_eval/src/management.rs +++ b/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" diff --git a/source/resource_estimator/src/counts.rs b/source/resource_estimator/src/counts.rs index 1cf8d8066b..e9bb4db2e3 100644 --- a/source/resource_estimator/src/counts.rs +++ b/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, + + /// 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, + /// 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) { + 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)>, 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, } } diff --git a/source/resource_estimator/src/counts/memory_compute.rs b/source/resource_estimator/src/counts/memory_compute.rs index e2df7d3fa5..60540bf50f 100644 --- a/source/resource_estimator/src/counts/memory_compute.rs +++ b/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, + next_id: usize, + in_use: usize, + max_in_use: usize, +} + +#[derive(Default)] +pub struct QubitPool { + free_list: Vec, + 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), LeastFrequentlyUsed(LeastFrequentlyUsedPriorityQueue), diff --git a/source/resource_estimator/src/counts/tests.rs b/source/resource_estimator/src/counts/tests.rs index f3fdc2a67f..92db45d7dd 100644 --- a/source/resource_estimator/src/counts/tests.rs +++ b/source/resource_estimator/src/counts/tests.rs @@ -3,6 +3,7 @@ use std::convert::Into; +use crate::system::LogicalResourceCounts; use expect_test::{Expect, expect}; use indoc::indoc; use miette::Report; @@ -14,7 +15,7 @@ use qsc::{ use super::LogicalCounter; -fn verify_logical_counts(source: &str, entry: Option<&str>, expect: &Expect) { +fn run_logical_counts(source: &str, entry: Option<&str>) -> LogicalResourceCounts { let source_map = SourceMap::new([("test".into(), source.into())], entry.map(Into::into)); let (std_id, store) = qsc::compile::package_store_with_stdlib(TargetCapabilityFlags::all()); @@ -40,9 +41,7 @@ fn verify_logical_counts(source: &str, entry: Option<&str>, expect: &Expect) { let mut out = GenericReceiver::new(&mut stdout); match interpreter.eval_entry_with_sim(&mut counter, &mut out) { - Ok(_) => { - expect.assert_debug_eq(&counter.logical_resources()); - } + Ok(_) => counter.logical_resources(), Err(err) => { for e in err { let report = Report::from(e); @@ -53,6 +52,11 @@ fn verify_logical_counts(source: &str, entry: Option<&str>, expect: &Expect) { } } +fn verify_logical_counts(source: &str, entry: Option<&str>, expect: &Expect) { + let logical_counts = run_logical_counts(source, entry); + expect.assert_debug_eq(&logical_counts); +} + #[test] fn gates_are_counted() { verify_logical_counts( @@ -427,3 +431,43 @@ fn post_selection_can_take_impossible_branch() { "#]], ); } + +#[test] +fn manual_memory_qubits_load_store() { + let counts = run_logical_counts( + indoc! {" + operation Main() : Unit { + use q = Qubit(); + Std.ResourceEstimation.MemoryQubitStore(q); + Std.ResourceEstimation.MemoryQubitLoad(q); + } + "}, + None, + ); + assert_eq!(counts.write_to_memory_count, Some(1)); + assert_eq!(counts.read_from_memory_count, Some(1)); + assert_eq!(counts.num_compute_qubits, Some(1)); + assert_eq!(counts.num_qubits, 2); +} + +#[test] +fn manual_memory_qubits_allocator() { + let counts = run_logical_counts( + indoc! {" + operation Main() : Unit { + Std.ResourceEstimation.AllocateMemoryQubits(); + use q1 = Qubit[5]; + Std.ResourceEstimation.AllocateComputeQubits(); + use q2 = Qubit[2]; + Std.ResourceEstimation.MemoryQubitLoad(q1[0]); + } + "}, + None, + ); + + // Initially allocated 5 memory, 2 compute qubits. + // Then moved a qubit from memory to compute. + // So, in total need 5 memory and 3 compute qubits. + assert_eq!(counts.num_compute_qubits, Some(3)); + assert_eq!(counts.num_qubits, 8); +}