Fix location emission

Author: driazati

.gitignore

@@ -271,3 +271,6 @@ gallery/how_to/work_with_microtvm/micro_tvmc.py
 
 # Printed TIR code on disk
 *.tir
+
+# GDB history file
+.gdb_history

src/printer/tir_text_printer_debug.cc

@@ -25,25 +25,64 @@
 
 #include "tir_text_printer_debug.h"
 
+#include <optional>
 #include <string>
 
 #include "text_printer.h"
 
 namespace tvm {
 namespace tir {
 
-std::string span_text(const Span& span) {
+std::optional<std::string> span_text(const Span& span) {
   if (!span.defined()) {
-    return "missing";
+    return std::nullopt;
+  }
+
+  std::string source("main.tir");
+  if (span->source_name.defined() && span->source_name->name.get()) {
+    source = span->source_name->name;
   }
-  std::string source("file");
   return source + ":" + std::to_string(span->line) + ":" + std::to_string(span->column);
 }
 
+template <typename ObjectPtr>
+void add_all_relevant_lines(const std::vector<std::tuple<const ObjectPtr*, size_t>>& data,
+                            size_t current_line, Doc* output) {
+  ICHECK(output) << "output must be a valid Doc";
+  for (const auto& item : data) {
+    if (std::get<1>(item) != current_line - 1) {
+      // Item is not relevant for this line, skip it
+      continue;
+    }
+
+    // Print out the item's span info if present
+    auto text = span_text(std::get<0>(item)->span);
+    if (text.has_value()) {
+      *output << *text;
+    } else {
+      *output << "missing";
+    }
+    *output << ", ";
+  }
+}
+
 Doc TIRTextPrinterDebug::NewLine() {
   current_line_ += 1;
 
-  return TIRTextPrinter::NewLine();
+  if (!show_spans_) {
+    return TIRTextPrinter::NewLine();
+  }
+
+  Doc output;
+
+  output << " [";
+
+  add_all_relevant_lines(exprs_by_line_, current_line_, &output);
+  add_all_relevant_lines(stmts_by_line_, current_line_, &output);
+
+  output << "]" << TIRTextPrinter::NewLine();
+
+  return output;
 }
 
 #define X(TypeName)                                               \

src/printer/tir_text_printer_debug.h

@@ -37,7 +37,8 @@ namespace tir {
 
 class TIRTextPrinterDebug : public TIRTextPrinter {
  public:
-  TIRTextPrinterDebug() : TIRTextPrinter(false, &meta_), current_line_(1) {}
+  explicit TIRTextPrinterDebug(bool show_spans)
+      : TIRTextPrinter(false, &meta_), current_line_(1), show_spans_(show_spans) {}
 
   std::vector<std::tuple<const PrimExprNode*, size_t>> GetExprsByLine() const {
     return exprs_by_line_;
@@ -61,6 +62,9 @@ class TIRTextPrinterDebug : public TIRTextPrinter {
   // Line that the printer is currently printing
   size_t current_line_;
 
+  // Whether to include spans relevant to each line before a newline or not
+  bool show_spans_;
+
   // Record of all stmts and exprs and their corresponding line
   std::vector<std::tuple<const StmtNode*, size_t>> stmts_by_line_;
   std::vector<std::tuple<const PrimExprNode*, size_t>> exprs_by_line_;

src/target/llvm/codegen_cpu.cc

@@ -186,11 +186,16 @@ void CodeGenCPU::Init(const std::string& module_name, LLVMTarget* llvm_target, b
 llvm::DISubprogram* CodeGenCPU::CreateDebugFunction(const PrimFunc& f) {
 #if TVM_LLVM_VERSION >= 50
   llvm::SmallVector<llvm::Metadata*, 4> paramTys;
+
+  paramTys.push_back(GetDebugType(f->ret_type));
+  for (const auto& param : f->params) {
+    paramTys.push_back(GetDebugType(GetType(param)));
+  }
+
   auto* DIFunctionTy = dbg_info_->di_builder_->createSubroutineType(
       dbg_info_->di_builder_->getOrCreateTypeArray(paramTys));
 
-  // TODO(driazati): add the right argument info to the function
-  bool local_to_unit = false;
+  bool local_to_unit = llvm::GlobalVariable::isLocalLinkage(llvm::GlobalValue::InternalLinkage);
 
 #if TVM_LLVM_VERSION >= 80
   auto SPFlags = llvm::DISubprogram::toSPFlags(local_to_unit, /*IsDefinition=*/true,
@@ -207,14 +212,15 @@ llvm::DISubprogram* CodeGenCPU::CreateDebugFunction(const PrimFunc& f) {
       /*Flags=*/llvm::DINode::FlagPrototyped, /*isOptimized=*/true);
 #endif
   return DIFunction;
+#else
+  return nullptr;
 #endif
 }
 
 void CodeGenCPU::AddFunction(const PrimFunc& f) {
 #if TVM_LLVM_VERSION >= 50
   di_subprogram_ = CreateDebugFunction(f);
 #endif
-
   EmitDebugLocation(f->span);
   CodeGenLLVM::AddFunction(f);
   if (f_tvm_register_system_symbol_ != nullptr) {
@@ -272,6 +278,9 @@ void CodeGenCPU::AddDebugInformation(PrimFunc f_tir, llvm::Function* f_llvm) {
 #endif
 }
 
+llvm::DIType* CodeGenCPU::GetDebugType(const Type& ty_tir) {
+  return GetDebugType(ty_tir, GetLLVMType(ty_tir));
+}
 llvm::DIType* CodeGenCPU::GetDebugType(const Type& ty_tir, llvm::Type* ty_llvm) {
   if (ty_llvm == t_void_) {
     return nullptr;
@@ -952,7 +961,6 @@ llvm::Value* CodeGenCPU::CreateCallPacked(const CallNode* op, bool use_string_lo
 }
 
 llvm::Value* CodeGenCPU::CreateCallTracePacked(const CallNode* op) {
-  EmitDebugLocation(op);
   ICHECK_EQ(op->args.size(), 6U);
   PackedCall pc = MakeCallPackedLowered(op->args, op->dtype, op->args[3].as<IntImmNode>()->value,
                                         op->args[4].as<IntImmNode>()->value, true);
@@ -1388,7 +1396,6 @@ void CodeGenCPU::AddStartupFunction() {
 }
 
 llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode* op) {
-  EmitDebugLocation(op);
   if (op->op.same_as(builtin::tvm_call_packed_lowered())) {
     return CreateCallPacked(op, true /* use_string_lookup */);
   } else if (op->op.same_as(builtin::tvm_call_trace_packed_lowered())) {

src/target/llvm/codegen_cpu.h

@@ -195,6 +195,7 @@ class CodeGenCPU : public CodeGenLLVM {
 
   // Get the DWARF type corresponding to the LLVM type |ty|. The current API in practice only
   // generates |int32|, and |int8*|.
+  llvm::DIType* GetDebugType(const Type& ty_tir);
   llvm::DIType* GetDebugType(const Type& ty_tir, llvm::Type* ty_llvm);
   // Adds the DWARF debug information for |function| to |dbg_info_|.
   void AddDebugInformation(PrimFunc f_tir, llvm::Function* f_llvm);