tract-core 0.23.9

//! Give a model's dataflow a batch axis, on axis 0 of every input and output, so
//! that a turn can seat several callers in one run.
//!
//! Every op the axis reaches is asked to host it, through the same
//! `TypedOp::change_axes` protocol declutter propagates axis changes with, and
//! the answer is accepted only if the op **carries the axis through** to all of
//! its outputs. That is a stricter contract than the protocol's own, which only
//! owes semantics at extent 1: an `EinSum` offered an extra leading axis on an
//! input takes it as a fresh contraction label, summing an axis of extent one
//! away, which is exact there and sums the seats together once the extent is the
//! batch. So the batched `EinSum` is built here instead, with the label on the
//! output as well, and an op that answers with anything but the axis carried
//! through is refused by name.
//!
//! An input the caller lists as **shared** is one every seat reads whole -- a
//! table, a language id. It keeps its own shape, and the ops that need their
//! operands to agree on rank read it through an `AxisOp::Add(0)`, so it
//! broadcasts across the turn.
//!
//! An input already sized by the batch symbol has that axis **moved** to the
//! front, as an interface change plus the inverse `AxisOp::Move` on the wire
//! behind it: an export putting the batch inside (`[T, B, H]` being the common
//! shape) becomes a graph edit declutter can absorb rather than a transpose the
//! runtime would pay per turn. Only axis 0 makes a seat's values a contiguous
//! run, which is what the laned runtime addresses and what stacking a turn's
//! answer back relies on.
//!
//! What this does not touch is state: a seat's own history is the laned
//! runtime's lane addressing, and a batch axis on the dataflow does not imply
//! one. Batchify gives you seats, lane addressing gives each seat its own past.

use std::collections::HashSet;

use crate::internal::*;
use crate::ops::array::{MultiBroadcastTo, TypedConcat};
use crate::ops::change_axes::AxisOp;
use crate::ops::einsum::EinSum;
use crate::ops::source::TypedSource;
use crate::transform::ModelTransform;

#[derive(Debug, Default, serde::Deserialize)]
pub struct BatchifyConfig {
    /// Symbol sizing the batch axis. Defaults to "B".
    pub symbol: Option<String>,
    /// Inputs every seat shares, by node name.
    #[serde(default)]
    pub shared: Option<Vec<String>>,
}

#[derive(Debug)]
pub struct Batchify(pub BatchifyConfig);

impl ModelTransform for Batchify {
    fn name(&self) -> StaticName {
        "batchify".into()
    }

    fn transform(&self, model: &mut TypedModel) -> TractResult<()> {
        let name = self.0.symbol.as_deref().unwrap_or("B");
        let symbol = model.symbols.sym(name);
        let shared = self.0.shared.as_deref().unwrap_or(&[]);
        *model = batchify(model, &symbol, shared)?;
        Ok(())
    }
}

/// What an op does with the batch axis arriving on some of its inputs: the op to
/// wire in its place, and the input slots that must read through an
/// `AxisOp::Add(0)` to agree on rank with the batched ones.
struct Hosted {
    op: Option<Box<dyn TypedOp>>,
    pad: TVec<usize>,
    /// Whether a padded slot must also be broadcast to the batch: an op whose
    /// operands have to agree axis by axis, rather than broadcast against each
    /// other, reads an axis of extent one as an axis of extent one.
    broadcast: bool,
}

