LTC->MLIR Debug Info support (#1922)

* LTC->MLIR Debug Info support * SW-95317 Propagate Lazy->Jit->MLIR scope name. * Enhance location information based on op names Currently, the location information attached to the ops just considers the filename, line number and column number. Attaching operation name would help identify the type of computation by just looking at the profile of execution. * Update locations logic; updated debug-info.py test * Use {scope}/{op_name} format to track names by default --------- Co-authored-by: Gleb Kazantaev <[email protected]> Co-authored-by: Mark Browning <[email protected]> Co-authored-by: Vimal Patel <[email protected]>
llvm · Aug 2, 2023 · fb52a73 · fb52a73
1 parent 4c24472
commit fb52a73
Show file tree

Hide file tree

Showing 6 changed files with 128 additions and 20 deletions.
diff --git a/python/torch_mlir/csrc/base_lazy_backend/mlir_lowering_context.cpp b/python/torch_mlir/csrc/base_lazy_backend/mlir_lowering_context.cpp
@@ -15,6 +15,7 @@
 #include <torch/csrc/jit/api/compilation_unit.h>
 #include <torch/csrc/jit/passes/refine_tuple_types.h>
 #include <torch/csrc/lazy/core/lazy_graph_executor.h>
+#include <torch/csrc/lazy/core/config.h>
 #include "torch-mlir-c/Registration.h"
 #include "torch-mlir-c/Transforms.h"
 #include "mlir-c/IR.h"
@@ -205,13 +206,46 @@ void TorchMlirLoweringContext::AssignOutputOp(
     const Output& output, torch::jit::Value* op) {
   PRINT_FUNCTION();
 
-  // TODO (antoniojkim): Do we need this?
-  // auto torch_mlir_node =
-  //     NodeCast<TorchMlirNode>(output.node, output.node->op());
-  // if (!torch_mlir_node->getPythonStacktrace().empty()) {
-  //   op->node()->s_(
-  //       c10::Symbol::attr("source"), torch_mlir_node->getPythonStacktrace());
-  // }
+  auto torch_mlir_node =
+      NodeCast<TorchMlirNode>(output.node, output.node->op());
+
+  std::vector<std::string> source_files, functions;
+  std::vector<int64_t> line_numbers;
+  const auto& metadata = torch_mlir_node->metadata();
+  const auto& frames = metadata.frame_info;
+  if (!frames.empty()) {
+    static std::vector<std::string> g_roots =
+      string_split(sys_util::GetEnvString("LTC_IR_DEBUG_ROOT_PATH", ""), ":");
+
+    std::for_each(frames.rbegin(), frames.rend(),
+      [&](const torch::lazy::SourceLocation& location) {
+        functions.push_back(location.function);
+        line_numbers.push_back(location.line);
+
+        std::string file_name = location.file;
+        for (const std::string& root : g_roots) {
+          if (startswith(file_name, root)) {
+            // location.file starts with root, strip it off
+            file_name = file_name.substr(root.size());
+            break;
+          }
+        }
+        source_files.push_back(file_name);
+    });
+
+    if (!source_files.empty()) {
+      op->node()->ss_(
+          c10::Symbol::attr("source_files"), source_files);
+      op->node()->ss_(
+          c10::Symbol::attr("functions"), functions);
+      op->node()->is_(
+          c10::Symbol::attr("line_numbers"), line_numbers);
+    }
+  }
+  auto scope = ::c10::Symbol::scope(metadata.scope);
+  op->node()->setScope(
+    c10::make_intrusive<torch::jit::Scope>()->push(scope));
+
   emitted_outputs_[output] = std::move(op);
 }
 
@@ -424,7 +458,11 @@ const std::string TorchMlirComputation::to_string() const {
     *ss_ptr << std::string(part.data, part.length);
   };
   std::stringstream ss;
-  mlirOperationPrint(mlirModuleGetOperation(module_op_), print_callback, &ss);
+
+  // Setup flags for MLIR serialization.
+  MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
+  mlirOpPrintingFlagsEnableDebugInfo(flags, FLAGS_torch_lazy_ir_debug, false);
+  mlirOperationPrintWithFlags(mlirModuleGetOperation(module_op_), flags, print_callback, &ss);
   return ss.str();
 }
 

diff --git a/python/torch_mlir/csrc/base_lazy_backend/utils/string_utils.h b/python/torch_mlir/csrc/base_lazy_backend/utils/string_utils.h
@@ -22,6 +22,24 @@ std::string string_join(const std::vector<T>& v, const std::string& delimiter) {
     return joined.str();
 }
 
