sc_neurocore_engine/partition.rs
1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Commercial license available
3// © Concepts 1996–2026 Miroslav Šotek. All rights reserved.
4// © Code 2020–2026 Miroslav Šotek. All rights reserved.
5// ORCID: 0009-0009-3560-0851
6// Contact: www.anulum.li | protoscience@anulum.li
7// SC-NeuroCore — Rust KL refinement for HierarchicalPartitioner (parity with Python _refine)
8
9//! Kernighan-Lin local refinement for the chiplet hierarchical
10//! partitioner. Mirrors `HierarchicalPartitioner._refine` step-for-step
11//! so the Rust output is bit-identical to the Python reference (the
12//! algorithm is fully deterministic given a CSR adjacency, a list of
13//! per-edge `|scc|` weights, vertex weights, an initial part_map and
14//! `(kl_iterations, correlation_penalty, n_parts)`).
15//!
16//! Inputs (CSR-style flat arrays):
17//! - adj_offsets: [V+1] — row pointers; row i covers
18//! adj_neighbours[adj_offsets[i] .. adj_offsets[i+1]]
19//! - adj_neighbours: [E_total] — neighbour vertex ids
20//! - adj_scc_abs: [E_total] — `|edge_scc(v, n)|` for each adjacency
21//! entry (already-pre-computed so the kernel does NO graph lookup)
22//! - vertex_weights: [V] — `vw` per Python `graph.vertex_weights.get(v, 1.0)`
23//! - part_map: [V] (mut) — initial partition id per vertex; the
24//! kernel writes the refined assignment back into the same slot
25//!
26//! Behaviour: same as `_refine` — for each KL iteration, walk every
27//! vertex of every partition, compute its full per-partition cost
28//! vector in ONE neighbour scan, and move the vertex to the
29//! best-gain target if the gain is strictly positive AND the source
30//! partition has more than one vertex remaining. Stop early if no
31//! moves happened in an iteration. Returns total moves performed.
32
33pub fn kl_refine(
34 adj_offsets: &[i64],
35 adj_neighbours: &[i32],
36 adj_scc_abs: &[f64],
37 vertex_weights: &[f64],
38 part_map: &mut [i32],
39 parts_concat: &[i32],
40 parts_offsets: &[i64],
41 n_parts: i32,
42 kl_iterations: i32,
43 correlation_penalty: f64,
44) -> u64 {
45 let n_parts_us = n_parts as usize;
46
47 // Seed per-partition vertex lists from the input concat array so
48 // the iteration order EXACTLY mirrors Python's
49 // for i, part in enumerate(partitions):
50 // for v in list(part): # snapshot of part NOW
51 // — this ordering is load-bearing for the algorithm output (the
52 // KL move order picks ties differently if vertex visit order
53 // differs; the Python reference is the ground truth). Building
54 // these from part_map alone (vertex-id order) gave 5/100
55 // disagreement at V=100 — see commit notes.
56 let mut parts: Vec<Vec<i32>> = vec![Vec::new(); n_parts_us];
57 for i in 0..n_parts_us {
58 let lo = parts_offsets[i] as usize;
59 let hi = parts_offsets[i + 1] as usize;
60 parts[i].extend_from_slice(&parts_concat[lo..hi]);
61 }
62
63 let mut weight_to: Vec<f64> = vec![0.0; n_parts_us];
64 let mut total_moves: u64 = 0;
65
66 for _ in 0..kl_iterations {
67 let mut improved = false;
68 for i in 0..n_parts_us {
69 // Snapshot the part list at entry — matches Python
70 // `for v in list(part)`. Subsequent moves into THIS
71 // part within this iteration won't be revisited until
72 // the next outer i sees them.
73 let snapshot: Vec<i32> = parts[i].clone();
74 for v_i32 in snapshot {
75 let v = v_i32 as usize;
76 if parts[i].len() <= 1 {
77 continue;
78 }
79 let cur_part = part_map[v];
80 if (cur_part as usize) != i {
81 // v was moved out earlier in this same iter — skip
82 // (Python's `for v in list(part)` would still visit
83 // it, but `part.remove(v)` would have raised since
84 // v is no longer in the list. Defensive guard.)
85 continue;
86 }
87 let vw = vertex_weights[v];
88
89 // Single-pass per-partition cost vector.
90 for w in weight_to.iter_mut() {
91 *w = 0.0;
92 }
93 let mut total_weight: f64 = 0.0;
94 let lo = adj_offsets[v] as usize;
95 let hi = adj_offsets[v + 1] as usize;
96 for k in lo..hi {
97 let n = adj_neighbours[k] as usize;
98 let scc = adj_scc_abs[k];
99 let contrib = vw * (1.0 + scc * correlation_penalty);
100 total_weight += contrib;
101 let tgt = part_map[n];
102 if (0..n_parts).contains(&tgt) {
103 weight_to[tgt as usize] += contrib;
104 }
105 }
106
107 let current_cost = total_weight - weight_to[i];
108 let mut best_target = i as i32;
109 let mut best_gain: f64 = 0.0;
110 for j in 0..n_parts_us {
111 if j == i {
112 continue;
113 }
114 let cand_cost = total_weight - weight_to[j];
115 let gain = current_cost - cand_cost;
116 if gain > best_gain {
117 best_gain = gain;
118 best_target = j as i32;
119 }
120 }
121
122 if best_target != (i as i32) && best_gain > 0.0 {
123 // Remove v from parts[i] (linear scan, O(|part_i|)
124 // — same complexity as Python's list.remove).
125 let pos = parts[i]
126 .iter()
127 .position(|&x| x == v_i32)
128 .expect("v in parts[i]");
129 parts[i].remove(pos);
130 parts[best_target as usize].push(v_i32);
131 part_map[v] = best_target;
132 total_moves += 1;
133 improved = true;
134 }
135 }
136 }
137 if !improved {
138 break;
139 }
140 }
141
142 total_moves
143}
144
145#[cfg(test)]
146mod tests {
147 use super::*;
148
149 #[test]
150 fn empty_graph_zero_moves() {
151 let part: &mut [i32] = &mut [];
152 // n_parts=2 → parts_offsets needs P+1=3 entries
153 let n = kl_refine(&[0], &[], &[], &[], part, &[], &[0, 0, 0], 2, 5, 0.5);
154 assert_eq!(n, 0);
155 }
156
157 #[test]
158 fn single_partition_no_moves() {
159 let mut pm = vec![0i32, 0, 0, 0];
160 let offsets = vec![0i64, 1, 2, 3, 4];
161 let neighbours = vec![1i32, 0, 3, 2];
162 let scc = vec![0.1, 0.1, 0.1, 0.1];
163 let vw = vec![1.0; 4];
164 let parts_concat = vec![0i32, 1, 2, 3];
165 let parts_offsets = vec![0i64, 4];
166 let n = kl_refine(
167 &offsets,
168 &neighbours,
169 &scc,
170 &vw,
171 &mut pm,
172 &parts_concat,
173 &parts_offsets,
174 1,
175 5,
176 0.5,
177 );
178 assert_eq!(n, 0);
179 assert_eq!(pm, vec![0, 0, 0, 0]);
180 }
181
182 #[test]
183 fn two_partitions_isolated_swap() {
184 let offsets = vec![0i64, 1, 2, 3, 4];
185 let neighbours = vec![1i32, 0, 3, 2];
186 let scc = vec![0.5, 0.5, 0.5, 0.5];
187 let vw = vec![1.0; 4];
188 let mut pm = vec![0i32, 1, 0, 1];
189 let parts_concat = vec![0i32, 2, 1, 3];
190 let parts_offsets = vec![0i64, 2, 4];
191 let n = kl_refine(
192 &offsets,
193 &neighbours,
194 &scc,
195 &vw,
196 &mut pm,
197 &parts_concat,
198 &parts_offsets,
199 2,
200 5,
201 0.5,
202 );
203 assert!(n > 0);
204 }
205
206 #[test]
207 fn part_size_one_vertex_pinned() {
208 let offsets = vec![0i64, 2, 3, 4];
209 let neighbours = vec![1i32, 2, 0, 0];
210 let scc = vec![0.1, 0.1, 0.1, 0.1];
211 let vw = vec![1.0; 3];
212 let mut pm = vec![0i32, 1, 1];
213 let parts_concat = vec![0i32, 1, 2];
214 let parts_offsets = vec![0i64, 1, 3];
215 let _ = kl_refine(
216 &offsets,
217 &neighbours,
218 &scc,
219 &vw,
220 &mut pm,
221 &parts_concat,
222 &parts_offsets,
223 2,
224 5,
225 0.5,
226 );
227 assert_eq!(pm[0], 0, "lone-vertex partition must stay populated");
228 }
229}