bb_ir/
peer_class.rs

1//! Peer-class metadata for partitioning by "what kind of Node does
2//! this op run on." [`PEER_CLASS_KEY`] tags `Output<PeerId>`
3//! producers; [`HOME_CLASS_KEY`] tags each NodeProto with the class
4//! of Node it runs on (stamped by `infer_peer_classes`).
5
6use crate::proto::onnx::{NodeProto, ValueInfoProto};
7
8/// Class of peer an `Output<PeerId>` producer yields. Stamped on
9/// `ValueInfoProto` (Graph::input) or `NodeProto` (frontend).
10pub const PEER_CLASS_KEY: &str = "ai.bytesandbrains.peer_class";
11
12/// Class of Node a `NodeProto` runs on. Stamped by
13/// `infer_peer_classes`; partition keys + cross-class detection
14/// consume it.
15pub const HOME_CLASS_KEY: &str = "ai.bytesandbrains.home_class";
16
17/// Sentinel for ops whose data inputs all originate on the local
18/// Node. Function inputs start here; `wire.Send` re-homes downstream.
19pub const SELF_CLASS: &str = "@self";
20
21/// Read `peer_class` off a `ValueInfoProto`. `None` if absent.
22pub fn peer_class_of_value_info(vi: &ValueInfoProto) -> Option<&str> {
23    vi.metadata_props
24        .iter()
25        .find(|p| p.key == PEER_CLASS_KEY)
26        .map(|p| p.value.as_str())
27}
28
29/// Read `peer_class` off a `NodeProto`. `None` if absent.
30pub fn peer_class_of_node(node: &NodeProto) -> Option<&str> {
31    node.metadata_props
32        .iter()
33        .find(|p| p.key == PEER_CLASS_KEY)
34        .map(|p| p.value.as_str())
35}
36
37/// Read `home_class` off a `NodeProto`. `None` for unvisited nodes;
38/// partition falls back to [`SELF_CLASS`].
39pub fn home_class_of_node(node: &NodeProto) -> Option<&str> {
40    node.metadata_props
41        .iter()
42        .find(|p| p.key == HOME_CLASS_KEY)
43        .map(|p| p.value.as_str())
44}
bb_ir/peer_class.rs

bb_ir/
peer_class.rs
