Example with experimental syntax

programming_examples/README.md

@@ -10,7 +10,14 @@
 
 # <ins>Programming Examples</ins>
 
-These programming examples are provided so that application programmers can learn how to leverage the IRON design flow with mlir-aie python bindings, and the mlir-aie intermediate representation directly to build applications targeting AI Engines. They are organized into the following directories:
+These programming examples are provided so that application programmers can learn how to leverage the IRON design flow with mlir-aie python bindings, and the mlir-aie intermediate representation directly to build applications targeting AI Engines. 
+
+Each IRON example has one or more implementations:
+* `aie2.py` - These are written using the original IRON syntax
+* `aie2_alt.py` - These are written using an alternate form of `runtime_sequence`, but the design is likely otherwise unchanged.
+* `aie2_iron.py` - These are written using an alternative IRON syntax
+
+They are organized into the following directories:
 
 ## [basic](./basic) 
 

programming_examples/basic/dma_transpose/Makefile

@@ -22,9 +22,17 @@ K ?= 32
 
 aie_py_src=aie2.py
 use_alt?=0
+use_iron?=0
 
 ifeq (${use_alt}, 1)
 aie_py_src=aie2_alt.py
+ifeq (${use_iron}, 1)
+$(error Cannot specify both alternative design and IRON)
+endif
+endif
+
+ifeq (${use_iron}, 1)
+aie_py_src=aie2_iron.py
 endif
 
 build/aie.mlir: ${srcdir}/${aie_py_src}
@@ -51,7 +59,7 @@ endif
 run: ${targetname}.exe build/final.xclbin
 	${powershell} ./$< -x build/final.xclbin -i build/insts.txt -k MLIR_AIE --M ${M} --K ${K}
 
-generate_access_map: ${srcdir}/aie2.py
+generate_access_map: ${srcdir}/${aie_py_src}
 	mkdir -p ${@D}
 	python3 $< --generate-access-map ${M} ${K}
 

programming_examples/basic/dma_transpose/aie2_alt.py

@@ -1,4 +1,4 @@
-# dma_transpose/aie2.py -*- Python -*-
+# dma_transpose/aie2_alt.py -*- Python -*-
 #
 # This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.

programming_examples/basic/dma_transpose/aie2_iron.py

@@ -0,0 +1,59 @@
+# dma_transpose/aie2_iron.py -*- Python -*-
+#
+# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#
+# (c) Copyright 2024 Advanced Micro Devices, Inc. or its affiliates
+import argparse
+import numpy as np
+import sys
+
+from aie.iron import ObjectFifo, Program, Runtime
+from aie.iron.device import NPU1Col1, AnyComputeTile
+from aie.iron.placers import SequentialPlacer
+from aie.helpers.taplib import TensorTiler2D
+
+
+def my_passthrough(M, K, generate_acccess_map=False):
+    tensor_ty = np.ndarray[(M, K), np.dtype[np.int32]]
+
+    tap_in = TensorTiler2D.simple_tiler((M, K), tile_col_major=True)[0]
+
+    if generate_acccess_map:
+        tap_in.visualize(file_path="iron_transpose_data.png", show_tile=False)
+        return
+
+    of_in = ObjectFifo(tensor_ty)
+    of_out = of_in.cons().forward(AnyComputeTile)
+
+    rt = Runtime()
+    with rt.sequence(tensor_ty, tensor_ty, tensor_ty) as (a_in, _, c_out):
+        rt.fill(of_in.prod(), a_in, tap_in)
+        rt.drain(of_out.cons(), c_out, wait=True)
+
+    my_program = Program(NPU1Col1(), rt)
+    module = my_program.resolve_program(SequentialPlacer())
+    print(module)
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser()
+    p.add_argument("dims", help="M K", type=int, nargs="*", default=[64, 64])
+    p.add_argument(
+        "--generate-access-map",
+        action="store_true",
+        help="Produce a file showing data access order",
+    )
+    args = p.parse_args()
+
+    if len(args.dims) != 2:
+        print(
+            "ERROR: Must provide either no dimensions or both M and K", file=sys.stderr
+        )
+        exit(-1)
+    my_passthrough(
+        M=args.dims[0],
+        K=args.dims[1],
+        generate_acccess_map=args.generate_access_map,
+    )

