[debugging] Use different linkage names for different functions

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: walter-erquinigo

src/cuda/tile/_exception.py

@@ -12,17 +12,25 @@
 
 @dataclass(eq=False, frozen=True)
 class FunctionDesc:
-    name: str | None
+    name: str | None  # None for lambdas
     filename: str
-    line: int
+    line: int  # 1-based
+    column: int  # 1-based
+    # If this FunctionDesc represents a concrete specialization of a source
+    # function other than the kernel entry point, this value will hold a
+    # unique identifier, which is used to distinguish distinct specialized
+    # functions in debug info.
+    specialization_id: str | None = None
+    # True for the FunctionDesc that represents the kernel entry point.
+    is_entry: bool = False
 
     def __str__(self):
-        return f"'{self.name}' @{self.filename}:{self.line}"
+        return f"'{self.name}' @{self.filename}:{self.line}:{self.column}"
 
     def short_str(self):
         if self.name is None:
             base_name = os.path.basename(self.filename)
-            return f"<lambda at {base_name}:{self.line}>"
+            return f"<lambda at {base_name}:{self.line}:{self.column}>"
         else:
             return f"<function {self.name}>"
 

src/cuda/tile/_ir/ir.py

@@ -40,6 +40,13 @@ def __init__(self, log_ir_on_error: bool, tileiras_version: BytecodeVersion):
         self.log_ir_on_error = log_ir_on_error
         self._aggregate_values: Dict[str, Any] = dict()
         self.tileiras_version: BytecodeVersion = tileiras_version
+        self._function_specialization_id_counter = itertools.count()
+
+    def next_function_specialization_id(self) -> str:
+        # Monotonic counter used as a unique id when creating concrete FunctionDescs
+        # to distinguish them in debug info. The monotonicity ensures that the
+        # emitted bytecode is deterministic for cache hits.
+        return f"s{next(self._function_specialization_id_counter)}"
 
     #  Make a Var with a unique name based on `name`.
     def make_var(self, name: str, loc: Loc) -> Var:

src/cuda/tile/_ir/ops.py

@@ -167,7 +167,7 @@ def format_var(var):
 
 @impl(hir_stubs.loop)
 async def loop_impl(body: hir.Block, iterable: Var):
-    from .._passes.hir2ir import dispatch_hir_block
+    from .._passes.hir2ir import dispatch_hir_block, retarget_loc
 
     scope = Scope.get_current()
     range_ty = require_optional_range_type(iterable)
@@ -202,7 +202,7 @@ async def loop_impl(body: hir.Block, iterable: Var):
 
     # Process the loop body
     loop_info = ControlFlowInfo(stored_locals)
-    body_loc = body.loc.with_call_site(scope.call_site)
+    body_loc = retarget_loc(body.loc, scope)
     with enter_nested_block(body_loc) as new_body, scope.change_loop_info(loop_info), \
             scope.local.enter_branch():
         # Define body variables. Not all of them will eventually be kept,
@@ -396,7 +396,7 @@ async def _flatten_branch(branch: hir.Block) -> Var | None:
 
 @impl(hir_stubs.if_else)
 async def if_else_impl(cond: Var, then_block: hir.Block, else_block: hir.Block) -> Var | None:
-    from .._passes.hir2ir import dispatch_hir_block
+    from .._passes.hir2ir import dispatch_hir_block, retarget_loc
 
     require_bool(cond)
     if cond.is_constant():
@@ -415,7 +415,7 @@ async def if_else_impl(cond: Var, then_block: hir.Block, else_block: hir.Block)
 
     # Convert the "then" branch from HIR to IR
     info = ControlFlowInfo(stored_locals)
-    then_loc = then_block.loc.with_call_site(scope.call_site)
+    then_loc = retarget_loc(then_block.loc, scope)
     with enter_nested_block(then_loc) as new_then_block, scope.change_if_else_info(info), \
             scope.local.enter_branch():
         await dispatch_hir_block(then_block)
@@ -427,7 +427,7 @@ async def if_else_impl(cond: Var, then_block: hir.Block, else_block: hir.Block)
     #        EndBranch
     #    <else_block>
     # This is to avoid the situation where none of the branches yield.
