tnc/tensornetwork/
partitioning.rs

1//! Functionality to partition composite tensors (i.e., tensor networks) into
2//! multiple tensor networks.
3use std::iter::zip;
4
5use itertools::Itertools;
6use kahypar::{partition, KaHyParContext};
7use rustc_hash::FxHashMap;
8
9use crate::tensornetwork::tensor::{CompositeTensor, LeafTensor};
10
11mod partition_config;
12pub use partition_config::PartitioningStrategy;
13
14/// The scale factor for the log weights used in the partitioning.
15/// This is used to convert the log weights to integer weights for KaHyPar.
16/// The current value allows for bond dimensions up to 144497 before two consecutive
17/// bond dimensions are rounded to the same weight. At the same time, it allows for
18/// the sum of ~150000 such bond dims before an i32 overflow occurs.
19const LOG_SCALE_FACTOR: f64 = 1e5;
20
21/// Partitions input tensor network using `KaHyPar` library.
22///
23/// Returns a `Vec<usize>` of length equal to the number of input tensors storing final partitioning results.
24/// The usize associated with each Tensor indicates its partionining.
25///
26/// # Arguments
27/// * `tensor_network` - [`Tensor`] to be partitionined
28/// * `k` - imbalance parameter for `KaHyPar`
29/// * `partition_strategy` - The strategy to pass to `KaHyPar`
30/// * `min` - if `true` performs `min_cut` to partition tensor network, if `false`, uses `max_cut`
31pub fn find_partitioning(
32    tensor_network: &CompositeTensor,
33    k: i32,
34    partitioning_strategy: PartitioningStrategy,
35    min: bool,
36) -> Vec<usize> {
37    if k == 1 {
38        return vec![0; tensor_network.len()];
39    }
40
41    let num_vertices = tensor_network.len() as u32;
42    let mut context = KaHyParContext::new();
43    partitioning_strategy.apply(&mut context);
44
45    let x = if min { 1 } else { -1 };
46
47    let imbalance = 0.03;
48    let mut objective = 0;
49    let mut hyperedge_weights = vec![];
50    let mut hyperedge_indices = vec![0];
51    let mut hyperedges = vec![];
52
53    let mut edge_dict = FxHashMap::default();
54    for (tensor_id, tensor) in tensor_network.tensors().iter().enumerate() {
55        let tensor = tensor
56            .as_leaf()
57            .expect("Partitioning currently only supports one level of nesting");
58        for (leg, dim) in tensor.edges() {
59            edge_dict
60                .entry(leg)
61                .and_modify(|entry| {
62                    hyperedges.push(*entry as u32);
63                    hyperedges.push(tensor_id as u32);
64                    hyperedge_indices.push(hyperedge_indices.last().unwrap() + 2);
65                    // Use log weights, because KaHyPar minimizes the sum of weights while we need the product.
66                    // Since it accepts only integer weights, we scale the log values before rounding.
67                    let weight = LOG_SCALE_FACTOR * (dim as f64).log2();
68                    hyperedge_weights.push(x * weight as i32);
69                    *entry = tensor_id;
70                })
71                .or_insert(tensor_id);
72        }
73    }
74
75    let mut partitioning = vec![-1; num_vertices as usize];
76    partition(
77        num_vertices,
78        hyperedge_weights.len() as u32,
79        imbalance,
80        k,
81        None,
82        Some(hyperedge_weights),
83        &hyperedge_indices,
84        hyperedges.as_slice(),
85        &mut objective,
86        &mut context,
87        &mut partitioning,
88    );
89    partitioning.iter().map(|e| *e as usize).collect()
90}
91
92/// Repeatedly partitions a tensor network to identify a communication scheme.
93/// Returns a `Vec<usize>` of length equal to the number of input tensors minus one, acts as a communication scheme.
94///
95/// # Arguments
96/// * `tensors` - &[(usize, `Tensor`)] to be partitioned. Each tuple contains the intermediate contraction cost and intermediate tensor for communication.
97/// * `k` - number of partitions
98/// * `partitioning_strategy` - The strategy to pass to `KaHyPar`
99/// * `min` - if `true` performs `min_cut` to partition tensor network, if `false`, uses `max_cut`
100pub fn communication_partitioning(
101    tensors: &[(usize, LeafTensor)],
102    k: i32,
103    imbalance: f64,
104    partitioning_strategy: PartitioningStrategy,
105    min: bool,
106) -> Vec<usize> {
107    assert!(k > 1, "Partitioning only valid for more than one process");
108    let num_vertices = tensors.len() as u32;
109    let mut context = KaHyParContext::new();
110    partitioning_strategy.apply(&mut context);
111
112    let x = if min { 1 } else { -1 };
113
114    let mut objective = 0;
115    let mut hyperedge_weights = vec![];
116
117    let mut hyperedge_indices = vec![0];
118    let mut hyperedges = vec![];
119
120    // Bidirectional mapping to a new index as KaHyPar indexes from 0.
121    // let mut edge_to_virtual_edge = FxHashMap::default();
122    // New index that starts from 0
123    // let mut edge_count = 0;
124    let mut edge_dict = FxHashMap::default();
125    for (tensor_id, (_, tensor)) in tensors.iter().enumerate() {
126        for (leg, dim) in tensor.edges() {
127            edge_dict
128                .entry(leg)
129                .and_modify(|entry| {
130                    hyperedges.push(*entry as u32);
131                    hyperedges.push(tensor_id as u32);
132                    hyperedge_indices.push(hyperedge_indices.last().unwrap() + 2);
133                    // Use log weights, because KaHyPar minimizes the sum of weights while we need the product.
134                    // Since it accepts only integer weights, we scale the log values before rounding.
135                    let weight = LOG_SCALE_FACTOR * (dim as f64).log2();
136                    hyperedge_weights.push(x * weight as i32);
137                    *entry = tensor_id;
138                })
139                .or_insert(tensor_id);
140        }
141    }
142
143    let mut partitioning = vec![-1; num_vertices as usize];
144    partition(
145        num_vertices,
146        hyperedge_weights.len() as u32,
147        imbalance,
148        k,
149        None,
150        Some(hyperedge_weights),
151        &hyperedge_indices,
152        hyperedges.as_slice(),
153        &mut objective,
154        &mut context,
155        &mut partitioning,
156    );
157
158    // partitioning
159    partitioning.iter().map(|e| *e as usize).collect()
160}
161
162/// Partitions the tensor network based on the `partitioning` vector that assigns
163/// each vector to a partition.
164pub fn partition_tensor_network(tn: CompositeTensor, partitioning: &[usize]) -> CompositeTensor {
165    let partition_ids = partitioning.iter().unique().copied().collect_vec();
166    let partition_dict =
167        zip(partition_ids.iter().copied(), 0..partition_ids.len()).collect::<FxHashMap<_, _>>();
168
169    let mut partitions = vec![CompositeTensor::default(); partition_ids.len()];
170    for (partition_id, tensor) in zip(partitioning, tn.tensors()) {
171        partitions[partition_dict[partition_id]].push_tensor(tensor.clone());
172    }
173    CompositeTensor::new(partitions)
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    use approx::assert_abs_diff_eq;
181    use rustc_hash::FxHashMap;
182
183    use crate::tensornetwork::partitioning::partition_config::PartitioningStrategy;
184    use crate::tensornetwork::tensor::{EdgeIndex, LeafTensor};
185
186    fn setup_complex() -> (CompositeTensor, FxHashMap<EdgeIndex, u64>) {
187        let bond_dims = FxHashMap::from_iter([
188            (0, 27),
189            (1, 18),
190            (2, 12),
191            (3, 15),
192            (4, 5),
193            (5, 3),
194            (6, 18),
195            (7, 22),
196            (8, 45),
197            (9, 65),
198            (10, 5),
199            (11, 17),
200        ]);
201        (
202            CompositeTensor::new(vec![
203                LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
204                LeafTensor::new_from_map(vec![0, 1, 3, 2], &bond_dims),
205                LeafTensor::new_from_map(vec![4, 5, 6], &bond_dims),
206                LeafTensor::new_from_map(vec![6, 8, 9], &bond_dims),
207                LeafTensor::new_from_map(vec![10, 8, 9], &bond_dims),
208                LeafTensor::new_from_map(vec![5, 1, 0], &bond_dims),
209            ]),
210            bond_dims,
211        )
212    }
213
214    #[test]
215    fn test_simple_partitioning() {
216        let (tn, bond_dims) = setup_complex();
217        let ref_tensor_1 = CompositeTensor::new(vec![
218            LeafTensor::new_from_map(vec![6, 8, 9], &bond_dims),
219            LeafTensor::new_from_map(vec![10, 8, 9], &bond_dims),
220        ]);
221        let ref_tensor_2 = CompositeTensor::new(vec![
222            LeafTensor::new_from_map(vec![0, 1, 3, 2], &bond_dims),
223            LeafTensor::new_from_map(vec![5, 1, 0], &bond_dims),
224        ]);
225        let ref_tensor_3 = CompositeTensor::new(vec![
226            LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
227            LeafTensor::new_from_map(vec![4, 5, 6], &bond_dims),
228        ]);
229        let partitioning = find_partitioning(&tn, 3, PartitioningStrategy::MinCut, true);
230        assert_eq!(partitioning, [2, 1, 2, 0, 0, 1]);
231        let partitioned_tn = partition_tensor_network(tn, partitioning.as_slice());
232        assert_eq!(partitioned_tn.tensors().len(), 3);
233
234        assert_abs_diff_eq!(partitioned_tn.tensor(2), ref_tensor_1.as_tensor());
235        assert_abs_diff_eq!(partitioned_tn.tensor(1), ref_tensor_2.as_tensor());
236        assert_abs_diff_eq!(partitioned_tn.tensor(0), ref_tensor_3.as_tensor());
237    }
238
239    #[test]
240    fn test_single_partition() {
241        let (tn, _) = setup_complex();
242        let partitioning = find_partitioning(&tn, 1, PartitioningStrategy::MinCut, true);
243        assert_eq!(partitioning, [0, 0, 0, 0, 0, 0]);
244    }
245}
tnc/tensornetwork/partitioning.rs

tnc/tensornetwork/
partitioning.rs
