Skip to content

Commit f59995b

Browse files
maksleventalmakslevental
committed
[WIP] use peano for kernels
fixes #637
1 parent 37bb7f1 commit f59995b

File tree

3 files changed

+66
-49
lines changed

3 files changed

+66
-49
lines changed

‎build_tools/ci/run_matmul_test.sh

+15-11
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,7 @@ function run_matmul_test() {
182182

183183
local amd_aie_install_path="${IREE_INSTALL_DIR}"
184184

185-
local vitis_path="${VITIS}"
185+
local vitis_path=""
186186

187187
local use_chess="false"
188188

@@ -540,16 +540,15 @@ run_matmul_test \
540540
# MLIR-AIR Matmul tests
541541
###################################################################
542542

543-
if [ -d "$VITIS" ]; then
544-
run_matmul_test \
545-
--name_prefix "ukern" \
546-
--lower_to_aie_pipeline "air" \
547-
--tile_pipeline "pad-pack" \
548-
--lhs_rhs_type "bf16" \
549-
--acc_type "f32" \
550-
--m "256" --k "256" --n "256" \
551-
--use_ukernel "1"
552-
fi
543+
run_matmul_test \
544+
--name_prefix "ukern" \
545+
--lower_to_aie_pipeline "air" \
546+
--tile_pipeline "pad-pack" \
547+
--lhs_rhs_type "bf16" \
548+
--acc_type "f32" \
549+
--m "256" --k "256" --n "256" \
550+
--vitis_path "${VITIS}" \
551+
--use_ukernel "1"
553552

554553
# Example of a run with a group of 2+ matmuls. Currently this test is passed
555554
# the flag '--num_repeat_runs 0" as there is currently an issue with the runtime if
@@ -720,6 +719,7 @@ if [ -d "$VITIS" ]; then
720719
--lhs_rhs_type "bf16" \
721720
--acc_type "f32" \
722721
--num_repeat_runs "2" \
722+
--vitis_path "${VITIS}" \
723723
--use_ukernel "1"
724724

725725
run_matmul_test_on_shapes ${bf16_ukernel_shapes_medium[@]} \
@@ -729,6 +729,7 @@ if [ -d "$VITIS" ]; then
729729
--lhs_rhs_type "bf16" \
730730
--acc_type "f32" \
731731
--num_repeat_runs "2" \
732+
--vitis_path "${VITIS}" \
732733
--use_ukernel "1"
733734
fi
734735

@@ -746,6 +747,7 @@ if [ -d "$VITIS" ]; then
746747
--n "32" \
747748
--k "32" \
748749
--use_chess "1" \
750+
--vitis_path "${VITIS}" \
749751
--num_repeat_runs "10"
750752

751753
run_matmul_test \
@@ -757,6 +759,7 @@ if [ -d "$VITIS" ]; then
757759
--k "64" \
758760
--use_chess "1" \
759761
--num_repeat_runs "10" \
762+
--vitis_path "${VITIS}" \
760763
--use_ukernel "1"
761764

762765
run_matmul_test \
@@ -769,6 +772,7 @@ if [ -d "$VITIS" ]; then
769772
--n "32" \
770773
--k "32" \
771774
--use_chess "1" \
775+
--vitis_path "${VITIS}" \
772776
--num_repeat_runs "10"
773777

774778
fi

‎build_tools/download_peano.sh

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#!/bin/bash
22

3-
RELEASE=19.0.0.2024082221+90abe71b
3+
RELEASE=19.0.0.2024083101+42158757
44
pip download llvm_aie==$RELEASE -f https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly
55
unzip llvm_aie*whl

‎compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/XCLBinGen.cpp

+50-37
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ namespace {
9696
FailureOr<std::string> getTargetDir(const std::string &npuVersion) {
9797
if (npuVersion == "npu1") return std::string{"target_aie_ml"};
9898
if (npuVersion == "npu4") return std::string{"target_aie2p"};
99-
llvm::errs() << "unsupported NPUVersion: " << npuVersion;
99+
llvm::errs() << "unsupported NPUVersion: " << npuVersion << "\n";
100100
return failure();
101101
}
102102

@@ -152,7 +152,8 @@ FailureOr<Path> findVitis(std::optional<Path> &vitisDir,
152152
return failure();
153153
}
154154
if (!std::filesystem::exists(licenseFile)) {
155-
llvm::errs() << "ERROR: license file" << licenseFile << " does not exist";
155+
llvm::errs() << "ERROR: license file" << licenseFile << " does not exist"
156+
<< "\n";
156157
return failure();
157158
}
158159
}
@@ -215,7 +216,7 @@ std::pair<std::string, std::vector<std::string>> makeChessArgs(
215216
archVersion = "21";
216217
modelDir = "aie2p";
217218
} else {
218-
llvm::errs() << "unsupported NPU version: " << npuVersion;
219+
llvm::errs() << "unsupported NPU version: " << npuVersion << "\n";
219220
llvm::report_fatal_error("unsupported NPU version");
220221
}
221222

@@ -475,7 +476,7 @@ static FailureOr<Path> assembleStringUsing(
475476
if (auto maybeErr = dumpStrToDisk(inputFileStr, inputFile.string());
476477
maybeErr.has_value()) {
477478
llvm::errs() << "Failed to dump to disk " << inputFile.string()
478-
<< " because: " << maybeErr;
479+
<< " because: " << maybeErr << "\n";
479480
return failure();
480481
}
481482

@@ -487,7 +488,8 @@ static FailureOr<Path> assembleStringUsing(
487488
}
488489
if (failed(assembler(inputFile.string(), outputFile.string(), extraArgs,
489490
workDir, toolDir, npuVersion, verbose))) {
490-
llvm::errs() << "Failed to assemble " << outputFileName << ".o";
491+
llvm::errs() << "Failed to assemble " << outputFileName << ".o"
492+
<< "\n";
491493
return failure();
492494
}
493495
return outputFile;
@@ -532,22 +534,31 @@ static LogicalResult generateCoreElfFiles(
532534
Path cwd = std::filesystem::current_path();
533535
FailureOr<Path> mmObjectFilePath;
534536
if (ukernel && (ukernel == "mm" || ukernel == "all")) {
535-
FailureOr<Path> maybeVitisDir = findVitis(vitisDir, npuVersion);
536-
if (failed(maybeVitisDir)) {
537-
llvm::errs() << "compiling ukernels currently requires chess (even if "
538-
"you're using peano)";
539-
return failure();
540-
}
541537
if (!std::filesystem::exists(cwd / "mm.o")) {
542-
mmObjectFilePath = assembleStringUsingChess(
543-
/*inputFileStr=*/_MM_CC,
544-
/*inputFileName=*/"mm.cc",
545-
/*outputFileName=*/"mm.o",
546-
/*outputDir=*/cwd,
547-
/*extraArgs*/ std::vector<std::string>{},
548-
/*workDir=*/tempDir,
549-
/*vitisDir=*/*maybeVitisDir,
550-
/*npuVersion*/ npuVersion, verbose);
538+
if (useChess) {
539+
if (verbose) llvm::outs() << "using chess for ukernel codegen\n";
540+
FailureOr<Path> maybeVitisDir = findVitis(vitisDir, npuVersion);
541+
mmObjectFilePath = assembleStringUsingChess(
542+
/*inputFileStr=*/_MM_CC,
543+
/*inputFileName=*/"mm.cc",
544+
/*outputFileName=*/"mm.o",
545+
/*outputDir=*/cwd,
546+
/*extraArgs*/ std::vector<std::string>{},
547+
/*workDir=*/tempDir,
548+
/*vitisDir=*/*maybeVitisDir,
549+
/*npuVersion*/ npuVersion, verbose);
550+
} else {
551+
if (verbose) llvm::outs() << "using peano for ukernel codegen\n";
552+
mmObjectFilePath = assembleStringUsingPeano(
553+
/*inputFileStr=*/_MM_CC,
554+
/*inputFileName=*/"mm.cc",
555+
/*outputFileName=*/"mm.o",
556+
/*outputDir=*/cwd,
557+
/*extraArgs*/ std::vector<std::string>{},
558+
/*workDir=*/tempDir,
559+
/*peanoDir=*/peanoDir,
560+
/*npuVersion*/ npuVersion, verbose);
561+
}
551562
if (failed(mmObjectFilePath)) return failure();
552563
} else {
553564
mmObjectFilePath = cwd / "mm.o";
@@ -579,13 +590,14 @@ static LogicalResult generateCoreElfFiles(
579590
{
580591
auto bcfOutput = openOutputFile(bcfPath.string(), &errorMessage);
581592
if (!bcfOutput) {
582-
llvm::errs() << "failed to open bcf file because: " << errorMessage;
593+
llvm::errs() << "failed to open bcf file because: " << errorMessage
594+
<< "\n";
583595
return failure();
584596
}
585597

586598
if (failed(mlir::iree_compiler::AMDAIE::AIETranslateToBCF(
587599
deviceOp, bcfOutput->os(), col, row))) {
588-
llvm::errs() << "Failed to generate BCF";
600+
llvm::errs() << "Failed to generate BCF\n";
589601
return failure();
590602
}
591603
bcfOutput->keep();
@@ -614,7 +626,7 @@ static LogicalResult generateCoreElfFiles(
614626
openOutputFile(ldscriptPath.string(), &errorMessage);
615627
if (!ldscriptOutput) {
616628
llvm::errs() << "Failed to open ldscript file because: "
617-
<< errorMessage;
629+
<< errorMessage << "\n";
618630
return failure();
619631
}
620632
if (failed(mlir::iree_compiler::AMDAIE::AIETranslateToLdScript(
@@ -654,7 +666,7 @@ static LogicalResult generateCDO(MLIRContext *context, AIE::DeviceOp deviceOp,
654666
deviceOp = *copy.getOps<AIE::DeviceOp>().begin();
655667
if (failed(mlir::iree_compiler::AMDAIE::AIETranslateToCDODirect(
656668
deviceOp, tempDir.string()))) {
657-
llvm::errs() << "failed to emit CDO";
669+
llvm::errs() << "failed to emit CDO\n";
658670
return failure();
659671
}
660672
copy->erase();
@@ -750,7 +762,7 @@ static LogicalResult generateXCLBin(
750762
dumpStrToDisk(memTopologyData, memTopologyJsonFile.string());
751763
maybeErr.has_value()) {
752764
llvm::errs() << "failed to dump to disk mem_topology.json because: "
753-
<< *maybeErr;
765+
<< *maybeErr << "\n";
754766
return failure();
755767
}
756768
}
@@ -797,7 +809,7 @@ static LogicalResult generateXCLBin(
797809
dumpStrToDisk(aiePartitionJsonData, aiePartitionJsonFile.string());
798810
maybeErr.has_value()) {
799811
llvm::errs() << "failed to dump to disk aie_partition.json because: "
800-
<< *maybeErr;
812+
<< *maybeErr << "\n";
801813
return failure();
802814
}
803815
}
@@ -816,7 +828,7 @@ static LogicalResult generateXCLBin(
816828
if (auto maybeErr = dumpStrToDisk(kernelStr, kernelsJsonFile.string());
817829
maybeErr.has_value()) {
818830
llvm::errs() << "failed to dump to disk kernels.json because: "
819-
<< *maybeErr;
831+
<< *maybeErr << "\n";
820832
return failure();
821833
}
822834
}
@@ -825,7 +837,8 @@ static LogicalResult generateXCLBin(
825837
{
826838
auto designBifOut = openOutputFile(designBifFile.string(), &errorMessage);
827839
if (!designBifOut) {
828-
llvm::errs() << "failed to open design.bif because: " << errorMessage;
840+
llvm::errs() << "failed to open design.bif because: " << errorMessage
841+
<< "\n";
829842
return failure();
830843
}
831844

@@ -868,7 +881,7 @@ static LogicalResult generateXCLBin(
868881
}
869882
if (iree_aie_bootgen_main(cstrings.size(),
870883
const_cast<const char **>(&cstrings[0]))) {
871-
llvm::errs() << "failed to execute bootgen";
884+
llvm::errs() << "failed to execute bootgen\n";
872885
return failure();
873886
}
874887
}
@@ -892,14 +905,14 @@ static LogicalResult generateXCLBin(
892905
"--force", "--input", *inputXclbin};
893906

894907
if (failed(runTool(xclbinutilBin.value().string(), inputFlags, verbose))) {
895-
llvm::errs() << "failed to execute xclbinutil";
908+
llvm::errs() << "failed to execute xclbinutil\n";
896909
return failure();
897910
}
898911
auto aieInputPartitionOut =
899912
openInputFile(aieInputPartitionJsonFile.string(), &errorMessage);
900913
if (!aieInputPartitionOut) {
901914
llvm::errs() << "failed to open aie_input_partition.json because: "
902-
<< errorMessage;
915+
<< errorMessage << "\n";
903916
return failure();
904917
}
905918
Expected<json::Value> aieInputPartitionOutValue =
@@ -913,7 +926,7 @@ static LogicalResult generateXCLBin(
913926
if (!aiePartitionOut) {
914927
llvm::errs() << "failed to open aie aie_input_partition.json for "
915928
"output because: "
916-
<< errorMessage;
929+
<< errorMessage << "\n";
917930
return failure();
918931
}
919932
llvm::Expected<llvm::json::Value> aiePartitionOutValue =
@@ -931,7 +944,7 @@ static LogicalResult generateXCLBin(
931944
maybeErr.has_value()) {
932945
llvm::errs()
933946
<< "failed to dump to disk aie_input_partition.json because: "
934-
<< errorMessage;
947+
<< errorMessage << "\n";
935948
return failure();
936949
}
937950
flags.insert(flags.end(), {"--input", *inputXclbin});
@@ -1040,14 +1053,14 @@ static LogicalResult generateUnifiedObject(
10401053
}
10411054

10421055
if (failed(pm.run(moduleOpCopy))) {
1043-
llvm::errs() << "Failed to lower to LLVM";
1056+
llvm::errs() << "Failed to lower to LLVM\n";
10441057
return failure();
10451058
}
10461059

10471060
llvm::LLVMContext llvmContext;
10481061
auto llvmModule = translateModuleToLLVMIR(moduleOpCopy, llvmContext);
10491062
if (!llvmModule) {
1050-
llvm::errs() << "Failed to translate module to LLVMIR";
1063+
llvm::errs() << "Failed to translate module to LLVMIR\n";
10511064
return failure();
10521065
}
10531066

@@ -1081,7 +1094,7 @@ static LogicalResult generateUnifiedObject(
10811094
if (auto maybeErr = dumpStrToDisk(inputLLStr, LLVMIRFile.string());
10821095
maybeErr.has_value()) {
10831096
llvm::errs() << "Failed to dump to disk input.ll"
1084-
<< " because: " << maybeErr;
1097+
<< " because: " << maybeErr << "\n";
10851098
return failure();
10861099
}
10871100
Path peanoOptBin = peanoDir / "bin" / "opt";
@@ -1096,7 +1109,7 @@ static LogicalResult generateUnifiedObject(
10961109
args.reserve(args.size() + peanoArgs.size());
10971110
args.insert(args.end(), peanoArgs.begin(), peanoArgs.end());
10981111
if (failed(runTool(peanoOptBin.string(), args, verbose))) {
1099-
llvm::errs() << "Failed to optimize ll with peano";
1112+
llvm::errs() << "Failed to optimize ll with peano\n";
11001113
return failure();
11011114
}
11021115

0 commit comments

Comments
 (0)