-    else_loc = else_block.loc.with_call_site(scope.call_site)
+    else_loc = retarget_loc(else_block.loc, scope)
     if len(info.jumps) == 0:
         info = ControlFlowInfo(())
         with enter_nested_block(else_loc) as new_else_block, scope.change_if_else_info(info), \

src/cuda/tile/_ir/scope.py

@@ -9,7 +9,7 @@
 from dataclasses import dataclass
 from typing import TypeVar, Generic
 
-from cuda.tile._exception import Loc, TileSyntaxError
+from cuda.tile._exception import Loc, FunctionDesc, TileSyntaxError
 from cuda.tile._ir import hir
 from cuda.tile._ir.hir import ResolvedName
 from cuda.tile._ir.ir import Operation, Var, IRContext
@@ -142,6 +142,9 @@ class Scope:
     call_site: Loc | None
     hir2ir_varmap: IntMap[Var]
     func_hir: hir.Function
+    # FunctionDesc that represents the concrete specialization of `func_hir`
+    # representing this scope.
+    concrete_func_desc: FunctionDesc
     context_stack: list[ContextManagerState] = dataclasses.field(default_factory=list)
     loop_context_stack_depth: int | None = None
 

src/cuda/tile/_ir2bytecode.py

@@ -4,6 +4,7 @@
 
 import functools
 import os
+import re
 from contextlib import contextmanager
 from typing import Dict, Tuple, Any, Optional
 import warnings
@@ -232,14 +233,56 @@ def encode_comparison(builder: bc.CodeBuilder, fn: str, lhs: bc.Value, rhs: bc.V
         raise TileInternalError(f'Unexpected dtype: {dtype}')
 
 
+def create_synthetic_linkage_name(func_desc: FunctionDesc) -> str:
+    # Build a synthetic linkage name for a helper function or lambda. Format:
+    #
+    #     <name>@<basename>:<line>:<column>_<specialization_id>
+    #
+    # By construction every FunctionDesc reaching this point has been
+    # concretized by hir2ir — only the kernel entry skips this path (it uses
+    # the externally-visible symbol instead), so the specialization_id is
+    # required here.
+    assert func_desc.specialization_id is not None, (
+        f"create_synthetic_linkage_name called on a FunctionDesc without a "
+        f"specialization_id: {func_desc}. hir2ir must concretize every "
+        f"non-entry function before bytecode generation."
+    )
+    base_name = os.path.basename(func_desc.filename) or "unknown"
+    stem = os.path.splitext(base_name)[0]
+    # Convert any non-alphanumeric chars to _.
+    stem = re.sub(r"[^A-Za-z0-9_]", "_", stem) or "anonymous"
+    func_part = func_desc.name if func_desc.name is not None else "lambda"
+    return (f"{func_part}@{stem}:{func_desc.line}:{func_desc.column}"
+            f"_{func_desc.specialization_id}")
+
+
 class DebugAttrMap:
-    def __init__(self, debug_attr_table: bc.DebugAttrTable, linkage_name: str, anonymize: bool):
+    def __init__(self,
+                 debug_attr_table: bc.DebugAttrTable,
+                 entry_symbol: str,
+                 anonymize: bool):
         self._subprogram_cache = {}
         self._debug_attr_table = debug_attr_table
-        self._linkage_name = linkage_name
+        self._entry_symbol = entry_symbol
         self._anonymize = anonymize
 
+    def _linkage_for(self, func_desc: FunctionDesc) -> str:
+        # The kernel entry point keeps the externally-visible symbol so it can
+        # be looked up by the loader. Every other function gets a per-function
+        # artificial linkage name.
+        if func_desc.is_entry:
+            return self._entry_symbol
+        return create_synthetic_linkage_name(func_desc)
+
     def get_subprogram(self, func_desc: FunctionDesc) -> bc.DebugAttrId:
+        # Every FunctionDesc reaching DI emission must satisfy: a function has
+        # no specialization_id iff it is the kernel entry. hir2ir leaves the
+        # entry's abstract desc as-is and concretizes everyone else; anything
+        # else is a bug.
+        assert func_desc.is_entry == (func_desc.specialization_id is None), (
+            f"FunctionDesc invariant violated: is_entry={func_desc.is_entry} "
+            f"but specialization_id={func_desc.specialization_id!r}: {func_desc}"
+        )
         try:
             return self._subprogram_cache[func_desc]
         except KeyError:
@@ -252,7 +295,7 @@ def get_subprogram(self, func_desc: FunctionDesc) -> bc.DebugAttrId:
             file=file_attr,
             line=func_desc.line,
             name="<lambda>" if func_desc.name is None else func_desc.name,
-            linkage_name=self._linkage_name,
+            linkage_name=self._linkage_for(func_desc),
             compile_unit=compile_unit_attr,
             scope_line=func_desc.line,
         )
