calyx_opt/analysis/
graph_coloring.rs

1use calyx_utils::{Idx, WeightGraph};
2use itertools::Itertools;
3use petgraph::algo;
4use std::{
5    collections::{BTreeMap, HashMap},
6    fmt::Display,
7    hash::Hash,
8};
9
10/// Defines a greedy graph coloring algorithm over a generic conflict graph.
11pub struct GraphColoring<T> {
12    graph: WeightGraph<T>,
13    // color_freq_map records similar information as `all_colors` does in the
14    // `color_greedy()` method, but `color_freq_map` stays alive after the
15    // function call, and doesn't get rid of colors once they reach the bound
16    color_freq_map: HashMap<Idx, i64>,
17}
18
19impl<T, C> From<C> for GraphColoring<T>
20where
21    T: Hash + Eq + Ord,
22    C: Iterator<Item = T>,
23{
24    fn from(nodes: C) -> Self {
25        let graph = WeightGraph::from(nodes);
26        GraphColoring {
27            graph,
28            color_freq_map: HashMap::new(),
29        }
30    }
31}
32
33impl<'a, T> GraphColoring<T>
34where
35    T: 'a + Eq + Hash + Clone + Ord,
36{
37    /// Add a conflict edge between `a` and `b`.
38    #[inline(always)]
39    pub fn insert_conflict(&mut self, a: &T, b: &T) {
40        self.graph.add_edge(a, b);
41    }
42
43    /// Add conflict edges between all given items.
44    pub fn insert_conflicts<C>(&mut self, items: C)
45    where
46        C: Iterator<Item = &'a T> + Clone,
47    {
48        self.graph.add_all_edges(items)
49    }
50
51    pub fn has_nodes(&self) -> bool {
52        self.graph.graph.node_count() > 0
53    }
54
55    /// increases the frequency of `idx` in `color_freq_map` by one
56    fn increase_freq(&mut self, idx: Idx) {
57        self.color_freq_map
58            .entry(idx)
59            .and_modify(|v| *v += 1)
60            .or_insert(1);
61    }
62
63    /// provides a hashmap that gives the sharing frequencies
64    pub fn get_share_freqs(&mut self) -> HashMap<i64, i64> {
65        let mut pdf: HashMap<i64, i64> = HashMap::new();
66        // hold total value so we know how much to divide by at the end
67        for value in self.color_freq_map.values() {
68            // in`pdf`, each key represents a possible number of times a cell
69            // is shared-- for example, x-- and the corresponding key represents
70            // how many cells in the new program were shared x times
71            pdf.entry(*value).and_modify(|v| *v += 1).or_insert(1);
72        }
73        pdf
74    }
75
76    /// Given an `ordering` of `T`s, find a mapping from nodes to `T`s such
77    /// that no node has a neighbor with the same `T`.
78    /// `keep_self_color` indicates whether to keep the mapping of the node to
79    /// itself in the returned HashMap (since nodes are "colors")
80    pub fn color_greedy(
81        &mut self,
82        bound: Option<i64>,
83        keep_self_color: bool,
84    ) -> HashMap<T, T> {
85        let mut all_colors: BTreeMap<Idx, i64> = BTreeMap::new();
86        let mut coloring: HashMap<Idx, Idx> = HashMap::new();
87        let always_share = bound.is_none();
88        // if we always_share is true, then we don't care about bound
89        let bound_if_exists = if always_share { 0 } else { bound.unwrap() };
90
91        // get strongly get components of graph
92        let sccs = algo::tarjan_scc(&self.graph.graph);
93        // sort strongly components from largest to smallest
94        for scc in sccs.into_iter().sorted_by(|a, b| b.len().cmp(&a.len())) {
95            // check if graph component is complete
96            let is_complete = scc.iter().all(|&idx| {
97                self.graph.graph.neighbors(idx).count() == scc.len() - 1
98            });
99            // if graph is complete, then every node needs a new color. so there's no reason to
100            // check neighbors
101            if is_complete {
102                let mut available_colors: Vec<_> =
103                    all_colors.keys().cloned().collect_vec();
104
105                // every node will need a different color
106                for nidx in scc.into_iter().sorted() {
107                    if !available_colors.is_empty() {
108                        let c = available_colors.remove(0);
109                        coloring.insert(nidx, c);
110                        self.increase_freq(c);
111                        if let Some(num_used) = all_colors.get_mut(&c) {
112                            *num_used += 1;
113                            if !always_share && *num_used == bound_if_exists {
114                                all_colors.remove(&c);
115                            }
116                        }
117                    } else {
118                        all_colors.insert(nidx, 1);
119                        coloring.insert(nidx, nidx);
120                        self.increase_freq(nidx);
121                        if !always_share && bound_if_exists == 1 {
122                            all_colors.remove(&nidx);
123                        }
124                    }
125                }
126            } else {
127                for nidx in scc.into_iter().sorted() {
128                    let mut available_colors = all_colors.clone();
129                    // search neighbors for used colors
130                    for item in self.graph.graph.neighbors(nidx) {
131                        // if the neighbor is already colored
132                        if coloring.contains_key(&item) {
133                            // remove it from the available colors
134                            available_colors.remove(&coloring[&item]);
135                        }
136                    }
137
138                    let color = available_colors.iter().next();
139                    match color {
140                        Some((c, _)) => {
141                            coloring.insert(nidx, *c);
142                            self.increase_freq(*c);
143                            if let Some(num_used) = all_colors.get_mut(c) {
144                                *num_used += 1;
145                                if !always_share && *num_used == bound_if_exists
146                                {
147                                    all_colors.remove(c);
148                                }
149                            }
150                        }
151                        None => {
152                            // use self as color if nothing else
153                            all_colors.insert(nidx, 1);
154                            coloring.insert(nidx, nidx);
155                            self.increase_freq(nidx);
156                            if !always_share && bound_if_exists == 1 {
157                                all_colors.remove(&nidx);
158                            }
159                        }
160                    };
161                }
162            }
163        }
164
165        let rev_map = self.graph.reverse_index();
166        coloring
167            .into_iter()
168            .map(|(n1, n2)| (rev_map[&n1].clone(), rev_map[&n2].clone()))
169            .filter(|(a, b)| (a != b) || keep_self_color)
170            .collect()
171    }
172
173    // Reverses a coloring by mapping color C -> vec of nodes colored C.
174    pub fn reverse_coloring(coloring: &HashMap<T, T>) -> HashMap<T, Vec<T>> {
175        let mut rev_coloring: HashMap<T, Vec<T>> = HashMap::new();
176        for (node, color) in coloring {
177            rev_coloring
178                .entry(color.clone())
179                .or_default()
180                .push(node.clone());
181        }
182        rev_coloring
183    }
184    pub fn welsh_powell_coloring(&self) -> HashMap<T, T> {
185        let mut coloring: HashMap<T, T> = HashMap::new();
186
187        let mut degree_ordering: Vec<&T> = self
188            .graph
189            .nodes()
190            .sorted()
191            .sorted_by(|a, b| self.graph.degree(b).cmp(&self.graph.degree(a)))
192            .collect();
193
194        let rev_map = self.graph.reverse_index();
195        while !degree_ordering.is_empty() {
196            let head = degree_ordering.remove(0);
197            // eprintln!("{}", self.graph.degree(head));
198            if !coloring.contains_key(head) {
199                coloring.insert(head.clone(), head.clone());
200                for &node in &degree_ordering {
201                    if coloring.contains_key(node) {
202                        continue;
203                    }
204                    if !self
205                        .graph
206                        .graph
207                        .neighbors(self.graph.index_map[node])
208                        .any(|x| coloring.get(&rev_map[&x]) == Some(head))
209                    {
210                        coloring.insert(node.clone(), head.clone());
211                    }
212                }
213            }
214        }
215
216        coloring
217    }
218}
219
220impl<T: Eq + Hash + ToString + Clone + Ord> Display for GraphColoring<T> {
221    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
222        self.graph.to_string().fmt(f)
223    }
224}
calyx_opt/analysis/graph_coloring.rs

calyx_opt/analysis/
graph_coloring.rs
