bb_ir/types/common_relations.rs
1//! Named `TypeRelation` slice constants for the common op shapes.
2//!
3//! Sharing one set of canonical slices across `bb-ops::backends`,
4//! `bb-derive`, and any third-party crate keeps op declarations
5//! one-liners and lets the TypeSolver work with stable
6//! pointer-equal constraint pools (the solver indexes relations by
7//! their `&'static` address). New shapes graft in alongside these
8//! without breaking the convention.
9
10use super::{relations::TypeRelation, PortRef};
11
12/// `[]` — explicit "no constraint" slot. Use this on ops whose type
13/// is attribute-driven (`Cast`), variadic (`Concat`, `Split`,
14/// `Gemm`), structural (`Reshape`'s output shape comes from an
15/// attribute), or whose I/O types are heterogeneous and lack a
16/// shared element-type bound (most role placeholders).
17pub static NO_RELATIONS: &[TypeRelation] = &[];
18
19/// Unary element-wise: `output.TypeNode == input.TypeNode`. Shape
20/// preserved, element type preserved. Use for `Neg`, `Abs`, `Sqrt`,
21/// `Exp`, `Log`, `Identity`, `Relu`, `Sigmoid`, `Tanh`, `Softmax`,
22/// `LeakyRelu`, `Gelu`, `GlobalAveragePool`.
23pub static ELEMENTWISE: &[TypeRelation] = &[TypeRelation::Elementwise {
24 input: PortRef::Input(0),
25 output: PortRef::Output(0),
26}];
27
28/// `input[0].ElementType == output[0].ElementType` without
29/// constraining shape. Use for ops that preserve element type but
30/// reshape (`Reshape`, `Transpose`, `Slice`, `Squeeze`, `Unsqueeze`).
31pub static UNARY_SAME_ELEMENT: &[TypeRelation] = &[TypeRelation::SameElementType(&[
32 PortRef::Input(0),
33 PortRef::Output(0),
34])];
35
36/// Binary broadcast arithmetic: `Add`, `Sub`, `Mul`, `Div`, `Pow`.
37/// Inputs + output share element type AND the output's shape is
38/// the broadcast of the two inputs.
39pub static BROADCAST_BINARY: &[TypeRelation] = &[
40 TypeRelation::SameElementType(&[PortRef::Input(0), PortRef::Input(1), PortRef::Output(0)]),
41 TypeRelation::BroadcastShape {
42 in0: PortRef::Input(0),
43 in1: PortRef::Input(1),
44 out: PortRef::Output(0),
45 },
46];
47
48/// `MatMul` / `Dot` — both inputs and the output share an element
49/// type; shape is matmul-specific and stays unconstrained until a
50/// `Custom` shape relation lands.
51pub static MATMUL_BINARY: &[TypeRelation] = &[TypeRelation::SameElementType(&[
52 PortRef::Input(0),
53 PortRef::Input(1),
54 PortRef::Output(0),
55])];
56
57/// `ReduceSum` / `ReduceMean` / `ReduceMax` / `ReduceMin` — same
58/// element type, reduced shape driven by `axes` / `keepdims`
59/// attributes.
60pub static REDUCE_AXIS: &[TypeRelation] = &[TypeRelation::ReduceOver {
61 input: PortRef::Input(0),
62 output: PortRef::Output(0),
63}];