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.

Prerequisites: Run the following to enter the development environment and install torch.

nix develop nix-amd-npu#iron-full
iron-fhs
pip install torch                  # only once

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.5 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()

  2. aie.core()

  3. aie.runtime_sequence()

  4. aie.objectfifo()

  5. aie.objectfifo.link()

  6. aie.objectfifo.acquire() and aie.objectfifo.release()

  7. aiex.npu.dma_memcpy_nd()

  8. aiex.npu.dma_wait()

Task 2 - Data Layouts and Loops

Sketch the data movement between main memory, shim tile, memory tile, and compute tile. Describe which mlir-aie operation is involved in each step.

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.