Passthrough Kernel for NPU2/Strix

programming_examples/basic/passthrough_kernel/Makefile

@@ -12,6 +12,7 @@ srcdir := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST))))
 
 include ${srcdir}/../../makefile-common
 
+device = npu
 targetname = passThroughKernel
 VPATH := ${srcdir}/../../../aie_kernels/generic
 data_size = 4096
@@ -24,27 +25,44 @@ all: build/final_${data_size}.xclbin
 
 build/aie2_lineBased_8b_${data_size}.mlir: ${srcdir}/aie2.py
 	mkdir -p ${@D}
-	python3 $< ${data_size} 0 > $@
+	python3 $< ${device} ${data_size} 0 > $@
 
 build/aie_trace__lineBased_8b_${data_size}.mlir: ${srcdir}/aie2.py
 	mkdir -p ${@D}
-	python3 $< ${data_size} ${trace_size} > $@
+	python3 $< ${device} ${data_size} ${trace_size} > $@
 
 build/passThrough.cc.o: passThrough.cc
 	mkdir -p ${@D}
+ifeq ($(device),npu)
 	cd ${@D} && ${PEANO_INSTALL_DIR}/bin/clang++ ${PEANOWRAP2_FLAGS} -DBIT_WIDTH=8 -c $< -o ${@F}
+else ifeq ($(device),npu2)
+	cd ${@D} && xchesscc_wrapper ${CHESSCCWRAP2P_FLAGS} -DBIT_WIDTH=8 -c $< -o ${@F}
+else
+	echo "Device type not supported"
+endif
 
 build/final_${data_size}.xclbin: build/aie2_lineBased_8b_${data_size}.mlir build/passThrough.cc.o
 	mkdir -p ${@D}
+ifeq ($(device),npu)
 	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_size}.txt $(<:%=../%)
+else
+	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host \
+		--xclbin-name=${@F} --npu-insts-name=insts_${data_size}.txt $(<:%=../%)
+endif
 
 build/final_trace_${data_size}.xclbin: build/aie2_lineBased_8b_${data_size}.mlir build/passThrough.cc.o
 	mkdir -p ${@D}
+ifeq ($(device),npu)
 	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_size}.txt $(<:%=../%)
+else
+	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host \
+		--xclbin-name=${@F} --npu-insts-name=insts_${data_size}.txt $(<:%=../%)
+endif
+
 
 ${targetname}_${data_size}.exe: ${srcdir}/test.cpp
 	rm -rf _build

programming_examples/basic/passthrough_kernel/aie2.py

@@ -16,11 +16,11 @@
 import aie.utils.trace as trace_utils
 
 
-def passthroughKernel(vector_size, trace_size):
+def passthroughKernel(dev, vector_size, trace_size):
     N = vector_size
     lineWidthInBytes = N // 4  # chop input in 4 sub-tensors
 
-    @device(AIEDevice.npu1_1col)
+    @device(dev)
     def device_body():
         # define types
         vector_ty = np.ndarray[(N,), np.dtype[np.uint8]]
@@ -85,13 +85,20 @@ def sequence(inTensor, outTensor, notUsed):
 
 
 try:
-    vector_size = int(sys.argv[1])
+    device_name = str(sys.argv[1])
+    if device_name == "npu":
+        dev = AIEDevice.npu1_1col
+    elif device_name == "npu2":
+        dev = AIEDevice.npu2
+    else:
+        raise ValueError("[ERROR] Device name {} is unknown".format(sys.argv[1]))
+    vector_size = int(sys.argv[2])
     if vector_size % 64 != 0 or vector_size < 512:
         print("Vector size must be a multiple of 64 and greater than or equal to 512")
         raise ValueError
-    trace_size = 0 if (len(sys.argv) != 3) else int(sys.argv[2])
+    trace_size = 0 if (len(sys.argv) != 4) else int(sys.argv[3])
 except ValueError:
     print("Argument has inappropriate value")
 with mlir_mod_ctx() as ctx:
