Merge branch 'main' into test_backend_dynamic

Author: negiyas

docs/AddCustomAccelerators.md

@@ -5,6 +5,7 @@ In general, onnx-mlir handles custom accelerators as pluggins which can be turne
 Besides this document, [NNPA accelerator](../src/Accelerators/NNPA) can be used as an example that has been deployed in onnx-mlir.
 
 ## 1. Code folder
+
 In onnx-mlir, all code for an accelerator should be put inside a separate folder under `src/Accelerators`. Thus, the first step to support an accelerator is to create a folder for it inside `src/Accelerators`.
 
 The folder name will be used as the accelerator name in onnx-mlir. In particular, it is used to
@@ -15,22 +16,25 @@ The folder name will be used as the accelerator name in onnx-mlir. In particular
 The folder content is flexible depending on each accelerator. However, we recomment to follow the same structure as the root folder of `onnx-mlir` as much as possbile. This helps maintain the consitency across the whole project.
 
 ### 1.1 Build accelerators in onnx-mlir
-To build accelerators in onnx-mlir, use the cmake variable `ONNX_MLIR_ACCELERATORS` when building onnx-mlir. `ONNX_MLIR_ACCELERATORS` accepts a comma-separated list of accelerator names. For example,
+
+To build accelerators in onnx-mlir, use the cmake variable `ONNX_MLIR_ACCELERATORS` when building onnx-mlir. `ONNX_MLIR_ACCELERATORS` accepts a semicolon-separated list of accelerator names. For example,
 ```bash
 $ cd build
-$ cmake .. -DONNX_MLIR_ACCELERATORS=accel1,accel2
+$ cmake .. -DONNX_MLIR_ACCELERATORS=accel1;accel2
 ```
 
 ### 1.2 Compile a model to run with selected accelerators.
-The compiler command `onnx-mlir` has an option, i.e. `--maccel`, to compile a model for selected accelerators. `--maccel` accepts a comma-separated list of accelerator names. For example,
+
+The compiler command `onnx-mlir` has an option, i.e. `--maccel`, to compile a model for selected accelerators. For each accelerator add a `--maccel=accel_name` entry. For example,
 
 ```bash
-$ onnx-mlir --maccel=accel1,accel2 model.onnx
+$ onnx-mlir --maccel=accel1 --maccel=accel2 model.onnx
 ```
 
 Only built accelerators can be used with `--maccel`.
 
 ### 1.3 Run passes related to selected accelerators.
