Skip to content

Commit

Permalink
[amd-aie] add entry_point names to kernels.json and flatbuffer (nod-a…
Browse files Browse the repository at this point in the history
…i#58)

Currently the kernel name is hard-coded as "MLIR-AIE". This PR changes
that to take the dispatch entry point function name and also adds that
to the flatbuffer so it can be read at runtime.
I tested this e2e in a sperate branch and will send PRs with the runtime
bits needed to get the matmul working through IREE HAL next.
  • Loading branch information
nirvedhmeshram authored Dec 20, 2023
1 parent 2d5b2e6 commit e3feed6
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 15 deletions.
54 changes: 39 additions & 15 deletions compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/AIETarget.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
#include "iree/compiler/Codegen/Dialect/IREECodegenDialect.h"
#include "iree/compiler/Utils/FlatbufferUtils.h"
#include "llvm/Bitcode/BitcodeWriter.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
Expand All @@ -39,6 +40,8 @@
#include "runtime/plugins/AMD-AIE/iree-amd-aie/schemas/xrt_executable_def_builder.h"
#include "runtime/plugins/AMD-AIE/iree-amd-aie/schemas/xrt_executable_def_reader.h"

#define DEBUG_TYPE "aie-target"

// Forward declaration of some translate methods from AIE. THis is done
// here since the headers in MLIR-AIE repo are in a place that is
// not where you would expect.
Expand Down Expand Up @@ -120,8 +123,8 @@ class AIETargetBackend final : public IREE::HAL::TargetBackend {
addConfig("target_arch", StringAttr::get(context, "chip-tbd"));
auto configAttr = b.getDictionaryAttr(configItems);
return IREE::HAL::ExecutableTargetAttr::get(
context, b.getStringAttr("amd-aie"), b.getStringAttr("elf"),
configAttr);
context, b.getStringAttr("amd-aie"),
b.getStringAttr("amdaie-xclbin-fb"), configAttr);
}

AMDAIEOptions options;
Expand Down Expand Up @@ -394,6 +397,7 @@ LogicalResult generateCoreElfFiles(ModuleOp moduleOp, Artifact &objFile,
LogicalResult generateXCLBin(MLIRContext *context, ModuleOp moduleOp,
std::string workDir, const AMDAIEOptions &options,
const XclBinGeneratorKit &toolkit,
ArrayRef<std::string> entryPointNames,
raw_ostream &xclBin) {
// This corresponds to `process_host_cgen`, which is listed as host
// compilation in aiecc.py... not sure we need this.
Expand Down Expand Up @@ -673,12 +677,22 @@ LogicalResult generateXCLBin(MLIRContext *context, ModuleOp moduleOp,
}
SmallVector<char, 0> kernelsJsonDataString;
llvm::raw_svector_ostream ostream(kernelsJsonDataString);
std::string kernels_json_data = R"(
std::string kernels_json_data = "";
for (auto entryPointName : entryPointNames) {
// this print is useful for non-IREE test benches as they need the
// Kernel names
if (toolkit.isVerbose()) {
llvm::errs() << "Entry point name: " << entryPointName << "\n";
}
LLVM_DEBUG(
{ llvm::errs() << "Entry point name: " << entryPointName << "\n"; });
kernels_json_data = kernels_json_data + R"(
{
"ps-kernels": {
"kernels": [
{
"name": "MLIR_AIE",
"name": ")" +
entryPointName + R"(",
"type": "dpu",
"extended-data": {
"subtype": "DPU",
Expand Down Expand Up @@ -731,6 +745,7 @@ LogicalResult generateXCLBin(MLIRContext *context, ModuleOp moduleOp,
}
}
)";
}
ostream << kernels_json_data;
kernelsJsonFile->write(kernelsJsonDataString);
kernelsJsonFile->close();
Expand Down Expand Up @@ -829,7 +844,7 @@ LogicalResult generateXCLBin(MLIRContext *context, ModuleOp moduleOp,
xclbinFile->keep();

if (!xclbinFile->readInto(xclBin)) {
return moduleOp.emitOpError("failed to get xlcbin bits");
return moduleOp.emitOpError("failed to get xclbin bits");
}
return success();
}
Expand Down Expand Up @@ -910,19 +925,30 @@ LogicalResult AIETargetBackend::serializeExecutable(
clonedModuleOp->erase();
}

// Create a flatbuffer containing (for now) lx6 instructions and xclbin.
FlatbufferBuilder builder;
iree_amd_aie_hal_xrt_ExecutableDef_start_as_root(builder);

auto ipuInstrsRef = builder.createInt32Vec(ipuInstrs);
iree_amd_aie_hal_xrt_ExecutableDef_asm_instrs_add(builder, ipuInstrsRef);

llvm::SmallVector<char, 0> xclbin;
llvm::raw_svector_ostream ostream(xclbin);
// collect names of kernels as they need to be in kernels.json
// generated by `generateXCLBin`
SmallVector<std::string> entryPointNames;
for (auto exportOp : variantOp.getExportOps()) {
entryPointNames.emplace_back(exportOp.getSymName());
}
if (failed(generateXCLBin(context, moduleOp, workDir.value(), getOptions(),
toolkit, ostream))) {
toolkit, entryPointNames, ostream))) {
return moduleOp.emitOpError() << "failed to generate XCLbin";
}
// Serialize the executable to flatbuffer format
FlatbufferBuilder builder;
iree_amd_aie_hal_xrt_ExecutableDef_start_as_root(builder);
auto entryPointsRef = builder.createStringVec(entryPointNames);

iree_amd_aie_hal_xrt_ExecutableDef_entry_points_add(builder, entryPointsRef);

iree_amd_aie_hal_xrt_AsmInstDef_vec_start(builder);
auto ipuInstrsVec = builder.createInt32Vec(ipuInstrs);
iree_amd_aie_hal_xrt_AsmInstDef_vec_push_create(builder, ipuInstrsVec);
auto ipuInstrsRef = iree_amd_aie_hal_xrt_AsmInstDef_vec_end(builder);
iree_amd_aie_hal_xrt_ExecutableDef_asm_instrs_add(builder, ipuInstrsRef);
llvm::StringRef xclbinStringView(xclbin.begin(), xclbin.size());
auto xclbinStringRef = builder.createString(xclbinStringView);
iree_amd_aie_hal_xrt_ExecutableDef_xclbin_add(builder, xclbinStringRef);
Expand All @@ -935,8 +961,6 @@ LogicalResult AIETargetBackend::serializeExecutable(
binaryOp.setMimeTypeAttr(
executableBuilder.getStringAttr("application/x-flatbuffers"));

// TODO(JamesNewling) We need to test that the above logic is correct,
// returning success here to enable runtime testing.
return success();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -298,4 +298,6 @@ LogicalResult XclBinGeneratorKit::runXclBinUtil(ArrayRef<std::string> flags,
return runCommand(cmdLine);
}

bool XclBinGeneratorKit::isVerbose() const { return verbose; }

} // namespace mlir::iree_compiler::AMDAIE
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,8 @@ class XclBinGeneratorKit {
LogicalResult runXclBinUtil(ArrayRef<std::string> flags, Artifact &input,
Artifact &output) const;

bool isVerbose() const;

private:
std::string peanoInstallDir;

Expand Down

0 comments on commit e3feed6

Please sign in to comment.