Merge pull request #1443 from IntelPython/improve/function_arg_mangling

Add the mangled_args property to kernel_api types.

Author: ZzEeKkAa

numba_dpex/core/debuginfo.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: 2020 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Implements a custom debug metadata generator class for numba-dpex kernels.
+"""
+
+from numba.core import debuginfo
+
+
+class DIBuilder(debuginfo.DIBuilder):
+    """Overrides Numba's default DIBuilder with numba-dpex-specific customizations."""
+
+    # pylint: disable=too-many-arguments
+    def mark_subprogram(self, function, qualname, argnames, argtypes, line):
+        """Sets DW_AT_name and DW_AT_linkagename tags for a kernel decorated function.
+
+        Numba generates a unique name for every function it compiles, but in
+        upstream Numba the unique name is not used as the "qualified" name of
+        the function. The behavior leads to a bug discovered in Numba-dpex when
+        a compiled function uses closure variables.
+        Refer (https://github.com/IntelPython/numba-dpex/issues/898).
+        To resolve the issue numba-dpex uses the unique_name as the qualified
+        name. Refer to
+        :class:`numba_dpex.core.passes.passes.QualNameDisambiguationLowering`.
+        However, doing so breaks setting GDB breakpoints based on function
+        name as the function name is no longer what is in the source, but what
+        is the unique name generated by Numba. To fix it, numba-dpex uses a
+        modified DISubprogram metadata generator. The name (DW_AT_name) tag is
+        set to the base function name, discarding the unique qualifier and
+        linkagename is set to an empty string.
+        """
+        name = qualname[0 : qualname.find("$")]  # noqa: E203
+        argmap = dict(zip(argnames, argtypes))
+
+        di_subp = self._add_subprogram(
+            name=name,
+            linkagename="",
+            line=line,
+            function=function,
+            argmap=argmap,
+        )
+        function.set_metadata("dbg", di_subp)

numba_dpex/core/types/kernel_api/atomic_ref.py

@@ -80,3 +80,13 @@ def cast_python_value(self, args):
         AtomicRefType throws a NotImplementedError.
         """
         raise NotImplementedError
+
+    @property
+    def mangling_args(self):
+        args = [
+            self.dtype,
+            self.memory_order,
+            self.memory_scope,
+            self.address_space,
+        ]
+        return self.__class__.__name__, args

numba_dpex/core/types/kernel_api/index_space_ids.py

@@ -31,6 +31,11 @@ def key(self):
     def cast_python_value(self, args):
         raise NotImplementedError
 
+    @property
+    def mangling_args(self):
+        args = [self.ndim]
+        return self.__class__.__name__, args
+
 
 class ItemType(types.Type):
     """Numba-dpex type corresponding to :class:`numba_dpex.kernel_api.Item`"""
@@ -53,6 +58,11 @@ def key(self):
         """Numba type specific overload"""
         return self._ndim
 
+    @property
+    def mangling_args(self):
+        args = [self.ndim]
+        return self.__class__.__name__, args
+
     def cast_python_value(self, args):
         raise NotImplementedError
 
@@ -78,5 +88,10 @@ def key(self):
         """Numba type specific overload"""
         return self._ndim
 
+    @property
+    def mangling_args(self):
+        args = [self.ndim]
+        return self.__class__.__name__, args
+
     def cast_python_value(self, args):
         raise NotImplementedError

numba_dpex/core/types/kernel_api/ranges.py

@@ -28,6 +28,11 @@ def ndim(self):
     def key(self):
         return self._ndim
 
+    @property
+    def mangling_args(self):
+        args = [self.ndim]
+        return self.__class__.__name__, args
+
 
 class NdRangeType(types.Type):
     """Numba-dpex type corresponding to
@@ -49,3 +54,8 @@ def ndim(self):
     @property
     def key(self):
         return self._ndim
+
+    @property
+    def mangling_args(self):
+        args = [self.ndim]
+        return self.__class__.__name__, args

numba_dpex/kernel_api_impl/spirv/target.py

@@ -11,7 +11,7 @@
 import dpnp
 from llvmlite import binding as ll
 from llvmlite import ir as llvmir
-from numba.core import cgutils, funcdesc
+from numba.core import cgutils
 from numba.core import types as nb_types
 from numba.core import typing
 from numba.core.base import BaseContext
@@ -21,8 +21,10 @@
 from numba.core.typing import cmathdecl, enumdecl
 
 from numba_dpex.core.datamodel.models import _init_kernel_data_model_manager
+from numba_dpex.core.debuginfo import DIBuilder as DpexDIbuilder
 from numba_dpex.core.types import IntEnumLiteral
 from numba_dpex.core.typing import dpnpdecl
+from numba_dpex.core.utils import itanium_mangler
 from numba_dpex.kernel_api.flag_enum import FlagEnum
 from numba_dpex.kernel_api.memory_enums import AddressSpace as address_space
 from numba_dpex.kernel_api_impl.spirv import printimpl
@@ -134,6 +136,7 @@ class SPIRVTargetContext(BaseContext):
 
     implement_powi_as_math_call = True
     allow_dynamic_globals = True
+    DIBuilder = DpexDIbuilder
 
     def __init__(self, typingctx, target=SPIRV_TARGET_NAME):
         super().__init__(typingctx, target)
@@ -188,7 +191,7 @@ def _generate_spir_kernel_wrapper(self, func, argtypes):
         wrapper_module = self._internal_codegen.create_empty_spirv_module(
             "dpex.kernel.wrapper"
         )
