Assignment 09: XDNA GEMM

This week you will perform a larger matrix multiplication on the NPU. For this you will write data-movement code in the MLIR-AIE dialect and add loops around your XDNA tensor kernel.

Data Layout and Data Movement

In the main memory, the matrices are stored in row-major order (in0: MK, in1: KN, and out: MN with M=256, N=128, and K=1024). During the data movement from L3 (main memory) to L1 (scratchpad), the matrices are tiled. The dimensions are split as follows:

  • M->apm with a=16, p=2, m=8,

  • N->bqn with b=8, q=2, n=8, and

  • K->crk with c=16, r=8, k=8.

This yields the views in0: apmcrk, in1: crkbqn, and out: apmbqn. During the data movement to the L1 scratchpad memory, the layout must be changed to in0: prmk, in1: rqkn, and out: pqmn. The dimensions a, b, and c are handled sequentially through loops on the compute tile. The DMAs move the corresponding tiles to the compute tile. Before the c loop, the output tile is zero-initialized. When writing the output tensor, its layout is changed back to a matrix layout (out: MN).

Task 0 - Setup

Copy your XDNA tensor kernel into the src/ directory and copy your verify() function to the driver. Set the maximum absolute error to 2 and the maximum relative error to 0.5.

Task 1 - MLIR-AIE operations

Give a brief summary of the following mlir-aie operations:

  1. aie.tile(): This operation creates an AIE tile in the AIE array

  2. aie.core(): This operation represents an AIEngine processor core belonging to a tile

  3. aie.runtime_sequence(): Instructions in this operation allow for runtime (re-)configuration of the AI Engine array, such as configuring data movement buffer descriptors. (e.g. data transfers between host and AIE array on the shims)

  4. aie.objectfifo(): Creates a circular buffer or channel between two tiles

  5. aie.objectfifo.link(): Links two objectFifos through an intermediary tile’s DMA (Direct Memory Access)

  6. aie.objectfifo.acquire() and aie.objectfifo.release(): Acquire/Release operation to lock and return objects of an ObjectFifo

  7. aiex.npu.dma_memcpy_nd(): An n-dimensional half DMA operator. Programs a DMA to access a memory memref with an access pattern specified by offsets, sizes and strides or static_offsets, static_sizes and static_strides

  8. aiex.npu.dma_wait(): blocks until the DMA referenced through symbol completes execution and issues a task-complete-token. e.g. aiex.npu.dma_wait { symbol = @out0 }

Task 2 - Data Layouts and Loops

Sketch the data movement between main memory, shim tile, memory tile, and compute tile.

  • data will be loaded from main memory to the shim tile, then to the memory tile, and finally to the compute tile where the matrix multiplication will be performed. The output will be moved back in the reverse order. Describe which mlir-aie operation is involved in each step.

  • aie.objectfifo will create a channel from shim to memory tile and a channel from memory tile to compute tile. both will be linked by aie.objectfifo.link

  • aiex.npu.dma_memcpy_nd() maps main memory to shim and aiex.npu.dma_wait will wait until writing is finished

Task 3 - Implementation

  1. Implement the data movement inside src/matmul.mlir and change the dimension sizes in src/driver.py. Replace the TODOs with the corresponding code.

  2. Verify your implementation by executing make run_matmul.

Task 4 - Performance

Change the data movement inside the MLIR code so that there is no blocking wait, i.e., there is always a data movement operation that can be issued (except for the last one).

Task 5 - Buffer Placement (optional)

Find the buffer placement operations inside the lowered MLIR code (build/matmul.mlir.prj/input_with_addresses.mlir). Describe how you would place the buffers to reduce bank conflicts and what changes would be needed to your XDNA tensor kernel.