vj-codegen 0.2.0

//! Splitting one design module into several for parallel JIT code generation.
//!
//! ORC's `LLJIT` (C API) compiles a module on whichever thread first looks up
//! one of its symbols, and a module is one compile unit, so a design lowered
//! into a single module is code-generated on one core. [`plan`] partitions
//! the module's defined functions into bins (connected components of the
//! direct-call graph, packed largest-first onto the least-loaded bin, so a
//! caller and its callees stay in one object and never pay a cross-object
//! call), and [`strip_to_bin`] turns a freshly parsed copy of the module into
//! the bin's slice: every function another bin owns loses its body (becoming
//! an external declaration the JIT resolves against the owning module) and,
//! for every bin but the first, externally visible globals become
//! declarations too, so each symbol is defined exactly once.
//!
//! Code generation is per function and deterministic, so the machine code of
//! each function is the same as in the unsplit module; only where it lives
//! changes. [`plan`] declines (`None`) whenever sharing could change
//! behaviour: aliases, ifuncs, appending/weak/common linkage, thread-local
//! globals, local-linkage functions, or a global initializer naming a mutable
//! local-linkage global. A mutable local-linkage global (a string-literal id
//! cache, a class static, a static local) is allowed because [`plan`] puts every
//! function that names one into the same bin, so the one copy that is ever
//! used is the only one with state; the unused copies in other bins are never
//! touched.

use std::collections::{HashMap, HashSet};

use inkwell::module::Module;
use llvm_sys::LLVMLinkage;
use llvm_sys::LLVMTypeKind;
use llvm_sys::core::{
    LLVMCountBasicBlocks, LLVMDeleteBasicBlock, LLVMGetFirstBasicBlock, LLVMGetFirstFunction,
    LLVMGetFirstGlobal, LLVMGetFirstGlobalAlias, LLVMGetFirstGlobalIFunc, LLVMGetFirstInstruction,
    LLVMGetInitializer, LLVMGetLinkage, LLVMGetNextBasicBlock, LLVMGetNextFunction,
    LLVMGetNextGlobal, LLVMGetNextInstruction, LLVMGetNumOperands, LLVMGetOperand, LLVMGetTypeKind,
    LLVMGetUndef, LLVMGetValueName2, LLVMInstructionEraseFromParent, LLVMIsAConstant,
    LLVMIsAGlobalValue, LLVMIsDeclaration, LLVMIsGlobalConstant, LLVMIsThreadLocal,
    LLVMReplaceAllUsesWith, LLVMSetInitializer, LLVMSetLinkage, LLVMTypeOf,
};
use llvm_sys::prelude::LLVMValueRef;

/// The functions each bin owns (defines); bin 0 also owns every global.
pub struct SplitPlan {
    pub bins: Vec<Vec<String>>,
    /// Mutable local-linkage globals too widely used to pin their users to
    /// one bin. [`strip_to_bin`] makes each one externally visible, defined
    /// once in bin 0 and declared elsewhere, so every bin shares one copy
    /// (through the GOT, like any external global).
    pub shared_globals: HashSet<String>,
}

fn value_name(v: LLVMValueRef) -> String {
    let mut len = 0usize;
    // SAFETY: `v` is a live named global value of the module the caller holds
    // alive; the returned pointer/length describe its name, valid until the
    // value is renamed or destroyed, and it is copied out immediately.
    let p = unsafe { LLVMGetValueName2(v, &mut len) };
    // SAFETY: `p` points at `len` readable bytes (see above).
    let bytes = unsafe { std::slice::from_raw_parts(p.cast::<u8>(), len) };
    String::from_utf8_lossy(bytes).into_owned()
}

fn functions(module: &Module) -> Vec<LLVMValueRef> {
    let mut out = Vec::new();
    // SAFETY: the module is alive for the whole walk and not mutated by it.
    let mut f = unsafe { LLVMGetFirstFunction(module.as_mut_ptr()) };
    while !f.is_null() {
        out.push(f);
        // SAFETY: `f` is a live function of the module.
        f = unsafe { LLVMGetNextFunction(f) };
    }
    out
}

/// Every instruction of `f`, in block order.
fn instructions(f: LLVMValueRef) -> Vec<LLVMValueRef> {
    let mut out = Vec::new();
    // SAFETY (whole walk): `f` is a live function with a body; the block and
    // instruction lists are not mutated while they are iterated.
    let mut bb = unsafe { LLVMGetFirstBasicBlock(f) };
    while !bb.is_null() {
        let mut i = unsafe { LLVMGetFirstInstruction(bb) };
        while !i.is_null() {
            out.push(i);
            i = unsafe { LLVMGetNextInstruction(i) };
        }
        bb = unsafe { LLVMGetNextBasicBlock(bb) };
    }
    out
}

