Getting Started with MLIR-AIE

To enable the experimental MLIR-AIE codegen, specify "aie" as the target in the dataflow.build function.

Currently, the supported target platforms include XDNA1 and XDNA2. By default, the target platform is set to XDNA1. To switch to XDNA2, please run:

export NPU2=1

Example

Vector addition

import os
import allo
from allo.ir.types import int32
import allo.dataflow as df
import numpy as np

def test_vector_scalar_add():
    # https://github.com/Xilinx/mlir-aie/tree/main/programming_examples/basic/vector_scalar_add
    Ty = int32
    M = 1024

    @df.region()
    def top(A: Ty[M], B: Ty[M]):
         @df.kernel(mapping=[1], args=[A, B])
         def core(local_A: Ty[M], local_B: Ty[M]):
             local_B[:] = allo.add(local_A, 1)

    A = np.random.randint(0, 100, M).astype(np.int32)
    if "MLIR_AIE_INSTALL_DIR" in os.environ:
        mod = df.build(top, target="aie")
        B = np.zeros(M).astype(np.int32)
        mod(A, B)
        np.testing.assert_allclose(B, A + 1)
        print("PASSED!")
    else:
        print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")

Matrix multiplication

import allo
from allo.ir.types import int32
import allo.dataflow as df
import numpy as np
from allo.memory import Layout

 S = Layout.Shard
 R = Layout.Replicate

def _test_gemm_1D():
    Ty = int32
    M, N, K = 16, 16, 16
    P0 = 2
    LyA = [S(0), R]

    @df.region()
    def top(A: Ty[M, K], B: Ty[K, N], C: Ty[M, N]):
         @df.kernel(mapping=[P0], args=[A, B, C])
         def gemm(local_A: Ty[M, K] @ LyA, local_B: Ty[K, N], local_C: Ty[M, N] @ LyA):
             local_C[:, :] = allo.matmul(local_A, local_B)

    mod = df.build(top, target="aie")
    A = np.random.randint(0, 64, (M, K)).astype(np.int32)
    B = np.random.randint(0, 64, (K, N)).astype(np.int32)
    C = np.zeros((M, N)).astype(np.int32)
    mod(A, B, C)
    np.testing.assert_allclose(C, A @ B, atol=1e-5)
    print("PASSED!")

Producer-consumer

import os
import allo
from allo.ir.types import int32, Stream
import allo.dataflow as df
import numpy as np

Ty = int32
M, N, K = 16, 16, 16


@df.region()
def top(A: Ty[M, N], B: Ty[M, N]):
    pipe: Stream[Ty, 4]

     @df.kernel(mapping=[1], args=[A])
     def producer(local_A: Ty[M, N]):
         for i, j in allo.grid(M, N):
             # load data
             out: Ty = local_A[i, j]
             # send data
             pipe.put(out)

     @df.kernel(mapping=[1], args=[B])
     def consumer(local_B: Ty[M, N]):
         for i, j in allo.grid(M, N):
             # receive data
             data = pipe.get()
             # computation
             local_B[i, j] = data + 1


def test_producer_consumer():
    A = np.random.randint(0, 64, (M, K)).astype(np.int32)
    B = np.zeros((M, N), dtype=np.int32)

    if "MLIR_AIE_INSTALL_DIR" in os.environ:
        mod = df.build(top, target="aie")
        mod(A, B)
        np.testing.assert_allclose(A + 1, B, atol=1e-5)
        print("Passed!")
    else:
        print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")

Learning Materials