Skip to content

Commit

Permalink
[WIP] use peano for kernels
Browse files Browse the repository at this point in the history 
fixes #637
  • Loading branch information
@makslevental
makslevental committed Aug 31, 2024
1 parent 37bb7f1 commit f59995b
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 49 deletions.
26 changes: 15 additions & 11 deletions build_tools/ci/run_matmul_test.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -540,16 +540,15 @@ 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" \
--vitis_path "${VITIS}" \
--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
Expand Down Expand Up @@ -720,6 +719,7 @@ if [ -d "$VITIS" ]; then
--lhs_rhs_type "bf16" \
--acc_type "f32" \
--num_repeat_runs "2" \
--vitis_path "${VITIS}" \
--use_ukernel "1"

run_matmul_test_on_shapes ${bf16_ukernel_shapes_medium[@]} \
Expand All @@ -729,6 +729,7 @@ if [ -d "$VITIS" ]; then
--lhs_rhs_type "bf16" \
--acc_type "f32" \
--num_repeat_runs "2" \
--vitis_path "${VITIS}" \
--use_ukernel "1"
fi

Expand All @@ -746,6 +747,7 @@ if [ -d "$VITIS" ]; then
--n "32" \
--k "32" \
--use_chess "1" \
--vitis_path "${VITIS}" \
--num_repeat_runs "10"

run_matmul_test \
Expand All @@ -757,6 +759,7 @@ if [ -d "$VITIS" ]; then
--k "64" \
--use_chess "1" \
--num_repeat_runs "10" \
--vitis_path "${VITIS}" \
--use_ukernel "1"

run_matmul_test \
Expand All @@ -769,6 +772,7 @@ if [ -d "$VITIS" ]; then
--n "32" \
--k "32" \
--use_chess "1" \
--vitis_path "${VITIS}" \
--num_repeat_runs "10"

fi
Expand Down
2 changes: 1 addition & 1 deletion build_tools/download_peano.sh
View file
Original file line number Diff line number Diff line change
@@ -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
87 changes: 50 additions & 37 deletions compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/XCLBinGen.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ namespace {
FailureOr<std::string> getTargetDir(const std::string &npuVersion) {
if (npuVersion == "npu1") return std::string{"target_aie_ml"};
if (npuVersion == "npu4") return std::string{"target_aie2p"};
llvm::errs() << "unsupported NPUVersion: " << npuVersion;
llvm::errs() << "unsupported NPUVersion: " << npuVersion << "\n";
return failure();
}

Expand Down Expand Up @@ -152,7 +152,8 @@ FailureOr<Path> findVitis(std::optional<Path> &vitisDir,
return failure();
}
if (!std::filesystem::exists(licenseFile)) {
llvm::errs() << "ERROR: license file" << licenseFile << " does not exist";
llvm::errs() << "ERROR: license file" << licenseFile << " does not exist"
<< "\n";
return failure();
}
}
Expand Down Expand Up @@ -215,7 +216,7 @@ std::pair<std::string, std::vector<std::string>> makeChessArgs(
archVersion = "21";
modelDir = "aie2p";
} else {
llvm::errs() << "unsupported NPU version: " << npuVersion;
llvm::errs() << "unsupported NPU version: " << npuVersion << "\n";
llvm::report_fatal_error("unsupported NPU version");
}