+inline std::vector<std::string> string_split(
+    const std::string& str,
+    const std::string& sep
+) {
+    std::vector<std::string> tokens;
+    std::size_t pos1 = str.find_first_not_of(sep);
+    while (pos1 != std::string::npos) {
+        std::size_t pos2 = str.find_first_of(sep, pos1);
+        if (pos2 == std::string::npos) {
+            tokens.push_back(str.substr(pos1));
+            pos1 = pos2;
+        } else {
+            tokens.push_back(str.substr(pos1, pos2 - pos1));
+            pos1 = str.find_first_not_of(sep, pos2 + 1);
+        }
+    }
+    return tokens;
+}
 
 /*
  * Returns true if str starts with prefix

diff --git a/python/torch_mlir/csrc/base_lazy_backend/utils/sys_utils.h b/python/torch_mlir/csrc/base_lazy_backend/utils/sys_utils.h
@@ -14,6 +14,14 @@ static T GetEnv(const std::string& name, const T& default_value = T(0)) {
   return T(std::atoi(env));
 }
 
+static std::string GetEnvString(const std::string& name, const std::string& default_value) {
+  const char* env = std::getenv(name.c_str());
+  if (!env) {
+    return default_value;
+  }
+  return std::string(env);
+}
+
 static bool GetEnvBool(const char* name, bool defval) {
   const char* env = std::getenv(name);
   if (env == nullptr) {

diff --git a/python/torch_mlir/csrc/reference_lazy_backend/reference_lazy_backend_pybind.cpp b/python/torch_mlir/csrc/reference_lazy_backend/reference_lazy_backend_pybind.cpp
@@ -8,6 +8,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "torch/csrc/jit/python/pybind.h"
+#include "torch/csrc/lazy/core/config.h"
 #include "torch/csrc/lazy/backend/backend_interface.h"
 
 #include <torch_mlir/csrc/base_lazy_backend/mlir_lowering_context.h>
@@ -25,6 +26,7 @@ namespace py = pybind11;
 
 namespace {
 bool verbose = sys_util::GetEnv("VERBOSE", false);
+bool ir_debug = sys_util::GetEnv("LTC_IR_DEBUG", false);
 
 struct NoGilSection {
   NoGilSection() : state(PyEval_SaveThread()) {}
@@ -52,6 +54,11 @@ void Initialize() {
   if (verbose) {
     std::cout << "MLIR LTC PyTorch Plugin Initialized." << std::endl;
   }
+
+  if (ir_debug) {
+      FLAGS_torch_lazy_ir_debug = true;
+      std::cout << "Enabled lazy tensor IR debugging." << std::endl;
+  }
 }
 
 /**

diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/torch_to_mlir_utils.cpp b/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/torch_to_mlir_utils.cpp
@@ -11,6 +11,8 @@
 #include "function_importer.h"
 #include "ivalue_importer.h"
 
+#include <c10/util/irange.h>
+
 #include <ATen/TensorUtils.h>
 #include <unordered_map>
 
@@ -407,15 +409,53 @@ MlirAttribute torch_mlir::importAttribute(MlirLocation loc,
 
 MlirLocation torch_mlir::getMlirLocationFromNode(MlirContext context,
                                                  torch::jit::Node *node) {
-  auto flc = node->sourceRange().file_line_col();
-  if (flc) {
+  MlirLocation loc = mlirLocationUnknownGet(context);
+
+  if (node->hasAttribute(c10::Symbol::attr("source_files"))) {
+    const auto &sourceFiles = node->ss(c10::Symbol::attr("source_files"));
+    const auto &lineNumbers = node->is(c10::Symbol::attr("line_numbers"));
+    const auto &functions = node->ss(c10::Symbol::attr("functions"));
+
+    // Chain a sequence of calls to construct single MlirLocation.
+    for (const auto i : c10::irange(sourceFiles.size())) {
+      MlirLocation newLoc = mlirLocationNameGet(
+          context, toMlirStringRef(functions[i]),
+          mlirLocationFileLineColGet(context, toMlirStringRef(sourceFiles[i]),
+                                     lineNumbers[i],
+                                     0 /* column is not available */
+                                     ));
+      loc = (i == 0 ? newLoc : mlirLocationCallSiteGet(newLoc, loc));
+    }
+    if (sourceFiles.size() == 1) {
+      // Somehow a callstack depth of 1...
+      // Disambiguate function name from scope name below.
+      loc = mlirLocationCallSiteGet(loc, mlirLocationUnknownGet(context));
+    }
+  } else if (auto flc = node->sourceRange().file_line_col()) {
     const std::string &file = std::get<0>(*flc);
     int line = std::get<1>(*flc);
     int col = std::get<2>(*flc);
-    return mlirLocationFileLineColGet(context, toMlirStringRef(file), line,
-                                      col);
+    loc = mlirLocationFileLineColGet(context, toMlirStringRef(file), line, col);
+  }
+
+  std::string locationName;
+  auto scopeName = node->scopeName();
+  if (!scopeName.empty()) {
+    locationName = scopeName;
+  }
+
+  if (const c10::FunctionSchema *schema = node->maybeSchema()) {
+    if (!locationName.empty()) {
+      locationName += "/";
+    }
+    locationName += schema->operator_name().name;
   }
-  return mlirLocationUnknownGet(context);
+
+  if (!locationName.empty()) {
+    loc = mlirLocationNameGet(context, toMlirStringRef(locationName), loc);
+  }
+
+  return loc;
 }
 
 std::vector<MlirType>

diff --git a/test/python/importer/jit_ir/node_import/debug-info.py b/test/python/importer/jit_ir/node_import/debug-info.py
@@ -17,14 +17,11 @@
 @mb.import_function
 @torch.jit.script
 def add3(t0, t1, t2):
-  # TODO: Checks for debug info are quite hard with the new trailing debug
-  # attribute print. See if this can be improved.
-  # CHECK: loc({{.*}}debug-info.py":[[# @LINE + 1]]
+  # CHECK-DAG: torch.aten.add.Tensor {{.*}} loc("aten::add"({{.*}}debug-info.py":[[# @LINE + 1]]
   intermediate = t0 + t1
-  # CHECK: loc({{.*}}debug-info.py":[[# @LINE + 1]]
-  final = intermediate + t2
-  return final
+  # CHECK-DAG: torch.aten.mul.Tensor {{.*}} loc("aten::mul"({{.*}}debug-info.py":[[# @LINE + 1]]
+  return intermediate * t2
 
 # Verify again with debug info present. Just checking that it makes it in there.
-mb.module.operation.print(enable_debug_info=True)
+mb.module.operation.print(enable_debug_info=True, use_local_scope=True)
 print()