[WIP] TensorTiler2D

programming_examples/basic/dma_transpose/Makefile

@@ -44,5 +44,9 @@ 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
+	mkdir -p ${@D}
+	python3 $< --generate-access-map ${M} ${K}
+
 clean:
 	rm -rf build _build inst ${targetname}.exe

programming_examples/basic/dma_transpose/README.md

@@ -15,11 +15,24 @@ This reference design can be run on a Ryzen™ AI NPU.
 In the [design](./aie2.py), a 2-D array in a row-major layout is read from external memory to `ComputeTile2` with a transposed layout,
 by using an implicit copy via the compute tile's Data Movement Accelerator (DMA). The data is read from and written to external memory through the Shim tile (`col`, 0).
 
+This data movement transformation can be visualized as a map which shows the order the data the data is streamed (e.g., in transposed layout):
+<p align="center">
+  <img
+    src="transpose_data.png">
+    <h3 align="center"> Visualization of the Transpose Data Transformation for M=32, K=16. 
+ </h3> 
+</p>
+
 The implicit copy is performed using the `object_fifo_link` operation that specifies how input data arriving via `of_in` should be sent further via `of_out` by specifically leveraging the compute tile's DMA. This operation and its functionality are described in more depth in [Section-2b](../../../programming_guide/section-2/section-2b/README.md/#object-fifo-link) of the programming guide.
 
 
 To compile and run the design for NPU:
-```
+```bash
 make
 make run
