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