From 31803e6ec70f4cfd7401242a6481d7498790fe14 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <Lunderberg@users.noreply.github.com>
Date: Fri, 22 Mar 2024 16:34:29 -0500
Subject: [PATCH] [LLVM] Lack of DWARF type is not an error (#16748)

Prior to this commit, the `CodeGenLLVM::GetDebugType` would raise an
exception if it could not convert the TIR data type to an equivalent
DWARF type for debug symbols.  This commit updates the behavior to
instead return `nullptr`, representing an unknown type in DWARF
---
 src/target/llvm/codegen_llvm.cc          | 25 ++++++++++++------------
 tests/python/tir-base/test_debug_info.py | 23 +++++++++++++++++++++-
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 8fe740dad197..938c18f19845 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -2229,19 +2229,18 @@ llvm::DIType* CodeGenLLVM::GetDebugType(const Type& ty_tir, llvm::Type* ty_llvm)
 
   } else if (auto* prim_type = ty_tir.as<PrimTypeNode>()) {
     DataType dtype = prim_type->dtype;
-    auto dwarf_type = [&]() -> llvm::dwarf::TypeKind {
-      if (dtype.is_bool()) {
-        return llvm::dwarf::DW_ATE_boolean;
-      } else if (dtype.is_float()) {
-        return llvm::dwarf::DW_ATE_float;
-      } else if (dtype.is_int()) {
-        return llvm::dwarf::DW_ATE_signed;
-      } else if (dtype.is_uint()) {
-        return llvm::dwarf::DW_ATE_unsigned;
-      } else {
-        LOG(FATAL) << "No DWARF representation for TIR type " << dtype;
-      }
-    }();
+    llvm::dwarf::TypeKind dwarf_type;
+    if (dtype.is_bool()) {
+      dwarf_type = llvm::dwarf::DW_ATE_boolean;
+    } else if (dtype.is_float()) {
+      dwarf_type = llvm::dwarf::DW_ATE_float;
+    } else if (dtype.is_int()) {
+      dwarf_type = llvm::dwarf::DW_ATE_signed;
+    } else if (dtype.is_uint()) {
+      dwarf_type = llvm::dwarf::DW_ATE_unsigned;
+    } else {
+      return nullptr;
+    }
 
     return dbg_info_->di_builder_->createBasicType(DLDataType2String(dtype),
                                                    dtype.bits() * dtype.lanes(), dwarf_type);
diff --git a/tests/python/tir-base/test_debug_info.py b/tests/python/tir-base/test_debug_info.py
index a94d4d74f2c8..7fc9bcf31633 100644
--- a/tests/python/tir-base/test_debug_info.py
+++ b/tests/python/tir-base/test_debug_info.py
@@ -19,7 +19,7 @@
 import tvm.testing
 from tvm import tir
 from tvm import relay
-from tvm.script import tir as T
+from tvm.script import tir as T, ir as I
 
 from typing import List, Dict
 import re
@@ -165,5 +165,26 @@ def test_llvm_ir_debug_accuracy():
     assert debug_line_no == 56
 
 
+def test_building_without_llvm_equivalent():
+    """A TIR PrimFunc may contain non-LLVM types
+
+    Types used in optimized kernels (e.g. "e4m3_float8") may not have
+    an equivalent in DWARF, or the mapping from TIR type to DWARF type
+    may not be defined.  If this occurs, the function should still be
+    able to be built.
+    """
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True)
+        def main(A_data: T.handle("e4m3_float8"), B_data: T.handle("e4m3_float8")):
+            A = T.decl_buffer(128, "e4m3_float8", data=A_data)
+            B = T.decl_buffer(128, "e4m3_float8", data=B_data)
+            for i in range(128):
+                B[i] = A[i]
+
+    tvm.target.codegen.build_module(Module, "llvm")
+
+
 if __name__ == "__main__":
     tvm.testing.main()