/// Batch axis on axis 0 of every input but the `shared` ones, sized by `batch`.
pub fn batchify(model: &TypedModel, batch: &Symbol, shared: &[String]) -> TractResult<TypedModel> {
    for name in shared {
        ensure!(
            model.input_outlets()?.iter().any(|o| &model.node(o.node).name == name),
            "{name} is listed as a shared input but is not a model input"
        );
    }
    // An export which kept the batch on one input and collapsed it to a leading
    // one on another -- an RNNT joint, whose decoder side carries the beam it is
    // batched over -- has the axis already, and the collapsed inputs join it
    // rather than introduce a second batch of their own.
    let mut join = false;
    for outlet in model.input_outlets()? {
        let name = &model.node(outlet.node).name;
        join |= !shared.contains(name) && carrier(model.outlet_fact(*outlet)?, batch).is_some();
    }
    let mut carried: HashSet<OutletId> = Default::default();
    let mut hosts: HashMap<usize, Hosted> = Default::default();
    for id in model.eval_order()? {
        let node = model.node(id);
        if let Some(source) = node.op_as::<TypedSource>() {
            if !shared.contains(&node.name)
                && carrier(&source.fact, batch).is_none()
                && !collapsed(&source.fact, batch, join)
            {
                carried.insert(id.into());
            }
            continue;
        }
        let batched: TVec<usize> = node
            .inputs
            .iter()
            .enumerate()
            .filter(|(_, input)| carried.contains(input))
            .map(|(slot, _)| slot)
            .collect();
        if batched.is_empty() {
            continue;
        }
        let mut hosted = host(model, node, &batched, &carried)
            .with_context(|| format!("Giving {node} a batch axis"))?;
        hosted.broadcast = node.op_is::<TypedConcat>();
        hosts.insert(id, hosted);
        carried.extend((0..node.outputs.len()).map(|slot| OutletId::new(id, slot)));
    }
    let wired = wire(model, batch, shared, join, &hosts)?;
    check(wired, batch)
}

/// The batch axis an input carries, if any: the first axis sized by an
/// expression the batch symbol is in.
fn carrier(fact: &TypedFact, batch: &Symbol) -> Option<usize> {
    fact.shape.iter().position(|dim| dim.symbols().contains(batch))
}

/// Whether an input's batch is the leading one an export collapsed it to: the
/// model carries the axis elsewhere, this input has no dim sized by it, and its
/// first axis is one. Such an input is resized rather than reshaped, so the
/// graph behind it hosts nothing.
fn collapsed(fact: &TypedFact, batch: &Symbol, join: bool) -> bool {
    join && carrier(fact, batch).is_none() && fact.shape.first() == Some(&1.to_dim())
}

/// Asks an op to host the batch axis on the inputs that carry it, and refuses an
/// answer that does not carry it through to every output.
fn host(
    model: &TypedModel,
    node: &TypedNode,
    batched: &[usize],
    carried: &HashSet<OutletId>,
) -> TractResult<Hosted> {
    if let Some(einsum) = node.op_as::<EinSum>() {
        return batched_einsum(einsum, batched);
    }
    if let Some(axis_op) = node.op_as::<AxisOp>() {
        return Ok(Hosted {
            op: Some(Box::new(shifted_axis_op(axis_op))),
            pad: tvec!(),
            broadcast: false,
        });
    }
    let change = AxisOp::Add(0);
    let consequence = node
        .op
        .change_axes(model, node, InOut::In(batched[0]), &change)?
        .context("the op takes no extra leading axis at all")?;
    for slot in 0..node.outputs.len() {
        ensure!(
            consequence.wire_changes.contains(&(InOut::Out(slot), change.clone())),
            "the op takes an extra leading axis on an input but does not carry it to output {slot}, \
             so the seats would be folded into one another"
        );
    }
    let pad = consequence
        .wire_changes
        .iter()
        .filter_map(|(io, _)| match io {
            InOut::In(slot) if !carried.contains(&node.inputs[*slot]) => Some(*slot),
            _ => None,
        })
        .collect();
    Ok(Hosted { op: consequence.substitute_op, pad, broadcast: false })
}

/// An `EinSum` batched by one label of its own, on axis 0 of the batched inputs
/// and of every output. Its `change_axes` would rather contract the axis away,
/// which is the same thing only while the extent is one.
fn batched_einsum(einsum: &EinSum, batched: &[usize]) -> TractResult<Hosted> {
    let label = einsum.axes.available_label();
    let mut axes = einsum.axes.clone().with_extra_axis(label, InOut::In(batched[0]), 0)?;
    for slot in &batched[1..] {
        axes = axes.with_extra_axis_occurency(label, InOut::In(*slot), 0)?;
    }
    for slot in 0..axes.output_count() {
        axes = axes.with_extra_axis_occurency(label, InOut::Out(slot), 0)?;
    }
    let op = EinSum { axes, ..einsum.clone() };
    Ok(Hosted { op: Some(Box::new(op)), pad: tvec!(), broadcast: false })
}