+
 Passes defined by an accelerator can be run or tested via `onnx-mlir-opt` command by using option `--maccel` which is similar to `--maccel` in `onnx-mlir` (See Sec. [1.2](#1.2-compile-a-model-to-run-with-selected-accelerators)). For example, to call a pass `--optimize-data-layout` defined by accelerator `accel1`:
 
 ```bash
@@ -41,6 +45,24 @@ Only built accelerators can be used with `--maccel`.
 
 ## 2. Code integration
 
+### 2.1 Macro
+
+Each accelerator is required to define a few macros. These needs to be included in [onnx_mlir::accel::Accelerator](../src/Accelerators/Accelerator.hpp). These macros are:
+
+1. `INSTRUMENTSTAGE_ENUM_<accel_name>`
+2. `INSTRUMENTSTAGE_CL_ENUM_<accel_name>`
+3. `PROFILEIR_CL_ENUM_<accel_name>`
+
+Replace `<accel_name>` with the name of the accelerator, for example if your accelerator is named `ACCEL1` use:
+
+```C
+#define INSTRUMENTSTAGE_ENUM_ACCEL1
+#define INSTRUMENTSTAGE_CL_ENUM_ACCEL1
+#define PROFILEIR_CL_ENUM_ACCEL1
+```
+
+### 2.2 Dialects and passes
+
 Writing code in MLIR typically involves desiging dialects and passes. So does supporting an accelerator. Thus, to integrate accelerator code into onnx-mlir is to register dialects and passes in onnx-mlir.
 
 We provide a base class [onnx_mlir::accel::Accelerator](../src/Accelerators/Accelerator.hpp) from which users can define an inherited class and write hooks to register dialects and passes.

src/Builder/FrontendDialectTransformer.cpp

@@ -51,6 +51,8 @@ SUPPRESS_WARNINGS_POP
 #include <array>
 #include <fstream>
 #include <map>
+#include <sstream>
+#include <string>
 #include <unordered_map>
 #include <vector>
 
@@ -450,6 +452,27 @@ class FrontendGenImpl {
     return attributes;
   }
 
+  // Generate a string vector from the dimParams option string
+  void getInputDimParamsMapFromOption(std::string optionStr,
+      std::map<int, std::string> &paramStrMap,
+      std::string &paramStrForAllArgs) {
+    std::stringstream paramStrStream(optionStr);
+    std::string dimParamStr;
+    while (std::getline(paramStrStream, dimParamStr, '|')) {
+      size_t pos = dimParamStr.find(':');
+      assert((pos > 0) && "invalid dimParams option string");
+      int idx = stoi(dimParamStr.substr(0, pos));
+      dimParamStr = dimParamStr.substr(pos + 1);
+      std::replace(dimParamStr.begin(), dimParamStr.end(), '=', ':');
+      if (idx < 0) // set all arguments
+        paramStrForAllArgs = dimParamStr;
+      else {
+        paramStrMap[idx] = dimParamStr;
+      }
+    }
+    return;
+  }
+
   /*!
    * An alternative graph importing procedure for importing ONNX subgraphs.
    * ONNX subgraphs, unlike the main computation graph, are imported as regions
@@ -490,6 +513,10 @@ class FrontendGenImpl {
     // See https://github.com/onnx/onnx/blob/main/docs/IR.md for more
     // information about dim_param.
     llvm::SmallVector<std::string, 4> inputDimParams, outputDimParams;
+    std::map<int, std::string> inputDimParamsFromOption;
+    std::string inputDimParamsFromOptionForAllArgs;
+    getInputDimParamsMapFromOption(options_.dimParams, inputDimParamsFromOption,
+        inputDimParamsFromOptionForAllArgs);
 
     // Import the input tensor types that are not constant and not initialized.
     int inputIndex = 0;
@@ -500,7 +527,16 @@ class FrontendGenImpl {
         std::string dimParams = "";
         Type argTy = ImportType(input.type(), &dimParams);
         argTy = modelInputShaper_.reshape(inputIndex, argTy);
-        if (!dimParams.empty())
+        // For each input tensor, use either all dimensions by the compiler
+        // option OR all dimensions in the original onnx model. Dimensions
+        // from the option and the model in a single input tensor are not
+        // merged.
+        if (inputDimParamsFromOption.find(inputIndex) !=
+            inputDimParamsFromOption.end())
+          inputDimParams.emplace_back(inputDimParamsFromOption[inputIndex]);
+        else if (!inputDimParamsFromOptionForAllArgs.empty())
+          inputDimParams.emplace_back(inputDimParamsFromOptionForAllArgs);
+        else if (!dimParams.empty())
           inputDimParams.emplace_back(dimParams);
 
         argTypes.emplace_back(argTy);

src/Builder/FrontendDialectTransformer.hpp

@@ -55,6 +55,18 @@ struct ImportOptions {
   //   - (arg0: tensor<3x4x5xf32>, arg1: tensor<10x5xf32>)
   //
   std::string shapeInformation = "";
+  // Custom onnx.dim_params attributes for the graph inputs for specifying
+  // relationship among their dynamic dimensions.
+  // Its format is 'input_id:dim_id=sym,dim_id=sym,...|input_id:
+  // dim_id=sym,dim_id=sym,...|input_id...'
+  // E.g. An ONNX model has two dynamic inputs
+  //   - (arg0: tensor<?x5xf32>, arg1: tensor<?x5xf32>)
+  // If we want to specify that the first unknown dimension of arg0 and the
+  // first unknown dimension of arg1 are the same, we can assign the two
+  // dimensions to the same symbol "batch" as follows.
+  //   - dimParams = '0:0=batch|1:0=batch'
+  //
+  std::string dimParams = "";
   // Directory to look for external data if any tensor has external
   // data location. If empty then external data is disabled.
   std::string externalDataDir = "";

src/Compiler/CompilerOptions.cpp

@@ -41,6 +41,7 @@ std::vector<std::string> onnxConstPropDisablePatterns; // common for both
 bool enableONNXHybridPass;                             // common for both
 std::vector<std::string> functionsToDecompose;         // common for both
 std::string opsForCall;                                // common for both
+bool disableKrnlOpFusion;                              // common for both
 EmissionTargetType emissionTarget;                     // onnx-mlir only
 bool invokeOnnxVersionConverter;                       // onnx-mlir only
 bool preserveLocations;                                // onnx-mlir only
@@ -51,6 +52,7 @@ bool preserveMLIR;                                     // onnx-mlir only
 bool useOnnxModelTypes;                                // onnx-mlir only
 int repeatOnnxTransform;                               // onnx-mlir only
 std::string shapeInformation;                          // onnx-mlir only
+std::string dimParams;                                 // onnx-mlir only
 ModelSize modelSize;                                   // onnx-mlir only
 bool storeConstantsToFile;                             // onnx-mlir only
 float constantsToFileTotalThreshold;                   // onnx-mlir only
@@ -82,6 +84,7 @@ std::vector<std::string> extraLibs;                    // onnx-mlir only
 ProfileIRs profileIR;                                  // onnx-mlir only
 OptReport optReport;                                   // onnx-mlir only
 bool useOldBufferization;                              // onnx-mlir only
+bool enableTiming;                                     // onnx-mlir only
 bool split_input_file;                                 // onnx-mlir-opt only
 bool verify_diagnostics;                               // onnx-mlir-opt only
 bool verify_passes;                                    // onnx-mlir-opt only
@@ -200,6 +203,13 @@ static llvm::cl::list<std::string, std::vector<std::string>>
         llvm::cl::location(functionsToDecompose),
         llvm::cl::cat(OnnxMlirCommonOptions));
 
+static llvm::cl::opt<bool, true> disableKrnlOpFusionOpt(
+    "disable-krnl-op-fusion",
+    llvm::cl::desc("disable op fusion in onnx-to-krnl pass (default=false)\n"
+                   "Set to 'true' if you want to disable fusion."),
+    llvm::cl::location(disableKrnlOpFusion), llvm::cl::init(false),
+    llvm::cl::cat(OnnxMlirCommonOptions));
+
 static llvm::cl::opt<bool, true> disableRecomposeOptionOpt("disable-recompose",
     llvm::cl::desc("Disable recomposition of ONNX operations."),
     llvm::cl::location(disableRecomposeOption), llvm::cl::init(false),
@@ -281,6 +291,23 @@ static llvm::cl::opt<std::string, true> shapeInformationOpt("shapeInformation",
     llvm::cl::value_desc("value"), llvm::cl::location(shapeInformation),
     llvm::cl::cat(OnnxMlirOptions));
 
+static llvm::cl::opt<std::string, true> dimParamsOpt("dimParams",
+    llvm::cl::desc(
+        "Custom onnx.dim_params attributes for the inputs of the ONNX model for"
+        "specifying relationship among dynamic dimensions of the inputs.\n"
+        "\"value\" is in the format of "
+        "\"INPUT_ID1:D1=S1,D2=S2,...,Dn=Sn|INPUT_ID2:D1=T1,D2=T2,...Dn=Tn|"
+        "...\" where \"INPUT_ID1, INPUT_ID2, ...\" are input indices "
+        "(starting from 0 or being -1 for all input indices), and\n"
+        "\"S1, S2, ...\" and \"T2, T2, ...\" are symbols to specify that same "
+        "symbols have the same value. "
+        "All dimensions of onnx.dim_params for a specified input index in "
+        "the original onnx model are cleared and repalced by this option. "
+        "onnx.dim_params for other input indices in the original onnx model "
+        "are not cleared"),
+    llvm::cl::value_desc("value"), llvm::cl::location(dimParams),
+    llvm::cl::cat(OnnxMlirOptions));
+
 // Default value is defined by the OnnxMlirEnvOptionName constant string
 // variable, but the default setting mechanism here cannot be used here as we
 // need to evaluate this value prior to the compiler options being set. Proper
@@ -546,6 +573,12 @@ static llvm::cl::opt<OptReport, true> optReportOpt("opt-report",
         clEnumVal(Simd, "Provide report on how SIMD is applied to ONNX ops.")),
     llvm::cl::init(OptReport::NoReport), llvm::cl::cat(OnnxMlirOptions));
 
+static llvm::cl::opt<bool, true> enable_timing("enable-timing",
+    llvm::cl::desc("Enable compile timing (default is false)\n"
+                   "Set to 'true' if you want to enable compile timing."),
+    llvm::cl::location(enableTiming), llvm::cl::init(false),
+    llvm::cl::cat(OnnxMlirOptions));
+
 // Options for onnx-mlir-opt only
 static llvm::cl::opt<bool, true> split_input_file_opt("split-input-file",
     llvm::cl::desc("Split the input file into pieces and process each "

src/Compiler/CompilerOptions.hpp

@@ -84,6 +84,7 @@ extern std::vector<std::string> onnxConstPropDisablePatterns; // common for both
 extern bool enableONNXHybridPass;                             // common for both
 extern std::vector<std::string> functionsToDecompose;         // common for both
 extern std::string opsForCall;                                // common for both
+extern bool disableKrnlOpFusion;                              // common for both
 extern EmissionTargetType emissionTarget;                     // onnx-mlir only
 extern bool invokeOnnxVersionConverter;                       // onnx-mlir only
 extern bool preserveLocations;                                // onnx-mlir only
@@ -94,6 +95,7 @@ extern bool preserveMLIR;                                     // onnx-mlir only
 extern bool useOnnxModelTypes;                                // onnx-mlir only
 extern int repeatOnnxTransform;                               // onnx-mlir only
 extern std::string shapeInformation;                          // onnx-mlir only
+extern std::string dimParams;                                 // onnx-mlir only
 extern ModelSize modelSize;                                   // onnx-mlir only
 extern bool storeConstantsToFile;                             // onnx-mlir only
 extern float constantsToFileTotalThreshold;                   // onnx-mlir only
@@ -125,6 +127,7 @@ extern std::vector<std::string> extraLibs;                    // onnx-mlir only
 extern ProfileIRs profileIR;                                  // onnx-mlir only
 extern OptReport optReport;                                   // onnx-mlir only
 extern bool useOldBufferization;                              // onnx-mlir only
+extern bool enableTiming;                                     // onnx-mlir only
 extern bool split_input_file;          // onnx-mlir-opt only
 extern bool verify_diagnostics;        // onnx-mlir-opt only
 extern bool verify_passes;             // onnx-mlir-opt only

src/Compiler/CompilerUtils.cpp

@@ -15,13 +15,16 @@
 #include "CompilerUtils.hpp"
 
 #include <fstream>
+#include <memory>
 #include <regex>
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Parser/Parser.h"
+#include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/FileUtilities.h"
+#include "mlir/Support/Timing.h"
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
@@ -35,6 +38,7 @@
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/ToolOutputFile.h"
+#include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 
 #include "src/Accelerators/Accelerator.hpp"
@@ -49,6 +53,8 @@
 using namespace mlir;
 using namespace onnx_mlir;
 
+mlir::DefaultTimingManager timingManager;
+mlir::TimingScope rootTimingScope;
 namespace onnx_mlir {
 
 // Make a function that forces preserving all files using the runtime arguments
@@ -327,6 +333,8 @@ std::string getTargetFilename(
 // Returns 0 on success, error code on failure.
 static int genLLVMBitcode(const mlir::OwningOpRef<ModuleOp> &module,
     std::string outputNameNoExt, std::string optimizedBitcodeNameWithExt) {
+  auto llvmTiming = rootTimingScope.nest(
+      "[onnx-mlir] Compiling MLIR module to LLVM Optimized Bitcode");
   std::error_code error;
 
   // Write bitcode to a file.
@@ -397,7 +405,8 @@ static int genLLVMBitcode(const mlir::OwningOpRef<ModuleOp> &module,
 // Return 0 on success, error code on failure.
 static int genModelObject(
     std::string bitcodeNameWithExt, std::string &modelObjNameWithExt) {
-
+  auto objectTiming =
+      rootTimingScope.nest("[onnx-mlir] Compiling LLVM Bitcode to Object File");
   std::string llcPath = getToolPath("llc");
   Command llvmToObj(/*exePath=*/llcPath);
   setXllcOption({"--code-model", modelSizeStr[modelSize]});
@@ -418,6 +427,8 @@ static int genModelObject(
 // Return 0 on success, error code on failure.
 static int genJniObject(const mlir::OwningOpRef<ModuleOp> &module,
     std::string jniSharedLibPath, std::string jniObjPath) {
+  auto jniTiming =
+      rootTimingScope.nest("[onnx-mlir] Compiling JNI Object File");
   Command ar(/*exePath=*/getToolPath("ar", true));
   int rc = ar.appendStr("x")
                // old version of ar does not support --output so comment out
@@ -436,7 +447,8 @@ static int genJniObject(const mlir::OwningOpRef<ModuleOp> &module,
 static int genSharedLib(std::string sharedLibNameWithExt,
     std::vector<std::string> opts, std::vector<std::string> objs,
     std::vector<std::string> libs, std::vector<std::string> libDirs) {
-
+  auto sharedLibTiming =
+      rootTimingScope.nest("[onnx-mlir] Linking Shared Library");
 #ifdef _WIN32
   std::vector<std::string> outputOpt = {"/Fe:" + sharedLibNameWithExt};
   // link has to be before libpath since they need to be passed through to the
@@ -486,6 +498,7 @@ static int genSharedLib(std::string sharedLibNameWithExt,
 // Return 0 on success, error code on failure.
 static int genJniJar(const mlir::OwningOpRef<ModuleOp> &module,
     std::string modelSharedLibPath, std::string modelJniJarPath) {
+  auto jniJarTiming = rootTimingScope.nest("[onnx-mlir] Creating JNI Jar");
   llvm::SmallString<8> libraryPath(getLibraryPath());
   llvm::sys::path::append(libraryPath, "javaruntime.jar");
   std::string javaRuntimeJarPath = llvm::StringRef(libraryPath).str();
@@ -623,6 +636,7 @@ int processInputFile(StringRef inputFilename, mlir::MLIRContext &context,
     options.useOnnxModelTypes = useOnnxModelTypes;
     options.invokeOnnxVersionConverter = invokeOnnxVersionConverter;
     options.shapeInformation = shapeInformation;
+    options.dimParams = dimParams;
     options.allowSorting = allowSorting;
     options.externalDataDir = dirName(inputFilename);
     options.functionsToDecompose.insert(options.functionsToDecompose.end(),
@@ -879,6 +893,9 @@ static int emitOutput(mlir::OwningOpRef<ModuleOp> &module,
 int compileModule(mlir::OwningOpRef<ModuleOp> &module,
     mlir::MLIRContext &context, std::string outputNameNoExt,
     EmissionTargetType emissionTarget) {
+  auto compileModuleTiming =
+      rootTimingScope.nest("[onnx-mlir] Compiling Module using MLIR");
+
   int rc = setupModule(module, context, outputNameNoExt);
   if (rc != CompilerSuccess)
     return rc;
@@ -904,10 +921,14 @@ int compileModule(mlir::OwningOpRef<ModuleOp> &module,
         heapLogFileame, reportHeapBefore, reportHeapAfter));
   }
   (void)mlir::applyPassManagerCLOptions(pm);
-  mlir::applyDefaultTimingPassManagerCLOptions(pm);
+
+  if (enableTiming) {
+    pm.enableTiming(compileModuleTiming);
+  }
 
   if (mlir::failed(pm.run(*module)))
     return CompilerFailure;
+  compileModuleTiming.stop();
   return emitOutput(module, context, outputNameNoExt, pm, emissionTarget);
 }
 

src/Compiler/CompilerUtils.hpp

@@ -18,13 +18,17 @@
 
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/OwningOpRef.h"
+#include "mlir/Support/Timing.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Path.h"
 
 #include <optional>
 #include <string>
 #include <vector>
 
+extern mlir::DefaultTimingManager timingManager;
+extern mlir::TimingScope rootTimingScope;
+
 namespace onnx_mlir {
 
 struct Command {