Section 2c - Data Layout Transformations¶
While the ObjectFifo primitive aims to reduce the complexity tied to data movement configuration on the AI Engine array, it also gives the user control over some advanced features of the underlying architecture. One such feature is the ability to do data layout transformations on the fly using the tile's dedicated hardware: the Direct Memory Access channels (DMAs). This is available on AIE-ML devices.
Tile DMAs interact directly with the memory modules of their tiles and are responsible for pushing and retrieving data to and from the AXI stream interconnect. When data is pushed onto the stream, the user can program the DMA's n-dimensional address generation scheme such that the data's layout when pushed may be different than how it is stored in the tile's local memory. In the same way, a user can also specify in what layout a DMA should store the data retrieved from the AXI stream.
DMA blocks contain buffer descriptor operations that summarize what data is being moved, from what offset, how much of it, and in what layout. These buffer descriptors are the AIE_DMABDOp operations in MLIR and have their own auto-generated Python binding (available under <MLIR_AIE_INSTALL_PATH>/python/aie/dialects/_aie_ops_gen.py after the repository is built):
def dma_bd
(
buffer,
*,
offset=None,
len=None,
dimensions=None,
bd_id=None,
next_bd_id=None,
loc=None,
ip=None
)
A data layout transformation is presented as a tuple of pairs, where each pair represents a size and a stride for a particular dimension of the data:
<size_2, stride_2>, while the last pair of the array gives the inner-most dimension's stride and size <size_0, stride_0>. All strides are expressed in multiples of the element width.
NOTE: Only for 4B data types the inner-most dimension's stride must be 1 by design.
Data layout transformations can be viewed as a way to specify to the hardware which location in the data to access next and as such it is possible to model the access pattern using a series of nested loops. For example, the transformation using the strides and sizes from above can be expressed as:
int *buffer;
for(int i = 0; i < size_2; i++)
for(int j = 0; j < size_1; j++)
for(int k = 0; k < size_0; k++)
// access/store element at/to buffer[ i * stride_2
// + j * stride_1
// + k * stride_0]
It is important to note that data layout transformations are interpreted differently depending on whether data is pushed onto or read from the AXI stream:
- when data are pushed to the AXI stream, the layout describes from where in memory the DMA reads the elements to push to the stream (where to get items from'');
- when data are read from the AXI stream, the data layout describes where in memory the DMA writes the elements that are arriving over the stream in-sequence (where toput arriving items'').
As a practical example, here is an access pattern that corresponds to alternating between even and odd elements every 8 elements in a 128 element buffer/stream:
which translates to:for(int i = 0; i < 8; i++) // size_2
for(int j = 0; j < 2; j++) // size_1
for(int k = 0; k < 8; k++) // size_0
// access/store element at/to index:
(
i * 16 // stride_2
+ j * 1 // stride_1
+ k * 2 // stride_0
)
Data Layout Transformations with the ObjectFifo¶
Remember that the ObjectFifo class constructor has two default-valued inputs: dimensionsToStream and dimensionsFromStreamPerConsumer.
class object_fifo:
def __init__(
self,
name,
producerTile,
consumerTiles,
depth,
datatype,
dimensionsToStream=None,
dimensionsFromStreamPerConsumer=None,
)
Our compiler directly lowers ObjectFifos that make use of the aforementioned data layout transformations to AIE_DMABDOp. You can use the dimensionsToStream input to describe in which order the producerTile's DMA should push the objects onto the stream. Similarly, the dimensionsFromStreamPerConsumer input describes to the DMAs of each individual tile in the consumerTiles in what layout to retrieve the objects from the stream.
NOTE: Data layout transformations are applied to individual ObjectFifo objects and cannot act across object boundaries.
As an example, the ObjectFifo in the code below contains objects with datatype <4x8xi8>. Using the dimensionsToStream input it performs a data layout transformation on the producer tile side that pushes elements from memory onto the stream as follows: For every even length-8 row, select the first three even-indexed elements.
A = tile(1, 1)
B = tile(1, 3)
of0 = object_fifo
(
"objfifo0",
A,
B,
3,
np.ndarray[(4, 8), np.dtype[np.int8]],
[
(2, 16),
(3, 2),
],
)
for(int i = 0; i < 2; i++) // size_1
for(int j = 0; j < 3; j++) // size_0
// access/store element at/to index:
(
i * 16 // stride_1
+ j * 2 // stride_0
)
Please see the to_stream_transformations and from_stream_transformations designs for end-to-end IRON examples of each pattern.
Other examples containing data layout transformations are available in the programming_examples. A few notable ones are matrix_vector_multiplication and matrix_multiplication_whole_array.
When using the implicit copy feature of the ObjectFifo for a join or distribute data movement pattern, data layout transformations are not supported on the output of a join pattern or on the input of a distribute one. The reasoning behind this decision is largely due to the complexity of the DMA program that is required to achieve a data layout transformation across the larger data tensor while ensuring race-free execution of the DMA buffer descriptor logic. Users may however program the DMA buffer descriptors themselves for specific designs; this is further detailed in this discussion.
Data Layout Transformations with the Runtime Sequence¶
The runtime sequence uses another representation of data layout transformations named Tensor Access Pattern, or tap, which is available in the taplib library on AIEs with IRON. An in-depth introduction to taplib is available here.
Runtime sequence operations such as fill() and drain() in IRON, or dma_wait and npu_dma_memcpy_nd in the AIE dialect, can optionally take a tap as input to change the access pattern to/from external memory on-the-fly. For more details on programming the runtime sequence please see the corresponding section.
Examples containing taps are available in the programming_examples. A few notable ones are transposes (specifically --strategy=dma) and row_wise_bias_add.