/// An `AxisOp` one axis further up, which is what hosting the batch on axis 0
/// asks of it: an export collapsed to one stream reaches the graph through an
/// `Add(0)` putting the batch back, and the ops the change_axes protocol
/// answers for would carry the seats to axis 1 instead, folding them into the
/// axis the collapse left behind.
fn shifted_axis_op(op: &AxisOp) -> AxisOp {
    match op {
        AxisOp::Add(a) => AxisOp::Add(a + 1),
        AxisOp::Rm(a) => AxisOp::Rm(a + 1),
        AxisOp::Move(from, to) => AxisOp::Move(from + 1, to + 1),
        AxisOp::Reshape(at, from, to) => AxisOp::Reshape(at + 1, from.clone(), to.clone()),
    }
}

/// The batchified model, rebuilt through its ops so that every fact comes from
/// Moves a `pulse.*_axes` property right by one wherever the interface it
/// indexes gained the batch axis. A negative entry marks an interface with no
/// streaming axis and stays as it is.
fn shift_pulse_axes(model: &mut TypedModel, key: &str, widened: &[bool]) -> TractResult<()> {
    let Some(axes) = model.properties.get(key) else { return Ok(()) };
    let mut axes: Vec<i64> = axes.cast_to::<i64>()?.try_as_plain_ram()?.as_slice::<i64>()?.to_vec();
    ensure!(
        axes.len() == widened.len(),
        "{key} holds {} axes for {} interfaces",
        axes.len(),
        widened.len()
    );
    for (axis, widened) in axes.iter_mut().zip(widened) {
        if *widened && *axis >= 0 {
            *axis += 1;
        }
    }
    model.properties.insert(key.to_string(), tensor1(&axes).into_arc_tensor());
    Ok(())
}

/// `output_facts` with the batch extent in place rather than from the facts the
/// unbatched model stored.
fn wire(
    model: &TypedModel,
    batch: &Symbol,
    shared: &[String],
    join: bool,
    hosts: &HashMap<usize, Hosted>,
) -> TractResult<TypedModel> {
    let mut target = TypedModel { symbols: model.symbols.clone(), ..TypedModel::default() };
    let mut mapping: HashMap<OutletId, OutletId> = Default::default();
    let mut interface: HashMap<OutletId, OutletId> = Default::default();
    for id in model.eval_order()? {
        let node = model.node(id);
        let wires = if let Some(source) = node.op_as::<TypedSource>() {
            let (source, interior) = batched_source(
                &mut target,
                node,
                source,
                batch,
                !shared.contains(&node.name),
                join,
            )?;
            interface.insert(id.into(), source);
            tvec!(interior)
        } else {
            let hosted = hosts.get(&id);
            let mut inputs: TVec<OutletId> =
                node.inputs.iter().map(|input| mapping[input]).collect();
            for slot in hosted.map(|h| &*h.pad).unwrap_or(&[]) {
                let mut wire = target.wire_node(
                    format!("{}.batchify.rank.{slot}", node.name),
                    AxisOp::Add(0),
                    &[inputs[*slot]],
                )?[0];
                if hosted.is_some_and(|h| h.broadcast) {
                    let mut shape: TVec<TDim> = target.outlet_fact(wire)?.shape.to_tvec();
                    shape[0] = batch.to_dim();
                    wire = target.wire_node(
                        format!("{}.batchify.broadcast.{slot}", node.name),
                        MultiBroadcastTo { shape: shape.into() },
                        &[wire],
                    )?[0];
                }
                inputs[*slot] = wire;
            }
            let op = hosted.and_then(|h| h.op.clone()).unwrap_or_else(|| node.op.clone());
            target
                .wire_node(&node.name, op, &inputs)
                .with_context(|| format!("Wiring {node} with a batch axis"))?
        };
        for (slot, wire) in wires.into_iter().enumerate() {
            let outlet = OutletId::new(id, slot);
            if let Some(label) = model.outlet_label(outlet) {
                target.set_outlet_label(wire, label.to_string())?;
            }
            mapping.insert(outlet, wire);
        }
    }
    let inputs: Vec<OutletId> = model.input_outlets()?.iter().map(|i| interface[i]).collect();
    let mut outputs: Vec<OutletId> = model.output_outlets()?.iter().map(|o| mapping[o]).collect();
    for (ix, output) in outputs.iter_mut().enumerate() {
        let fact = target.outlet_fact(*output)?.clone();
        if let Some(axis) = carrier(&fact, batch).filter(|axis| *axis > 0) {
            let label = target.outlet_label(*output).map(|l| l.to_string());
            *output = target.wire_node(
                format!("batchify.move.output.{ix}"),
                AxisOp::Move(axis, 0),
                &[*output],
            )?[0];
            if let Some(label) = label {
                target.set_outlet_label(*output, label)?;
            }
        }
    }
    target.set_input_outlets(&inputs)?;
    target.select_output_outlets(&outputs)?;
    // A batched model is the same model, so it keeps what the graph says about
    // itself -- its name, and the pulse properties a streaming caller reads to
    // find the delay and the streaming axes. The axes among them are indices
    // into shapes this transform has just widened.
    target.properties.clone_from(&model.properties);
    if !join {
        let batched_inputs: Vec<bool> = model
            .input_outlets()?
            .iter()
            .map(|o| !shared.contains(&model.node(o.node).name))
            .collect();
        shift_pulse_axes(&mut target, "pulse.input_axes", &batched_inputs)?;
        let batched_outputs = vec![true; model.output_outlets()?.len()];
        shift_pulse_axes(&mut target, "pulse.output_axes", &batched_outputs)?;
    }
    Ok(target)
}

