Assignment 10: Using the whole NPU

This week you will perform the matrix multiplication on the whole NPU. For this, you will adapt the data-movement code and execute your XDNA tensor kernel on all compute tiles.

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 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 first tiled and then broadcast along the columns or rows. The dimensions are split as follows:

  • M->axpm with a=2, x=8, p=2, m=8,

  • N->byqn with b=2, y=4, q=2, n=8, and

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

This yields the views in0: axpmcrk, in1: crkbyqn, and out: axpmbyqn. 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 tiles. The dimensions x and y are handled spatially by distributing them across the compute-tile columns and rows, respectively. The DMAs move the corresponding tiles from the memory tile to the compute tiles, broadcasting the in0 tiles along the columns and the in1 tiles along the rows, so that each column receives a different in0 tile and each row receives a different in1 tile. Before the c loop, the output tile is zero-initialized. When writing the out tiles, the four out tiles produced by the compute tiles of one column are joined into the intermediate layout ypqmn. The dimension y is realized by giving each row’s L1L2 FIFO a different write offset into the joined L2 buffer. Note that the layout is changed to ypmqn when reading the data from L2 to the stream. When writing the output tiles from the memory tile to the main memory, the layout is changed to a matrix layout (out: MN).

Task 1 - Setup of the Whole NPU

Create variables for each tile (shim, memory, and compute). Adapt the fused ab-loop size and duplicate the core function for each compute tile. Set the FIFO queue suffixes inside each core function so that they match its column and row: replace the placeholder suffixes _0 (for in0/in1) and _0_0 (for out) with the actual tile coordinates, yielding in0_<col>, in1_<row>, and out_<col>_<row>. Row indices in the FIFO suffixes are zero-based and counted from the first compute-tile row (so aie.tile(_, 2) → row index 0, aie.tile(_, 3) → row index 1, etc.). For example, for aie.tile(7, 3) the suffixes are @in0_L2L1_7, @in1_L2L1_1, and @out_L1L2_7_1.

Task 2 - Broadcasting the Inputs

Extend the existing L2L1 input FIFO queues by adding the additional consumer tiles along row 0 (in0) or column 0 (in1). Create the needed L2L1 and L3L2 input FIFO queues for the remaining columns and rows. Note that only four in1_L3L2 FIFO queues are needed. Their placement on the shim tiles can be freely chosen. Adapt the input dma_memcpy_nd operations to their corresponding FIFO queue. Note that each shim tile has 16 buffer descriptors; consequently, FIFO queues on different shim tiles may reuse the same BD id.

Task 3 - Writing the Output

Create the needed L1L2 and L2L3 output FIFO queues and join the L1L2 FIFO queues along a column to the corresponding L2L3 FIFO queue. Adapt the output dma_memcpy_nd operations to their corresponding FIFO queue and change their sizes, strides, and offsets. Note that the dma_memcpy_nd operations in each shim tile receive the output tensor tiles in the layout aby(p*m)(q*n), and x is handled spatially across the columns and therefore does not appear in this view. The dimensions b and y could be fused in this view, but we keep them separate.

Task 4 - Testing

Test your implementation by running make run_matmul.

Task 5 - Selection of Spatial Dimensions (optional)

Explain how the data transfer would change if the roles of the dimensions y and x were switched, i.e., the dimensions were split as follows:

  • M->aypm with a=4, y=4, p=2, m=8,

  • N->bxqn with b=1, x=8, q=2, n=8, and

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

In which cases should in0 be broadcast along the rows, and in which cases should in1 be broadcast along the rows? When is no difference in performance expected?

Group Specific Component

At the end of the course, your group will work on a project of your own choosing.

Prepare a pitch for at least two possible project ideas: one targeting the XDNA NPU and one targeting GPUs with cuTile. Create at least two slides per project idea, covering

  • a brief introduction,

  • the problem formulation,

  • the proposed solution, and

  • the expected results or insights.

You will present both project ideas on 17.06.2026, with about 5 minutes for each idea.

You do not have to hand in the slides. Hand in a short description (2-3 sentences) about each project idea.