diff --git a/CMakeLists.txt b/CMakeLists.txt index 3a5c39e..9db14ef 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -30,4 +30,59 @@ raptor_ensure_symlink( "${CMAKE_CURRENT_SOURCE_DIR}/test/PIM" ) +# Patch onnx-mlir sources for PIM accelerator support. +# Each patch searches for a context-aware anchor string rather than relying on +# line numbers, so that moderate upstream changes are tolerated. +function(raptor_apply_patch file_path anchor replacement description) + file(READ "${file_path}" contents) + + # Already applied – replacement text is present + string(FIND "${contents}" "${replacement}" already_applied_pos) + if(NOT already_applied_pos EQUAL -1) + message(STATUS "Patch already applied: ${description}") + return() + endif() + + # Anchor must exist for the patch to be applicable + string(FIND "${contents}" "${anchor}" anchor_pos) + if(anchor_pos EQUAL -1) + message(FATAL_ERROR + "Patch anchor not found – onnx-mlir may have changed.\n" + " Patch : ${description}\n" + " File : ${file_path}\n" + " Anchor: ${anchor}" + ) + endif() + + string(REPLACE "${anchor}" "${replacement}" patched "${contents}") + file(WRITE "${file_path}" "${patched}") + message(STATUS "Patch applied: ${description}") +endfunction() + +set(ONNX_MLIR_DIR "${CMAKE_CURRENT_SOURCE_DIR}/onnx-mlir") + +# 1) Disable the absl dependency +raptor_apply_patch( + "${ONNX_MLIR_DIR}/CMakeLists.txt" + "find_package(absl REQUIRED)" + "#find_package(absl REQUIRED)" + "Disable find_package(absl)" +) + +# 2) Register PIM compiler options alongside NNPA +raptor_apply_patch( + "${ONNX_MLIR_DIR}/src/Accelerators/Accelerator.hpp" + "#include \"src/Accelerators/NNPA/Compiler/NNPACompilerOptions.hpp\"" + "#include \"src/Accelerators/NNPA/Compiler/NNPACompilerOptions.hpp\"\n#include \"src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp\"" + "Add PIM compiler options include" +) + +# 3) Short-circuit output emission for the PIM accelerator +raptor_apply_patch( + "${ONNX_MLIR_DIR}/src/Compiler/CompilerUtils.cpp" + "switch (emissionTarget) {\n case EmitObj: {" + "if (llvm::is_contained(maccel, accel::Accelerator::Kind::PIM))\n return CompilerSuccess;\n switch (emissionTarget) {\n case EmitObj: {" + "Skip output emission for PIM accelerator" +) + add_subdirectory(onnx-mlir) diff --git a/src/PIM/CMakeLists.txt b/src/PIM/CMakeLists.txt index 03d8cc9..6a704c7 100644 --- a/src/PIM/CMakeLists.txt +++ b/src/PIM/CMakeLists.txt @@ -18,8 +18,9 @@ add_subdirectory(Common) add_onnx_mlir_library(OMPIMAccel PimAccelerator.cpp Transforms/PimBufferizationPass.cpp - Pass/MessagePass.cpp Pass/CountInstructionPass.cpp + Pass/EmitPimJsonPass.cpp + Pass/MessagePass.cpp EXCLUDE_FROM_OM_LIBS diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index 16438d7..605028e 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -433,7 +433,7 @@ std::string getMemorySizeAsString(size_t size) { return std::to_string(size) + " Bytes"; } -int compileModuleToPIMJSON(const OwningOpRef& moduleOpRef, std::string& outputDirPath) { +OnnxMlirCompilerErrorCodes compileModuleToPIMJSON(const OwningOpRef& moduleOpRef, std::string& outputDirPath) { ModuleOp moduleOp = moduleOpRef.get(); if (pimEmissionTarget != EmitPimCodegen) { diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp index bddd179..695b0a2 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -17,40 +17,44 @@ namespace onnx_mlir { -llvm::cl::opt pimEmissionTarget( - llvm::cl::desc("[Optional] Choose PIM-related target to emit " - "(once selected it will cancel the other targets):"), - llvm::cl::values(clEnumVal(EmitSpatial, "Lower model to spatial IR")), - llvm::cl::values(clEnumVal(EmitPim, "Lower model to PIM IR")), - llvm::cl::values( - clEnumVal(EmitPimBufferized, "Lower model to PIM IR and bufferize it")), - llvm::cl::values(clEnumVal(EmitPimCodegen, "Lower model to PIM IR and " - "generate code for PIM")), - llvm::cl::init(EmitPimCodegen), llvm::cl::cat(OnnxMlirOptions)); +llvm::cl::opt pimOutputDir("pim-output-dir", + llvm::cl::desc("Directory where pim json code will be emitted"), + llvm::cl::init("pim"), + llvm::cl::cat(OnnxMlirOptions)); -llvm::cl::opt pimOnlyCodegen("pim-only-codegen", - llvm::cl::desc("Only generate code for PIM (assume input is already in " - "bufferized PIM IR)"), - llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions)); +llvm::cl::opt pimEmissionTarget( + llvm::cl::desc("[Optional] Choose PIM-related target to emit (once selected it will cancel the other targets):"), + llvm::cl::values(clEnumVal(EmitSpatial, "Lower model to spatial IR")), + llvm::cl::values(clEnumVal(EmitPim, "Lower model to PIM IR")), + llvm::cl::values(clEnumVal(EmitPimBufferized, "Lower model to PIM IR and bufferize it")), + llvm::cl::values(clEnumVal(EmitPimCodegen, "Lower model to PIM IR and generate code for PIM")), + llvm::cl::init(EmitPimCodegen), + llvm::cl::cat(OnnxMlirOptions)); + +llvm::cl::opt + pimOnlyCodegen("pim-only-codegen", + llvm::cl::desc("Only generate code for PIM (assume input is already in bufferized PIM IR)"), + llvm::cl::init(false), + llvm::cl::cat(OnnxMlirOptions)); llvm::cl::opt useExperimentalConvImpl("use-experimental-conv-impl", - llvm::cl::desc("Use experimental implementation for convolution"), - llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions)); + llvm::cl::desc("Use experimental implementation for convolution"), + llvm::cl::init(false), + llvm::cl::cat(OnnxMlirOptions)); -llvm::cl::opt crossbarSize("crossbar-size", - llvm::cl::desc("Width and heigth of a single crossbar"), llvm::cl::init(2)); +llvm::cl::opt + crossbarSize("crossbar-size", llvm::cl::desc("Width and heigth of a single crossbar"), llvm::cl::init(2)); -llvm::cl::opt crossbarCountInCore("crossbar-count", - llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(2)); +llvm::cl::opt + crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(2)); llvm::cl::opt coresCount("core-count", - llvm::cl::desc("Number of cores in the chip. `-1` to use the minimum " - "amount of cores."), - llvm::cl::init(-1)); + llvm::cl::desc("Number of cores in the chip. `-1` to use the minimum amount of cores."), + llvm::cl::init(-1)); -llvm::cl::opt ignoreConcatError("ignore-concat-error", - llvm::cl::desc( - "Ignore ConcatOp corner case: do not assert and do a simplification"), - llvm::cl::init(false)); +llvm::cl::opt + ignoreConcatError("ignore-concat-error", + llvm::cl::desc("Ignore ConcatOp corner case: do not assert and do a simplification"), + llvm::cl::init(false)); -} // namespace onnx_mlir \ No newline at end of file +} // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimCompilerOptions.hpp b/src/PIM/Compiler/PimCompilerOptions.hpp index 8b2a51b..053b625 100644 --- a/src/PIM/Compiler/PimCompilerOptions.hpp +++ b/src/PIM/Compiler/PimCompilerOptions.hpp @@ -21,7 +21,8 @@ typedef enum { } PimEmissionTargetType; extern llvm::cl::OptionCategory OnnxMlirOptions; -extern llvm::cl::opt pimEmissionTarget; +extern llvm::cl::opt pimOutputDir; +extern llvm::cl::opt pimEmissionTarget; extern llvm::cl::opt pimOnlyCodegen; extern llvm::cl::opt useExperimentalConvImpl; diff --git a/src/PIM/Compiler/PimCompilerUtils.hpp b/src/PIM/Compiler/PimCompilerUtils.hpp index 76581be..077d566 100644 --- a/src/PIM/Compiler/PimCompilerUtils.hpp +++ b/src/PIM/Compiler/PimCompilerUtils.hpp @@ -13,7 +13,7 @@ void addPassesPim(mlir::OwningOpRef& module, EmissionTargetType& emissionTarget, std::string outputNameNoExt); -int compileModuleToPIMJSON(const mlir::OwningOpRef& moduleOpRef, - std::string& outputDirName); +OnnxMlirCompilerErrorCodes compileModuleToPIMJSON(const mlir::OwningOpRef& moduleOpRef, + std::string& outputDirName); } // namespace onnx_mlir diff --git a/src/PIM/Pass/EmitPimJsonPass.cpp b/src/PIM/Pass/EmitPimJsonPass.cpp new file mode 100644 index 0000000..1b857c2 --- /dev/null +++ b/src/PIM/Pass/EmitPimJsonPass.cpp @@ -0,0 +1,45 @@ +#include "mlir/Pass/Pass.h" + +#include + +#include "Compiler/PimCompilerOptions.hpp" +#include "Compiler/PimCompilerUtils.hpp" + +using namespace mlir; + +namespace onnx_mlir { + +namespace { + +struct EmitPimJsonPass : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(EmitPimJsonPass); + StringRef getArgument() const override { return "emit-pim-json-pass"; } + StringRef getDescription() const override { return "Emit json code for the pim simulators"; } + + EmitPimJsonPass() {} + EmitPimJsonPass(const EmitPimJsonPass& pass) {} + + void runOnOperation() final { + ModuleOp moduleOp = getOperation(); + + std::filesystem::path pimDir(pimOutputDir.data()); + + std::error_code error_code; + std::filesystem::create_directories(pimDir, error_code); + if (error_code) { + moduleOp.emitError("Failed to create PIM output directory: " + error_code.message()); + signalPassFailure(); + return; + } + + int compiler_error_code = compileModuleToPIMJSON(moduleOp, pimOutputDir); + if (compiler_error_code != CompilerSuccess) + signalPassFailure(); + } +}; + +} // namespace + +std::unique_ptr createEmitPimJsonPass() { return std::make_unique(); } + +} // namespace onnx_mlir diff --git a/src/PIM/Pass/PimPasses.hpp b/src/PIM/Pass/PimPasses.hpp index 9b1c0f1..f5a320f 100644 --- a/src/PIM/Pass/PimPasses.hpp +++ b/src/PIM/Pass/PimPasses.hpp @@ -15,6 +15,8 @@ std::unique_ptr createSpatialToPIMPass(); std::unique_ptr createBufferizePimPass(); +std::unique_ptr createEmitPimJsonPass(); + std::unique_ptr createMessagePass(std::string message); std::unique_ptr createCountInstructionPass(); diff --git a/src/PIM/PimAccelerator.cpp b/src/PIM/PimAccelerator.cpp index 8c5bcd4..13ce3a6 100644 --- a/src/PIM/PimAccelerator.cpp +++ b/src/PIM/PimAccelerator.cpp @@ -8,7 +8,7 @@ #include "mlir/Dialect/Tosa/IR/TosaOps.h" #include "mlir/Dialect/Tosa/Transforms/Passes.h" #include "mlir/IR/BuiltinTypes.h" -#include "mlir/Transforms/Passes.h" + #include "llvm/Support/Debug.h" #include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp" @@ -19,40 +19,40 @@ #include "src/Accelerators/PIM/Pass/PimPasses.hpp" #include "src/Accelerators/PIM/PimAccelerator.hpp" -#include - #define DEBUG_TYPE "PimAccelerator" namespace onnx_mlir { namespace accel { -Accelerator *createPIM() { return PimAccelerator::getInstance(); } +Accelerator* createPIM() { return PimAccelerator::getInstance(); } -PimAccelerator *PimAccelerator::instance = nullptr; +PimAccelerator* PimAccelerator::instance = nullptr; -PimAccelerator *PimAccelerator::getInstance() { +PimAccelerator* PimAccelerator::getInstance() { if (instance == nullptr) instance = new PimAccelerator(); return instance; } -PimAccelerator::PimAccelerator() : Accelerator(Accelerator::Kind::PIM) { +PimAccelerator::PimAccelerator() +: Accelerator(Kind::PIM) { LLVM_DEBUG(llvm::dbgs() << "Creating a PIM accelerator\n"); acceleratorTargets.push_back(this); -}; +} PimAccelerator::~PimAccelerator() { delete instance; } uint64_t PimAccelerator::getVersionNumber() const { return 0x000001; } -void PimAccelerator::addPasses(mlir::OwningOpRef &module, - mlir::PassManager &pm, onnx_mlir::EmissionTargetType &emissionTarget, - std::string outputNameNoExt) const { +void PimAccelerator::addPasses(OwningOpRef& module, + PassManager& pm, + EmissionTargetType& emissionTarget, + std::string outputNameNoExt) const { LLVM_DEBUG(llvm::dbgs() << "Adding passes for PIM accelerator\n"); addPassesPim(module, pm, emissionTarget, outputNameNoExt); } -void PimAccelerator::registerDialects(mlir::DialectRegistry ®istry) const { +void PimAccelerator::registerDialects(DialectRegistry& registry) const { LLVM_DEBUG(llvm::dbgs() << "Registering dialects for PIM accelerator\n"); registry.insert(); registry.insert(); @@ -61,8 +61,7 @@ void PimAccelerator::registerDialects(mlir::DialectRegistry ®istry) const { registry.insert(); tensor::registerBufferizableOpInterfaceExternalModels(registry); arith::registerBufferizableOpInterfaceExternalModels(registry); - mlir::bufferization::func_ext::registerBufferizableOpInterfaceExternalModels( - registry); + bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(registry); spatial::registerBufferizableOpInterfaceExternalModels(registry); spatial::registerONNXBufferizableOpInterfaceExternalModels(registry); pim::registerBufferizableOpInterfaceExternalModels(registry); @@ -70,11 +69,11 @@ void PimAccelerator::registerDialects(mlir::DialectRegistry ®istry) const { void PimAccelerator::registerPasses(int optLevel) const { LLVM_DEBUG(llvm::dbgs() << "Registering passes for PIM accelerator\n"); - // Register here all the passes that could be used - mlir::registerPass(createONNXToSpatialPass); - mlir::registerPass(createSpatialToGraphvizPass); - mlir::registerPass(createSpatialToPIMPass); - mlir::registerPass(createBufferizePimPass); + registerPass(createONNXToSpatialPass); + registerPass(createSpatialToGraphvizPass); + registerPass(createSpatialToPIMPass); + registerPass(createBufferizePimPass); + registerPass(createEmitPimJsonPass); } void PimAccelerator::configurePasses() const { @@ -82,27 +81,26 @@ void PimAccelerator::configurePasses() const { // TODO: This does nothing for now. } -mlir::MemRefType PimAccelerator::convertTensorTypeToMemRefType( - const mlir::TensorType tensorType) const { +MemRefType PimAccelerator::convertTensorTypeToMemRefType(const TensorType tensorType) const { // Do not convert tensor types to memref types. return nullptr; } -void PimAccelerator::conversionTargetONNXToKrnl( - mlir::ConversionTarget &target) const { +void PimAccelerator::conversionTargetONNXToKrnl(ConversionTarget& target) const { target.addLegalDialect(); } -void PimAccelerator::rewritePatternONNXToKrnl(mlir::RewritePatternSet &patterns, - mlir::TypeConverter &typeConverter, mlir::MLIRContext *ctx) const { +void PimAccelerator::rewritePatternONNXToKrnl(RewritePatternSet& patterns, + TypeConverter& typeConverter, + MLIRContext* ctx) const { // TODO: Add patterns for conversion } -void PimAccelerator::conversionTargetKrnlToLLVM( - mlir::ConversionTarget &target) const {} +void PimAccelerator::conversionTargetKrnlToLLVM(ConversionTarget& target) const {} -void PimAccelerator::rewritePatternKrnlToLLVM(mlir::RewritePatternSet &patterns, - mlir::LLVMTypeConverter &typeConverter, mlir::MLIRContext *ctx) const { +void PimAccelerator::rewritePatternKrnlToLLVM(RewritePatternSet& patterns, + LLVMTypeConverter& typeConverter, + MLIRContext* ctx) const { // We should not need this, since we offload it all to PIM. } diff --git a/validation/validate.py b/validation/validate.py index efaf037..cea876a 100644 --- a/validation/validate.py +++ b/validation/validate.py @@ -43,7 +43,7 @@ def build_dump_ranges(config_path, outputs_descriptor): def run_pim_simulator(simulator_dir, pim_dir, output_bin_path, dump_ranges): subprocess.run( - ["cargo", "run", "--package", "pim-simulator", "--bin", "pim-simulator", "--", + ["cargo", "run", "--release", "--package", "pim-simulator", "--bin", "pim-simulator", "--", "-f", str(pim_dir), "-o", str(output_bin_path), "-d", dump_ranges], cwd=simulator_dir, check=True )