I finished a timing tool that allows the flag “--enable-timing” to start a detailed timing report for the entire onnx-mlir execution process, including setup in main() in onnx-mlir.cpp, input file processing before compilation, the passes run in the pass manager, and finally, the lower level compilation/generation stages (LLVM, Object, Shared Library, JNI). This tool will print the detailed execution report as a tree of all the timed stages to the terminal with the (wall time) and (relative percentage of program execution) for each timed stage and the total execution time of the onnx-mlir compiler. I used the DefaultTimingManager from llvm-project “timing.cpp” and timing scopes to time and name essential code sections where work is being done.

Signed-off-by: Paramvir Sran <[email protected]>

Author: ParamvirSran

CMakeLists.txt

@@ -7,7 +7,7 @@ project(onnx-mlir)
 
 option(ONNX_MLIR_BUILD_TESTS "Build ONNX-MLIR test executables. If OFF, just generate build targets." ON)
 option(ONNX_MLIR_CCACHE_BUILD "Set to ON for a ccache enabled build." OFF)
-option(ONNX_MLIR_ENABLE_STABLEHLO "Enable StableHLO support." ON)
+option(ONNX_MLIR_ENABLE_STABLEHLO "Enable StableHLO support." OFF)
 option(ONNX_MLIR_DECOMP_ONNX_CONVTRANSPOSE "Enable ONNXConvTransposeOp decomposition." ON)
 option(ONNX_MLIR_ENABLE_WERROR "Enable warnings as errors." OFF)
 option(ONNX_MLIR_SUPPRESS_THIRD_PARTY_WARNINGS "Suppress warning in third_party code." ON)

docs/devcontainer-example/devcontainer.json