/// Extra code-generation cost, in instruction equivalents, of one wide
/// (over 128 bits) integer instruction. Instruction selection legalizes it into
/// a straight-line sequence of word operations, but only operations whose
/// result depends on a runtime value in every word are expensive: a select, an
/// add or compare, a shift by a variable amount, a phi. Measured on UVM, the
/// packer's `set_packed_bits` (9 such ops of 32k..256k bits) takes 1 s to
/// compile, while `uvm_string_to_bits` (5420 constant shifts, zero-extends and
/// ors of 65k..262k bits) takes 1.5 s. Counting bare instructions put all the
/// packer functions, which sit together in definition order, into one bin
/// with three to five times the compile time of the others.
fn wide_op_cost(inst: LLVMValueRef) -> usize {
    use llvm_sys::LLVMOpcode::*;
    // A compare's result is i1: its operand carries the width.
    // SAFETY: `inst` is a live instruction; operand 0 exists for the opcodes
    // that reach it (ICmp has two operands).
    let opcode = unsafe { llvm_sys::core::LLVMGetInstructionOpcode(inst) };
    let width = if opcode == LLVMICmp {
        int_width(unsafe { LLVMGetOperand(inst, 0) })
    } else {
        int_width(inst)
    };
    if width <= 128 {
        return 0;
    }
    let words = width.div_ceil(64);
    let variable = |k: u32| {
        // SAFETY: operand index is below the operand count for the opcodes
        // that call this (binary operators have two).
        unsafe { LLVMIsAConstant(LLVMGetOperand(inst, k)).is_null() }
    };
    match opcode {
        LLVMShl | LLVMLShr | LLVMAShr if variable(1) => words * 8,
        LLVMAdd | LLVMSub | LLVMMul | LLVMUDiv | LLVMSDiv | LLVMURem | LLVMSRem | LLVMICmp
        | LLVMSelect | LLVMPHI => words * 8,
        _ => words / 32,
    }
}

/// Bit width of the integer type of `v`, or 0.
fn int_width(v: LLVMValueRef) -> usize {
    // SAFETY: `v` is a live value.
    let t = unsafe { LLVMTypeOf(v) };
    // SAFETY: `t` is its type.
    if unsafe { LLVMGetTypeKind(t) } == LLVMTypeKind::LLVMIntegerTypeKind {
        // SAFETY: `t` is an integer type.
        unsafe { llvm_sys::core::LLVMGetIntTypeWidth(t) as usize }
    } else {
        0
    }
}

fn find(parent: &mut [usize], mut x: usize) -> usize {
    while parent[x] != x {
        parent[x] = parent[parent[x]];
        x = parent[x];
    }
    x
}

/// Modules below this many IR instructions never have a giant component cut
/// (the extra cross-object calls only pay off when code generation is long).
const GIANT_SPLIT_MIN_INSTRS: usize = 500_000;

/// A mutable local global named by more functions than this is shared as an
/// external global rather than pinning its users to one bin.
const SHARED_GLOBAL_USERS: usize = 2;

/// A giant component is cut into about this many pieces per bin: the slow
/// functions of a design (UVM's packer family) sit next to each other in
/// definition order, and whole-bin-sized pieces would put them all in one bin.
const PIECES_PER_BIN: usize = 4;

/// Depth bound for walking constant operand trees.
const CONST_WALK_DEPTH: u32 = 8;

/// Pushes the index of every mutable local global the constant (or global
/// value) `v` names, looking through constant expressions and aggregates.
fn collect_mutable_globals(
    v: LLVMValueRef,
    mutable_local: &HashMap<LLVMValueRef, usize>,
    out: &mut Vec<usize>,
    depth: u32,
) {
    if let Some(&gi) = mutable_local.get(&v) {
        out.push(gi);
        return;
    }
    // SAFETY: `v` is a live value; the casts and operand reads are read-only.
    unsafe {
        if depth >= CONST_WALK_DEPTH
            || LLVMIsAConstant(v).is_null()
            || !LLVMIsAGlobalValue(v).is_null()
        {
            return;
        }
        for k in 0..LLVMGetNumOperands(v).max(0) as u32 {
            collect_mutable_globals(LLVMGetOperand(v, k), mutable_local, out, depth + 1);
        }
    }
}

