Enum StageKind 

pub enum StageKind {
    Load {
        tensor: TensorId,
        pattern: AccessPattern,
        num_blocks: usize,
        multicast: Option<GroupId>,
    },
    Store {
        tensor: TensorId,
        pattern: AccessPattern,
        sub_blocks: (usize, usize),
    },
    Recv {
        group: GroupId,
    },
    AllReduce {
        pool: ReducePool,
        tensor: TensorId,
        combine: Box<ComputeOp>,
    },
    Compute(ComputeOp),
    Scalar(ScalarBody),
}

Variants

Load

Fields

tensor: TensorId

pattern: AccessPattern

num_blocks: usize

multicast: Option<GroupId>

This load is the SENDER of a multicast group: after landing a block it multicasts it to the group.

Store

Fields

tensor: TensorId

pattern: AccessPattern

sub_blocks: (usize, usize)

Recv

The receiver side of a multicast group: its output buffer is filled by the group’s sender.

Fields

group: GroupId

AllReduce

CROSS-CHIP REDUCTION. Every chip of group holds a PARTIAL of tensor — the full output shape, summed over only its slice of the contraction axis. This stage combines them, and until it runs the tensor does not hold the answer.

★ ITS ARITHMETIC IS ALREADY IN THE CLOSED ALGEBRA. An all-reduce’s combine is an N-ary elementwise fold — for Sum, Add(Add(op0, op1), op2)… — which ComputeOp::Generic expresses exactly. So combine is a Generic built from the pool and the operand count, and every consumer downstream (the scheduler, the evaluator, the cost model) sees an ordinary compute op rather than a new form. What is genuinely NEW here is the TOPOLOGY, not the arithmetic: this stage’s operands live on different chips, which is what the certificate’s UnreducedShard obligation is about.

★ AND THE MEMBERSHIP IS NOT STORED, because it cannot then disagree with the truth. The chips taking part are exactly the sites of the buffers in ins — Group is the multicast RECTANGLE abstraction and a rectangle of (x, y) cannot name a set of chips, so widening it here would stretch a type to carry something it was not built for. A derived answer beats a field that can drift.

Fields

pool: ReducePool

tensor: TensorId

combine: Box<ComputeOp>

Compute(ComputeOp)

Scalar(ScalarBody)

Implementations

impl StageKind

pub fn remap_buffers(&mut self, f: &dyn Fn(BufferId) -> BufferId)

Rename every buffer the stage’s KIND names beyond its ports — an indirect address’s index buffer, a matmul’s spill buffer, a bias. A pass that renumbers buffers maps these with ins/outs, or the stage silently points at somebody else’s buffer.

pub fn role(&self) -> SiteRole

Trait Implementations

impl Clone for StageKind

fn clone(&self) -> StageKind

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for StageKind

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<'de> Deserialize<'de> for StageKind

fn deserialize<__D>(__deserializer: __D) -> Result<Self, __D::Error>

where __D: Deserializer<'de>,

Deserialize this value from the given Serde deserializer.

impl Hash for StageKind

fn hash<__H: Hasher>(&self, state: &mut __H)

Feeds this value into the given Hasher.

fn hash_slice<H>(data: &[Self], state: &mut H)

where H: Hasher, Self: Sized,

Feeds a slice of this type into the given Hasher.

impl PartialEq for StageKind

fn eq(&self, other: &StageKind) -> bool

Tests for self and other values to be equal, and is used by ==.

fn ne(&self, other: &Rhs) -> bool

Tests for !=. The default implementation is almost always sufficient, and should not be overridden without very good reason.

impl Serialize for StageKind

fn serialize<__S>(&self, __serializer: __S) -> Result<__S::Ok, __S::Error>

where __S: Serializer,

Serialize this value into the given Serde serializer.

impl Eq for StageKind

impl StructuralPartialEq for StageKind