use std::collections::{BTreeMap, BTreeSet, HashMap};
use crate::native::cfg::BodyBlock;
use crate::native::ir::{LlType, LlValue};
#[derive(Clone, Debug, PartialEq, Eq)]
pub(in crate::native) enum NetworkClass {
Uniform,
WholeVsPart(LlType),
ReinterpretMix,
Unclassified,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(in crate::native) struct PointerNetwork {
pub(in crate::native) members: Vec<String>,
pub(in crate::native) pointees: Vec<LlType>,
pub(in crate::native) class: NetworkClass,
}
#[derive(Default)]
struct UnionFind {
parent: HashMap<String, String>,
}
impl UnionFind {
fn find(&mut self, x: &str) -> String {
let mut root = x.to_string();
while let Some(p) = self.parent.get(&root) {
if p == &root {
break;
}
root = p.clone();
}
let mut cur = x.to_string();
while cur != root {
let next = self
.parent
.get(&cur)
.cloned()
.unwrap_or_else(|| cur.clone());
self.parent.insert(cur.clone(), root.clone());
if next == cur {
break;
}
cur = next;
}
root
}
fn union(&mut self, a: &str, b: &str) {
self.parent
.entry(a.to_string())
.or_insert_with(|| a.to_string());
self.parent
.entry(b.to_string())
.or_insert_with(|| b.to_string());
let ra = self.find(a);
let rb = self.find(b);
if ra != rb {
self.parent.insert(ra, rb);
}
}
}
fn peel_scalar(ty: &LlType) -> &LlType {
match ty {
LlType::Vector(e, _) | LlType::Array(e, _) => peel_scalar(e),
other => other,
}
}
fn is_primitive_scalar(ty: &LlType) -> bool {
matches!(
ty,
LlType::Float | LlType::Half | LlType::BFloat | LlType::Int(_)
)
}
fn classify(pointees: &[LlType]) -> NetworkClass {
if pointees.len() <= 1 {
return NetworkClass::Uniform;
}
let mut scalars: Vec<&LlType> = Vec::new();
for p in pointees {
let s = peel_scalar(p);
if !is_primitive_scalar(s) {
return NetworkClass::Unclassified;
}
if !scalars.contains(&s) {
scalars.push(s);
}
}
match scalars.as_slice() {
[only] => NetworkClass::WholeVsPart((*only).clone()),
_ => NetworkClass::ReinterpretMix,
}
}
fn pointer_edges(blocks: &[BodyBlock]) -> Vec<(String, String)> {
let mut edges = Vec::new();
for block in blocks {
if let Some(carrier) = &block.typed {
pointer_edges_from_carrier(carrier, &mut edges);
}
}
edges
}
fn pointer_edges_from_carrier(
carrier: &crate::native::tir::TirBlock,
edges: &mut Vec<(String, String)>,
) {
for inst in &carrier.insts {
let Some(name) = &inst.result else {
continue;
};
if inst.opcode == "phi" {
if let Some((ty, incomings)) = &inst.phi_incoming {
if matches!(ty, LlType::Ptr(_)) {
for (value, _pred) in incomings {
if let LlValue::Local(inc) = value {
edges.push((name.clone(), inc.clone()));
}
}
}
}
} else if inst.opcode == "select" && inst.operands.len() == 3 {
let arms = [&inst.operands[1], &inst.operands[2]];
if arms.iter().all(|arm| is_ptr_operand(arm)) {
for arm in arms {
if let Some(tv) = arm.as_typed_value() {
if let LlValue::Local(local) = tv.value {
edges.push((name.clone(), local));
}
}
}
}
} else if inst.opcode == "bitcast" {
if let Some((src, dst)) = &inst.bitcast {
if matches!(src.ty, LlType::Ptr(_)) && dst.trim_start().starts_with("ptr") {
if let LlValue::Local(local) = &src.value {
edges.push((name.clone(), local.clone()));
}
}
}
}
}
}
fn is_ptr_operand(op: &crate::native::tir::TirOperand) -> bool {
op.as_typed_value()
.is_some_and(|tv| matches!(tv.ty, LlType::Ptr(_)))
}
fn build_components(blocks: &[BodyBlock]) -> Vec<Vec<String>> {
let edges = pointer_edges(blocks);
let mut uf = UnionFind::default();
let mut nodes: BTreeSet<String> = BTreeSet::new();
for (a, b) in &edges {
uf.union(a, b);
nodes.insert(a.clone());
nodes.insert(b.clone());
}
let mut groups: BTreeMap<String, Vec<String>> = BTreeMap::new();
for n in &nodes {
let root = uf.find(n);
groups.entry(root).or_default().push(n.clone());
}
groups
.into_values()
.map(|mut members| {
members.sort();
members.dedup();
members
})
.collect()
}
fn census<'a, F, I>(members: Vec<String>, types_of: F) -> PointerNetwork
where
F: Fn(&str) -> I,
I: IntoIterator<Item = &'a LlType>,
{
let mut distinct: Vec<LlType> = Vec::new();
for m in &members {
for p in types_of(m) {
if !distinct.contains(p) {
distinct.push(p.clone());
}
}
}
distinct.sort_by_key(|t| format!("{t:?}"));
let class = classify(&distinct);
PointerNetwork {
members,
pointees: distinct,
class,
}
}
pub(in crate::native) fn analyze_pointer_networks(
blocks: &[BodyBlock],
pointees: &HashMap<String, LlType>,
) -> Vec<PointerNetwork> {
build_components(blocks)
.into_iter()
.map(|members| census(members, |m| pointees.get(m)))
.collect()
}
pub(in crate::native) fn analyze_networks_by_access(blocks: &[BodyBlock]) -> Vec<PointerNetwork> {
let access = access_pointees(blocks);
build_components(blocks)
.into_iter()
.map(|members| census(members, |m| access.get(m).into_iter().flatten()))
.collect()
}
fn access_pointees(blocks: &[BodyBlock]) -> HashMap<String, Vec<LlType>> {
use crate::native::emitter::helpers::gep_pointee;
let mut map: HashMap<String, Vec<LlType>> = HashMap::new();
let mut record = |name: &str, ty: LlType| {
let entry = map.entry(name.to_string()).or_default();
if !entry.contains(&ty) {
entry.push(ty);
}
};
for block in blocks {
if let Some(carrier) = &block.typed {
{
for inst in &carrier.insts {
if inst.opcode == "load" {
if let (Some(ty), Some(ptr)) = (
&inst.result_ty,
inst.operands.first().and_then(local_operand),
) {
record(&ptr, ty.clone());
}
} else if inst.opcode == "store" {
if let (Some(val), Some(ptr)) = (
inst.operands.first().and_then(|o| o.as_typed_value()),
inst.operands.get(1).and_then(local_operand),
) {
record(&ptr, val.ty);
}
} else if inst.opcode == "getelementptr" {
if let Some(gep) = &inst.gep {
if let LlValue::Local(base) = &gep.base.value {
record(base, gep.source_ty.clone());
}
if let (Some(result), Ok(pointee)) =
(&inst.result, gep_pointee(&gep.source_ty, &gep.indices))
{
record(result, pointee);
}
}
}
}
}
}
}
map
}
fn local_operand(op: &crate::native::tir::TirOperand) -> Option<String> {
match op.as_typed_value()?.value {
LlValue::Local(name) => Some(name),
_ => None,
}
}
pub(in crate::native) fn array_indexed_scalar_bases(blocks: &[BodyBlock]) -> BTreeSet<String> {
use crate::native::emitter::helpers::const_index;
use crate::native::ir::LlGep;
let mut set = BTreeSet::new();
let mut consider = |gep: &LlGep| {
let LlValue::Local(base) = &gep.base.value else {
return;
};
if !is_primitive_scalar(&gep.source_ty) || gep.indices.len() != 1 {
return;
}
if const_index(gep.indices.first()) != Some(0) {
set.insert(base.clone());
}
};
for block in blocks {
if let Some(carrier) = &block.typed {
for inst in &carrier.insts {
if let Some(gep) = &inst.gep {
consider(gep);
}
}
}
}
set
}
#[cfg(test)]
mod tests {
use super::*;
fn block(lines: &[&str]) -> BodyBlock {
let mut lines: Vec<String> = lines.iter().map(|l| l.to_string()).collect();
lines.push("ret void".to_string());
let typed = crate::native::tir::lower_block_carrier("entry", &lines, &HashMap::new());
assert!(typed.is_some(), "fixture must lower: {lines:?}");
BodyBlock {
name: "entry".to_string(),
role: crate::native::cfg::BlockRole::Normal,
typed,
}
}
fn net_containing<'a>(nets: &'a [PointerNetwork], member: &str) -> &'a PointerNetwork {
nets.iter()
.find(|n| n.members.iter().any(|m| m == member))
.expect("network with member")
}
#[test]
fn peels_vector_and_array_to_scalar() {
assert_eq!(
peel_scalar(&LlType::Vector(Box::new(LlType::Float), 4)),
&LlType::Float
);
assert_eq!(
peel_scalar(&LlType::Array(Box::new(LlType::Half), 8)),
&LlType::Half
);
assert_eq!(peel_scalar(&LlType::Float), &LlType::Float);
}
#[test]
fn classifies_whole_vs_part_as_finest_scalar() {
let c = classify(&[LlType::Float, LlType::Vector(Box::new(LlType::Float), 4)]);
assert_eq!(c, NetworkClass::WholeVsPart(LlType::Float));
}
#[test]
fn classifies_same_scalar_across_two_vector_widths() {
let c = classify(&[
LlType::Vector(Box::new(LlType::Float), 2),
LlType::Vector(Box::new(LlType::Float), 4),
]);
assert_eq!(c, NetworkClass::WholeVsPart(LlType::Float));
}
#[test]
fn classifies_reinterpret_mix_for_different_scalars() {
assert_eq!(
classify(&[LlType::Float, LlType::Int(32)]),
NetworkClass::ReinterpretMix
);
}
#[test]
fn classifies_struct_mix_as_unclassified() {
let c = classify(&[
LlType::Float,
LlType::Struct(vec![LlType::Half, LlType::Half]),
]);
assert_eq!(c, NetworkClass::Unclassified);
}
#[test]
fn single_pointee_is_uniform() {
assert_eq!(classify(&[LlType::Float]), NetworkClass::Uniform);
assert_eq!(classify(&[]), NetworkClass::Uniform);
}
#[test]
fn groups_mixed_granularity_loop_phi_as_whole_vs_part() {
let blocks = [block(&[
"%144 = phi ptr [ %148, %loop ], [ %142, %entry ]",
"%148 = getelementptr <4 x float>, ptr %144, i64 %125",
"%142 = getelementptr float, ptr %129, i64 %140",
"%147 = load <4 x float>, ptr %144",
])];
let mut pointees = HashMap::new();
pointees.insert(
"%144".to_string(),
LlType::Vector(Box::new(LlType::Float), 4),
);
pointees.insert(
"%148".to_string(),
LlType::Vector(Box::new(LlType::Float), 4),
);
pointees.insert("%142".to_string(), LlType::Float);
let nets = analyze_pointer_networks(&blocks, &pointees);
let net = net_containing(&nets, "%144");
assert!(net.members.contains(&"%148".to_string()));
assert!(net.members.contains(&"%142".to_string()));
assert_eq!(net.class, NetworkClass::WholeVsPart(LlType::Float));
}
#[test]
fn access_census_classifies_05_phi_as_whole_vs_part() {
let blocks = [block(&[
"%144 = phi ptr [ %148, %loop ], [ %142, %entry ]",
"%148 = getelementptr <4 x float>, ptr %144, i64 %125",
"%142 = getelementptr float, ptr %129, i64 %140",
"%147 = load <4 x float>, ptr %144",
])];
let nets = analyze_networks_by_access(&blocks);
let net = net_containing(&nets, "%144");
assert!(net.members.contains(&"%148".to_string()));
assert!(net.members.contains(&"%142".to_string()));
assert_eq!(net.class, NetworkClass::WholeVsPart(LlType::Float));
}
#[test]
fn access_census_follows_bitcast_alias_to_scalar_arm() {
let blocks = [block(&[
"%144 = phi ptr addrspace(1) [ %148, %loop ], [ %142, %entry ]",
"%148 = getelementptr <4 x float>, ptr addrspace(1) %144, i64 %125",
"%141 = getelementptr float, ptr addrspace(1) %129, i64 %140",
"%142 = bitcast ptr addrspace(1) %141 to ptr addrspace(1)",
"%147 = load <4 x float>, ptr addrspace(1) %144",
])];
let nets = analyze_networks_by_access(&blocks);
let net = net_containing(&nets, "%144");
assert!(net.members.contains(&"%141".to_string()));
assert!(net.members.contains(&"%142".to_string()));
assert_eq!(net.class, NetworkClass::WholeVsPart(LlType::Float));
}
#[test]
fn access_census_uniform_when_all_scalar() {
let blocks = [block(&[
"%p = phi ptr [ %q, %loop ], [ %r, %entry ]",
"%q = getelementptr float, ptr %p, i64 1",
"store float %v, ptr %r",
])];
let nets = analyze_networks_by_access(&blocks);
let net = net_containing(&nets, "%p");
assert_eq!(net.class, NetworkClass::Uniform);
assert_eq!(net.pointees, vec![LlType::Float]);
}
#[test]
fn groups_select_arms_into_one_network() {
let blocks = [block(&[
"%m = select i1 %c, ptr addrspace(1) %pa, ptr addrspace(1) %pb",
])];
let mut pointees = HashMap::new();
pointees.insert("%pa".to_string(), LlType::Float);
pointees.insert("%pb".to_string(), LlType::Float);
let nets = analyze_pointer_networks(&blocks, &pointees);
let net = net_containing(&nets, "%m");
assert!(net.members.contains(&"%pa".to_string()));
assert!(net.members.contains(&"%pb".to_string()));
assert_eq!(net.class, NetworkClass::Uniform);
}
#[test]
fn transitively_fuses_chained_phis() {
let blocks = [block(&[
"%a = phi ptr [ %b, %x ], [ %c, %y ]",
"%d = phi ptr [ %a, %z ], [ %e, %w ]",
])];
let pointees = HashMap::new();
let nets = analyze_pointer_networks(&blocks, &pointees);
let net = net_containing(&nets, "%e");
for name in ["%a", "%b", "%c", "%d", "%e"] {
assert!(net.members.contains(&name.to_string()), "missing {name}");
}
assert_eq!(
nets.iter()
.filter(|n| n.members.contains(&"%a".to_string()))
.count(),
1
);
}
#[test]
fn disjoint_networks_stay_separate() {
let blocks = [block(&[
"%a = phi ptr [ %b, %x ], [ %c, %y ]",
"%p = phi ptr [ %q, %x ], [ %r, %y ]",
])];
let pointees = HashMap::new();
let nets = analyze_pointer_networks(&blocks, &pointees);
let na = net_containing(&nets, "%a");
let np = net_containing(&nets, "%p");
assert!(!na.members.contains(&"%q".to_string()));
assert!(!np.members.contains(&"%b".to_string()));
}
#[test]
fn array_indexed_scalar_bases_flags_nonzero_and_dynamic_scalar_steps() {
let blocks = [block(&[
"%r1 = getelementptr float, ptr %p, i64 1",
"%r2 = getelementptr float, ptr %q, i64 %i",
])];
let flagged = array_indexed_scalar_bases(&blocks);
assert!(flagged.contains("%p"));
assert!(flagged.contains("%q"));
}
#[test]
fn array_indexed_scalar_bases_ignores_identity_and_aggregate_steps() {
let blocks = [block(&[
"%r1 = getelementptr float, ptr %p, i64 0",
"%r2 = getelementptr [8 x float], ptr %q, i64 0, i64 %i",
"%r3 = getelementptr %struct.S, ptr %s, i64 0, i32 2",
])];
let flagged = array_indexed_scalar_bases(&blocks);
assert!(!flagged.contains("%p"));
assert!(!flagged.contains("%q"));
assert!(!flagged.contains("%s"));
}
}