plotnik_lib/emit/
layout.rs

1//! Cache-aligned instruction layout.
2//!
3//! Uses Pettis-Hansen inspired greedy chain extraction to place
4//! hot paths contiguously and avoid cache line straddling.
5
6use std::collections::{BTreeMap, HashSet};
7
8use crate::bytecode::{InstructionIR, Label, LayoutResult};
9
10const CACHE_LINE: usize = 64;
11const STEP_SIZE: usize = 8;
12
13/// Successor graph for layout analysis.
14struct Graph {
15    /// label -> list of successor labels
16    successors: BTreeMap<Label, Vec<Label>>,
17    /// label -> list of predecessor labels
18    predecessors: BTreeMap<Label, Vec<Label>>,
19}
20
21impl Graph {
22    fn build(instructions: &[InstructionIR]) -> Self {
23        let mut successors: BTreeMap<Label, Vec<Label>> = BTreeMap::new();
24        let mut predecessors: BTreeMap<Label, Vec<Label>> = BTreeMap::new();
25
26        for instr in instructions {
27            let label = instr.label();
28            successors.entry(label).or_default();
29
30            for succ in instr.successors() {
31                successors.entry(label).or_default().push(succ);
32                predecessors.entry(succ).or_default().push(label);
33            }
34        }
35
36        Self {
37            successors,
38            predecessors,
39        }
40    }
41
42    fn successors(&self, label: Label) -> &[Label] {
43        self.successors
44            .get(&label)
45            .map(|v| v.as_slice())
46            .unwrap_or(&[])
47    }
48
49    fn predecessor_count(&self, label: Label) -> usize {
50        self.predecessors.get(&label).map(|v| v.len()).unwrap_or(0)
51    }
52}
53
54/// Cache-aligned layout strategy.
55pub struct CacheAligned;
56
57impl CacheAligned {
58    /// Compute layout for instructions with given entry points.
59    ///
60    /// Returns mapping from labels to step IDs and total step count.
61    pub fn layout(instructions: &[InstructionIR], entries: &[Label]) -> LayoutResult {
62        if instructions.is_empty() {
63            return LayoutResult::empty();
64        }
65
66        let graph = Graph::build(instructions);
67        let label_to_instr: BTreeMap<Label, &InstructionIR> =
68            instructions.iter().map(|i| (i.label(), i)).collect();
69
70        let chains = extract_chains(&graph, instructions, entries);
71        let ordered = order_chains(chains, entries);
72
73        assign_step_ids(ordered, &label_to_instr)
74    }
75}
76
77/// Extract linear chains from the control flow graph.
78fn extract_chains(
79    graph: &Graph,
80    instructions: &[InstructionIR],
81    entries: &[Label],
82) -> Vec<Vec<Label>> {
83    let mut visited = HashSet::new();
84    let mut chains = Vec::new();
85
86    // Start with entry points (hot paths)
87    for &entry in entries {
88        if visited.contains(&entry) {
89            continue;
90        }
91        chains.push(build_chain(entry, graph, &mut visited));
92    }
93
94    // Then remaining unvisited instructions
95    for instr in instructions {
96        let label = instr.label();
97        if visited.contains(&label) {
98            continue;
99        }
100        chains.push(build_chain(label, graph, &mut visited));
101    }
102
103    chains
104}
105
106/// Build a single chain starting from a label.
107///
108/// Extends the chain while there's a single unvisited successor with a single predecessor.
109fn build_chain(start: Label, graph: &Graph, visited: &mut HashSet<Label>) -> Vec<Label> {
110    let mut chain = vec![start];
111    visited.insert(start);
112
113    let mut current = start;
114    while let [next] = graph.successors(current)
115        && !visited.contains(next)
116        && graph.predecessor_count(*next) == 1
117    {
118        chain.push(*next);
119        visited.insert(*next);
120        current = *next;
121    }
122
123    chain
124}
125
126/// Order chains: entries first, then by size (larger = hotter assumption).
127fn order_chains(mut chains: Vec<Vec<Label>>, entries: &[Label]) -> Vec<Vec<Label>> {
128    let entry_set: HashSet<Label> = entries.iter().copied().collect();
129
130    // Partition into entry chains and non-entry chains
131    let (mut entry_chains, mut other_chains): (Vec<_>, Vec<_>) =
132        chains.drain(..).partition(|chain| {
133            chain
134                .first()
135                .map(|l| entry_set.contains(l))
136                .unwrap_or(false)
137        });
138
139    // Sort other chains by size (descending) for better locality
140    other_chains.sort_by_key(|chain| std::cmp::Reverse(chain.len()));
141
142    // Entry chains first, then others
143    entry_chains.extend(other_chains);
144    entry_chains
145}
146
147/// Assign step IDs with cache line awareness.
148fn assign_step_ids(
149    chains: Vec<Vec<Label>>,
150    label_to_instr: &BTreeMap<Label, &InstructionIR>,
151) -> LayoutResult {
152    let mut mapping = BTreeMap::new();
153
154    let mut current_step = 0u16;
155    let mut current_offset = 0usize; // Byte offset for cache alignment
156
157    for chain in chains {
158        for label in chain {
159            let Some(instr) = label_to_instr.get(&label) else {
160                continue;
161            };
162            let size = instr.size();
163
164            // Cache line alignment for large instructions
165            if size >= 48 {
166                let line_offset = current_offset % CACHE_LINE;
167                if line_offset + size > CACHE_LINE {
168                    // Would straddle cache line - pad to next line
169                    let padding_bytes = CACHE_LINE - line_offset;
170                    let padding_steps = (padding_bytes / STEP_SIZE) as u16;
171                    current_step += padding_steps;
172                    current_offset += padding_bytes;
173                }
174            }
175
176            mapping.insert(label, current_step);
177            let step_count = (size / STEP_SIZE) as u16;
178            current_step += step_count;
179            current_offset += size;
180        }
181    }
182
183    LayoutResult::new(mapping, current_step)
184}
plotnik_lib/emit/layout.rs

plotnik_lib/emit/
layout.rs