-        wrappername = func.name.replace("dpex_fn", "dpex_kernel")
+        wrappername = func.name + ("dpex_kernel")
         argtys = list(arginfo.argument_types)
         fnty = llvmir.FunctionType(
             llvmir.IntType(32),
@@ -319,12 +322,9 @@ def target_data(self):
 
     def mangler(self, name, types, *, abi_tags=(), uid=None):
         """
-        Generates a name for a function by appending \"dpex_fn\" to the
-        name of the function before calling Numba's default function name
-        mangler."""
-        return funcdesc.default_mangler(
-            name + "dpex_fn", types, abi_tags=abi_tags, uid=uid
-        )
+        Generates a mangled function name using numba_dpex's itanium mangler.
+        """
+        return itanium_mangler.mangle(name, types, abi_tags=abi_tags, uid=uid)
 
     def prepare_spir_kernel(self, func, argtypes):
         """Generates a wrapper function with \"spir_kernel\" calling conv that

numba_dpex/tests/debugging/test_backtraces.py

@@ -25,7 +25,5 @@ def test_backtrace(app):
 
     app.backtrace()
 
-    app.expect(
-        r"#0.*__main__::func_sum.* at simple_dpex_func.py:12", with_eol=True
-    )
-    app.expect(r"#1.*__main__::kernel_sum", with_eol=True)
+    app.expect(r"#0.*func_sum.* at simple_dpex_func.py:12", with_eol=True)
+    app.expect(r"#1.*kernel_sum", with_eol=True)

numba_dpex/tests/debugging/test_breakpoints.py

@@ -56,11 +56,6 @@ def test_device_func_breakpoint(
     SAT-4449
     """
 
-    if api == "numba-ndpx-kernel" and breakpoint != "side-by-side.py:15":
-        pytest.skip(
-            "Breakpoint by function name not working for numba-dpex."
-        )  # TODO: https://github.com/IntelPython/numba-dpex/issues/1242
-
     app.breakpoint(breakpoint, condition=condition)
     app.run(f"side-by-side.py --api={api}")
     app.expect_hit_breakpoint("side-by-side.py:15")

numba_dpex/tests/debugging/test_info.py

@@ -104,7 +104,7 @@ def test_info_functions(app):
 
     app.info_functions("data_parallel_sum")
 
-    app.child.expect(r"11:\s+[a-z 0-9\*]+__main__::data_parallel_sum")
+    app.child.expect(r"11:\s+[a-z 0-9\*]+data_parallel_sum")
 
 
 @pytest.mark.parametrize(

numba_dpex/tests/debugging/test_stepping.py

@@ -31,7 +31,7 @@ def test_next(app: gdb):
     )
     # checking that we did not step in
     app.next()
-    app.expect(r"in _ZN8__main__21kernel_sum_", with_eol=True)
+    app.expect(r"in _ZN8__main__14kernel_sum_", with_eol=True)
 
 
 def test_step(app: gdb):
@@ -44,7 +44,7 @@ def test_step(app: gdb):
     )
     app.set_scheduler_lock()
     app.step()
-    app.expect(r"__main__::func_sum.* at simple_dpex_func.py:12", with_eol=True)
+    app.expect(r"func_sum.* at simple_dpex_func.py:12", with_eol=True)
     app.expect(r"12\s+result = a_in_func \+ b_in_func", with_eol=True)
     app.step()
     app.expect(

numba_dpex/tests/test_debuginfo.py

@@ -126,8 +126,8 @@ def data_parallel_sum(item, a, b, c):
         c[i] = func_sum(a[i], b[i])
 
     ir_tags = [
-        r'\!DISubprogram\(name: ".*func_sum\$?\d*"',
-        r'\!DISubprogram\(name: ".*data_parallel_sum\$?\d*"',
+        r'\!DISubprogram\(name: ".*func_sum*"',
+        r'\!DISubprogram\(name: ".*data_parallel_sum*"',
     ]
 
     sig = (itemty, f32arrty, f32arrty, f32arrty)
@@ -156,8 +156,8 @@ def data_parallel_sum(item, a, b, c):
         c[i] = func_sum(a[i], b[i])
 
     ir_tags = [
-        r'\!DISubprogram\(name: ".*func_sum\$?\d*"',
-        r'\!DISubprogram\(name: ".*data_parallel_sum\$\d*"',
+        r'\!DISubprogram\(name: ".*func_sum*"',
+        r'\!DISubprogram\(name: ".*data_parallel_sum"',
     ]
 
     sig = (itemty, f32arrty, f32arrty, f32arrty)
@@ -176,6 +176,8 @@ def data_parallel_sum(item, a, b, c):
 
 
 def test_debuginfo_DISubprogram_linkageName():
+    """Tests to check that the linkagename tag is not set by numba-dpex."""
+
     def foo(item, a, b):
         i = item.get_id(0)
         b[i] = a[i]
@@ -190,7 +192,9 @@ def foo(item, a, b):
     kernel_ir = kcres.library.get_llvm_str()
 
     for tag in ir_tags:
-        assert make_check(kernel_ir, tag)
+        # Ensure that linkagename (DW_AT_linkagename) tag is not present for
+        # the DISubprogram attribute.
+        assert not make_check(kernel_ir, tag)
 
 
 def test_debuginfo_DICompileUnit_language_and_producer():