/// Whether the constant `v` names any mutable local global.
fn constant_names_global(
    v: LLVMValueRef,
    mutable_local: &HashMap<LLVMValueRef, usize>,
    depth: u32,
) -> bool {
    let mut hits = Vec::new();
    collect_mutable_globals(v, mutable_local, &mut hits, depth);
    !hits.is_empty()
}

/// `None`, after naming why under `VJ_TIMING` (a profiling aid).
fn decline(reason: &str) -> Option<SplitPlan> {
    if crate::timing::enabled() {
        eprintln!("[timing] jit split declined: {reason}");
    }
    None
}

/// Partitions the module's defined functions into at most `nbins` non-empty
/// bins, or `None` when the module must stay whole (see the module docs) or
/// would not fill two bins.
pub fn plan(module: &Module, nbins: usize) -> Option<SplitPlan> {
    let m = module.as_mut_ptr();
    // SAFETY: `m` is the live module `module` wraps; these are read-only
    // accessors on it.
    if unsafe { !LLVMGetFirstGlobalAlias(m).is_null() || !LLVMGetFirstGlobalIFunc(m).is_null() } {
        return decline("alias or ifunc present");
    }
    // SAFETY: as above.
    let mut g = unsafe { LLVMGetFirstGlobal(m) };
    // Mutable local globals (string-literal id caches, class statics, static
    // locals): sound to split only if every function touching one lands in
    // the same bin, so each is unioned with its users below.
    let mut mutable_local: HashMap<LLVMValueRef, usize> = HashMap::new();
    while !g.is_null() {
        // SAFETY: `g` is a live global variable of the module.
        let (link, tls, constant) = unsafe {
            (
                LLVMGetLinkage(g),
                LLVMIsThreadLocal(g) != 0,
                LLVMIsGlobalConstant(g) != 0,
            )
        };
        let ok = match link {
            LLVMLinkage::LLVMExternalLinkage => true,
            LLVMLinkage::LLVMInternalLinkage | LLVMLinkage::LLVMPrivateLinkage => {
                if !constant && !tls {
                    let n = mutable_local.len();
                    mutable_local.insert(g, n);
                }
                true
            }
            _ => false,
        };
        if tls || !ok {
            return decline("a global is thread-local or has unsupported linkage/mutability");
        }
        // SAFETY: `g` is live.
        g = unsafe { LLVMGetNextGlobal(g) };
    }

    let mut defined: Vec<LLVMValueRef> = Vec::new();
    for f in functions(module) {
        // SAFETY: `f` is a live function of the module.
        if unsafe { LLVMIsDeclaration(f) } != 0 {
            continue;
        }
        // SAFETY: as above.
        if unsafe { LLVMGetLinkage(f) } != LLVMLinkage::LLVMExternalLinkage {
            return decline("a defined function has non-external linkage");
        }
        defined.push(f);
    }
    if defined.len() < 2 {
        return decline("fewer than 2 defined functions");
    }
    let index: HashMap<LLVMValueRef, usize> =
        defined.iter().enumerate().map(|(i, &f)| (f, i)).collect();

    let nf = defined.len();
    // A global initializer naming a mutable local global would make the
    // global's users depend on another definition: keep the module whole.
    if !mutable_local.is_empty() {
        // SAFETY: `m` is live; read-only walk of the globals.
        let mut g = unsafe { LLVMGetFirstGlobal(m) };
        while !g.is_null() {
            // SAFETY: `g` is a live global variable.
            let init = unsafe { LLVMGetInitializer(g) };
            if !init.is_null() && constant_names_global(init, &mutable_local, 0) {
                return decline("a global initializer references a mutable local global");
            }
            // SAFETY: as above.
            g = unsafe { LLVMGetNextGlobal(g) };
        }
    }
    let mut parent: Vec<usize> = (0..nf + mutable_local.len()).collect();
    // Only the mutable-global links: the atoms a giant component may be cut
    // along (functions sharing a mutable global never separate).
    let mut atom: Vec<usize> = (0..nf + mutable_local.len()).collect();
    let mut size = vec![0usize; nf];
    // The distinct mutable local globals each function names.
    let mut fn_globals: Vec<Vec<usize>> = vec![Vec::new(); nf];
    for (i, &f) in defined.iter().enumerate() {
        let insts = instructions(f);
        size[i] = insts.len();
        for &inst in &insts {
            size[i] += wide_op_cost(inst);
        }
        for inst in insts {
            // SAFETY: `inst` is a live instruction; operand indices are below
            // its operand count.
            let n = unsafe { LLVMGetNumOperands(inst) };
            for k in 0..n.max(0) as u32 {
                // SAFETY: as above.
                let op = unsafe { LLVMGetOperand(inst, k) };
                if !mutable_local.is_empty() {
                    collect_mutable_globals(op, &mutable_local, &mut fn_globals[i], 0);
                }
                if let Some(&j) = index.get(&op) {
                    let (a, b) = (find(&mut parent, i), find(&mut parent, j));
                    if a != b {
                        parent[b] = a;
                    }
                }
            }
        }
    }
    // A mutable local global named by many functions would chain them (and
    // through them the whole design) into one bin: share it as an external
    // global instead. The rest pin their few users to one bin.
    let mut users = vec![0usize; mutable_local.len()];
    for gs in &mut fn_globals {
        gs.sort_unstable();
        gs.dedup();
        for &gi in gs.iter() {
            users[gi] += 1;
        }
    }
    for (i, gs) in fn_globals.iter().enumerate() {
        for &gi in gs {
            if users[gi] > SHARED_GLOBAL_USERS {
                continue;
            }
            let (a, b) = (find(&mut parent, i), find(&mut parent, nf + gi));
            if a != b {
                parent[b] = a;
            }
            let (a, b) = (find(&mut atom, i), find(&mut atom, nf + gi));
            if a != b {
                atom[b] = a;
            }
        }
    }
    let mut shared_globals: HashSet<String> = HashSet::new();
    for (&g, &gi) in &mutable_local {
        if users[gi] > SHARED_GLOBAL_USERS {
            shared_globals.insert(value_name(g));
        }
    }
    // Components in first-function order, so the plan is deterministic.
    let mut comp_of_root: HashMap<usize, usize> = HashMap::new();
    let mut comps: Vec<(usize, Vec<usize>)> = Vec::new();
    for (i, &sz) in size.iter().enumerate() {
        let r = find(&mut parent, i);
        let c = *comp_of_root.entry(r).or_insert_with(|| {
            comps.push((0, Vec::new()));
            comps.len() - 1
        });
        comps[c].0 += sz;
        comps[c].1.push(i);
    }
    // A call graph that is one giant connected component (a UVM design's
    // virtual dispatch links every class) cannot fill the bins by components
    // alone, which leaves code generation on one core. Cut such a component
    // in definition order (a class's methods are lowered together, so the
    // pieces keep most calls local) into pieces of about a bin's share; the
    // calls that cross pieces become cross-object calls the JIT links.
    // Small modules keep whole components.
    let total: usize = size.iter().sum();
    if total >= GIANT_SPLIT_MIN_INSTRS {
        let target = total.div_ceil(nbins.max(1) * PIECES_PER_BIN);
        let mut c = 0;
        while c < comps.len() {
            if comps[c].0 <= target * 3 / 2 {
                c += 1;
                continue;
            }
            let (_, fns) = std::mem::take(&mut comps[c]);
            // Atoms in first-function order, each with its members.
            let mut atoms: Vec<(usize, Vec<usize>)> = Vec::new();
            let mut atom_of_root: HashMap<usize, usize> = HashMap::new();
            for &i in &fns {
                let r = find(&mut atom, i);
                let a = *atom_of_root.entry(r).or_insert_with(|| {
                    atoms.push((0, Vec::new()));
                    atoms.len() - 1
                });
                atoms[a].0 += size[i];
                atoms[a].1.push(i);
            }
            if crate::timing::enabled() {
                let mut a: Vec<(usize, usize)> = atoms.iter().map(|x| (x.0, x.1.len())).collect();
                a.sort_unstable_by_key(|x| std::cmp::Reverse(x.0));
                a.truncate(5);
                eprintln!(
                    "[timing] jit split: cut component of {} atoms, target {target}, largest atoms {a:?}",
                    atoms.len()
                );
            }
            let mut pieces: Vec<(usize, Vec<usize>)> = vec![(0, Vec::new())];
            for (sz, members) in atoms {
                // A heavy atom (one huge function) gets a piece of its own, so
                // it balances against the others instead of dragging along
                // whatever small atoms happened to precede it.
                let heavy = sz >= target / 2;
                if pieces
                    .last()
                    .is_some_and(|p| p.0 >= target || (heavy && p.0 > 0))
                {
                    pieces.push((0, Vec::new()));
                }
                if let Some(p) = pieces.last_mut() {
                    p.0 += sz;
                    p.1.extend(members);
                }
                if heavy {
                    pieces.push((0, Vec::new()));
                }
            }
            let mut it = pieces.into_iter();
            comps[c] = it.next().unwrap_or_default();
            comps.extend(it);
            c += 1;
        }
        comps.retain(|c| !c.1.is_empty());
    }
    // Largest first (stable: ties keep first-function order), each onto the
    // least-loaded bin (lowest index on ties).
    let mut order: Vec<usize> = (0..comps.len()).collect();
    order.sort_by_key(|&c| std::cmp::Reverse(comps[c].0));
    if crate::timing::enabled() {
        let top: Vec<(usize, usize)> = order
            .iter()
            .take(5)
            .map(|&c| (comps[c].0, comps[c].1.len()))
            .collect();
        eprintln!(
            "[timing] jit split: {} components of {} functions; largest 5 (instrs, functions): {top:?}",
            comps.len(),
            defined.len()
        );
    }
    let nbins = nbins.min(comps.len());
    if nbins < 2 {
        return decline("fewer than 2 call-graph components");
    }
    let mut load = vec![0usize; nbins];
    let mut bins: Vec<Vec<usize>> = vec![Vec::new(); nbins];
    for c in order {
        let b = (0..nbins).min_by_key(|&b| (load[b], b)).unwrap_or(0);
        load[b] += comps[c].0;
        bins[b].extend(&comps[c].1);
    }
    if crate::timing::enabled() {
        eprintln!("[timing] jit split: bin loads (weighted instrs) {load:?}");
    }
    let bins = bins
        .into_iter()
        .map(|mut ids| {
            ids.sort_unstable();
            ids.into_iter().map(|i| value_name(defined[i])).collect()
        })
        .collect();
    Some(SplitPlan {
        bins,
        shared_globals,
    })
}