@@ -435,7 +478,8 @@ def generate_bytecode_for_kernel(func_body: Block,
                                 num_worker_warps_per_cta=num_worker_warps)
 
     param_type_ids = [typeid(writer.type_table, p.get_type()) for p in func_body.params]
-    debug_attr_map = DebugAttrMap(writer.debug_attr_table, symbol, anonymize_debug_attr)
+    debug_attr_map = DebugAttrMap(writer.debug_attr_table, symbol,
+                                  anonymize=anonymize_debug_attr)
     func_debug_attr = debug_attr_map.get_debugattr(func_body.loc)
 
     with writer.function(name=symbol,

src/cuda/tile/_passes/ast2hir.py

@@ -52,7 +52,7 @@ def get_function_hir(pyfunc: Callable, entry_point: bool) -> hir.Function:
     assert len(mod.body[0].body) == 1
     func_def = mod.body[0].body[0]
     assert isinstance(func_def, ast.FunctionDef)
-    _fix_line_numbers(func_def, first_line)
+    _fix_line_and_column_numbers(func_def, first_line)
 
     func_globals = dict(pyfunc.__builtins__)
     func_globals.update(pyfunc.__globals__)
@@ -62,7 +62,8 @@ def get_function_hir(pyfunc: Callable, entry_point: bool) -> hir.Function:
             func_globals[name] = cell.cell_contents
 
     filename = inspect.getfile(pyfunc)
-    desc = FunctionDesc(func_def.name, filename, first_line)
+    desc = FunctionDesc(func_def.name, filename, first_line, func_def.col_offset + 1,
+                        is_entry=entry_point)
     local_names, _, _ = ast_get_all_local_names(func_def)
     ctx = _Context(filename, first_line, desc, func_globals, local_names, entry_point)
     signature = inspect.signature(pyfunc)
@@ -73,9 +74,9 @@ def get_function_hir(pyfunc: Callable, entry_point: bool) -> hir.Function:
     return ret
 
 
-# Translate the 1-based line numbers of the chunk we passed to the AST parser
-# to the original 1-based line numbers in the file.
-def _fix_line_numbers(tree: ast.AST, first_line: int):
+# Translate the 1-based line and 0-based column numbers of the chunk we passed to the
+# AST parser to the original line and column numbers in the file.
+def _fix_line_and_column_numbers(tree: ast.AST, first_line: int):
     for node in ast.walk(tree):
         if hasattr(node, "lineno"):
             # Why -2?
@@ -846,9 +847,8 @@ def _pass_stmt(stmt: ast.Pass, ctx: _Context) -> None:
 def _make_closure(node: ast.FunctionDef | ast.Lambda, ctx: _Context) -> hir.Value:
     signature, default_exprs = _signature_from_ast_arguments(node.args)
     default_values = tuple(_expr(x, ctx) for x in default_exprs)
-    line_no = node.lineno
     name = None if isinstance(node, ast.Lambda) else node.name
-    desc = FunctionDesc(name, ctx.filename, line_no)
+    desc = FunctionDesc(name, ctx.filename, node.lineno, node.col_offset + 1)
     new_locals, new_globals, _ = ast_get_all_local_names(node)
     local_names = (ctx.local_names - new_globals) | new_locals
     new_ctx = _Context(ctx.filename, ctx.first_line, desc, ctx.frozen_globals,

src/cuda/tile/_passes/hir2ir.py

@@ -11,7 +11,8 @@
 from .ast2hir import get_function_hir
 from .. import TileTypeError
 from .._coroutine_util import resume_after, run_coroutine
-from .._exception import Loc, TileSyntaxError, TileInternalError, TileError, TileRecursionError
+from .._exception import Loc, FunctionDesc, TileSyntaxError, TileInternalError, TileError, \
+    TileRecursionError
 from .._execution import is_stub
 from .._ir import hir, ir
 from .._ir.ir import Var, IRContext, BoundMethodValue, ClosureValue, TupleValue
@@ -35,10 +36,11 @@ def hir2ir(func_hir: hir.Function,
     run_coroutine(_hir2ir_coroutine(func_hir, param_aggregate_vars, ir_ctx))
 
 
-async def _hir2ir_coroutine(func_hir: hir.Function,
-                            param_aggregate_vars: Sequence[ir.Var],
-                            ir_ctx: IRContext):
-    scope = _create_scope(func_hir, ir_ctx, call_site=None, parent_scopes=())
+async def _hir2ir_coroutine(
+    func_hir: hir.Function, param_aggregate_vars: Sequence[ir.Var], ir_ctx: IRContext
+):
+    scope = _create_scope(func_hir, ir_ctx, call_site=None, parent_scopes=(),
+                          concrete_func_desc=func_hir.desc)
     for local_idx, var in zip(func_hir.param_local_indices, param_aggregate_vars, strict=True):
         scope.local[local_idx] = var
 
@@ -56,9 +58,28 @@ async def _hir2ir_coroutine(func_hir: hir.Function,
 
 
 def _create_scope(func_hir: hir.Function, ir_ctx: IRContext, call_site: Loc | None,
-                  parent_scopes: tuple[LocalScope, ...]) -> Scope:
+                  parent_scopes: tuple[LocalScope, ...],
+                  concrete_func_desc: FunctionDesc) -> Scope:
     local_scope = LocalScope(func_hir.local_names, ir_ctx)
-    return Scope(parent_scopes + (local_scope,), None, None, call_site, IntMap(), func_hir)
+    return Scope(parent_scopes + (local_scope,), None, None, call_site, IntMap(), func_hir,
+                 concrete_func_desc=concrete_func_desc)
+
+
+def _concretize_func_desc(func_hir: hir.Function, ir_ctx: IRContext) -> FunctionDesc:
+    # Mint a fresh FunctionDesc whose `specialization_id` makes its synthesized
+    # linkage name unique across every inlining.
+    return dataclasses.replace(func_hir.desc,
+                               specialization_id=ir_ctx.next_function_specialization_id())
+
+
+def retarget_loc(loc: Loc, scope: Scope) -> 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.
+    new_function = loc.function
+    if loc.function is scope.func_hir.desc:
+        new_function = scope.concrete_func_desc
+    return dataclasses.replace(loc, function=new_function, call_site=scope.call_site)
 
 
 async def dispatch_hir_block(block: hir.Block, cur_builder: ir.Builder | None = None):
@@ -72,7 +93,7 @@ async def _dispatch_hir_block_inner(block: hir.Block, builder: ir.Builder):
     try:
         scope = Scope.get_current()
         for cursor, call in enumerate(block.calls):
-            loc = call.loc.with_call_site(scope.call_site)
+            loc = retarget_loc(call.loc, scope)
             with _wrap_exceptions(loc), builder.change_loc(loc):
                 await _dispatch_call(call, scope)
             if builder.is_terminated:
@@ -81,7 +102,7 @@ async def _dispatch_hir_block_inner(block: hir.Block, builder: ir.Builder):
                 return
         cursor = len(block.calls)
 
-        loc = block.jump_loc.with_call_site(scope.call_site)
+        loc = retarget_loc(block.jump_loc, scope)
         with _wrap_exceptions(loc), builder.change_loc(loc):
             _dispatch_hir_jump(block, scope)
     except Exception:
@@ -155,9 +176,13 @@ async def _call_user_defined(callee_hir: hir.Function,
             raise TileSyntaxError("Variadic keyword parameters in user-defined"
                                   " functions are not supported")
 
-    # Activate a fresh Scope.
+    # Activate a fresh Scope. Each inlining gets its own concretized
+    # FunctionDesc so that DI never merges two specializations whose generated
+    # IR might differ.
     new_scope = _create_scope(callee_hir, builder.ir_ctx, call_site=builder.loc,
-                              parent_scopes=parent_scopes)
+                              parent_scopes=parent_scopes,
+                              concrete_func_desc=_concretize_func_desc(callee_hir,
+                                                                       builder.ir_ctx))
     with new_scope.make_current():
         # Call store_var() to bind arguments to parameters.
         for arg, local_idx, param_loc in zip(arg_list, callee_hir.param_local_indices,