/// A source of the batchified model: a batched one carries the batch on axis 0,
/// by gaining the axis or by having the one it already carries moved there, and
/// the wire behind it moves it back so that the graph sees the layout it was
/// built for.
fn batched_source(
    target: &mut TypedModel,
    node: &TypedNode,
    source: &TypedSource,
    batch: &Symbol,
    batched: bool,
    join: bool,
) -> TractResult<(OutletId, OutletId)> {
    let carrier = carrier(&source.fact, batch);
    if !batched || carrier == Some(0) {
        let wire = target.wire_node(&node.name, source.clone(), &[])?[0];
        return Ok((wire, wire));
    }
    if collapsed(&source.fact, batch, join) {
        // The seats go where the collapse left its one, so the interface and the
        // interior are the same wire, one axis wider.
        let mut fact = source.fact.clone();
        let mut shape = fact.shape.to_tvec();
        shape[0] = batch.to_dim();
        fact.shape = shape.into();
        let wire = target.wire_node(&node.name, TypedSource::new(fact), &[])?[0];
        return Ok((wire, wire));
    }
    let mut shape = source.fact.shape.to_tvec();
    let mut restore = None;
    match carrier {
        Some(axis) => {
            let dim = shape.remove(axis);
            shape.insert(0, dim);
            restore = Some(AxisOp::Move(0, axis));
        }
        None => shape.insert(0, batch.to_dim()),
    }
    let mut fact = source.fact.clone();
    fact.shape = shape.into();
    let wire = target.wire_node(&node.name, TypedSource::new(fact), &[])?[0];
    match restore {
        Some(restore) => {
            let interior =
                target.wire_node(format!("{}.batchify.move", node.name), restore, &[wire])?[0];
            Ok((wire, interior))
        }
        None => Ok((wire, wire)),
    }
}

/// Fails on a model the laned runtime would refuse anyway, naming what it found:
/// a model answering nothing per seat has no turn to run.
fn check(model: TypedModel, batch: &Symbol) -> TractResult<TypedModel> {
    let batch = batch.to_dim();
    let mut shared = vec![];
    for (ix, outlet) in model.output_outlets()?.iter().enumerate() {
        let fact = model.outlet_fact(*outlet)?;
        if fact.shape.first() == Some(&batch) {
            return Ok(model);
        }
        shared.push(format!("output {ix} is {:?}", fact.shape));
    }
    bail!("No output carries the batch axis on axis 0: {}", shared.join(", "))
}

register_model_transform!("batchify", BatchifyConfig, |config| Ok(Box::new(Batchify(config))));