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};
use vyre_primitives::bitset::bitset_words;

/// One summary-update step for a single function.
///
/// Bit `i` of `summary_out` is set iff any of: the AST-derived bit
/// was set, the callgraph-aggregated bit from callee summaries was
/// set, or the cached summary from the previous scan was set. That
/// is a three-way OR  -  same kernel shape as `escape`, different
/// buffer names.
#[must_use]
pub fn summarize_function(
    fn_ast_in: &str,
    callgraph_in: &str,
    cached_summary_in: &str,
    summary_out: &str,
) -> Program {
    summarize_function_with_count(fn_ast_in, callgraph_in, cached_summary_in, summary_out, 64)
}

/// Version that takes the bit-lane count explicitly.
#[must_use]
pub fn summarize_function_with_count(
    fn_ast_in: &str,
    callgraph_in: &str,
    cached_summary_in: &str,
    summary_out: &str,
    bit_count: u32,
) -> Program {
    let words = bitset_words(bit_count).max(1);
    let w = Expr::InvocationId { axis: 0 };

    let body = vec![
        Node::let_bind("ast", Expr::load(fn_ast_in, w.clone())),
        Node::let_bind("cg", Expr::load(callgraph_in, w.clone())),
        Node::let_bind("cached", Expr::load(cached_summary_in, w.clone())),
        Node::store(
            summary_out,
            w.clone(),
            Expr::bitor(
                Expr::bitor(Expr::var("ast"), Expr::var("cg")),
                Expr::var("cached"),
            ),
        ),
    ];

    let buffers = vec![
        BufferDecl::storage(fn_ast_in, 0, BufferAccess::ReadOnly, DataType::U32).with_count(words),
        BufferDecl::storage(callgraph_in, 1, BufferAccess::ReadOnly, DataType::U32)
            .with_count(words),
        BufferDecl::storage(cached_summary_in, 2, BufferAccess::ReadOnly, DataType::U32)
            .with_count(words),
        BufferDecl::storage(summary_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,
            )]),
        }],
    )
}