@@ -3,14 +3,15 @@
   "customizations": {
     "vscode": {
       "settings": {
+        "extensions.verifySignature": false,
         "terminal.integrated.defaultProfile.linux": "bash",
         "cmake.mergedCompileCommands": "${workspaceFolder}/compile_commands.json",
         "cmakeExplorer.suiteDelimiter": ".",
         "cmakeExplorer.debugConfig": "(gdb) Launch",
         "editor.formatOnSave": true,
         "[cpp]": {
           "editor.defaultFormatter": "xaver.clang-format",
-          "editor.tabSize": 2
+          "editor.tabSize":2 
         },
         "clangd.path": "/usr/bin/clangd-15",
         "clangd.arguments": [

src/Compiler/CompilerOptions.cpp

@@ -82,6 +82,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
@@ -546,6 +547,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 pass timing (default is false)\n"
+                   "Set to 'true' if you want to enable pass timings."),
+    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 "
@@ -579,6 +586,7 @@ static llvm::cl::opt<bool, true> useOldBufferizationOpt("use-old-bufferization",
     llvm::cl::location(useOldBufferization), llvm::cl::init(true),
     llvm::cl::cat(OnnxMlirOptions));
 
+
 // Configuration states associated with certain options.
 // For example, when maccel is specified, NNPA can register
 // dependent libdnn.

src/Compiler/CompilerOptions.hpp

@@ -125,6 +125,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

@@ -14,14 +14,13 @@
 
 #include "CompilerUtils.hpp"
 
-#include <fstream>
-#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,7 +34,11 @@
 #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 <fstream>
+#include <memory>
+#include <regex>
 
 #include "src/Accelerators/Accelerator.hpp"
 #include "src/Builder/FrontendDialectTransformer.hpp"
@@ -49,6 +52,8 @@
 using namespace mlir;
 using namespace onnx_mlir;
 
+extern mlir::DefaultTimingManager timingManager;
+
 namespace onnx_mlir {
 
 // Make a function that forces preserving all files using the runtime arguments
@@ -146,8 +151,7 @@ int Command::exec(std::string wdir) const {
     llvm::errs() << llvm::join(argsRef, " ") << "\n"
                  << "Error message: " << errMsg << "\n"
                  << "Program path: " << _path << "\n"
-                 << "Command execution failed."
-                 << "\n";
+                 << "Command execution failed." << "\n";
     return rc;
   }
 
@@ -327,6 +331,10 @@ 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 rootScope = timingManager.getRootScope();
+  auto llvmScope =
+      rootScope.nest("[onnx-mlir] Compiling to LLVM Optimized Bitcode");
+
   std::error_code error;
 
   // Write bitcode to a file.
@@ -397,7 +405,9 @@ 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 rootScope = timingManager.getRootScope();
+  auto objectScope =
+      rootScope.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 +428,9 @@ 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 rootScope = timingManager.getRootScope();
+  auto jniScope = rootScope.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 +449,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 rootScope = timingManager.getRootScope();
+  auto libraryScope = rootScope.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 +500,9 @@ 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 rootScope = timingManager.getRootScope();
+  auto jniJarScope = rootScope.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();
@@ -508,6 +525,7 @@ static int genJniJar(const mlir::OwningOpRef<ModuleOp> &module,
 // Return 0 on success, error code on failure
 static int compileModuleToObject(const mlir::OwningOpRef<ModuleOp> &module,
     std::string outputNameWithoutExt, std::string &objectNameWithExt) {
+
   std::string bitcodeNameWithExt = outputNameWithoutExt + ".bc";
   int rc = genLLVMBitcode(module, outputNameWithoutExt, bitcodeNameWithExt);
   if (rc != CompilerSuccess)
@@ -522,6 +540,7 @@ static int compileModuleToObject(const mlir::OwningOpRef<ModuleOp> &module,
 static int compileModuleToSharedLibrary(
     const mlir::OwningOpRef<ModuleOp> &module, std::string outputNameNoExt,
     std::string &libNameWithExt) {
+
   std::string modelObjNameWithExt;
   int rc = compileModuleToObject(module, outputNameNoExt, modelObjNameWithExt);
   if (rc != CompilerSuccess)
@@ -879,6 +898,9 @@ static int emitOutput(mlir::OwningOpRef<ModuleOp> &module,
 int compileModule(mlir::OwningOpRef<ModuleOp> &module,
     mlir::MLIRContext &context, std::string outputNameNoExt,
     EmissionTargetType emissionTarget) {
+  auto rootScope = timingManager.getRootScope();
+  auto compileModuleScope = rootScope.nest("[onnx-mlir] Compile Module Setup");
+
   int rc = setupModule(module, context, outputNameNoExt);
   if (rc != CompilerSuccess)
     return rc;
@@ -904,7 +926,11 @@ int compileModule(mlir::OwningOpRef<ModuleOp> &module,
         heapLogFileame, reportHeapBefore, reportHeapAfter));
   }
   (void)mlir::applyPassManagerCLOptions(pm);
-  mlir::applyDefaultTimingPassManagerCLOptions(pm);
+
+  if (enableTiming) {
+    pm.enableTiming(rootScope);
+  }
+  compileModuleScope.stop();
 
   if (mlir::failed(pm.run(*module)))
     return CompilerFailure;

src/Dialect/ONNX/ElementsAttr/ElementsAttrBuilder.cpp

@@ -22,7 +22,6 @@
 
 #include <algorithm>
 #include <numeric>
-
 using namespace mlir;
 
 namespace onnx_mlir {
@@ -217,6 +216,8 @@ ElementsAttr ElementsAttrBuilder::combine(ElementsAttr lhs, ElementsAttr rhs,
     WideNum *dst = dstNums.data();
     const WideNum *lhsSrc = lhsNums.get().data();
     const WideNum *rhsSrc = rhsNums.get().data();
+    
+    
     for (auto &idxoffs :
         StridesRange<2>(combinedShape, {xpLhsStrides, xpRhsStrides})) {
       dst[idxoffs.flattenedIndex] =

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"
@@ -23,8 +23,9 @@
 
 using namespace onnx_mlir;
 
-int main(int argc, char *argv[]) {
+mlir::DefaultTimingManager timingManager;
 
+int main(int argc, char *argv[]) {
   // Register MLIR command line options.
   mlir::registerAsmPrinterCLOptions();
   mlir::registerMLIRContextCLOptions();
@@ -44,9 +45,14 @@ 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
+  mlir::applyDefaultTimingManagerCLOptions(timingManager);
+  timingManager.setEnabled(enableTiming);
+  auto rootScope = timingManager.getRootScope();
+  auto setupScope = rootScope.nest("[onnx-mlir] Preparing for Compilation");
+
   // Special handling of outputBaseName to derive output filename.
   // outputBaseName must specify a file, so ignore invalid values
   // such as ".", "..", "./", "/.", etc.
@@ -71,7 +77,8 @@ int main(int argc, char *argv[]) {
     LLVM_DEBUG(llvm::dbgs() << "multithreading is disabled\n");
   }
   loadDialects(context);
-
+  setupScope.stop();
+  auto inputFileScope = rootScope.nest("[onnx-mlir] Processing of Input File");
   mlir::OwningOpRef<mlir::ModuleOp> module;
   std::string errorMessage;
   int rc = processInputFile(inputFilename, context, module, &errorMessage);
@@ -80,6 +87,6 @@ int main(int argc, char *argv[]) {
       llvm::errs() << errorMessage << "\n";
     return 1;
   }
-
+  inputFileScope.stop();
   return compileModule(module, context, outputBaseName, emissionTarget);
 }