/// Reduces a freshly parsed copy of the planned module to bin `bin`: bodies of
/// functions outside `owned` are dropped (they stay as declarations), and
/// outside bin 0 so are the definitions of externally visible globals.
pub fn strip_to_bin(
    module: &Module,
    bin: usize,
    owned: &HashSet<String>,
    shared_globals: &HashSet<String>,
) {
    let m = module.as_mut_ptr();
    if !shared_globals.is_empty() {
        // SAFETY: `m` is the live module; only linkage of listed globals is
        // changed, before the (bin != 0) initializer clearing below.
        let mut g = unsafe { LLVMGetFirstGlobal(m) };
        while !g.is_null() {
            if shared_globals.contains(&value_name(g)) {
                // SAFETY: `g` is a live global variable.
                unsafe { LLVMSetLinkage(g, LLVMLinkage::LLVMExternalLinkage) };
            }
            // SAFETY: as above.
            g = unsafe { LLVMGetNextGlobal(g) };
        }
    }
    for f in functions(module) {
        // SAFETY: `f` is a live function of the module.
        if unsafe { LLVMIsDeclaration(f) } != 0 || owned.contains(&value_name(f)) {
            continue;
        }
        let insts = instructions(f);
        for &i in &insts {
            // SAFETY: `i` is a live instruction. Its result is replaced by
            // `undef` first so that erasing the instructions in any order
            // leaves no dangling use (loops and phis make the order unsafe).
            unsafe {
                let ty = LLVMTypeOf(i);
                if LLVMGetTypeKind(ty) != LLVMTypeKind::LLVMVoidTypeKind {
                    LLVMReplaceAllUsesWith(i, LLVMGetUndef(ty));
                }
            }
        }
        for i in insts {
            // SAFETY: `i` is live, has no remaining uses, and is erased once.
            unsafe { LLVMInstructionEraseFromParent(i) };
        }
        // SAFETY: `f` is live; every instruction (so every use of a block) is
        // gone, so deleting the blocks front to back is sound.
        unsafe {
            while LLVMCountBasicBlocks(f) > 0 {
                LLVMDeleteBasicBlock(LLVMGetFirstBasicBlock(f));
            }
        }
    }
    if bin == 0 {
        return;
    }
    // SAFETY: `m` is the live module; globals are only re-initialised here.
    let mut g = unsafe { LLVMGetFirstGlobal(m) };
    while !g.is_null() {
        // SAFETY: `g` is a live global variable. `plan` admitted only external
        // (visible, defined once in bin 0), constant local, or mutable local
        // globals (whose users all share one bin, so a duplicate is never
        // touched), and
        // clearing an external global's initializer turns it into the
        // declaration of bin 0's definition.
        unsafe {
            if LLVMIsDeclaration(g) == 0 && LLVMGetLinkage(g) == LLVMLinkage::LLVMExternalLinkage {
                LLVMSetInitializer(g, std::ptr::null_mut());
            }
            g = LLVMGetNextGlobal(g);
        }
    }
}