#![forbid(unsafe_code)]
use vyre_foundation::ir::{validate, Expr, Node, Program};
use vyre_primitives::fixpoint::persistent_fixpoint::{
persistent_fixpoint, persistent_fixpoint_grid,
};
use vyre_primitives::graph::persistent_bfs::{
persistent_bfs, persistent_bfs_batch, persistent_bfs_batch_with_density,
persistent_bfs_with_density,
};
use vyre_primitives::graph::program_graph::ProgramGraphShape;
use vyre_primitives::parsing::ast_cse_structural_hash::ast_cse_structural_hash_program;
use vyre_primitives::reduce::workgroup_tree::{
workgroup_max_f32, workgroup_max_u32, workgroup_min_f32, workgroup_min_u32, workgroup_sum_f32,
workgroup_sum_u32,
};
fn transfer_body() -> Vec<Node> {
vec![Node::if_then(
Expr::lt(Expr::InvocationId { axis: 0 }, Expr::u32(4)),
vec![Node::store(
"next",
Expr::InvocationId { axis: 0 },
Expr::load("current", Expr::InvocationId { axis: 0 }),
)],
)]
}
fn synchronizing_programs(max_iters: u32) -> Vec<(&'static str, Program)> {
let shape = ProgramGraphShape::new(64, 256);
vec![
(
"persistent_bfs",
persistent_bfs(shape, "frontier_in", "frontier_out", u32::MAX, max_iters),
),
(
"persistent_bfs_with_density",
persistent_bfs_with_density(
shape,
"frontier_in",
"frontier_out",
"density_active",
u32::MAX,
max_iters,
),
),
(
"persistent_bfs_batch",
persistent_bfs_batch(
shape,
"frontier_in",
"frontier_out",
"changed",
"converged",
4,
u32::MAX,
max_iters,
),
),
(
"persistent_bfs_batch_with_density",
persistent_bfs_batch_with_density(
shape,
"frontier_in",
"frontier_out",
"changed",
"converged",
"density_active",
4,
u32::MAX,
max_iters,
),
),
(
"persistent_fixpoint",
persistent_fixpoint(transfer_body(), "current", "next", "changed", 4, max_iters),
),
(
"persistent_fixpoint_grid",
persistent_fixpoint_grid(transfer_body(), "current", "next", "changed", 4, max_iters),
),
(
"workgroup_sum_f32",
workgroup_sum_f32("values", "out", 1024, 256),
),
(
"workgroup_sum_u32",
workgroup_sum_u32("values", "out", 1024, 256),
),
(
"workgroup_max_f32",
workgroup_max_f32("values", "out", 1024, 256),
),
(
"workgroup_max_u32",
workgroup_max_u32("values", "out", 1024, 256),
),
(
"workgroup_min_f32",
workgroup_min_f32("values", "out", 1024, 256),
),
(
"workgroup_min_u32",
workgroup_min_u32("values", "out", 1024, 256),
),
(
"ast_cse_structural_hash_program",
ast_cse_structural_hash_program(64, 128),
),
]
}
fn has_barrier(nodes: &[Node]) -> bool {
nodes.iter().any(|node| match node {
Node::Barrier { .. } => true,
Node::If {
then, otherwise, ..
} => has_barrier(then) || has_barrier(otherwise),
Node::Loop { body, .. } => has_barrier(body),
Node::Block(body) => has_barrier(body),
Node::Region { body, .. } => has_barrier(body),
_ => false,
})
}
fn synchronizing_loop_bodies<'a>(nodes: &'a [Node], found: &mut Vec<&'a [Node]>) {
for node in nodes {
match node {
Node::Loop { body, .. } => {
if has_barrier(body) {
found.push(body.as_slice());
}
synchronizing_loop_bodies(body, found);
}
Node::If {
then, otherwise, ..
} => {
synchronizing_loop_bodies(then, found);
synchronizing_loop_bodies(otherwise, found);
}
Node::Block(body) => synchronizing_loop_bodies(body, found),
Node::Region { body, .. } => synchronizing_loop_bodies(body, found),
_ => {}
}
}
}
fn messages(program: &Program) -> Vec<String> {
validate(program)
.iter()
.map(|error| error.message().to_string())
.collect()
}
#[test]
fn every_shipped_synchronizing_loop_program_validates() {
let mut broken: Vec<(&str, Vec<String>)> = Vec::new();
for (label, program) in synchronizing_programs(8) {
let found = messages(&program);
if !found.is_empty() {
broken.push((label, found));
}
}
assert_eq!(
broken,
Vec::new(),
"shipped programs must validate; each entry is (builder, errors)"
);
}
#[test]
fn assorted_iteration_counts_stay_valid() {
for max_iters in [0_u32, 1, 2, 7, 4096] {
for (label, program) in synchronizing_programs(max_iters) {
assert_eq!(
messages(&program),
Vec::<String>::new(),
"{label} must validate at max_iters {max_iters}"
);
}
}
}
#[test]
fn exactly_the_known_programs_contain_synchronizing_loops() {
let mut with_sync_loop: Vec<&str> = Vec::new();
for (label, program) in synchronizing_programs(8) {
let mut bodies: Vec<&[Node]> = Vec::new();
synchronizing_loop_bodies(program.entry(), &mut bodies);
if !bodies.is_empty() {
with_sync_loop.push(label);
}
}
with_sync_loop.sort_unstable();
assert_eq!(
with_sync_loop,
vec![
"persistent_bfs_batch",
"persistent_bfs_batch_with_density",
"persistent_bfs_with_density",
"persistent_fixpoint",
],
"the set of programs governed by the back-edge rule changed"
);
}
#[test]
fn exit_proof_and_exit_free_bodies_are_where_they_were_measured() {
let mut exit_proof: Vec<String> = Vec::new();
let mut exit_free: Vec<String> = Vec::new();
for (label, program) in synchronizing_programs(8) {
let mut bodies: Vec<&[Node]> = Vec::new();
synchronizing_loop_bodies(program.entry(), &mut bodies);
for (index, body) in bodies.iter().enumerate() {
let entry = format!("{label} body {index}");
if matches!(body.last(), Some(Node::Barrier { .. })) {
exit_proof.push(entry);
} else {
exit_free.push(entry);
}
}
}
exit_proof.sort_unstable();
exit_free.sort_unstable();
assert_eq!(
exit_proof,
vec![
"persistent_bfs_batch body 0".to_string(),
"persistent_fixpoint body 0".to_string(),
],
"the set of bodies ending in an unconditional barrier changed"
);
assert_eq!(
exit_free,
vec![
"persistent_bfs_batch_with_density body 0".to_string(),
"persistent_bfs_with_density body 0".to_string(),
],
"the set of bodies that are exit-free but not exit-proof changed"
);
}
#[test]
fn a_body_with_an_early_exit_ends_with_a_barrier() {
let mut exits: Vec<(String, usize, bool)> = Vec::new();
for (label, program) in synchronizing_programs(8) {
let mut bodies: Vec<&[Node]> = Vec::new();
synchronizing_loop_bodies(program.entry(), &mut bodies);
for (index, body) in bodies.iter().enumerate() {
let returns = returns_at_any_depth(body);
let exit_proof = matches!(body.last(), Some(Node::Barrier { .. }));
assert!(
returns == 0 || exit_proof,
"{label} body {index} holds {returns} early exit(s) and does NOT \
end with an unconditional barrier, so an invocation can leave \
after the body's last barrier while its siblings take the back \
edge"
);
exits.push((format!("{label} body {index}"), returns, exit_proof));
}
}
exits.sort_by(|left, right| left.0.cmp(&right.0));
assert_eq!(
exits,
vec![
("persistent_bfs_batch body 0".to_string(), 0, true),
(
"persistent_bfs_batch_with_density body 0".to_string(),
0,
false
),
("persistent_bfs_with_density body 0".to_string(), 0, false),
("persistent_fixpoint body 0".to_string(), 1, true),
],
"measured (body, early exits, ends with barrier) changed"
);
}
fn returns_at_any_depth(nodes: &[Node]) -> usize {
nodes
.iter()
.map(|node| match node {
Node::Return => 1,
Node::If {
then, otherwise, ..
} => returns_at_any_depth(then) + returns_at_any_depth(otherwise),
Node::Loop { body, .. } => returns_at_any_depth(body),
Node::Block(body) => returns_at_any_depth(body),
Node::Region { body, .. } => returns_at_any_depth(body),
_ => 0,
})
.sum()
}