+```
+
+To generate a data visualization of the transpose (like that above), run:
+```bash
+make generate_access_map
 ```

programming_examples/basic/dma_transpose/aie2.py

@@ -5,27 +5,32 @@
 # 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.dialects.aie import *
 from aie.dialects.aiex import *
 from aie.extras.context import mlir_mod_ctx
 from aie.helpers.dialects.ext.scf import _for as range_
+from aie.helpers.tensortiler.tensortiler2d import TensorTile
 
-N = 4096
-M = 64
-K = 64
 
-if len(sys.argv) == 3:
-    M = int(sys.argv[1])
-    K = int(sys.argv[2])
-    N = M * K
+def my_passthrough(M, K, N, generate_acccess_map=False):
+    tensor_ty = np.ndarray[(M, K), np.dtype[np.int32]]
+    data_transform = TensorTile(
+        tensor_height=M,
+        tensor_width=K,
+        sizes=[1, 1, K, M],
+        strides=[1, 1, 1, K],
+        offset=0,
+    )
+    if generate_acccess_map:
+        data_transform.visualize(
+            plot_access_count=False, file_path="transpose_data.png"
+        )
+        return
 
-tensor_ty = np.ndarray[(M, K), np.dtype[np.int32]]
-
-
-def my_passthrough():
     with mlir_mod_ctx() as ctx:
 
         @device(AIEDevice.npu1_1col)
@@ -56,8 +61,7 @@ def sequence(A, B, C):
                     metadata=of_in,
                     bd_id=1,
                     mem=A,
-                    sizes=[1, 1, K, M],
-                    strides=[1, 1, 1, K],
+                    tensor_tile=data_transform,
                     issue_token=True,
                 )
                 npu_dma_memcpy_nd(metadata=of_out, bd_id=0, mem=C, sizes=[1, 1, 1, N])
@@ -66,4 +70,24 @@ def sequence(A, B, C):
     print(ctx.module)
 
 
-my_passthrough()
+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],
+        N=args.dims[0] * args.dims[1],
+        generate_acccess_map=args.generate_access_map,
+    )

programming_examples/basic/matrix_multiplication/whole_array/tiling_scratch.py

@@ -0,0 +1,92 @@
+import numpy as np
+
+from aie.helpers.tensortiler.tensortiler2d import TensorTiler2D, TensorTile
+from util import construct_test
+
+
+# RUN: %python %s | FileCheck %s
+def ceildiv(a, b):
+    return -(a // -b)
+
+
+def run_checks(n_aie_cols, n_aie_rows, M, N, K, m, n, k):
+    tb_max_n_rows = 4
+    tb_n_rows = tb_max_n_rows // 2
+
+    # Define tilers
+    c_tiler = TensorTiler2D(M, N, m * n_aie_rows, n)
+    c_iter = c_tiler.tile_iter(
+        tile_repeat_step_horizontal=n_aie_cols, iter_step=tb_n_rows
+    )
+
+    for tb in range(ceildiv(M // m // n_aie_rows, tb_max_n_rows)):
+        for pingpong in [0, 1]:
+            row_base = tb * tb_max_n_rows + pingpong * tb_max_n_rows // 2
+            tb_n_rows = min([tb_max_n_rows // 2, M // m // n_aie_rows - row_base])
+            print(tb_n_rows)
+            if tb_n_rows <= 0:
+                # for small input sizes, we may not even need a "pong" iteration
+                break
+
+            for col in range(n_aie_cols):
+                C_row_offset = row_base * m * n_aie_rows * N
+                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_aie_rows, n]
+                C_strides = [m * n_aie_rows * N, n * n_aie_cols, N, 1]
+                expected_c_tile = TensorTile(
+                    M, N, offset=C_offset, sizes=C_sizes, strides=C_strides
+                )
+
+                c_tile = next(c_iter)
+                if c_tile != expected_c_tile:
+                    # equivalence for tensor tile checks offset, size, stride
+                    # but there may be different but equivalent transformations
+
+                    reference_access, reference_count = expected_c_tile.access_tensors()
+                    c_access, c_count = c_tile.access_tensors()
+
+                    """
+                    assert (reference_access == c_access).all(), (
+                        f"C access orders do not match. "
+                        f"Expected ({expected_c_tile}), got ({c_tile})"
+                    )
+                    assert (reference_count == c_count).all()
+                    """
+                    print(f"Expected: {expected_c_tile}")
+                    print(f"Actual: {c_tile}")
+
+
+def matrix_whole_array_tiling_sweep():
+    n_aie_cols_sweep = [1, 2, 4]  # TODO: when partial, add 3
+    n_aie_rows_sweep = [1, 2, 4]  # TODO: when partial, add 3
+    M_sweep = range(512, 4096, 512)
+    K_sweep = range(512, 4096, 512)
+    N_sweep = range(512, 4096, 512)
+    m_sweep = [16, 32, 64]
+    n_sweep = [16, 32, 64]
+    k_sweep = [16, 32, 64]
+
+    for n_aie_cols in n_aie_cols_sweep:
+        for n_aie_rows in n_aie_rows_sweep:
+            for M in M_sweep:
+                for N in N_sweep:
+                    for K in K_sweep:
+                        for m in m_sweep:
+                            for n in n_sweep:
+                                for k in k_sweep:
+                                    run_checks(
+                                        n_aie_cols=n_aie_cols,
+                                        n_aie_rows=n_aie_rows,
+                                        M=M,
+                                        N=N,
+                                        K=K,
+                                        m=m,
+                                        k=k,
+                                        n=n,
+                                    )
+                                    return
+
+
+if __name__ == "__main__":
+    matrix_whole_array_tiling_sweep()

programming_examples/basic/matrix_scalar_add/aie2.py

@@ -12,6 +12,7 @@
 from aie.dialects.aiex import *
 from aie.extras.context import mlir_mod_ctx
 from aie.helpers.dialects.ext.scf import _for as range_
+from aie.helpers.tensortiler.tensortiler2d import TensorTiler2D
 
 # Size of the entire image
 IMAGE_HEIGHT = 16
@@ -68,23 +69,24 @@ def core_body():
                 of_out1.release(ObjectFifoPort.Produce, 1)
 
         # To/from AIE-array data movement
+        tiler = TensorTiler2D(IMAGE_HEIGHT, IMAGE_WIDTH, TILE_HEIGHT, TILE_WIDTH)
+        t = next(tiler.tile_iter())  # Only transfer one (first) tile of data
+
         @runtime_sequence(tile_ty, tile_ty, tile_ty)
         def sequence(inTensor, notUsed, outTensor):
             npu_dma_memcpy_nd(
                 metadata=of_in1,
                 bd_id=1,
                 mem=inTensor,
-                sizes=[1, 1, TILE_HEIGHT, TILE_WIDTH],
-                strides=[1, 1, IMAGE_WIDTH, 1],
+                tensor_tile=t,
                 issue_token=True,
             )
 
             npu_dma_memcpy_nd(
                 metadata=of_out1,
                 bd_id=0,
                 mem=outTensor,
-                sizes=[1, 1, TILE_HEIGHT, TILE_WIDTH],
-                strides=[1, 1, IMAGE_WIDTH, 1],
+                tensor_tile=t,
             )
             dma_wait(of_in1, of_out1)
 

programming_examples/basic/row_wise_bias_add/aie2.py

@@ -11,6 +11,7 @@
 from aie.dialects.aiex import *
 from aie.extras.context import mlir_mod_ctx
 from aie.helpers.dialects.ext.scf import _for as range_
+from aie.helpers.tensortiler.tensortiler2d import TensorTiler2D
 
 
 def row_wise_bias_add(M, N, m, n):
@@ -48,28 +49,32 @@ def core_body():
                         in_fifo.release(ObjectFifoPort.Consume, 1)
                     bias_fifo.release(ObjectFifoPort.Consume, 1)
 
+        tiler = TensorTiler2D(M, N, m, n, tensor_col_major=True)
+        t = next(
+            tiler.tile_iter(tile_group_height=M // m, tile_group_width=N // n)
+        )  # Transfer all tiles at once
+        bias_tiler = TensorTiler2D(1, N, 1, n)
+        bias_t = next(bias_tiler.tile_iter(tile_group_width=N // n))
+
         @runtime_sequence(tensor_ty, bias_ty, tensor_ty)
         def sequence(inp, bias, out):
             npu_dma_memcpy_nd(
                 metadata=in_fifo,
                 bd_id=0,
                 mem=inp,
-                sizes=[1, N // n, M, n],
-                strides=[0, n, N, 1],
+                tensor_tile=t,
             )
             npu_dma_memcpy_nd(
                 metadata=bias_fifo,
                 bd_id=1,
                 mem=bias,
-                sizes=[1, 1, N // n, n],
-                strides=[0, 0, n, 1],
+                tensor_tile=bias_t,
             )
             npu_dma_memcpy_nd(
                 metadata=out_fifo,
                 bd_id=2,
                 mem=out,
-                sizes=[1, N // n, M, n],
-                strides=[0, n, N, 1],
+                tensor_tile=t,
             )
             # of_out will only complete after of_in completes, so we just wait on of_out instead of both
             dma_wait(out_fifo)

programming_examples/basic/tiling_exploration/per_tile/Makefile

@@ -0,0 +1,39 @@
+##===- Makefile -----------------------------------------------------------===##
+# 
+# This file 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
+#
+# Copyright (C) 2024, Advanced Micro Devices, Inc.
+# 
+##===----------------------------------------------------------------------===##
+
+srcdir := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST))))
+
+include ${srcdir}/../../../makefile-common
+
+tensor_height = 32
+tensor_width = 32
+tile_height = 4
+tile_width = 4
+data_str=${tensor_height}_${tensor_width}_${tile_height}_${tile_width}
+
+.PHONY: all template clean
+
+all: build/final_${data_str}.xclbin
+
+build/aie_${data_str}.mlir: ${srcdir}/aie2.py
+	mkdir -p ${@D}
+	python3 $< --tensor-height ${tensor_height} --tensor-width ${tensor_width} --tile-height ${tile_height} --tile-width ${tile_width} > $@
+
+build/final_${data_str}.xclbin: build/aie_${data_str}.mlir
+	mkdir -p ${@D}
+	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host \
+		--no-xchesscc --no-xbridge \
+		--xclbin-name=${@F} --npu-insts-name=insts_${data_str}.txt $(<:%=../%)
+
+run: build/final_${data_str}.xclbin build/insts_${data_str}.txt
+	${powershell} python3 ${srcdir}/test.py -x build/final_${data_str}.xclbin -i build/insts_${data_str}.txt -k MLIR_AIE --tensor-height ${tensor_height} --tensor-width ${tensor_width} --tile-height ${tile_height} --tile-width ${tile_width}
+
+clean:
+	rm -rf build

programming_examples/basic/tiling_exploration/per_tile/README.md

@@ -0,0 +1,31 @@
+<!---//===- README.md -----------------------------------------*- Markdown -*-===//
+//
+// 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
+//
+// Copyright (C) 2024, Advanced Micro Devices, Inc.
+// 
+//===----------------------------------------------------------------------===//-->
+
+# Tiling Exploration
+
+This IRON design flow example, called "Tiling Exploration", demonstrates how data may be `tiled` on input/output. This is a common data transformation pattern, and this example is meant to be interactive.
+
+## Source Files Overview
+
+TODO
+
+## Design Overview
+
+TODO
+
+## Design Component Details
+
+### AIE Array Structural Design
+
+TODO
+
+## Usage
+
+TODO