Skip to main content

bb_dsl/
syscalls.rs

1//! DSL-side syscall helpers — record canonical `NodeProto`s into a
2//! `Graph`. The runtime-side dispatch impls live in `bb-ops`; the
3//! two sides agree on stable `(domain, op_type)` string constants
4//! re-exported from `bb_ir::syscall_ids`.
5
6use crate::graph::Graph;
7use crate::output::Output;
8use bb_ir::proto::onnx::NodeProto;
9
10/// Canonical `(domain, op_type)` string constants the DSL helpers
11/// stamp onto recorded `NodeProto`s. The strings live in
12/// `bb_ir::syscall_ids` so the DSL + compiler + runtime cite one
13/// declaration; this module re-exports them under shorter local
14/// aliases for the helper bodies below.
15pub mod ids {
16    pub use bb_ir::syscall_ids::{
17        OP_GATE_DISPATCH as GATE_DISPATCH_OP, OP_PASS_THROUGH as PASS_THROUGH_OP, SYSCALL_DOMAIN,
18    };
19}
20
21/// `(domain, op_type)` registration key for the `ai.bytesandbrains.address_book`
22/// custom-op family. Matches the strings registered in
23/// `bb-ops/src/syscalls/peers/{insert,insert_many,lookup}.rs` via
24/// `register_op!`.
25const ADDRESS_BOOK_DOMAIN: &str = "ai.bytesandbrains.address_book";
26const INSERT_MANY_OP: &str = "InsertMany";
27const LOOKUP_OP: &str = "Lookup";
28
29/// Record a `PassThrough` syscall NodeProto into a `Graph`.
30///
31/// The framework's structural identity op - threads a value through
32/// a partition without doing any compute. Authors reach for it when
33/// a partition needs a non-wire node (e.g. a receiver class that
34/// only forwards values it receives over the wire). The recorded
35/// NodeProto's home class is inferred by the compiler from the
36/// input's home.
37pub fn pass_through(g: &mut Graph, input: Output) -> Output {
38    let out_name = g.next_site_name();
39    g.push_node(NodeProto {
40        op_type: ids::PASS_THROUGH_OP.into(),
41        domain: ids::SYSCALL_DOMAIN.into(),
42        input: vec![input.name],
43        output: vec![out_name.clone()],
44        ..Default::default()
45    });
46    g.declare_value_info(&out_name, input.type_node);
47    Output::new(out_name, input.type_node)
48}
49
50/// Record an `AddressBook::InsertMany(peer, addresses)` custom-op
51/// NodeProto into a `Graph`. New peer creates an entry with
52/// `ref_count = 1`; known peer dedupe-appends every address without
53/// touching `ref_count`. Empty `addresses` vec surfaces as a
54/// dispatch-time `OpError`.
55pub fn address_book_insert_many(g: &mut Graph, peer: Output, addresses: Output) -> Output {
56    let out_name = g.next_site_name();
57    g.push_node(NodeProto {
58        op_type: INSERT_MANY_OP.into(),
59        domain: ADDRESS_BOOK_DOMAIN.into(),
60        input: vec![peer.name, addresses.name],
61        output: vec![out_name.clone()],
62        ..Default::default()
63    });
64    g.declare_value_info(&out_name, &bb_ir::types::TYPE_TRIGGER);
65    Output::new(out_name, &bb_ir::types::TYPE_TRIGGER)
66}
67
68/// Record an `AddressBook::Lookup(peer)` custom-op NodeProto into a
69/// `Graph`. Output carries the full ordered `TYPE_ADDRESS_VEC`;
70/// callers that need a single address pick one downstream. Unknown
71/// or empty-address peer surfaces as a dispatch-time `OpError`.
72pub fn address_book_lookup(g: &mut Graph, peer: Output) -> Output {
73    let out_name = g.next_site_name();
74    g.push_node(NodeProto {
75        op_type: LOOKUP_OP.into(),
76        domain: ADDRESS_BOOK_DOMAIN.into(),
77        input: vec![peer.name],
78        output: vec![out_name.clone()],
79        ..Default::default()
80    });
81    g.declare_value_info(&out_name, &bb_ir::types::TYPE_ADDRESS_VEC);
82    Output::new(out_name, &bb_ir::types::TYPE_ADDRESS_VEC)
83}
84
85/// Record a `GateDispatch` syscall NodeProto into a `Graph` - a
86/// multi-edge synchronization barrier.
87pub fn gate_dispatch(g: &mut Graph, inputs: &[Output]) -> Output {
88    let out_name = g.next_site_name();
89    g.push_node(NodeProto {
90        op_type: ids::GATE_DISPATCH_OP.into(),
91        domain: ids::SYSCALL_DOMAIN.into(),
92        input: inputs.iter().map(|o| o.name.clone()).collect(),
93        output: vec![out_name.clone()],
94        ..Default::default()
95    });
96    g.declare_value_info(&out_name, &bb_ir::types::TYPE_BYTES);
97    Output::new(out_name, &bb_ir::types::TYPE_BYTES)
98}