Add runtime tracing support to passes under IREE's PassManager usage. (…

…#5160)

Thanks to [this tip](https://discord.com/channels/689900678990135345/760577505840463893/772892255106105365) way back on Discord :D

Sample instrumented trace: 
![image](https://user-images.githubusercontent.com/4010439/111682242-59dd5180-87e1-11eb-8c16-cbb8ed4a7ef7.png)

That trace is measuring a debug compile of this model/flag combination: https://github.com/google/iree/blob/789689dd12a7d17292e2371945d72e56447f53f6/iree/test/e2e/models/CMakeLists.txt#L54-L65

iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.cpp

@@ -28,6 +28,7 @@
 #include "iree/compiler/Dialect/VM/Target/CallingConventionUtils.h"
 #include "iree/compiler/Dialect/VM/Transforms/Passes.h"
 #include "iree/compiler/Utils/FlatbufferUtils.h"
+#include "iree/compiler/Utils/TracingUtils.h"
 #include "iree/schemas/bytecode_module_def_builder.h"
 #include "iree/schemas/bytecode_module_def_json_printer.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -137,6 +138,7 @@ static LogicalResult canonicalizeModule(BytecodeTargetOptions targetOptions,
 
   PassManager passManager(context);
   mlir::applyPassManagerCLOptions(passManager);
+  passManager.addInstrumentation(std::make_unique<PassTracing>());
   auto &modulePasses = passManager.nest<IREE::VM::ModuleOp>();
 
   if (targetOptions.optimize) {

iree/compiler/Translation/BUILD

@@ -37,6 +37,7 @@ cc_library(
         "//iree/compiler/Dialect/VM/Conversion/StandardToVM",
         "//iree/compiler/Dialect/VM/Target/Bytecode",
         "//iree/compiler/Dialect/VM/Transforms",
+        "//iree/compiler/Utils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",

iree/compiler/Translation/CMakeLists.txt

@@ -45,6 +45,7 @@ iree_cc_library(
     iree::compiler::Dialect::VM::Conversion::StandardToVM
     iree::compiler::Dialect::VM::Target::Bytecode
     iree::compiler::Dialect::VM::Transforms
+    iree::compiler::Utils
     ${IREE_VM_CONDITIONAL_TARGETS}
   PUBLIC
 )

iree/compiler/Translation/IREEVM.cpp

@@ -21,6 +21,7 @@
 #include "iree/compiler/Dialect/IREE/Transforms/Passes.h"
 #include "iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.h"
 #include "iree/compiler/Dialect/VM/Transforms/Passes.h"
+#include "iree/compiler/Utils/TracingUtils.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Translation.h"
@@ -77,6 +78,7 @@ static BindingOptions getBindingOptionsFromFlags() {
 static LogicalResult convertToFlowModule(ModuleOp moduleOp) {
   PassManager passManager(moduleOp.getContext());
   mlir::applyPassManagerCLOptions(passManager);
+  passManager.addInstrumentation(std::make_unique<PassTracing>());
   IREE::Flow::buildInputTransformPassPipeline(passManager);
   IREE::Flow::buildFlowTransformPassPipeline(passManager);
   if (failed(passManager.run(moduleOp))) {
@@ -92,6 +94,7 @@ static LogicalResult convertToHALModule(
     ModuleOp moduleOp, IREE::HAL::TargetOptions executableOptions) {
   PassManager passManager(moduleOp.getContext());
   mlir::applyPassManagerCLOptions(passManager);
+  passManager.addInstrumentation(std::make_unique<PassTracing>());
   IREE::HAL::buildHALTransformPassPipeline(passManager, executableOptions);
   if (failed(passManager.run(moduleOp))) {
     return moduleOp.emitError()
@@ -107,6 +110,7 @@ static LogicalResult convertToVMModule(ModuleOp moduleOp,
                                        IREE::VM::TargetOptions targetOptions) {
   PassManager passManager(moduleOp.getContext());
   mlir::applyPassManagerCLOptions(passManager);
+  passManager.addInstrumentation(std::make_unique<PassTracing>());
   IREE::VM::buildVMTransformPassPipeline(passManager, targetOptions);
   if (failed(passManager.run(moduleOp))) {
     return moduleOp.emitError()
@@ -155,6 +159,7 @@ static LogicalResult translateFromMLIRToVM(
     IREE::VM::TargetOptions targetOptions) {
   PassManager passManager(moduleOp.getContext());
   mlir::applyPassManagerCLOptions(passManager);
+  passManager.addInstrumentation(std::make_unique<PassTracing>());
   buildIREEVMTransformPassPipeline(bindingOptions, executableOptions,
                                    targetOptions, passManager);
   if (failed(passManager.run(moduleOp))) {

iree/compiler/Utils/BUILD

@@ -30,9 +30,11 @@ cc_library(
         "FlatbufferUtils.h",
         "GraphUtils.h",
         "PatternUtils.h",
+        "TracingUtils.h",
     ],
     deps = [
         "//iree/base:flatcc",
+        "//iree/base:tracing",
         "//iree/compiler/Dialect/IREE/IR",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",

iree/compiler/Utils/CMakeLists.txt

@@ -17,6 +17,7 @@ iree_cc_library(
     "FlatbufferUtils.h"
     "GraphUtils.h"
     "PatternUtils.h"
+    "TracingUtils.h"
   SRCS
     "FlatbufferUtils.cpp"
     "GraphUtils.cpp"
@@ -29,6 +30,7 @@ iree_cc_library(
     MLIRTransformUtils
     MLIRTransforms
     iree::base::flatcc
+    iree::base::tracing
     iree::compiler::Dialect::IREE::IR
   PUBLIC
 )

iree/compiler/Utils/TracingUtils.h

@@ -0,0 +1,57 @@
+// Copyright 2021 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef IREE_COMPILER_UTILS_TRACINGUTILS_H_
+#define IREE_COMPILER_UTILS_TRACINGUTILS_H_
+
+#include "iree/base/tracing.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassInstrumentation.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+// Instruments passes using IREE's runtime tracing support.
+//
+// Usage:
+//   passManager.addInstrumentation(std::make_unique<PassTracing>());
+struct PassTracing : public PassInstrumentation {
+  PassTracing() {}
+  ~PassTracing() override = default;
+
+#if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
+  // Note: we could also trace pipelines and analyses.
+
+  void runBeforePass(Pass *pass, Operation *op) override {
+    std::string passName = pass->getName().str();
+    IREE_TRACE_ZONE_BEGIN_NAMED_DYNAMIC(z0, passName.data(), passName.size());
+    passTraceZonesStack.push_back(z0);
+  }
+  void runAfterPass(Pass *pass, Operation *op) override {
+    IREE_TRACE_ZONE_END(passTraceZonesStack.back());
+    passTraceZonesStack.pop_back();
+  }
+  void runAfterPassFailed(Pass *pass, Operation *op) override {
+    IREE_TRACE_ZONE_END(passTraceZonesStack.back());
+    passTraceZonesStack.pop_back();
+  }
+#endif  // IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
+
+  llvm::SmallVector<iree_zone_id_t, 8> passTraceZonesStack;
+};
+
+}  // namespace iree_compiler
+}  // namespace mlir
+
+#endif  // IREE_COMPILER_UTILS_TRACINGUTILS_H_