weirflow 0.1.0

GPU-first dataflow analysis primitives for Vyre and Santh compiler pipelines.
Documentation
use std::sync::Arc;

use super::OP_ID;
use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Ident, Node, Program};

/// Build a single-dispatch Program that OR-merges a direct call-edge
/// bitset with the transitive-closure bitset produced by
/// [`crate::points_to::andersen_points_to`] into the final
/// callgraph edge bitset.
///
/// `direct_edges_in` and `indirect_sites_in` are read-only bitsets
/// over `node_count` call-site lanes. `pts_closure_in` is the
/// points-to closure for each indirect call-site. `callgraph_out` is
/// the final bitset written lane-by-lane.
///
/// `node_count` is the number of call sites (one bit per site).
#[must_use]
pub fn callgraph_build(
    direct_edges_in: &str,
    indirect_sites_in: &str,
    pts_closure_in: &str,
    callgraph_out: &str,
) -> Program {
    callgraph_build_with_count(
        direct_edges_in,
        indirect_sites_in,
        pts_closure_in,
        callgraph_out,
        4,
    )
}

/// Version that takes the number of call-site lanes explicitly.
#[must_use]
pub fn callgraph_build_with_count(
    direct_edges_in: &str,
    indirect_sites_in: &str,
    pts_closure_in: &str,
    callgraph_out: &str,
    node_count: u32,
) -> Program {
    let domain = crate::graph_layout::LinearDomain::new(node_count);
    let words = domain.bitset_words();
    let w = Expr::InvocationId { axis: 0 };

    let body = vec![
        Node::let_bind("direct", Expr::load(direct_edges_in, w.clone())),
        Node::let_bind(
            "indirect",
            Expr::bitand(
                Expr::load(indirect_sites_in, w.clone()),
                Expr::load(pts_closure_in, w.clone()),
            ),
        ),
        Node::store(
            callgraph_out,
            w.clone(),
            Expr::bitor(Expr::var("direct"), Expr::var("indirect")),
        ),
    ];

    let buffers = vec![
        BufferDecl::storage(direct_edges_in, 0, BufferAccess::ReadOnly, DataType::U32)
            .with_count(words),
        BufferDecl::storage(indirect_sites_in, 1, BufferAccess::ReadOnly, DataType::U32)
            .with_count(words),
        BufferDecl::storage(pts_closure_in, 2, BufferAccess::ReadOnly, DataType::U32)
            .with_count(words),
        BufferDecl::storage(callgraph_out, 3, BufferAccess::ReadWrite, DataType::U32)
            .with_count(words),
    ];

    Program::wrapped(
        buffers,
        [256, 1, 1],
        vec![Node::Region {
            generator: Ident::from(OP_ID),
            source_region: None,
            body: Arc::new(vec![Node::if_then(
                Expr::lt(w.clone(), Expr::u32(words)),
                body,
            )]),
        }],
    )
}