Set unknown loc for dummy return in main function

Signed-off-by: Ziheng Deng <zihengd@nvidia.com>

Author: gunnersdeng

src/cuda/tile/_ir2bytecode.py

@@ -330,7 +330,7 @@ def get_subprogram(self, func_desc: FunctionDesc) -> bc.DebugAttrId:
         return ret
 
     def get_debugattr(self, loc: Loc) -> bc.DebugAttrId:
-        if self._anonymize:
+        if self._anonymize or loc.is_unknown():
             return bc.MISSING_DEBUG_ATTR_ID
 
         subprogram = self.get_subprogram(loc.function)

src/cuda/tile/_passes/ast2hir.py

@@ -1000,7 +1000,8 @@ def _ast2hir(func_def: ast.FunctionDef | ast.Lambda, ctx: _Context) -> hir.Block
             _stmt_list(func_def.body, ctx)
             # Add a Return jump to the root block if it doesn't have one
             if root_block.jump is None:
-                ctx.set_block_jump(hir.Jump.RETURN)
+                with ctx.change_loc(Loc.unknown()):
+                    ctx.set_block_jump(hir.Jump.RETURN)
         elif isinstance(func_def, ast.FunctionDef):
             # To enable early returns in a helper function, wrap the body in a loop.
             # Thus, we can use "break" to implement the return statement.

src/cuda/tile/_passes/hir2ir.py

@@ -75,6 +75,9 @@ def _concretize_func_desc(func_hir: hir.Function, ir_ctx: IRContext) -> Function
 
 
 def retarget_loc(loc: Loc, scope: Scope) -> Loc:
+    if loc.is_unknown():
+        return loc
+
     # Splice in the scope's call site and, if this loc belongs to the function
     # currently being inlined, swap in the per-specialization FunctionDesc so
     # emitted ops carry the right debug info.

test/test_function_debuginfo.py

@@ -7,11 +7,13 @@
 import pytest
 
 import cuda.tile as ct
+import cuda.tile._bytecode as bc
 from cuda.tile._bytecode.basic import StringTable
 from cuda.tile._bytecode.debug_info import DebugAttrTable
 from cuda.tile._cext import CallingConvention
 from cuda.tile._compile import compile_tile
-from cuda.tile._exception import FunctionDesc
+from cuda.tile._exception import FunctionDesc, Loc
+from cuda.tile._ir.ops import Return
 from cuda.tile._ir2bytecode import DebugAttrMap, create_synthetic_linkage_name
 from cuda.tile.compilation import ArrayConstraint, KernelSignature
 
@@ -202,6 +204,34 @@ def kernel(x):
     assert entry.line == kernel_line and entry.line != body_line
 
 
+def test_dummy_return_sets_unknown_loc():
+    def kernel(x):
+        t = ct.load(x, (0,), (1,))
+        ct.store(x, (1,), t)
+
+    body = _compile(kernel)
+    return_ops = [op for op in body.traverse() if isinstance(op, Return)]
+    assert len(return_ops) == 1, "expected one Return op"
+    assert return_ops[0].loc.is_unknown()
+
+
+def test_explicit_return_keeps_loc():
+    def kernel(x):
+        t = ct.load(x, (0,), (1,))
+        ct.store(x, (1,), t)
+        return None
+
+    body = _compile(kernel)
+    return_line, _ = _line_col(kernel, "return None")
+    return_ops = [op for op in body.traverse() if isinstance(op, Return)]
+    assert len(return_ops) == 1, "expected one Return op"
+    return_op = return_ops[0]
+    assert return_op, "expected a Return op"
+    assert return_op.loc.line == return_line, (return_op.loc.line, return_line)
+    assert return_op.loc.function is not None
+    assert return_op.loc.function.name == "kernel"
+
+
 def test_linkage_names_appear_in_emitted_bytecode():
     # End-to-end smoke check: every linkage name we compute should also be
     # serialized into the bytecode's string table. This catches regressions
@@ -258,3 +288,8 @@ def compile_once():
         )
 
     assert compile_once() == compile_once()
+
+
+def test_unknown_loc_maps_to_missing_debug_attr():
+    m = DebugAttrMap(DebugAttrTable(StringTable()), entry_symbol="kern", anonymize=False)
+    assert m.get_debugattr(Loc.unknown()) == bc.MISSING_DEBUG_ATTR_ID