filament

import "primitives/reshape.fil";

/// The abstract interface for all pipelined conv2D modules generated by Aetherling
comp Conv2dAbstract<'G: 1>(
    in[N]: ['G, 'G+1] 8,
) -> (
    out[N]: ['G+L, 'G+L+1] 8,
) with {
    /// The number of inputs the module accepts. This is always a factor of 16.
    some N where N > 0, N <= 16, 16 % N == 0;
    /// The latency of the module
    some L where L > 0;
} {
    // This is UNBOUND and we're using the fact that Filament has a global namespace to use whatever
    // implementation is available.
    C := new Conv2dWrapper;

    // Execute and forward
    c := C<'G>(in{0..C::N});
    out{0..C::N} = c.O{0..C::N};

    L := C::L;
    N := C::N;
}

/// The module that interfaces with the outside world which always provides all
/// 16 inputs at the same time and exactly for one cycle.
comp Wrapper<'G: II>(
    valid_up: interface['G],
    I[16]: ['G, 'G+1] 8
) -> (
    O[16]: ['G+L, 'G+L+1] 8
) with {
    /// The latency of the module
    some L where L > 0;
    /// The initiation interval of the module
    some II where II > 0;
} {
    // The convolution module.
    C := new Conv2dAbstract;
    let Chunks = 16 / C::N;

    // The serializer produces C::N chunks of inputs for the convolution module.
    ser := new Serialize[8, Chunks, C::N]<'G>(I{0..16});

    // Track the outputs of the convolution for deserialization.
    // The liveness this because we send N inputs every cycle and receive them back in L cycles.
    bundle co[Chunks][C::N]: for<i> ['G+i+C::L, 'G+i+C::L+1] 8;

    // Send the chunks to the conv and collect the outputs
    for i in 0..Chunks {
        // The start time of the chunk inputs
        let Base = i * C::N;
        // Send C::N inputs to the Convolution
        c := C<'G+i>(ser.out{i}{0..C::N});
        // Assign the outputs to the bundle
        co{i}{0..C::N} = c.out{0..C::N};
    }

    if Chunks > 1 {
        // The deserializer can start collecting the outputs after the latency of the convolution elapses.
        // It will collect C::N outputs every cycle and produce them at the end.
        de := new Deserialize[8, Chunks, C::N]<'G+C::L>(co{0..Chunks}{0..C::N});
        O{0..16} = de.out{0..16};
    } else {
        // No need to deserialize if there is only one chunk.
        O{0..16} = co{0..Chunks}{0..C::N};
    }

    // The output latency is determined by the Deserializer.
    L := C::L+16/C::N-1;

    // The II is limited by the speed at which the convolution can process the inputs.
    II := 16/C::N;
}

// The main component has the same interface for each design
comp main<'G: II>(
    go: interface['G],
    I_0: ['G, 'G+1] 8,
    I_1: ['G, 'G+1] 8,
    I_2: ['G, 'G+1] 8,
    I_3: ['G, 'G+1] 8,
    I_4: ['G, 'G+1] 8,
    I_5: ['G, 'G+1] 8,
    I_6: ['G, 'G+1] 8,
    I_7: ['G, 'G+1] 8,
    I_8: ['G, 'G+1] 8,
    I_9: ['G, 'G+1] 8,
    I_10: ['G, 'G+1] 8,
    I_11: ['G, 'G+1] 8,
    I_12: ['G, 'G+1] 8,
    I_13: ['G, 'G+1] 8,
    I_14: ['G, 'G+1] 8,
    I_15: ['G, 'G+1] 8,
) -> (
    O_0: ['G+L, 'G+L+1] 8,
    O_1: ['G+L, 'G+L+1] 8,
    O_2: ['G+L, 'G+L+1] 8,
    O_3: ['G+L, 'G+L+1] 8,
    O_4: ['G+L, 'G+L+1] 8,
    O_5: ['G+L, 'G+L+1] 8,
    O_6: ['G+L, 'G+L+1] 8,
    O_7: ['G+L, 'G+L+1] 8,
    O_8: ['G+L, 'G+L+1] 8,
    O_9: ['G+L, 'G+L+1] 8,
    O_10: ['G+L, 'G+L+1] 8,
    O_11: ['G+L, 'G+L+1] 8,
    O_12: ['G+L, 'G+L+1] 8,
    O_13: ['G+L, 'G+L+1] 8,
    O_14: ['G+L, 'G+L+1] 8,
    O_15: ['G+L, 'G+L+1] 8,
) with {
    // The initiation interval of the module
    some II where II > 0;
    // The latency of the module
    some L where L > 0;
} {
    // Initialize the input bundle
    bundle I[16]: ['G, 'G+1] 8;
    I{0} = I_0; I{1} = I_1; I{2} = I_2; I{3} = I_3;
    I{4} = I_4; I{5} = I_5; I{6} = I_6; I{7} = I_7;
    I{8} = I_8; I{9} = I_9; I{10} = I_10; I{11} = I_11;
    I{12} = I_12; I{13} = I_13; I{14} = I_14; I{15} = I_15;

    // The wrapper that interfaces with the outside world
    w := new Wrapper<'G>(I{0..16});

    // Connect the ouputs
    O_0 = w.O{0}; O_1 = w.O{1}; O_2 = w.O{2}; O_3 = w.O{3};
    O_4 = w.O{4}; O_5 = w.O{5}; O_6 = w.O{6}; O_7 = w.O{7};
    O_8 = w.O{8}; O_9 = w.O{9}; O_10 = w.O{10}; O_11 = w.O{11};
    O_12 = w.O{12}; O_13 = w.O{13}; O_14 = w.O{14}; O_15 = w.O{15};

    L := W::L;
    II := W::II;
}