-    passthroughKernel(vector_size, trace_size)
+    passthroughKernel(dev, vector_size, trace_size)
     print(ctx.module)

programming_examples/basic/passthrough_kernel/run_makefile.lit

@@ -3,6 +3,8 @@
 //
 // REQUIRES: ryzen_ai, peano 
 //
+// RUN: mkdir -p test
+// RUN: cd test
 // RUN: make -f %S/Makefile clean
 // RUN: make -f %S/Makefile 
 // RUN: %run_on_npu make -f %S/Makefile run | FileCheck %s

programming_examples/basic/passthrough_kernel/run_strix_makefile.lit

@@ -0,0 +1,10 @@
+// (c) Copyright 2024 Advanced Micro Devices, Inc.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// REQUIRES: ryzen_ai, chess
+//
+// RUN: mkdir -p test_stx
+// RUN: cd test_stx
+// RUN: make -f %S/Makefile clean
+// RUN: make -f %S/Makefile device=npu2 
+// RUN: %run_on_npu2 make -f %S/Makefile run device=npu2 | FileCheck %s

programming_examples/lit.cfg.py

@@ -46,6 +46,7 @@
 llvm_config.with_environment("PYTHONPATH", os.path.join(config.aie_obj_root, "python"))
 
 run_on_npu = "echo"
+run_on_npu2 = "echo"
 xrt_flags = ""
 
 # Not using run_on_board anymore, need more specific per-platform commands
@@ -140,9 +141,14 @@
             if len(m.groups()) == 3:
                 print("\tmodel:", m.group(3))
             config.available_features.add("ryzen_ai")
-            run_on_npu = (
-                f"flock /tmp/npu.lock {config.aie_src_root}/utils/run_on_npu.sh"
-            )
+            if str(m.group(3)) == "npu1":
+                run_on_npu = (
+                    f"flock /tmp/npu.lock {config.aie_src_root}/utils/run_on_npu.sh"
+                )
+            if str(m.group(3)) == "npu4":
+                run_on_npu2 = (
+                    f"flock /tmp/npu.lock {config.aie_src_root}/utils/run_on_npu.sh"
+                )
             break
     except:
         print("Failed to run xrt-smi")
@@ -151,6 +157,7 @@
     print("xrt not found")
 
 config.substitutions.append(("%run_on_npu", run_on_npu))
+config.substitutions.append(("%run_on_npu2", run_on_npu2))
 config.substitutions.append(("%xrt_flags", xrt_flags))
 config.substitutions.append(("%XRT_DIR", config.xrt_dir))
 config.environment["XRT_HACK_UNSECURE_LOADING_XCLBIN"] = "1"

programming_examples/lit.site.cfg.py.in

@@ -69,6 +69,8 @@ if lit.util.pythonize_bool("@AIETools_AIE_FOUND@"):
     config.vitis_components.append("AIE")
 if lit.util.pythonize_bool("@AIETools_AIE2_FOUND@"):
     config.vitis_components.append("AIE2")
+if lit.util.pythonize_bool("@AIETools_AIE2P_FOUND@"):
+    config.vitis_components.append("AIE2P")
 
 # Support substitution of the tools_dir with user parameters. This is
 # used when we can't determine the tool dir at configuration time.

programming_examples/makefile-common

@@ -11,6 +11,7 @@ CHESS_FLAGS = -P ${AIE_INCLUDE_DIR}
 
 CHESSCCWRAP1_FLAGS = aie -I ${AIETOOLS_DIR}/include 
 CHESSCCWRAP2_FLAGS = aie2 -I ${AIETOOLS_DIR}/include
+CHESSCCWRAP2P_FLAGS = aie2p -I ${AIETOOLS_DIR}/include 
 PEANOWRAP2_FLAGS = -O2 -v -std=c++20 --target=aie2-none-unknown-elf ${WARNING_FLAGS} -DNDEBUG -I ${AIETOOLS_DIR}/include 
 
 TEST_POWERSHELL := $(shell command -v powershell.exe >/dev/null 2>&1 && echo yes || echo no)