1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

use crate::internal::*;
use tract_core::ops::binary::*;
use tract_core::ops::logic::Iff;
use tract_pulse_opl::ops::Delay;

register_all!(UnaryOp: pulsify_un, TypedBinOp: pulsify_bin, Iff: pulsify_iff);

fn sync_inputs(
    node: &TypedNode,
    target: &mut PulsedModel,
    mapping: &HashMap<OutletId, OutletId>,
) -> TractResult<TVec<OutletId>> {
    let delay = node
        .inputs
        .iter()
        .map(|input| target.outlet_fact(mapping[input]).unwrap().delay)
        .max()
        .unwrap();
    let mut inputs = tvec!();
    for input in &node.inputs {
        let mut input = mapping[input];
        let fact = target.outlet_fact(input)?.clone();
        if fact.delay < delay {
            let add_delay = delay - fact.delay;
            input = target.wire_node(
                format!("{}.Delay", &*node.name),
                Delay::new(fact.axis, &fact.into(), add_delay, 0),
                &[input],
            )?[0];
        }
        inputs.push(input);
    }
    Ok(inputs)
}

fn pulsify_bin(
    op: &TypedBinOp,
    _source: &TypedModel,
    node: &TypedNode,
    target: &mut PulsedModel,
    mapping: &HashMap<OutletId, OutletId>,
    _pulse: usize,
) -> TractResult<TVec<OutletId>> {
    let inputs = &*sync_inputs(node, target, mapping)?;
    target.wire_node(&*node.name, op.clone(), &inputs)
}

impl PulsedOp for TypedBinOp {
    fn pulsed_output_facts(&self, inputs: &[&PulsedFact]) -> TractResult<TVec<PulsedFact>> {
        let mut fact = inputs[0].clone();
        fact.datum_type = self.0.result_datum_type(inputs[0].datum_type, inputs[1].datum_type)?;
        fact.shape = tract_core::broadcast::multi_broadcast(&[&inputs[0].shape, &inputs[1].shape])
            .ok_or_else(|| {
                format_err!("Can not broadcast: {:?} and {:?}", inputs[0].shape, inputs[1].shape)
            })?;
        Ok(tvec!(fact))
    }

    as_op!();
    pulsed_op_to_typed_op!();
}

fn pulsify_un(
    op: &UnaryOp,
    _source: &TypedModel,
    node: &TypedNode,
    target: &mut PulsedModel,
    mapping: &HashMap<OutletId, OutletId>,
    _pulse: usize,
) -> TractResult<TVec<OutletId>> {
    let input = mapping[&node.inputs[0]];
    target.wire_node(&*node.name, op.clone(), &[input])
}

impl PulsedOp for UnaryOp {
    fn pulsed_output_facts(&self, inputs: &[&PulsedFact]) -> TractResult<TVec<PulsedFact>> {
        let mut fact = inputs[0].clone();
        fact.datum_type = self.mini_op.result_datum_type(self.a.datum_type(), fact.datum_type)?;
        fact.shape = tract_core::broadcast::multi_broadcast(&[
            &inputs[0].shape,
            &self.a.shape().iter().map(|d| d.into()).collect(),
        ])
        .unwrap();
        Ok(tvec!(fact))
    }

    as_op!();
    pulsed_op_to_typed_op!();
}

fn pulsify_iff(
    op: &Iff,
    _source: &TypedModel,
    node: &TypedNode,
    target: &mut PulsedModel,
    mapping: &HashMap<OutletId, OutletId>,
    _pulse: usize,
) -> TractResult<TVec<OutletId>> {
    let inputs = &*sync_inputs(node, target, mapping)?;
    target.wire_node(&*node.name, op.clone(), &inputs)
}

impl PulsedOp for Iff {
    fn pulsed_output_facts(&self, inputs: &[&PulsedFact]) -> TractResult<TVec<PulsedFact>> {
        let mut fact = inputs[0].clone();
        fact.datum_type = inputs[1].datum_type;
        fact.shape = tract_core::broadcast::multi_broadcast(&[
            &inputs[0].shape,
            &inputs[1].shape,
            &inputs[2].shape,
        ])
        .ok_or_else(|| {
            format_err!(
                "Can not broadcast: {:?}, {:?}, {:?}",
                inputs[0].shape,
                inputs[1].shape,
                inputs[2].shape
            )
        })?;
        Ok(tvec!(fact))
    }

    as_op!();
    pulsed_op_to_typed_op!();
}