From 18c881556dd6b9150d362093ee239ca98c2c3b07 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram <96096277+nirvedhmeshram@users.noreply.github.com> Date: Sat, 15 Jun 2024 14:51:27 -0500 Subject: [PATCH] Modify aie2xclbin tool to allow merging kernels in existing XCLBIN (#1561) --- tools/aie2xclbin/XCLBinGen.cpp | 83 ++++++++++++++++++++++++++------- tools/aie2xclbin/XCLBinGen.h | 5 +- tools/aie2xclbin/aie2xclbin.cpp | 7 ++- 3 files changed, 74 insertions(+), 21 deletions(-) diff --git a/tools/aie2xclbin/XCLBinGen.cpp b/tools/aie2xclbin/XCLBinGen.cpp index 03946b250b..a36fad41a1 100644 --- a/tools/aie2xclbin/XCLBinGen.cpp +++ b/tools/aie2xclbin/XCLBinGen.cpp @@ -438,7 +438,8 @@ static json::Object makeKernelJSON(std::string name, std::string id, static LogicalResult generateXCLBin(MLIRContext *context, ModuleOp moduleOp, XCLBinGenConfig &TK, - const StringRef &Output) { + const StringRef &Output, + const StringRef &inputXclbin = "") { std::string errorMessage; // Create mem_topology.json. SmallString<64> memTopologyJsonFile(TK.TempDir); @@ -507,7 +508,8 @@ static LogicalResult generateXCLBin(MLIRContext *context, ModuleOp moduleOp, "type": "PRIMARY", "pdi_id": "0x01", "dpu_kernel_ids": [ - "0x901" + ")" + TK.XCLBinKernelID + + R"(" ], "pre_cdo_groups": [ "0xC1" @@ -581,22 +583,66 @@ static LogicalResult generateXCLBin(MLIRContext *context, ModuleOp moduleOp, if (runTool(bootgenBin, flags, TK.Verbose) != 0) return moduleOp.emitOpError("failed to execute bootgen"); } - + SmallVector flags; // Execute the xclbinutil command. + std::string memArg = "MEM_TOPOLOGY:JSON:" + std::string(memTopologyJsonFile); + std::string partArg = + "AIE_PARTITION:JSON:" + std::string(aiePartitionJsonFile); { - std::string memArg = - "MEM_TOPOLOGY:JSON:" + std::string(memTopologyJsonFile); - std::string partArg = - "AIE_PARTITION:JSON:" + std::string(aiePartitionJsonFile); - SmallVector flags{"--add-replace-section", - memArg, - "--add-kernel", - std::string(kernelsJsonFile), - "--add-replace-section", - partArg, - "--force", - "--output", - std::string(Output)}; + + if (!inputXclbin.empty()) { + // Create aie_partition.json. + SmallString<64> aieInputPartitionJsonFile(TK.TempDir); + sys::path::append(aieInputPartitionJsonFile, "aie_input_partition.json"); + + std::string inputPartArg = + "AIE_PARTITION:JSON:" + std::string(aieInputPartitionJsonFile); + SmallVector inputFlags{"--dump-section", inputPartArg, + "--force", "--input", + std::string(inputXclbin)}; + if (auto xclbinutil = sys::findProgramByName("xclbinutil")) { + if (runTool(*xclbinutil, inputFlags, TK.Verbose) != 0) + return moduleOp.emitOpError("failed to execute xclbinutil"); + } else { + return moduleOp.emitOpError("could not find xclbinutil"); + } + auto aieInputPartitionOut = + openInputFile(aieInputPartitionJsonFile, &errorMessage); + if (!aieInputPartitionOut) + return moduleOp.emitOpError(errorMessage); + Expected aieInputPartitionOutValue = + llvm::json::parse(aieInputPartitionOut->getBuffer()); + json::Array *aieInputPartionPDIs; + aieInputPartionPDIs = aieInputPartitionOutValue->getAsObject() + ->getObject("aie_partition") + ->getArray("PDIs"); + auto aiePartitionOut = openInputFile(aiePartitionJsonFile, &errorMessage); + if (!aiePartitionOut) + return moduleOp.emitOpError(errorMessage); + llvm::Expected aiePartitionOutValue = + llvm::json::parse(aiePartitionOut->getBuffer()); + json::Array *aiePartionPDIs; + aiePartionPDIs = aiePartitionOutValue->getAsObject() + ->getObject("aie_partition") + ->getArray("PDIs"); + aieInputPartionPDIs->insert(aieInputPartionPDIs->end(), + aiePartionPDIs->begin(), + aiePartionPDIs->end()); + // rewrite aie partion json file + auto aiePartitionJsonOut = + openOutputFile(aiePartitionJsonFile, &errorMessage); + if (!aiePartitionJsonOut) + return moduleOp.emitOpError(errorMessage); + aiePartitionJsonOut->os() << formatv("{0:2}", *aieInputPartitionOutValue); + aiePartitionJsonOut->keep(); + flags.insert(flags.end(), {"--input", std::string(inputXclbin)}); + + } else { + flags.insert(flags.end(), {"--add-replace-section", memArg}); + } + flags.insert(flags.end(), {"--add-kernel", std::string(kernelsJsonFile), + "--add-replace-section", partArg, "--force", + "--output", std::string(Output)}); if (auto xclbinutil = sys::findProgramByName("xclbinutil")) { if (runTool(*xclbinutil, flags, TK.Verbose) != 0) @@ -825,7 +871,8 @@ static LogicalResult generateUnifiedObject(MLIRContext *context, LogicalResult xilinx::aie2xclbin(MLIRContext *ctx, ModuleOp moduleOp, XCLBinGenConfig &TK, StringRef OutputNPU, - StringRef OutputXCLBin) { + StringRef OutputXCLBin, + StringRef InputXCLBin = "") { PassManager pm(ctx, moduleOp.getOperationName()); applyConfigToPassManager(TK, pm); @@ -886,7 +933,7 @@ LogicalResult xilinx::aie2xclbin(MLIRContext *ctx, ModuleOp moduleOp, if (failed(generateCDO(ctx, moduleOp, TK))) return moduleOp.emitOpError("Failed to generate CDO"); - if (failed(generateXCLBin(ctx, moduleOp, TK, OutputXCLBin))) + if (failed(generateXCLBin(ctx, moduleOp, TK, OutputXCLBin, InputXCLBin))) return moduleOp.emitOpError("Failed to generate XCLBin"); return success(); diff --git a/tools/aie2xclbin/XCLBinGen.h b/tools/aie2xclbin/XCLBinGen.h index 0ab910e668..a7e115dcf3 100644 --- a/tools/aie2xclbin/XCLBinGen.h +++ b/tools/aie2xclbin/XCLBinGen.h @@ -41,7 +41,8 @@ struct XCLBinGenConfig { void findVitis(XCLBinGenConfig &TK); mlir::LogicalResult aie2xclbin(mlir::MLIRContext *ctx, mlir::ModuleOp moduleOp, - XCLBinGenConfig &TK, llvm::StringRef OutputNPU, - llvm::StringRef OutputXCLBin); + XCLBinGenConfig &TK, llvm::StringRef outputNPU, + llvm::StringRef outputXCLBin, + llvm::StringRef inputXCLBin); } // namespace xilinx diff --git a/tools/aie2xclbin/aie2xclbin.cpp b/tools/aie2xclbin/aie2xclbin.cpp index b03881fdf9..dbd7339aa4 100644 --- a/tools/aie2xclbin/aie2xclbin.cpp +++ b/tools/aie2xclbin/aie2xclbin.cpp @@ -109,6 +109,11 @@ cl::opt cl::desc("Output xclbin filename for CDO/XCLBIN target"), cl::init("final.xclbin"), cl::cat(AIE2XCLBinCat)); +cl::opt inputXCLBinName( + "input-xclbin-name", + cl::desc("input xclbin filename on which new xclbin is merged into"), + cl::init(""), cl::cat(AIE2XCLBinCat)); + cl::opt XCLBinKernelName("xclbin-kernel-name", cl::desc("Kernel name in xclbin file"), cl::init("MLIR_AIE"), @@ -215,7 +220,7 @@ int main(int argc, char *argv[]) { return 1; if (failed(aie2xclbin(&ctx, *owning, TK, NPUInstsName.getValue(), - XCLBinName.getValue()))) + XCLBinName.getValue(), inputXCLBinName.getValue()))) return 1; return 0;