Skip to main content

tnc/contractionpath/
contraction_cost.rs

1//! Different methods to compute the computational and memory cost of contraction
2//! paths.
3
4use num_complex::Complex64;
5
6use crate::{
7    contractionpath::{ContractionPath, SimplePathRef},
8    tensornetwork::tensor::LeafTensor,
9    tensornetwork::tensor::Tensor,
10};
11
12/// Returns Schroedinger contraction time complexity of contracting two
13/// [`LeafTensor`] objects. Considers cost of complex operations.
14///
15/// # Examples
16/// ```
17/// # use tnc::tensornetwork::tensor::LeafTensor;
18/// # use tnc::contractionpath::contraction_cost::contract_cost_tensors;
19/// # use rustc_hash::FxHashMap;
20/// let bond_dims = FxHashMap::from_iter([(0, 5),(1, 7), (2, 9), (3, 11), (4, 13)]);
21/// let tensor1 = LeafTensor::new_from_map(vec![0, 1, 2], &bond_dims);
22/// let tensor2 = LeafTensor::new_from_map(vec![2, 3, 4], &bond_dims);
23/// // result = [0, 1, 2, 3, 4] // cost of (9-1)*54*5005 = 350350;
24/// assert_eq!(contract_cost_tensors(&tensor1, &tensor2), 350350.);
25/// ```
26pub fn contract_cost_tensors(t_1: &LeafTensor, t_2: &LeafTensor) -> f64 {
27    let final_dims = t_1 ^ t_2;
28    let shared_dims = t_1 & t_2;
29
30    let single_loop_cost = shared_dims.size();
31    (single_loop_cost - 1f64).mul_add(2f64, single_loop_cost * 6f64) * final_dims.size()
32}
33
34/// Returns Schroedinger contraction time complexity of contracting two
35/// [`LeafTensor`] objects. Naive op cost, does not consider costs of multiplication.
36///
37/// # Examples
38/// ```
39/// # use tnc::tensornetwork::tensor::LeafTensor;
40/// # use tnc::contractionpath::contraction_cost::contract_op_cost_tensors;
41/// # use rustc_hash::FxHashMap;
42/// let bond_dims = FxHashMap::from_iter([(0, 5),(1, 7), (2, 9), (3, 11), (4, 13)]);
43/// let tensor1 = LeafTensor::new_from_map(vec![0, 1, 2], &bond_dims);
44/// let tensor2 = LeafTensor::new_from_map(vec![2, 3, 4], &bond_dims);
45/// // result = [0, 1, 2, 3, 4] // cost of 5*7*9*11*13 = 45045;
46/// assert_eq!(contract_op_cost_tensors(&tensor1, &tensor2), 45045.);
47/// ```
48#[inline]
49pub fn contract_op_cost_tensors(t_1: &LeafTensor, t_2: &LeafTensor) -> f64 {
50    let all_dims = t_1 | t_2;
51    all_dims.size()
52}
53
54/// Returns Schroedinger contraction space complexity of contracting two
55/// [`LeafTensor`] objects.
56///
57/// # Examples
58/// ```
59/// # use tnc::tensornetwork::tensor::LeafTensor;
60/// # use tnc::contractionpath::contraction_cost::contract_size_tensors;
61/// # use rustc_hash::FxHashMap;
62/// let bond_dims = FxHashMap::from_iter([(0, 5),(1, 7), (2, 9), (3, 11), (4, 13)]);
63/// let tensor1 = LeafTensor::new_from_map(vec![0, 1, 2], &bond_dims); // 315 entries
64/// let tensor2 = LeafTensor::new_from_map(vec![2, 3, 4], &bond_dims); // 1287 entries
65/// // result = [0, 1, 3, 4] //  5005 entries -> total 6607 entries
66/// assert_eq!(contract_size_tensors(&tensor1, &tensor2), 6607.);
67/// ```
68#[inline]
69pub fn contract_size_tensors(t_1: &LeafTensor, t_2: &LeafTensor) -> f64 {
70    let diff = t_1 ^ t_2;
71    diff.size() + t_1.size() + t_2.size()
72}
73
74/// Returns contraction space complexity in bytes of contracting two [`Tensor`]
75/// objects.
76///
77/// # Examples
78/// ```
79/// # use tnc::tensornetwork::tensor::LeafTensor;
80/// # use tnc::contractionpath::contraction_cost::contract_size_tensors_bytes;
81/// # use rustc_hash::FxHashMap;
82/// let bond_dims = FxHashMap::from_iter([(0, 5),(1, 7), (2, 9), (3, 11), (4, 13)]);
83/// let tensor1 = LeafTensor::new_from_map(vec![0, 1, 2], &bond_dims); // 315 entries
84/// let tensor2 = LeafTensor::new_from_map(vec![2, 3, 4], &bond_dims); // 1287 entries
85/// // result = [0, 1, 3, 4] //  5005 entries -> total 6607 entries
86/// assert_eq!(contract_size_tensors_bytes(&tensor1, &tensor2), 6607. * 16.);
87/// ```
88#[inline]
89pub fn contract_size_tensors_bytes(i: &LeafTensor, j: &LeafTensor) -> f64 {
90    contract_size_tensors(i, j) * std::mem::size_of::<Complex64>() as f64
91}
92
93/// Returns Schroedinger contraction time and space complexity of fully contracting
94/// the input tensors.
95///
96/// # Arguments
97/// * `inputs` - Tensors to contract
98/// * `contract_path`  - Contraction order (replace path)
99/// * `only_count_ops` - If `true`, ignores cost of complex multiplication and addition and only counts number of operations
100#[inline]
101pub fn contract_path_cost(
102    inputs: &[Tensor],
103    contract_path: &ContractionPath,
104    only_count_ops: bool,
105) -> (f64, f64) {
106    let cost_function = if only_count_ops {
107        contract_op_cost_tensors
108    } else {
109        contract_cost_tensors
110    };
111    contract_path_custom_cost(inputs, contract_path, cost_function, contract_size_tensors)
112}
113
114/// Returns Schroedinger contraction time and space complexity of fully contracting
115/// the input tensors.
116///
117/// # Arguments
118/// * `inputs` - Tensors to contract
119/// * `contract_path`  - Contraction order (replace path)
120/// * `cost_function` - Function to calculate cost of contracting two tensors
121fn contract_path_custom_cost(
122    inputs: &[Tensor],
123    contract_path: &ContractionPath,
124    cost_function: fn(&LeafTensor, &LeafTensor) -> f64,
125    size_function: fn(&LeafTensor, &LeafTensor) -> f64,
126) -> (f64, f64) {
127    let mut op_cost = 0f64;
128    let mut mem_cost = 0f64;
129    let mut inputs = inputs.to_vec();
130
131    for (i, path) in &contract_path.nested {
132        let composite = inputs[*i].as_composite().unwrap();
133        let costs =
134            contract_path_custom_cost(composite.tensors(), path, cost_function, size_function);
135        op_cost += costs.0;
136        mem_cost = mem_cost.max(costs.1);
137        inputs[*i] = composite.external_tensor().into();
138    }
139
140    for (i, j) in &contract_path.toplevel {
141        let ti = inputs[*i].as_leaf().unwrap();
142        let tj = inputs[*j].as_leaf().unwrap();
143        op_cost += cost_function(ti, tj);
144        let tij = ti ^ tj;
145        let new_mem_cost = size_function(ti, tj);
146        mem_cost = mem_cost.max(new_mem_cost);
147        inputs[*i] = tij.into();
148    }
149
150    (op_cost, mem_cost)
151}
152
153/// Returns Schroedinger contraction time complexity using the critical path metric
154/// and using the sum metric. Additionally returns the space complexity.
155#[inline]
156pub fn communication_path_op_costs(
157    inputs: &[LeafTensor],
158    contract_path: SimplePathRef,
159    only_count_ops: bool,
160    tensor_cost: Option<&[f64]>,
161) -> ((f64, f64), f64) {
162    let (parallel_cost, _) =
163        communication_path_cost(inputs, contract_path, only_count_ops, true, tensor_cost);
164    let (serial_cost, mem_cost) =
165        communication_path_cost(inputs, contract_path, only_count_ops, false, tensor_cost);
166    ((parallel_cost, serial_cost), mem_cost)
167}
168
169/// Returns Schroedinger contraction time and space complexity of fully contracting
170/// the input tensors assuming all operations occur in parallel.
171///
172/// # Arguments
173/// * `inputs` - Tensors to contract
174/// * `contract_path`  - Contraction order (replace path)
175/// * `only_count_ops` - If `true`, ignores cost of complex multiplication and addition and only counts number of operations
176/// * `only_circital_path` - If `true`, only counts the cost along the critical path, otherwise the sum of all costs
177/// * `tensor_costs` - Initial cost for each tensor
178pub fn communication_path_cost(
179    inputs: &[LeafTensor],
180    contract_path: SimplePathRef,
181    only_count_ops: bool,
182    only_critical_path: bool,
183    tensor_cost: Option<&[f64]>,
184) -> (f64, f64) {
185    let cost_function = if only_count_ops {
186        contract_op_cost_tensors
187    } else {
188        contract_cost_tensors
189    };
190    let tensor_cost = if let Some(tensor_cost) = tensor_cost {
191        assert_eq!(inputs.len(), tensor_cost.len());
192        tensor_cost
193    } else {
194        &vec![0f64; inputs.len()]
195    };
196    if inputs.len() == 1 {
197        return (tensor_cost[0], tensor_cost[0]);
198    }
199
200    communication_path_custom_cost(
201        inputs,
202        contract_path,
203        cost_function,
204        only_critical_path,
205        tensor_cost,
206    )
207}
208
209/// Returns Schroedinger contraction time and space complexity of fully contracting
210/// the input tensors assuming all operations occur in parallel.
211///
212/// # Arguments
213/// * `inputs` - Tensors to contract
214/// * `contract_path`  - Contraction order (replace path)
215/// * `cost_function` - Function to calculate cost of contracting two tensors
216/// * `tensor_costs` - Initial cost for each tensor
217fn communication_path_custom_cost(
218    inputs: &[LeafTensor],
219    contract_path: SimplePathRef,
220    cost_function: fn(&LeafTensor, &LeafTensor) -> f64,
221    only_critical_path: bool,
222    tensor_cost: &[f64],
223) -> (f64, f64) {
224    let mut op_cost = 0f64;
225    let mut mem_cost = 0f64;
226    let mut inputs = inputs.to_vec();
227    let mut tensor_cost = tensor_cost.to_vec();
228
229    for &(i, j) in contract_path {
230        let ij = &inputs[i] ^ &inputs[j];
231        let new_mem_cost = contract_size_tensors(&inputs[i], &inputs[j]);
232        mem_cost = mem_cost.max(new_mem_cost);
233
234        op_cost = if only_critical_path {
235            cost_function(&inputs[i], &inputs[j]) + tensor_cost[i].max(tensor_cost[j])
236        } else {
237            cost_function(&inputs[i], &inputs[j]) + tensor_cost[i] + tensor_cost[j]
238        };
239        tensor_cost[i] = op_cost;
240        inputs[i] = ij;
241    }
242
243    (op_cost, mem_cost)
244}
245
246/// Computes the max memory requirements for contracting the tensor network using the
247/// given path. Uses `memory_estimator` to compute the memory required to contract
248/// two tensors.
249///
250/// Candidates for `memory_estimator` are e.g.:
251/// - [`contract_size_tensors`]
252/// - [`contract_size_tensors_bytes`]
253#[inline]
254pub fn compute_memory_requirements(
255    inputs: &[Tensor],
256    contract_path: &ContractionPath,
257    memory_estimator: fn(&LeafTensor, &LeafTensor) -> f64,
258) -> f64 {
259    fn id(_: &LeafTensor, _: &LeafTensor) -> f64 {
260        0.0
261    }
262    let (_, mem) = contract_path_custom_cost(inputs, contract_path, id, memory_estimator);
263    mem
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    use rustc_hash::FxHashMap;
271
272    use crate::path;
273    use crate::tensornetwork::tensor::CompositeTensor;
274
275    fn setup_simple() -> CompositeTensor {
276        let bond_dims =
277            FxHashMap::from_iter([(0, 5), (1, 2), (2, 6), (3, 8), (4, 1), (5, 3), (6, 4)]);
278        CompositeTensor::new(vec![
279            LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
280            LeafTensor::new_from_map(vec![0, 1, 3, 2], &bond_dims),
281            LeafTensor::new_from_map(vec![4, 5, 6], &bond_dims),
282        ])
283    }
284
285    fn setup_complex() -> CompositeTensor {
286        let bond_dims = FxHashMap::from_iter([
287            (0, 5),
288            (1, 2),
289            (2, 6),
290            (3, 8),
291            (4, 1),
292            (5, 3),
293            (6, 4),
294            (7, 3),
295            (8, 2),
296            (9, 2),
297        ]);
298        let t1_tensors = vec![
299            LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
300            LeafTensor::new_from_map(vec![0, 1, 3, 2], &bond_dims),
301            LeafTensor::new_from_map(vec![4, 5, 6], &bond_dims),
302        ];
303        let t1 = CompositeTensor::new(t1_tensors);
304
305        let t2_tensors = vec![
306            LeafTensor::new_from_map(vec![5, 6, 8], &bond_dims),
307            LeafTensor::new_from_map(vec![7, 8, 9], &bond_dims),
308        ];
309        let t2 = CompositeTensor::new(t2_tensors);
310        CompositeTensor::new(vec![t1, t2])
311    }
312
313    fn setup_parallel() -> Vec<LeafTensor> {
314        let bond_dims =
315            FxHashMap::from_iter([(0, 5), (1, 2), (2, 6), (3, 8), (4, 1), (5, 3), (6, 4)]);
316        vec![
317            LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
318            LeafTensor::new_from_map(vec![0, 1, 3, 2], &bond_dims),
319            LeafTensor::new_from_map(vec![4, 5, 6], &bond_dims),
320            LeafTensor::new_from_map(vec![5, 6], &bond_dims),
321        ]
322    }
323
324    #[test]
325    fn test_contract_path_cost() {
326        let tn = setup_simple();
327        let (op_cost, mem_cost) = contract_path_cost(tn.tensors(), &path![(0, 1), (0, 2)], false);
328        assert_eq!(op_cost, 4540.);
329        assert_eq!(mem_cost, 538.);
330        let (op_cost, mem_cost) = contract_path_cost(tn.tensors(), &path![(0, 2), (0, 1)], false);
331        assert_eq!(op_cost, 49296.);
332        assert_eq!(mem_cost, 1176.);
333    }
334
335    #[test]
336    fn test_contract_complex_path_cost() {
337        let tn = setup_complex();
338        let (op_cost, mem_cost) = contract_path_cost(
339            tn.tensors(),
340            &path![{(0, [(0, 1), (0, 2)]), (1, [(0, 1)])}, (0, 1)],
341            false,
342        );
343        assert_eq!(op_cost, 11188.);
344        assert_eq!(mem_cost, 538.);
345    }
346
347    #[test]
348    fn test_contract_path_cost_only_ops() {
349        let tn = setup_simple();
350        let (op_cost, mem_cost) = contract_path_cost(tn.tensors(), &path![(0, 1), (0, 2)], true);
351        assert_eq!(op_cost, 600.);
352        assert_eq!(mem_cost, 538.);
353        let (op_cost, mem_cost) = contract_path_cost(tn.tensors(), &path![(0, 2), (0, 1)], true);
354        assert_eq!(op_cost, 6336.);
355        assert_eq!(mem_cost, 1176.);
356    }
357
358    #[test]
359    fn test_contract_path_complex_cost_only_ops() {
360        let tn = setup_complex();
361        let (op_cost, mem_cost) = contract_path_cost(
362            tn.tensors(),
363            &path![{(0, [(0, 1), (0, 2)]), (1, [(0, 1)])}, (0, 1)],
364            true,
365        );
366        assert_eq!(op_cost, 1464.);
367        assert_eq!(mem_cost, 538.);
368    }
369
370    #[test]
371    fn test_communication_path_cost_only_ops() {
372        let tensors = setup_parallel();
373        let (op_cost, mem_cost) =
374            communication_path_cost(&tensors, &[(0, 1), (2, 3), (0, 2)], true, true, None);
375        assert_eq!(op_cost, 490.);
376        assert_eq!(mem_cost, 538.);
377    }
378
379    #[test]
380    fn test_communication_path_cost() {
381        let tensors = setup_parallel();
382        let (op_cost, mem_cost) =
383            communication_path_cost(&tensors, &[(0, 1), (2, 3), (0, 1)], false, true, None);
384        assert_eq!(op_cost, 7564.);
385        assert_eq!(mem_cost, 538.);
386    }
387
388    #[test]
389    fn test_communication_path_cost_only_ops_with_partition_cost() {
390        let tensors = setup_parallel();
391        let tensor_cost = vec![20., 30., 80., 10.];
392        let (op_cost, mem_cost) = communication_path_cost(
393            &tensors,
394            &[(0, 1), (2, 3), (0, 2)],
395            true,
396            true,
397            Some(&tensor_cost),
398        );
399        assert_eq!(op_cost, 520.);
400        assert_eq!(mem_cost, 538.);
401    }
402
403    #[test]
404    fn test_communication_path_cost_with_partition_cost() {
405        let tensors = setup_parallel();
406        let tensor_cost = vec![20., 30., 80., 10.];
407        let (op_cost, mem_cost) = communication_path_cost(
408            &tensors,
409            &[(0, 1), (2, 3), (0, 1)],
410            false,
411            true,
412            Some(&tensor_cost),
413        );
414        assert_eq!(op_cost, 7594.);
415        assert_eq!(mem_cost, 538.);
416    }
417}