diff --git a/src/Accelerators/Accelerator.cpp b/src/Accelerators/Accelerator.cpp index 47d2e5182d2..a1a7becdb42 100644 --- a/src/Accelerators/Accelerator.cpp +++ b/src/Accelerators/Accelerator.cpp @@ -15,17 +15,13 @@ //===----------------------------------------------------------------------===// #include "src/Accelerators/Accelerator.hpp" -#include -#include namespace onnx_mlir { namespace accel { -std::vector Accelerator::acceleratorTargets; +llvm::SmallVector Accelerator::acceleratorTargets; -Accelerator::~Accelerator() {} - -std::vector Accelerator::getAcceleratorList() { +const llvm::SmallVectorImpl &Accelerator::getAccelerators() { return acceleratorTargets; } diff --git a/src/Accelerators/Accelerator.hpp b/src/Accelerators/Accelerator.hpp index 60225b770e8..2fe48701375 100644 --- a/src/Accelerators/Accelerator.hpp +++ b/src/Accelerators/Accelerator.hpp @@ -17,39 +17,53 @@ #include "include/onnx-mlir/Compiler/OMCompilerTypes.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/Pass/PassManager.h" -#include - -extern bool InitAccelerators(); +#include "llvm/ADT/SmallVector.h" namespace onnx_mlir { namespace accel { +// Note: this function is generated by cmake. +extern bool initAccelerators(); + class Accelerator { public: /// Kinds of accelerators. + // clang-format off enum class Kind { -#include "src/Accelerators/AcceleratorKinds.hpp" + // Note: AcceleratorKinds.inc is generated by cmake. + #include "src/Accelerators/AcceleratorKinds.inc" }; + // clang-format on Accelerator(Kind kind) : kind(kind) {} - virtual ~Accelerator(); + virtual ~Accelerator() = default; /// Getter for the kind of this accelerator. Kind getKind() const { return kind; } - static std::vector getAcceleratorList(); + /// Returns the set of accelerators available. + static const llvm::SmallVectorImpl &getAccelerators(); + + /// Returns whether the accelerator is active. virtual bool isActive() const = 0; + + /// Load the MLIR dialects necessary to generate code for an accelerator. virtual void getOrLoadDialects(mlir::MLIRContext &context) const = 0; + + /// Add the transformations necessary to support the accelerator. virtual void addPasses(mlir::OwningOpRef &module, mlir::PassManager &pm, onnx_mlir::EmissionTargetType &emissionTarget) const = 0; + + /// Register the MLIR dialects required to support an accelerator. virtual void registerDialects(mlir::DialectRegistry ®istry) const = 0; + + /// Initialize the transformation passes required to generate code for an + /// accelerator. virtual void initPasses(int optLevel) const = 0; protected: - // static llvm::SmallPtrSet acceleratorTargets; - static std::vector acceleratorTargets; - + static llvm::SmallVector acceleratorTargets; /// Kind of accelerator. Kind kind; }; diff --git a/src/Accelerators/CMakeLists.txt b/src/Accelerators/CMakeLists.txt index a8c96f43f73..9fe27f5cbd2 100644 --- a/src/Accelerators/CMakeLists.txt +++ b/src/Accelerators/CMakeLists.txt @@ -1,43 +1,47 @@ # SPDX-License-Identifier: Apache-2.0 + set(ACCEL_LIST "") -file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/AcceleratorKinds.hpp "// Enumeration of accelerators \n") + +# Generate the enumerators used to identify the kind of accelerators. +set(AcceleratorKinds "${CMAKE_CURRENT_BINARY_DIR}/AcceleratorKinds.inc") +file(WRITE "${AcceleratorKinds}" "// Enumeration of accelerators.\n") if (ACCELERATORS_TO_BUILD) foreach(t ${ACCELERATORS_TO_BUILD}) message(DEBUG "Targeting ${t}") add_subdirectory(${t}) list(APPEND ACCEL_LIST "${t}Accel") - file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/AcceleratorKinds.hpp " ${t},\n") + file(APPEND "${AcceleratorKinds}" " ${t},\n") endforeach(t) -else (ACCELERATORS_TO_BUILD) endif (ACCELERATORS_TO_BUILD) -file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/AcceleratorKinds.hpp " NONE\n") +file(APPEND "${AcceleratorKinds}" " NONE\n") -file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "bool InitAccelerators() {\n") -file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/AcceleratorOptions.hpp " clEnumValN(\n") +# Generate the enumerators for the --maccel command line option. +set(AcceleratorOptions "${CMAKE_CURRENT_BINARY_DIR}/AcceleratorOptions.inc") +file(REMOVE "${AcceleratorOptions}") if (ACCELERATORS_TO_BUILD) foreach(t ${ACCELERATORS_TO_BUILD}) - file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "extern void create${t}();\n") - file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "create${t}();\n") - file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/AcceleratorOptions.hpp " accel::Accelerator::Kind::${t}, \"${t}\", \"Some Accelerator\"),\n clEnumValN(\n") - endforeach(t) - file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "return true;\n") -else (ACCELERATORS_TO_BUILD) - file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "return false;\n") + file(APPEND "${AcceleratorOptions}" " clEnumValN(accel::Accelerator::Kind::${t}, \"${t}\", \"${t} accelerator\"),\n") + endforeach(t) endif (ACCELERATORS_TO_BUILD) -file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "}\n") -file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/AcceleratorOptions.hpp " accel::Accelerator::Kind::NONE, \"NONE\", \"Some Accelerator\")\n") +file(APPEND "${AcceleratorOptions}" " clEnumValN(accel::Accelerator::Kind::NONE, \"NONE\", \"No accelerator\")\n") -file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "bool InitAccelerators() {\n") +# Generate the function used to initialize the accelerators. +set(InitAccelerators "${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp") +file(WRITE "${InitAccelerators}" "namespace onnx_mlir {\n" + "namespace accel {\n" + "bool initAccelerators() {\n") if (ACCELERATORS_TO_BUILD) foreach(t ${ACCELERATORS_TO_BUILD}) - file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "extern void create${t}();\n") - file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "create${t}();\n") + file(APPEND "${InitAccelerators}" " extern void create${t}();\n") + file(APPEND "${InitAccelerators}" " create${t}();\n") endforeach(t) - file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "return true;\n") + file(APPEND "${InitAccelerators}" " return true;\n") else (ACCELERATORS_TO_BUILD) - file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "return false;\n") + file(APPEND "${InitAccelerators}" " return false;\n") endif (ACCELERATORS_TO_BUILD) -file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/InitAccelerators.cpp "}\n") +file(APPEND "${InitAccelerators}" "}\n" + "} // namespace accel\n" + "} // namespace onnx_mlir\n") add_onnx_mlir_library(InitAccelerators InitAccelerators.cpp @@ -50,10 +54,15 @@ add_onnx_mlir_library(InitAccelerators add_onnx_mlir_library(Accelerator Accelerator.cpp + + EXCLUDE_FROM_OM_LIBS + DEPENDS MLIRIR + INCLUDE_DIRS PUBLIC ${ONNX_MLIR_SRC_ROOT}/include + LINK_LIBS PUBLIC LLVMSupport ) diff --git a/src/Accelerators/NNPA/NNPAAccelerator.cpp b/src/Accelerators/NNPA/NNPAAccelerator.cpp index 1ca8b5fcf71..09ba6a5f83f 100644 --- a/src/Accelerators/NNPA/NNPAAccelerator.cpp +++ b/src/Accelerators/NNPA/NNPAAccelerator.cpp @@ -27,14 +27,15 @@ #define DEBUG_TYPE "NNPAAccelerator" extern llvm::cl::OptionCategory OMNNPAPassOptions; -onnx_mlir::accel::nnpa::NNPAAccelerator *pnnpa = nullptr; - -void createNNPA() { pnnpa = new onnx_mlir::accel::nnpa::NNPAAccelerator; } namespace onnx_mlir { extern llvm::cl::list maccel; namespace accel { + +nnpa::NNPAAccelerator *pnnpa = nullptr; +void createNNPA() { pnnpa = new nnpa::NNPAAccelerator(); } + namespace nnpa { NNPAAccelerator::NNPAAccelerator() : Accelerator(Accelerator::Kind::NNPA) { @@ -42,13 +43,17 @@ NNPAAccelerator::NNPAAccelerator() : Accelerator(Accelerator::Kind::NNPA) { if (!initialized) { initialized = true; - // getAcceleratorList()->push_back(this); acceleratorTargets.push_back(this); } }; +NNPAAccelerator::~NNPAAccelerator() { + assert(initialized && "Expecting the accelerator to be initialized"); + delete pnnpa; +} + bool NNPAAccelerator::isActive() const { - if (initialized || llvm ::any_of(maccel, [](Accelerator::Kind kind) { + if (initialized || llvm::any_of(maccel, [](Accelerator::Kind kind) { return kind == Accelerator::Kind::NNPA; })) { LLVM_DEBUG(llvm::dbgs() << "NNPA accelerator is active\n"); @@ -60,7 +65,7 @@ bool NNPAAccelerator::isActive() const { } void NNPAAccelerator::getOrLoadDialects(mlir::MLIRContext &context) const { - // Load our Dialect in this MLIR Context. + // Load our dialects in this MLIR Context. context.getOrLoadDialect(); context.getOrLoadDialect(); } diff --git a/src/Accelerators/NNPA/NNPAAccelerator.hpp b/src/Accelerators/NNPA/NNPAAccelerator.hpp index 3e2b384fc43..e5be0856e23 100644 --- a/src/Accelerators/NNPA/NNPAAccelerator.hpp +++ b/src/Accelerators/NNPA/NNPAAccelerator.hpp @@ -26,6 +26,13 @@ class NNPAAccelerator final : public Accelerator { public: NNPAAccelerator(); + ~NNPAAccelerator(); + + /// Define classof to be able to use isa<>, cast<>, dyn_cast<>, etc. + static bool classof(const Accelerator *accel) { + return accel->getKind() == Accelerator::Kind::NNPA; + } + static bool classof(const NNPAAccelerator *) { return true; } bool isActive() const final; virtual void getOrLoadDialects(mlir::MLIRContext &context) const final; diff --git a/src/Compiler/CMakeLists.txt b/src/Compiler/CMakeLists.txt index 0c8c21480bf..d727d80c96b 100644 --- a/src/Compiler/CMakeLists.txt +++ b/src/Compiler/CMakeLists.txt @@ -81,6 +81,7 @@ add_onnx_mlir_library(CompilerUtils LINK_LIBS PUBLIC ${OMLibs} CompilerPasses + Accelerator InitAccelerators # Link LLVM libraries necessary to query which target architectures are configured. diff --git a/src/Compiler/CompilerOptions.cpp b/src/Compiler/CompilerOptions.cpp index a907054628e..2d47caf8f16 100644 --- a/src/Compiler/CompilerOptions.cpp +++ b/src/Compiler/CompilerOptions.cpp @@ -84,9 +84,12 @@ llvm::cl::opt march("march", llvm::cl::list maccel("maccel", llvm::cl::desc("Specify an accelerator to generate code for"), + // clang-format off llvm::cl::values( -#include "src/Accelerators/AcceleratorOptions.hpp" - ), + /// Note: AcceleratorOptions is generated by cmake. + #include "src/Accelerators/AcceleratorOptions.inc" + ), + // clang-format on llvm::cl::cat(OnnxMlirOptions), llvm::cl::ValueRequired); llvm::cl::opt VerboseOutput("v", llvm::cl::desc("Use verbose output"), diff --git a/src/Compiler/CompilerUtils.cpp b/src/Compiler/CompilerUtils.cpp index c3b2e39e4bc..ee560cd71bf 100644 --- a/src/Compiler/CompilerUtils.cpp +++ b/src/Compiler/CompilerUtils.cpp @@ -699,14 +699,14 @@ int compileModule(mlir::OwningOpRef &module, setupModule(module, context, outputBaseName); mlir::PassManager pm(&context, mlir::OpPassManager::Nesting::Implicit); - // Initialize accelerator if required + // Initialize accelerator(s) if required. if (!maccel.empty()) { - InitAccelerators(); - for (auto accel : onnx_mlir::accel::Accelerator::getAcceleratorList()) { - if (accel->isActive()) { - accel->getOrLoadDialects(context); - accel->addPasses(module, pm, emissionTarget); - } + onnx_mlir::accel::initAccelerators(); + for (auto *accel : onnx_mlir::accel::Accelerator::getAccelerators()) { + if (!accel->isActive()) + continue; + accel->getOrLoadDialects(context); + accel->addPasses(module, pm, emissionTarget); } } else addPasses(module, pm, emissionTarget); diff --git a/src/Tools/onnx-mlir-opt/CMakeLists.txt b/src/Tools/onnx-mlir-opt/CMakeLists.txt index e280ee07e95..32eea8e7475 100644 --- a/src/Tools/onnx-mlir-opt/CMakeLists.txt +++ b/src/Tools/onnx-mlir-opt/CMakeLists.txt @@ -8,6 +8,7 @@ add_onnx_mlir_executable(onnx-mlir-opt LINK_LIBS PRIVATE ${OMLibs} CompilerOptions + Accelerator InitAccelerators MLIRAffineTransforms MLIRLinalgTransforms diff --git a/src/Tools/onnx-mlir-opt/onnx-mlir-opt.cpp b/src/Tools/onnx-mlir-opt/onnx-mlir-opt.cpp index 0fac0c8d969..cf144d483d5 100644 --- a/src/Tools/onnx-mlir-opt/onnx-mlir-opt.cpp +++ b/src/Tools/onnx-mlir-opt/onnx-mlir-opt.cpp @@ -82,9 +82,6 @@ void scanAndSetOptLevel(int argc, char **argv) { } int main(int argc, char **argv) { - // Initialize accelerators if they exist. - bool hasAccelerators = InitAccelerators(); - mlir::DialectRegistry registry; registry.insert(); registry.insert(); @@ -99,9 +96,12 @@ int main(int argc, char **argv) { registry.insert(); registry.insert(); + // Initialize accelerators if they exist. + bool hasAccelerators = onnx_mlir::accel::initAccelerators(); + // Register dialects for accelerators. if (hasAccelerators) - for (auto accel : onnx_mlir::accel::Accelerator::getAcceleratorList()) + for (auto *accel : onnx_mlir::accel::Accelerator::getAccelerators()) if (accel->isActive()) accel->registerDialects(registry); @@ -118,9 +118,10 @@ int main(int argc, char **argv) { onnx_mlir::initOMPasses(OptimizationLevel); onnx_mlir::initMLIRPasses(); + // Initialize passes for accelerators. if (hasAccelerators) - for (auto accel : onnx_mlir::accel::Accelerator::getAcceleratorList()) + for (auto *accel : onnx_mlir::accel::Accelerator::getAccelerators()) if (accel->isActive()) accel->initPasses(OptimizationLevel); diff --git a/src/Tools/onnx-mlir-reduce/CMakeLists.txt b/src/Tools/onnx-mlir-reduce/CMakeLists.txt index 9fbd3f8a8a9..6f40575a99c 100644 --- a/src/Tools/onnx-mlir-reduce/CMakeLists.txt +++ b/src/Tools/onnx-mlir-reduce/CMakeLists.txt @@ -7,6 +7,8 @@ add_onnx_mlir_executable(onnx-mlir-reduce LINK_LIBS PRIVATE ${OMLibs} + Accelerator + InitAccelerators MLIRAffineTransforms MLIRLinalgTransforms MLIRReduceLib diff --git a/src/Tools/onnx-mlir-reduce/onnx-mlir-reduce.cpp b/src/Tools/onnx-mlir-reduce/onnx-mlir-reduce.cpp index 02654e1b2b9..96970be0171 100644 --- a/src/Tools/onnx-mlir-reduce/onnx-mlir-reduce.cpp +++ b/src/Tools/onnx-mlir-reduce/onnx-mlir-reduce.cpp @@ -19,6 +19,7 @@ #include "mlir/InitAllDialects.h" #include "mlir/InitAllPasses.h" #include "mlir/Tools/mlir-reduce/MlirReduceMain.h" +#include "src/Accelerators/Accelerator.hpp" #include "src/Dialect/Krnl/KrnlOps.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -37,6 +38,12 @@ static void registerDialects(DialectRegistry ®istry) { registry.insert(); registry.insert(); + + // Initialize and register dialects used by accelerators. + if (onnx_mlir::accel::initAccelerators()) + for (auto *accel : onnx_mlir::accel::Accelerator::getAccelerators()) + if (accel->isActive()) + accel->registerDialects(registry); } int main(int argc, char **argv) {