[WIP] use peano for kernels

fixes #637
nod-ai · Aug 31, 2024 · 4b0d88b · 4b0d88b
1 parent 37bb7f1
commit 4b0d88b
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 27 deletions.
diff --git a/build_tools/ci/run_matmul_test.sh b/build_tools/ci/run_matmul_test.sh
@@ -182,7 +182,7 @@ function run_matmul_test() {
 
   local amd_aie_install_path="${IREE_INSTALL_DIR}"
 
-  local vitis_path="${VITIS}"
+  local vitis_path=""
 
   local use_chess="false"
 
@@ -540,16 +540,14 @@ run_matmul_test \
 # MLIR-AIR Matmul tests
 ###################################################################
 
-if [ -d "$VITIS" ]; then
-  run_matmul_test \
-      --name_prefix "ukern" \
-      --lower_to_aie_pipeline "air" \
-      --tile_pipeline "pad-pack" \
-      --lhs_rhs_type "bf16" \
-      --acc_type "f32" \
-      --m "256"  --k "256" --n "256" \
-      --use_ukernel "1"
-fi
+run_matmul_test \
+    --name_prefix "ukern" \
+    --lower_to_aie_pipeline "air" \
+    --tile_pipeline "pad-pack" \
+    --lhs_rhs_type "bf16" \
+    --acc_type "f32" \
+    --m "256"  --k "256" --n "256" \
+    --use_ukernel "1"
 
 # Example of a run with a group of 2+ matmuls. Currently this test is passed
 # the flag '--num_repeat_runs 0" as there is currently an issue with the runtime if

diff --git a/build_tools/download_peano.sh b/build_tools/download_peano.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
 
-RELEASE=19.0.0.2024082221+90abe71b
+RELEASE=19.0.0.2024083101+42158757
 pip download llvm_aie==$RELEASE -f https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly
 unzip llvm_aie*whl
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/XCLBinGen.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/XCLBinGen.cpp
@@ -532,22 +532,31 @@ static LogicalResult generateCoreElfFiles(
     Path cwd = std::filesystem::current_path();
     FailureOr<Path> mmObjectFilePath;
     if (ukernel && (ukernel == "mm" || ukernel == "all")) {
-      FailureOr<Path> maybeVitisDir = findVitis(vitisDir, npuVersion);
-      if (failed(maybeVitisDir)) {
-        llvm::errs() << "compiling ukernels currently requires chess (even if "
-                        "you're using peano)";
-        return failure();
-      }
       if (!std::filesystem::exists(cwd / "mm.o")) {
-        mmObjectFilePath = assembleStringUsingChess(
-            /*inputFileStr=*/_MM_CC,
-            /*inputFileName=*/"mm.cc",
-            /*outputFileName=*/"mm.o",
-            /*outputDir=*/cwd,
-            /*extraArgs*/ std::vector<std::string>{},
-            /*workDir=*/tempDir,
-            /*vitisDir=*/*maybeVitisDir,
-            /*npuVersion*/ npuVersion, verbose);
+        FailureOr<Path> maybeVitisDir = findVitis(vitisDir, npuVersion);
+        if (succeeded(maybeVitisDir)) {
+          if (verbose) llvm::outs() << "using chess for ukernel codegen\n";
+          mmObjectFilePath = assembleStringUsingChess(
+              /*inputFileStr=*/_MM_CC,
+              /*inputFileName=*/"mm.cc",
+              /*outputFileName=*/"mm.o",
+              /*outputDir=*/cwd,
+              /*extraArgs*/ std::vector<std::string>{},
+              /*workDir=*/tempDir,
+              /*vitisDir=*/*maybeVitisDir,
+              /*npuVersion*/ npuVersion, verbose);
+        } else {
+          if (verbose) llvm::outs() << "using peano for ukernel codegen\n";
+          mmObjectFilePath = assembleStringUsingPeano(
+              /*inputFileStr=*/_MM_CC,
+              /*inputFileName=*/"mm.cc",
+              /*outputFileName=*/"mm.o",
+              /*outputDir=*/cwd,
+              /*extraArgs*/ std::vector<std::string>{},
+              /*workDir=*/tempDir,
+              /*peanoDir=*/peanoDir,
+              /*npuVersion*/ npuVersion, verbose);
+        }
         if (failed(mmObjectFilePath)) return failure();
       } else {
         mmObjectFilePath = cwd / "mm.o";