To Stream Data Layout Transformations

In the to_stream.py design we first bring 24xi32 data from external memory to L2 memory (i.e., a Mem tile) with of_in0. We then use of_in1 to forward the data from the MemTile to my_worker. Two FIFOs then move the output data from the Worker: - first to L2 via of_out1, applying a data layout transformation as the data is pushed onto the AXI stream by the Worker tile's DMA, - then to external memory via of_out0 as 24xi32 tensors. All FIFOs use double buffers.

# Dataflow with ObjectFifos
# Input
of_in0 = ObjectFifo(tile24_ty, name="in0")
of_in1 = of_in0.cons().forward(name="in1", obj_type=tile24_ty)

# Output
of_out1 = ObjectFifo(tile24_ty, name="out1", dims_to_stream=[(8, 1), (3, 8)])
of_out0 = of_out1.cons().forward(name="out0", obj_type=tile24_ty)

The process on the Worker acquires one object from of_in1 to consume and one object from of_out1 to produce into. It then reads the value of the input object and loads it into the output one before releasing both objects.

The data layout transformation dims_to_stream=[(8, 1), (3, 8)] expresses the access pattern in which the Worker will push the data from the 24xi32 tensor to the stream. This access pattern can also be expressed with for loops as follows:

for i in range(8):
    for j in range(3):
        # read data at index
        index = (i * 1 + j * 8)

If we imagine the 24-element wide tensor as 3 rows of 8 elements, the transformation above describes a column-major access pattern.

The design is wrapped in @iron.jit, so a single command JIT-compiles and runs it on the attached NPU:

make run                              # builds + runs on the NPU (devicename={npu,npu2})
make emit-mlir                        # writes the lowered MLIR to build/aie.mlir without touching the NPU

The # To/from AIE-array data movement section of the design code is described in detail in Section 2d.