Expand Down Expand Up @@ -475,7 +476,7 @@ static FailureOr<Path> assembleStringUsing(
if (auto maybeErr = dumpStrToDisk(inputFileStr, inputFile.string());
maybeErr.has_value()) {
llvm::errs() << "Failed to dump to disk " << inputFile.string()
<< " because: " << maybeErr;
<< " because: " << maybeErr << "\n";
return failure();
}

Expand All @@ -487,7 +488,8 @@ static FailureOr<Path> assembleStringUsing(
}
if (failed(assembler(inputFile.string(), outputFile.string(), extraArgs,
workDir, toolDir, npuVersion, verbose))) {
llvm::errs() << "Failed to assemble " << outputFileName << ".o";
llvm::errs() << "Failed to assemble " << outputFileName << ".o"
<< "\n";
return failure();
}
return outputFile;
Expand Down Expand Up @@ -532,22 +534,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);
if (useChess) {
if (verbose) llvm::outs() << "using chess for ukernel codegen\n";
FailureOr<Path> maybeVitisDir = findVitis(vitisDir, npuVersion);
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";
Expand Down Expand Up @@ -579,13 +590,14 @@ static LogicalResult generateCoreElfFiles(
{
auto bcfOutput = openOutputFile(bcfPath.string(), &errorMessage);
if (!bcfOutput) {
llvm::errs() << "failed to open bcf file because: " << errorMessage;
llvm::errs() << "failed to open bcf file because: " << errorMessage
<< "\n";
return failure();
}

if (failed(mlir::iree_compiler::AMDAIE::AIETranslateToBCF(
deviceOp, bcfOutput->os(), col, row))) {
llvm::errs() << "Failed to generate BCF";
llvm::errs() << "Failed to generate BCF\n";
return failure();
}
bcfOutput->keep();
Expand Down Expand Up @@ -614,7 +626,7 @@ static LogicalResult generateCoreElfFiles(
openOutputFile(ldscriptPath.string(), &errorMessage);
if (!ldscriptOutput) {
llvm::errs() << "Failed to open ldscript file because: "
<< errorMessage;
<< errorMessage << "\n";
return failure();
}
if (failed(mlir::iree_compiler::AMDAIE::AIETranslateToLdScript(
Expand Down Expand Up @@ -654,7 +666,7 @@ static LogicalResult generateCDO(MLIRContext *context, AIE::DeviceOp deviceOp,
deviceOp = *copy.getOps<AIE::DeviceOp>().begin();
if (failed(mlir::iree_compiler::AMDAIE::AIETranslateToCDODirect(
deviceOp, tempDir.string()))) {
llvm::errs() << "failed to emit CDO";
llvm::errs() << "failed to emit CDO\n";
return failure();
}
copy->erase();
Expand Down Expand Up @@ -750,7 +762,7 @@ static LogicalResult generateXCLBin(
dumpStrToDisk(memTopologyData, memTopologyJsonFile.string());
maybeErr.has_value()) {
llvm::errs() << "failed to dump to disk mem_topology.json because: "
<< *maybeErr;
<< *maybeErr << "\n";
return failure();
}
}
Expand Down Expand Up @@ -797,7 +809,7 @@ static LogicalResult generateXCLBin(
dumpStrToDisk(aiePartitionJsonData, aiePartitionJsonFile.string());
maybeErr.has_value()) {
llvm::errs() << "failed to dump to disk aie_partition.json because: "
<< *maybeErr;
<< *maybeErr << "\n";
return failure();
}
}
Expand All @@ -816,7 +828,7 @@ static LogicalResult generateXCLBin(
if (auto maybeErr = dumpStrToDisk(kernelStr, kernelsJsonFile.string());
maybeErr.has_value()) {
llvm::errs() << "failed to dump to disk kernels.json because: "
<< *maybeErr;
<< *maybeErr << "\n";
return failure();
}
}
Expand All @@ -825,7 +837,8 @@ static LogicalResult generateXCLBin(
{
auto designBifOut = openOutputFile(designBifFile.string(), &errorMessage);
if (!designBifOut) {
llvm::errs() << "failed to open design.bif because: " << errorMessage;
llvm::errs() << "failed to open design.bif because: " << errorMessage
<< "\n";
return failure();
}

Expand Down Expand Up @@ -868,7 +881,7 @@ static LogicalResult generateXCLBin(
}
if (iree_aie_bootgen_main(cstrings.size(),
const_cast<const char **>(&cstrings[0]))) {
llvm::errs() << "failed to execute bootgen";
llvm::errs() << "failed to execute bootgen\n";
return failure();
}
}
Expand All @@ -892,14 +905,14 @@ static LogicalResult generateXCLBin(
"--force", "--input", *inputXclbin};

if (failed(runTool(xclbinutilBin.value().string(), inputFlags, verbose))) {
llvm::errs() << "failed to execute xclbinutil";
llvm::errs() << "failed to execute xclbinutil\n";
return failure();
}
auto aieInputPartitionOut =
openInputFile(aieInputPartitionJsonFile.string(), &errorMessage);
if (!aieInputPartitionOut) {
llvm::errs() << "failed to open aie_input_partition.json because: "
<< errorMessage;
<< errorMessage << "\n";
return failure();
}
Expected<json::Value> aieInputPartitionOutValue =
Expand All @@ -913,7 +926,7 @@ static LogicalResult generateXCLBin(
if (!aiePartitionOut) {
llvm::errs() << "failed to open aie aie_input_partition.json for "
"output because: "
<< errorMessage;
<< errorMessage << "\n";
return failure();
}
llvm::Expected<llvm::json::Value> aiePartitionOutValue =
Expand All @@ -931,7 +944,7 @@ static LogicalResult generateXCLBin(
maybeErr.has_value()) {
llvm::errs()
<< "failed to dump to disk aie_input_partition.json because: "
<< errorMessage;
<< errorMessage << "\n";
return failure();
}
flags.insert(flags.end(), {"--input", *inputXclbin});
Expand Down Expand Up @@ -1040,14 +1053,14 @@ static LogicalResult generateUnifiedObject(
}

if (failed(pm.run(moduleOpCopy))) {
llvm::errs() << "Failed to lower to LLVM";
llvm::errs() << "Failed to lower to LLVM\n";
return failure();
}

llvm::LLVMContext llvmContext;
auto llvmModule = translateModuleToLLVMIR(moduleOpCopy, llvmContext);
if (!llvmModule) {
llvm::errs() << "Failed to translate module to LLVMIR";
llvm::errs() << "Failed to translate module to LLVMIR\n";
return failure();
}

Expand Down Expand Up @@ -1081,7 +1094,7 @@ static LogicalResult generateUnifiedObject(
if (auto maybeErr = dumpStrToDisk(inputLLStr, LLVMIRFile.string());
maybeErr.has_value()) {
llvm::errs() << "Failed to dump to disk input.ll"
<< " because: " << maybeErr;
<< " because: " << maybeErr << "\n";
return failure();
}
Path peanoOptBin = peanoDir / "bin" / "opt";
Expand All @@ -1096,7 +1109,7 @@ static LogicalResult generateUnifiedObject(
args.reserve(args.size() + peanoArgs.size());
args.insert(args.end(), peanoArgs.begin(), peanoArgs.end());
if (failed(runTool(peanoOptBin.string(), args, verbose))) {
llvm::errs() << "Failed to optimize ll with peano";
llvm::errs() << "Failed to optimize ll with peano\n";
return failure();
}

Expand Down

0 comments on commit f59995b

Please sign in to comment.