Misc. improvements to Accelerator code. (llvm#1281)

This PR contains several changes/improvements to the Accelerator support code: - use llvm::SmallVector rather than std::vector for accelerator list - move initAccelerators() in the onnx_mlir::accel namespace (to avoid polluting the global namespace) - add documentation to the public member function in the Accelerator class - delete the nnpa pointers when ~NNPAAccelerator is invoked - move createNNPA in the onnx_mlir::accel namespace - initialize accelerator dialects also for onn-mlir-reduce Signed-off-by: Ettore Tiotto etiotto@ca.ibm.com
nod-ai · Mar 30, 2022 · 993f296 · 993f296
1 parent f454310
commit 993f296
Show file tree

Hide file tree

Showing 12 changed files with 102 additions and 56 deletions.
diff --git a/src/Accelerators/Accelerator.cpp b/src/Accelerators/Accelerator.cpp
@@ -15,17 +15,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "src/Accelerators/Accelerator.hpp"
-#include <iostream>
-#include <vector>
 
 namespace onnx_mlir {
 namespace accel {
 
-std::vector<Accelerator *> Accelerator::acceleratorTargets;
+llvm::SmallVector<Accelerator *, 4> Accelerator::acceleratorTargets;
 
-Accelerator::~Accelerator() {}
-
-std::vector<Accelerator *> Accelerator::getAcceleratorList() {
+const llvm::SmallVectorImpl<Accelerator *> &Accelerator::getAccelerators() {
   return acceleratorTargets;
 }
 

diff --git 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 <vector>
-
-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<Accelerator *> getAcceleratorList();
+  /// Returns the set of accelerators available.
+  static const llvm::SmallVectorImpl<Accelerator *> &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<mlir::ModuleOp> &module,
       mlir::PassManager &pm,
       onnx_mlir::EmissionTargetType &emissionTarget) const = 0;
+
+  /// Register the MLIR dialects required to support an accelerator.
   virtual void registerDialects(mlir::DialectRegistry &registry) const = 0;
+
+  /// Initialize the transformation passes required to generate code for an
+  /// accelerator.
   virtual void initPasses(int optLevel) const = 0;
 
 protected:
-  // static llvm::SmallPtrSet<Accelerator *, 2> acceleratorTargets;
-  static std::vector<Accelerator *> acceleratorTargets;
-
+  static llvm::SmallVector<Accelerator *, 4> acceleratorTargets;
   /// Kind of accelerator.
   Kind kind;
 };

diff --git 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
@@ -27,28 +27,33 @@
 #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<onnx_mlir::accel::Accelerator::Kind> maccel;
 
 namespace accel {
+
+nnpa::NNPAAccelerator *pnnpa = nullptr;
+void createNNPA() { pnnpa = new nnpa::NNPAAccelerator(); }
+
 namespace nnpa {
 
 NNPAAccelerator::NNPAAccelerator() : Accelerator(Accelerator::Kind::NNPA) {
   LLVM_DEBUG(llvm::dbgs() << "Initializing NNPA accelerator\n");
 
   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<zhigh::ZHighDialect>();
   context.getOrLoadDialect<zlow::ZLowDialect>();
 }

diff --git 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
@@ -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
@@ -84,9 +84,12 @@ llvm::cl::opt<std::string> march("march",
 
 llvm::cl::list<accel::Accelerator::Kind> 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<bool> VerboseOutput("v", llvm::cl::desc("Use verbose output"),

diff --git a/src/Compiler/CompilerUtils.cpp b/src/Compiler/CompilerUtils.cpp
@@ -699,14 +699,14 @@ int compileModule(mlir::OwningOpRef<ModuleOp> &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
@@ -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
@@ -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<mlir::linalg::LinalgDialect>();
   registry.insert<mlir::AffineDialect>();
@@ -99,9 +96,12 @@ int main(int argc, char **argv) {
   registry.insert<mlir::ONNXDialect>();
   registry.insert<mlir::KrnlOpsDialect>();
 
+  // 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
@@ -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
@@ -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 &registry) {
 
   registry.insert<mlir::ONNXDialect>();
   registry.insert<mlir::KrnlOpsDialect>();
+
+  // 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) {