programming_examples/basic/dma_transpose/run_makefile_iron.lit

@@ -0,0 +1,11 @@
+// (c) Copyright 2024 Advanced Micro Devices, Inc.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// REQUIRES: ryzen_ai, peano 
+//
+// RUN: mkdir -p iron_test
+// RUN: cd iron_test
+// RUN: make -f %S/Makefile clean
+// RUN: env use_iron=1 make -f %S/Makefile 
+// RUN: %run_on_npu make -f %S/Makefile run | FileCheck %s
+// CHECK: PASS!

programming_examples/basic/matrix_multiplication/cascade/aie2_alt.py

@@ -7,12 +7,12 @@
 import argparse
 from ml_dtypes import bfloat16
 import numpy as np
-import sys
 
 from aie.extras.context import mlir_mod_ctx
 from aie.dialects.aie import *
 from aie.dialects.aiex import *
 from aie.helpers.dialects.ext.scf import _for as range_
+from aie.helpers.taplib import TensorAccessPattern, TensorAccessSequence
 
 dtype_map = {
     "bf16": bfloat16,
@@ -42,9 +42,15 @@ def main():
         "--dtype_out", type=str, choices=["bf16", "i16", "f32", "i32"], default="i32"
     )
     argparser.add_argument("--trace_size", type=int, default=0)
+    argparser.add_argument(
+        "--generate-taps",
+        action="store_true",
+        help="Generate TensorAccessPatterns, a Python object to represent each data transfer"
+        "of the input/output matrices. These objects can be used for visualization.",
+    )
     args = argparser.parse_args()
     with mlir_mod_ctx() as ctx:
