bb_ir/
wire_shape.rs

1//! Canonical Send / Recv NodeProto shape — the contract every
2//! wire-IR-touching pass + runtime gate agrees on.
3//!
4//! Wire ops carry a specific NodeProto shape. Centralizing the
5//! contract here keeps the DSL emitters, compiler passes, and
6//! runtime gates from drifting on string literals or attribute
7//! layouts.
8//!
9//! This module is the **single declarative description** of that
10//! shape. The DSL emits to it; the compiler passes mutate within
11//! it; the runtime consumes it. [`crate::verify::wire_shape`]
12//! checks a `ModelProto` against the contract.
13//!
14//! ## The shape
15//!
16//! ### Send
17//!
18//! ```text
19//! NodeProto {
20//!     op_type: "Send",
21//!     domain: "ai.bytesandbrains.wire",
22//!     input: [payload_0, payload_1, ..., payload_{N-1}, peer],
23//!     output: [handle],
24//!     attribute: [
25//!         (name: "peer", type: BYTES, t: <PeerId.to_bytes() multihash>),
26//!         (name: "dest_suffix.{name}", type: BYTES, t: <multiaddr-bytes>),
27//!         (name: "deadline_ns", type: INT, i: <i64-ns>),  // optional
28//!     ],
29//!     metadata_props: [
30//!         ("ai.bytesandbrains.wire.wire_id", "<u64>"),
31//!         ("ai.bytesandbrains.wire_transport", "data" | "trigger_only"),
32//!         ("ai.bytesandbrains.batch_group_id", "<u32>"),
33//!         ("ai.bytesandbrains.dest_site_name.{name}", "<recv-site-name>"),
34//!     ],
35//! }
36//! ```
37//!
38//! ### Recv
39//!
40//! ```text
41//! NodeProto {
42//!     op_type: "Recv",
43//!     domain: "ai.bytesandbrains.wire",
44//!     input: [],
45//!     output: [received_0, received_1, ..., received_{N-1}, sender],
46//!     metadata_props: [
47//!         ("ai.bytesandbrains.wire.wire_id", "<u64 matching paired Send>"),
48//!         ("ai.bytesandbrains.wire_transport", "data" | "trigger_only"),
49//!     ],
50//! }
51//! ```
52//!
53//! `wire_id` pairs the Send/Recv halves across the cut.
54//! `wire_transport` tells the runtime whether each fill carries a
55//! payload (`data`) or is firing-signal-only (`trigger_only`).
56//!
57//! ## Key invariants the contract pins
58//!
59//! - **`ATTR_PEER` is bytes, not i64.** The peer attribute on a
60//!   Send carries the PeerId's canonical multihash byte form
61//!   (`PeerId::to_bytes()`), NOT a `u64`-collapsed identity hash.
62//!   The runtime gates parse via
63//!   `PeerId::from_bytes(&attr.t)`.
64//!
65//! - **`wire_id` is the pairing token.** The DSL `Graph::wire`
66//!   mints a monotonic u64 and stamps it on BOTH halves; the
67//!   compiler's `discover_wire_edges` pair Send/Recv by it.
68//!
69//! - **`wire_transport` lives on the NodeProto itself.**
70//!   `analyze_wire_edges` stamps it onto
71//!   `partition.functions[0].node[i].metadata_props` in place, so
72//!   downstream passes and the runtime read a single source of
73//!   truth.
74//!
75//! - **`SlotFill.type_hash` is populated from the sender side's
76//!   `T::HASH`.** Receivers dispatch wire bytes via
77//!   `if envelope.fill.type_hash == T::HASH { T::deserialize(&fill.payload) }`.
78
79/// Wire-op domain. All Send / Recv NodeProtos live under here.
80pub const WIRE_DOMAIN: &str = "ai.bytesandbrains.wire";
81
82/// op_type for a Send node.
83pub const OP_SEND: &str = "Send";
84
85/// op_type for a Recv node.
86pub const OP_RECV: &str = "Recv";
87
88/// Attribute key on a Send NodeProto carrying the destination peer
89/// as **multihash bytes** (`PeerId::to_bytes()`), i.e. the
90/// `attribute.t` (bytes) field, not `attribute.i` (i64).
91///
92/// Aliased from [`crate::syscall_ids::ATTR_PEER`] so consumers can
93/// reach for `bb_ir::wire_shape::ATTR_PEER` or `bb_ir::keys::ATTR_PEER`
94/// — they're the same key, kept in one place.
95pub use crate::syscall_ids::ATTR_PEER;
96
97/// Attribute key prefix for per-fill multiaddr destination
98/// suffixes. Full key is `format!("{DEST_SUFFIX_ATTR_PREFIX}{slot_name}")`.
99pub use crate::keys::DEST_SUFFIX_ATTR_PREFIX;
100
101/// Attribute key for the optional static deadline (in nanoseconds
102/// since the reference clock epoch) stamped by
103/// `insert_async_deadlines`.
104pub use crate::syscall_ids::ATTR_DEADLINE_NS;
105
106/// `metadata_props` key carrying the wire-pairing token.
107pub use crate::keys::WIRE_ID_KEY;
108
109/// `metadata_props` key carrying the data-vs-trigger-only
110/// classification.
111pub use crate::keys::WIRE_TRANSPORT_KEY;
112
113use crate::proto::onnx::{attribute_proto, AttributeProto, StringStringEntryProto};
114
115/// Value of [`WIRE_TRANSPORT_KEY`] for full-payload edges.
116pub use crate::keys::WIRE_TRANSPORT_DATA;
117
118/// Value of [`WIRE_TRANSPORT_KEY`] for trigger-only edges.
119pub use crate::keys::WIRE_TRANSPORT_TRIGGER_ONLY;
120
121/// `metadata_props` key prefix for per-fill recv-site names. Full
122/// key is `format!("{DEST_SITE_NAME_PREFIX}{slot_name}")`.
123pub use crate::keys::DEST_SITE_NAME_PREFIX;
124
125/// Return `true` if the NodeProto is a `wire.Send`.
126pub fn is_send(node: &crate::proto::onnx::NodeProto) -> bool {
127    node.op_type == OP_SEND && node.domain == WIRE_DOMAIN
128}
129
130/// Return `true` if the NodeProto is a `wire.Recv`.
131pub fn is_recv(node: &crate::proto::onnx::NodeProto) -> bool {
132    node.op_type == OP_RECV && node.domain == WIRE_DOMAIN
133}
134
135/// Read the wire_id metadata stamp from a Send or Recv node.
136/// Returns `None` if the key is absent or non-numeric.
137pub fn read_wire_id(node: &crate::proto::onnx::NodeProto) -> Option<u64> {
138    node.metadata_props
139        .iter()
140        .find(|p| p.key == WIRE_ID_KEY)
141        .and_then(|p| p.value.parse::<u64>().ok())
142}
143
144/// Read the destination peer's multihash bytes from a Send / gate
145/// NodeProto. Returns `None` if the attribute is absent or carries
146/// no byte content. The byte payload lives on `attribute.s` per the
147/// ONNX convention for raw bytes (paired with
148/// `AttributeType::String`); callers reconstruct the PeerId via
149/// `PeerId::from_bytes(read_peer_bytes(node)?)`.
150pub fn read_peer_bytes(node: &crate::proto::onnx::NodeProto) -> Option<&[u8]> {
151    let attr = node.attribute.iter().find(|a| a.name == ATTR_PEER)?;
152    if attr.s.is_empty() {
153        None
154    } else {
155        Some(&attr.s)
156    }
157}
158
159/// Stamp the destination peer onto a Send / gate NodeProto's
160/// `attribute.s` (bytes) using the canonical multihash form. Used
161/// by the compiler's gate-insertion passes and any pass synthesizing
162/// new Send NodeProtos.
163pub fn stamp_peer_bytes(node: &mut crate::proto::onnx::NodeProto, peer_bytes: Vec<u8>) {
164    let attr_type = attribute_proto::AttributeType::String as i32;
165    if let Some(existing) = node.attribute.iter_mut().find(|a| a.name == ATTR_PEER) {
166        existing.s = peer_bytes;
167        existing.r#type = attr_type;
168        existing.i = 0;
169    } else {
170        node.attribute.push(AttributeProto {
171            name: ATTR_PEER.to_string(),
172            s: peer_bytes,
173            r#type: attr_type,
174            ..Default::default()
175        });
176    }
177}
178
179/// Stamp the [`WIRE_TRANSPORT_KEY`] classification onto a wire-op
180/// NodeProto's `metadata_props` (idempotent: replaces an existing
181/// value).
182pub fn stamp_wire_transport(node: &mut crate::proto::onnx::NodeProto, transport: &str) {
183    if let Some(entry) = node
184        .metadata_props
185        .iter_mut()
186        .find(|p| p.key == WIRE_TRANSPORT_KEY)
187    {
188        entry.value = transport.to_string();
189    } else {
190        node.metadata_props.push(StringStringEntryProto {
191            key: WIRE_TRANSPORT_KEY.to_string(),
192            value: transport.to_string(),
193        });
194    }
195}
196
bb_ir/wire_shape.rs

bb_ir/
wire_shape.rs
