Merge branch 'main' into pr_extend_dynamic_backend_test

Author: negiyas

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
@@ -83,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
@@ -201,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),
@@ -564,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
@@ -126,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();
@@ -880,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;
@@ -905,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 {

src/Conversion/ONNXToKrnl/CMakeLists.txt

@@ -80,6 +80,7 @@ add_onnx_mlir_library(OMONNXToKrnl
 
   LINK_LIBS PUBLIC
   OMAccelerator
+  OMCompilerOptions
   OMONNXConversionCommon
   OMONNXOps
   OMSupport

src/Conversion/ONNXToKrnl/Math/Elementwise.cpp

@@ -17,6 +17,7 @@
 
 #include "llvm/Support/Debug.h"
 
+#include "src/Compiler/CompilerOptions.hpp"
 #include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
 #include "src/Dialect/Krnl/DialectBuilder.hpp"
 #include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
@@ -1884,6 +1885,9 @@ bool OpFusionHelper::areInputsValidForFusion(
 // A successor op (user) is fusible if it is the only user, it is in the
 // fusible elementwise op list, and its inputs are valid for fusion.
 void OpFusionHelper::findFusibleOps() {
+  // Direct return if fusion is disabled
+  if (disableKrnlOpFusion)
+    return;
   Operation *defOp = rootOp;
   while (defOp->hasOneUse()) {
     // the possible ONNX Ops.

src/onnx-mlir.cpp

@@ -11,9 +11,9 @@
 // Implements main for onnx-mlir driver.
 //===----------------------------------------------------------------------===//
 
-#include <iostream>
 #include <regex>
 
+#include "mlir/Support/Timing.h"
 #include "src/Compiler/CompilerOptions.hpp"
 #include "src/Compiler/CompilerUtils.hpp"
 #include "src/Version/Version.hpp"
@@ -24,12 +24,10 @@
 using namespace onnx_mlir;
 
 int main(int argc, char *argv[]) {
-
   // Register MLIR command line options.
   mlir::registerAsmPrinterCLOptions();
   mlir::registerMLIRContextCLOptions();
   mlir::registerPassManagerCLOptions();
-  mlir::registerDefaultTimingManagerCLOptions();
   mlir::registerAsmPrinterCLOptions();
 
   llvm::cl::SetVersionPrinter(getVersionPrinter);
@@ -44,9 +42,13 @@ int main(int argc, char *argv[]) {
     llvm::errs() << "Failed to parse options\n";
     return 1;
   }
-
   initCompilerConfig();
 
+  // Timing manager reporting enabled via "--enable-timing" compiler flag
+  timingManager.setEnabled(enableTiming);
+  rootTimingScope = timingManager.getRootScope();
+  auto setupTiming = rootTimingScope.nest("[onnx-mlir] Loading Dialects");
+
   // Special handling of outputBaseName to derive output filename.
   // outputBaseName must specify a file, so ignore invalid values
   // such as ".", "..", "./", "/.", etc.
@@ -71,7 +73,9 @@ int main(int argc, char *argv[]) {
     LLVM_DEBUG(llvm::dbgs() << "multithreading is disabled\n");
   }
   loadDialects(context);
-
+  setupTiming.stop();
+  auto inputFileTiming =
+      rootTimingScope.nest("[onnx-mlir] Importing Input Model to MLIR");
   mlir::OwningOpRef<mlir::ModuleOp> module;
   std::string errorMessage;
   int rc = processInputFile(inputFilename, context, module, &errorMessage);
@@ -80,6 +84,6 @@ int main(int argc, char *argv[]) {
       llvm::errs() << errorMessage << "\n";
     return 1;
   }
-
+  inputFileTiming.stop();
   return compileModule(module, context, outputBaseName, emissionTarget);
 }