-        my_matmul(
+        maybe_taps = my_matmul(
             args.M,
             args.K,
             args.N,
@@ -55,16 +61,32 @@ def main():
             args.dtype_in,
             args.dtype_out,
             args.trace_size,
+            args.generate_taps,
         )
         # print(ctx.module.operation.verify())
         print(ctx.module)
 
+    if args.generate_taps:
+        return maybe_taps
+
 
 def ceildiv(a, b):
     return (a + b - 1) // b
 
 
-def my_matmul(M, K, N, m, k, n, n_aie_cols, dtype_in_str, dtype_out_str, trace_size):
+def my_matmul(
+    M,
+    K,
+    N,
+    m,
+    k,
+    n,
+    n_aie_cols,
+    dtype_in_str,
+    dtype_out_str,
+    trace_size,
+    generate_taps=False,
+):
 
     n_aie_rows = 4
     n_aie_cores = n_aie_rows * n_aie_cols
@@ -128,6 +150,12 @@ def my_matmul(M, K, N, m, k, n, n_aie_cols, dtype_in_str, dtype_out_str, trace_s
     elif n_aie_cols == 4:
         dev = AIEDevice.npu1_4col
 
+    # These will hold TensorAccessPattern objects that represent the runtime
+    # npu_dma_memcpy_nd operations of this design. They are only used if generate_taps is true
+    A_taps = []
+    B_taps = []
+    C_taps = []
+
     @device(dev)
     def device_body():
         A_l2_ty = np.ndarray[(m * k * n_A_tiles_per_shim,), np.dtype[dtype_in]]
@@ -349,60 +377,75 @@ def sequence(A, B, C):
                 for col in range(n_aie_cols):
                     C_col_offset = col * n
                     C_offset = C_col_offset + C_row_offset
-
+                    C_sizes = [tb_n_rows, N // n // n_aie_cols, m, n]
+                    C_strides = [m * N, n * n_aie_cols, N, 1]
+                    C_tap = TensorAccessPattern(
+                        (M, N), C_offset, sizes=C_sizes, strides=C_strides
+                    )
                     c_task = shim_dma_single_bd_task(
                         C_l2l3_fifos[col],
                         C,
-                        offset=C_offset,
-                        sizes=[tb_n_rows, N // n // n_aie_cols, m, n],
-                        strides=[m * N, n * n_aie_cols, N, 1],
+                        tap=C_tap,
                         issue_token=True,
                     )
                     dma_start_task(c_task)
                     out_tasks.append(c_task)
+                    C_taps.append(C_tap)
 
                     for tile_row in range(tb_n_rows):
                         A_block_offset = ((tb * tb_max_n_rows) + tile_row) * m * K
                         A_row_offset = col * n_A_tiles_per_shim * k
                         A_offset = A_block_offset + A_row_offset
+                        A_sizes = [
+                            N // n // n_aie_cols,
+                            K // k // n_aie_rows,
+                            m * n_A_tiles_per_shim,
+                            k,
+                        ]
+                        A_strides = [0, k * n_aie_rows, K, 1]
+                        A_tap = TensorAccessPattern(
+                            (M, K), A_offset, sizes=A_sizes, strides=A_strides
+                        )
                         B_col_offset = col * n
                         a_task = shim_dma_single_bd_task(
                             A_l3l2_fifos[col],
                             A,
-                            offset=A_offset,
-                            sizes=[
-                                N // n // n_aie_cols,
-                                K // k // n_aie_rows,
-                                m * n_A_tiles_per_shim,
-                                k,
-                            ],
-                            strides=[0, k * n_aie_rows, K, 1],
+                            tap=A_tap,
                         )
                         dma_start_task(a_task)
                         in_tasks.append(a_task)
-
+                        A_taps.append(A_tap)
+
+                        B_sizes = [
+                            N // n // n_aie_cols,
+                            K // k // n_aie_rows,
+                            k * n_aie_rows,
+                            n,
+                        ]
+                        B_strides = [n * n_aie_cols, k * n_aie_rows * N, N, 1]
+                        B_tap = TensorAccessPattern(
+                            (K, N), B_col_offset, sizes=B_sizes, strides=B_strides
+                        )
                         b_task = shim_dma_single_bd_task(
-                            B_l3l2_fifos[col],
-                            B,
-                            offset=B_col_offset,
-                            sizes=[
-                                N // n // n_aie_cols,
-                                K // k // n_aie_rows,
-                                k * n_aie_rows,
-                                n,
-                            ],
-                            strides=[n * n_aie_cols, k * n_aie_rows * N, N, 1],
+                            B_l3l2_fifos[col], B, tap=B_tap
                         )
                         dma_start_task(b_task)
                         in_tasks.append(b_task)
+                        B_taps.append(B_tap)
                 dma_await_task(*out_tasks)
                 out_tasks = []
                 dma_free_task(*in_tasks)
                 in_tasks = []
 
+    if generate_taps:
+        # If generate taps is true, return a representation of tensor access patterns
+        # representing all the npu_dma_memcpy_nd runtime sequence operations per input/ouput tensor.
+        return (
+            TensorAccessSequence.from_taps(A_taps),
+            TensorAccessSequence.from_taps(B_taps),
+            TensorAccessSequence.from_taps(C_taps),
+        )
+
 
 if __name__ == "__main__":
     main()
-else:
-    print("Not meant to be imported")
-    sys.exit(1)

programming_examples/basic/matrix_multiplication/matrix_vector/Makefile

@@ -21,10 +21,19 @@ k=32
 
 kernels=mv_${m}x${k}
 use_alt?=0
+use_iron?=0
 
 ifeq (${use_alt}, 1)
 aie_py_src=aie2_alt.py
+ifeq (${use_iron}, 1)
+$(error Cannot specify both alternative design and IRON)
 endif
+endif
+
+ifeq (${use_iron}, 1)
+aie_py_src=aie2_iron.py
+endif
+
 
 SELF_DIR := $(dir $(lastword $(MAKEFILE_LIST)))
 include ${SELF_DIR}../makefile-common