Function fifo_buffer

python/tvm/topi/nn/fifo_buffer.py:29–188

FIFO buffer to enable computation reuse in CNNs with sliding indow input Compute equivalent of .. code-block:: python concat(buffer, data, axis=axis) .slice_axis(axis=axis, begin=data.shape[axis], end=data.shape[axis]+buffer.sha

(data, buffer, axis)

@tvm.te.tag_scope(tag=tag.INJECTIVE + ",fifo_buffer")
def fifo_buffer(data, buffer, axis):
    """
    FIFO buffer to enable computation reuse in CNNs with sliding indow input

    Compute equivalent of

    .. code-block:: python

        concat(buffer, data, axis=axis)
        .slice_axis(axis=axis,
                    begin=data.shape[axis],
                    end=data.shape[axis]+buffer.shape[axis])

    Useful for

    * Encoding explicit re-use of computation in convolution ops operated on a sliding window input
    * Implementing a FIFO queue to cache intermediate results, e.g. as in Fast WaveNet.

    Parameters
    ----------
    data : tvm.te.Tensor
        The input data
    buffer : tvm.te.Tensor
        Previous value of the FIFO buffer
    axis : int
        Specify which axis should be used for buffering

    Returns
    -------
    result : tvm.te.Tensor
        Updated value for the buffer
    """
    assert len(data.shape) == len(buffer.shape), (
        f"buffer and data must have same number of dimensions, "
        f"buffer.shape = {buffer.shape}, data.shape = {data.shape}"
    )
    assert len(buffer.shape) >= 1, "Zero-dimension tensor not supported"
    assert 0 <= axis < len(buffer.shape), "buffer axis out of range"
    for i in range(len(data.shape)):
        if i == axis:
            assert int(str(data.shape[i])) <= int(str(buffer.shape[i]))
        else:
            assert int(str(data.shape[i])) == int(str(buffer.shape[i]))

    buflen = buffer.shape[axis]
    data_size = data.shape[axis]

    # Explicitly write out formula up to 4D, and then use concat+slice combo for 5D and higher
    if len(buffer.shape) == 1:
        return te.compute(
            buffer.shape,
            lambda i: tvm.tirx.if_then_else(
                i < buflen - data_size, buffer[i + data_size], data[i - buflen + data_size]
            ),
            name="new_buffer",
        )
    if len(buffer.shape) == 2:
        if axis == 0: