[Frontend] Use full module and qualname when mangling non-kernel function names (#7025)

Without this, calling two different Python jit functions with the same
name will cause a name collision.

Author: Mogball

python/test/gluon/test_frontend.py

@@ -307,25 +307,25 @@ def anchor(x):
 @filecheck_test
 @gluon.jit
 def test_warp_specialize():
-    # CHECK-LABEL: tt.func public @test_warp_specialize
+    # CHECK-LABEL: test_warp_specialize
     # CHECK-NEXT:    [[A:%.*]] = tt.make_range {end = 1 : i32, start = 0 : i32}
     # CHECK-NEXT:    [[B:%.*]] = tt.make_range {end = 2 : i32, start = 0 : i32}
     # CHECK-NEXT:    [[C:%.*]] = tt.make_range {end = 4 : i32, start = 0 : i32}
     # CHECK-NEXT:    [[OUTS:%.*]]:3 = ttg.warp_specialize([[A]], [[B]], [[C]]) {{.*}}requestedRegisters = array<i32: 24, 48>
     # CHECK-NEXT:    default {
-    # CHECK-NEXT:      [[RESULTS:%.*]]:3 = tt.call @"warp_specialize_default{{.*}}"([[A]], [[B]], [[C]])
+    # CHECK-NEXT:      [[RESULTS:%.*]]:3 = tt.call @{{.*}}warp_specialize_default{{.*}}([[A]], [[B]], [[C]])
     # CHECK-NEXT:      warp_yield [[RESULTS]]#0, [[RESULTS]]#1, [[RESULTS]]#2
     # CHECK-NEXT:    }
     # CHECK-NEXT:    partition0(%arg0: tensor<1xi32>, %arg1: tensor<2xi32>, %arg2: tensor<4xi32>) num_warps(4) {
-    # CHECK-NEXT:      call @"warp_specialize_worker0{{.*}}"(%arg0, %arg1, %arg2)
+    # CHECK-NEXT:      call @{{.*}}warp_specialize_worker0{{.*}}(%arg0, %arg1, %arg2)
     # CHECK-NEXT:      warp_return
     # CHECK-NEXT:    }
     # CHECK-NEXT:    partition1(%arg0: tensor<1xi32>, %arg1: tensor<2xi32>, %arg2: tensor<4xi32>) num_warps(4) {
-    # CHECK-NEXT:      call @"warp_specialize_worker1{{.*}}"(%arg0, %arg1, %arg2)
+    # CHECK-NEXT:      call @{{.*}}warp_specialize_worker1{{.*}}(%arg0, %arg1, %arg2)
     # CHECK-NEXT:      warp_return
     # CHECK-NEXT:    }
-    # CHECK-NEXT:    call @anchor{{.*}}([[OUTS]]#0)
-    # CHECK-NEXT:    call @"anchor{{.*}}"([[OUTS]]#1, [[OUTS]]#2)
+    # CHECK-NEXT:    call @{{.*}}anchor{{.*}}([[OUTS]]#0)
+    # CHECK-NEXT:    call @{{.*}}anchor{{.*}}([[OUTS]]#1, [[OUTS]]#2)
     pair = Pair(tl.arange(0, 1), tl.arange(0, 2))
     a, b = ttgl.warp_specialize((pair, tl.arange(0, 4)), warp_specialize_default,
                                 [warp_specialize_worker0, warp_specialize_worker1], [4, 4], [24, 48])
@@ -584,10 +584,10 @@ def kernel():
 module {
   tt.func public @kernel() attributes {noinline = false} {
     %0 = ttg.local_alloc : () -> !ttg.memdesc<32x32xi32, #shared, #smem, mutable>
-    tt.call @"smem_and_layout_user__MDi32S32_32SLSSS_1_1_1_constexpr[1]_constexpr[0]____SSSLAS[32, 32]ASMD__(1,)cconstexpr_SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=(constexpr_1_ ,constexpr_0_), ctas_per_cga=None, cta_split_num=None, cta_order=None)_"(%0) : (!ttg.memdesc<32x32xi32, #shared, #smem, mutable>) -> ()
+    tt.call @"test_frontend.smem_and_layout_user__MDi32S32_32SLSSS_1_1_1_constexpr[1]_constexpr[0]____SSSLAS[32, 32]ASMD__(1,)cconstexpr_SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=(constexpr_1_ ,constexpr_0_), ctas_per_cga=None, cta_split_num=None, cta_order=None)_"(%0) : (!ttg.memdesc<32x32xi32, #shared, #smem, mutable>) -> ()
     tt.return
   }
-  tt.func private @"smem_and_layout_user__MDi32S32_32SLSSS_1_1_1_constexpr[1]_constexpr[0]____SSSLAS[32, 32]ASMD__(1,)cconstexpr_SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=(constexpr_1_ ,constexpr_0_), ctas_per_cga=None, cta_split_num=None, cta_order=None)_"(%arg0: !ttg.memdesc<32x32xi32, #shared, #smem, mutable>) attributes {noinline = false} {
+  tt.func private @"test_frontend.smem_and_layout_user__MDi32S32_32SLSSS_1_1_1_constexpr[1]_constexpr[0]____SSSLAS[32, 32]ASMD__(1,)cconstexpr_SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=(constexpr_1_ ,constexpr_0_), ctas_per_cga=None, cta_split_num=None, cta_order=None)_"(%arg0: !ttg.memdesc<32x32xi32, #shared, #smem, mutable>) attributes {noinline = false} {
     tt.return
   }
 }

python/test/unit/language/test_frontend.py

@@ -42,7 +42,7 @@ def test_assign_attribute():
     scalar = 11
     pair = Pair(tl.arange(0, 4), scalar)
     # CHECK: %c42_i32 = arith.constant 42 : i32
-    # CHECK-NEXT: call @"anchor{{.*}}"([[RANGE]], %c42_i32)
+    # CHECK-NEXT: call @{{.*}}anchor{{.*}}([[RANGE]], %c42_i32)
     pair.second = 42
     anchor(pair)
 
@@ -58,7 +58,7 @@ def test_augassign_attribute():
     # CHECK: %c42_i32 = arith.constant 42 : i32
     # CHECK: [[VALUE:%.*]] = arith.addi %c11_i32, %c42_i32
     pair.second += 42
-    # CHECK-NEXT: call @"anchor{{.*}}"([[RANGE]], [[VALUE]])
+    # CHECK-NEXT: call @{{.*}}anchor{{.*}}([[RANGE]], [[VALUE]])
     anchor(pair)
 
 
@@ -69,12 +69,12 @@ def test_jit_method():
     # CHECK: %c11_i32 = arith.constant 11 : i32
     # CHECK: [[RANGE:%.*]] = tt.make_range {end = 4 : i32, start = 0 : i32}
     scalar = 11
-    # CHECK: [[V:%.*]]:2 = tt.call @"unpack{{.*}}"([[RANGE]], %c11_i32)
+    # CHECK: [[V:%.*]]:2 = tt.call @{{.*}}unpack{{.*}}([[RANGE]], %c11_i32)
     pair = Pair(tl.arange(0, 4), scalar)
     a, b = pair.unpack()
-    # CHECK: call @anchor{{.*}}([[V]]#0)
+    # CHECK: call @{{.*}}anchor{{.*}}([[V]]#0)
     anchor(a)
-    # CHECK: call @anchor{{.*}}([[V]]#1)
+    # CHECK: call @{{.*}}anchor{{.*}}([[V]]#1)
     anchor(b)
 
 
@@ -95,10 +95,10 @@ def test_aggregate_initializers():
     # CHECK-LABEL: test_aggregate_initializers
     value = TypeWithBuiltinInitializer()
     # CHECK: [[RANGE:%.*]] = tt.make_range {end = 4 : i32, start = 0 : i32}
-    # CHECK: call @"anchor{{.*}}"([[RANGE]])
+    # CHECK: call @{{.*}}anchor{{.*}}([[RANGE]])
     anchor(value)
     # CHECK: [[RANGE:%.*]] = tt.make_range {end = 8 : i32, start = 4 : i32}
-    # CHECK: call @"anchor{{.*}}"([[RANGE]])
+    # CHECK: call @{{.*}}anchor{{.*}}([[RANGE]])
     value.modify(tl.arange(4, 8))
     anchor(value)
 
@@ -118,11 +118,11 @@ def list_of_functions_constexpr(arg, fns: tl.constexpr):
 @triton.jit
 def test_list_of_functions():
     # CHECK-LABEL: test_list_of_functions
-    # CHECK: call @"list_of_functions_constexpr{{.*}}cJITFunction(test_frontend:anchor){{.*}}cJITFunction(test_frontend:forward)"
+    # CHECK: call @{{.*}}list_of_functions_constexpr{{.*}}cJITFunction(test_frontend:anchor){{.*}}cJITFunction(test_frontend:forward)
 
-    # CHECK-LABEL: tt.func private @"list_of_functions_constexpr
-    # CHECK-NEXT: call @anchor
-    # CHECK-NEXT: call @forward
+    # CHECK: tt.func private @{{.*}}list_of_functions_constexpr
+    # CHECK-NEXT: call @{{.*}}anchor
+    # CHECK-NEXT: call @{{.*}}forward
     list_of_functions_constexpr(tl.arange(0, 4), [anchor, forward])
 
 
@@ -138,11 +138,34 @@ def test_call_in_loop():
     # CHECK-LABEL: test_call_in_loop
     acc = 0
     # CHECK: scf.for
-    # CHECK:   call @accumulate
+    # CHECK:   call @{{.*}}accumulate
     for i in range(10):
         acc = accumulate(acc, i)
 
 
+@tl.core._aggregate
+class FunctionParent:
+
+    @triton.jit
+    def function_with_name():
+        pass
+
+
+@triton.jit
+def function_with_name():
+    pass
+
+
+@filecheck_test
+@triton.jit
+def test_function_name_mangling():
+    # CHECK-LABEL: test_function_name_mangling
+    # CHECK: call @test_frontend.function_with_name
+    # CHECK: call @test_frontend.FunctionParent.function_with_name
+    function_with_name()
+    FunctionParent.function_with_name()
+
+
 @tl.core._aggregate
 class AggregateWithConstexpr:
     a: tl.tensor
@@ -166,10 +189,10 @@ def add_rhs_constexpr(agg):
 @triton.jit
 def test_aggregate_with_constexpr():
     # CHECK-LABEL: test_aggregate_with_constexpr
-    # CHECK: tt.call @"add_rhs_constexpr__test_frontend.AggregateWithConstexpr<i32S4S, constexpr[42]>
+    # CHECK: tt.call @"test_frontend.add_rhs_constexpr__test_frontend.AggregateWithConstexpr<i32S4S, constexpr[42]>
     agg = AggregateWithConstexpr.create(tl.arange(0, 4))
     add_rhs_constexpr(agg)
 
-    # CHECK: tt.func private @"add_rhs_constexpr__test_frontend.AggregateWithConstexpr<i32S4S, constexpr[42]>
+    # CHECK: tt.func private @"test_frontend.add_rhs_constexpr__test_frontend.AggregateWithConstexpr<i32S4S, constexpr[42]>
     # CHECK: %cst = arith.constant dense<42> : tensor<4xi32>
     # CHECK: arith.addi %arg0, %cst : tensor<4xi32>

python/test/unit/language/test_module.py

@@ -0,0 +1,6 @@
+import triton
+
+
+@triton.jit
+def function_with_name():
+    pass

python/test/unit/test_filecheck.py

@@ -17,7 +17,7 @@ def test_kernel():
         # CHECK-LABEL: test_kernel
         scalar = 42
         # CHECK: %c42_i32 = arith.constant 42 : i32
-        # CHECK-NEXT: call @anchor{{.*}}(%c42_i32) : (i32) -> ()
+        # CHECK-NEXT: call @{{.*}}anchor{{.*}}(%c42_i32) : (i32) -> ()
         anchor(scalar)
 
     run_filecheck_test(test_kernel)

python/triton/compiler/code_generator.py

@@ -13,7 +13,7 @@
 from .._C.libtriton import ir, gluon_ir
 from ..language import constexpr, semantic, str_to_ty, tensor
 from ..language.core import _unwrap_if_constexpr, base_value, base_type
-from ..runtime.jit import get_jit_fn_file_line
+from ..runtime.jit import get_jit_fn_file_line, get_full_name
 # ideally we wouldn't need any runtime component
 from ..runtime import JITFunction
 from .._utils import find_paths_if, get_iterable_path, set_iterable_path
@@ -315,6 +315,7 @@ def __init__(self, context, prototype, gscope, function_name, jit_fn: JITFunctio
         self.jit_fn = jit_fn
         # TODO: we currently generate illegal names for non-kernel functions involving constexprs!
         if is_kernel:
+            function_name = function_name[function_name.rfind('.') + 1:]
             function_name = check_identifier_legality(function_name, "function")
         self.function_name = function_name
         self.is_kernel = is_kernel
@@ -1200,7 +1201,7 @@ def call_JitFunction(self, fn: JITFunction, args, kwargs):
         args_path = find_paths_if(args, lambda _, x: not _is_constexpr(x))
         args_val = [get_iterable_path(args, path) for path in args_path]
         # mangle
-        fn_name = mangle_fn(fn.__name__, [arg.type for arg in args_val], args_cst)
+        fn_name = mangle_fn(get_full_name(fn), [arg.type for arg in args_val], args_cst)
         # generate function def if necessary
         if not self.module.has_function(fn_name):
             gscope = fn.__globals__

python/triton/runtime/jit.py

@@ -485,6 +485,10 @@ def dynamic_func({", ".join(list(map(arg, sig.parameters.items())) + ["**options
     type_canonicalisation_dict[v] = v
 
 
+def get_full_name(fn):
+    return f"{fn.__module__}.{fn.__qualname__}"
+
+
 @dataclass
 class JitFunctionInfo:
     module: ModuleType
@@ -511,7 +515,7 @@ def _call_hook(
         if not hook:
             return None
 
-        name = self.fn.__name__
+        name = get_full_name(self.fn)
         module = self.fn.__module__
         arg_reprs = ", ".join([f"{param.name}: {ty}" for param, ty in zip(self.params, key[1])])
         repr = f"{name}[num_warps={options.num_warps}, num_ctas={options.num_ctas}, num_stages={options.num_stages}, enable_fp_fusion={options.enable_fp_fusion}, launch_cooperative_grid={options.launch_cooperative_grid}]({arg_reprs})"
@@ -653,7 +657,7 @@ def __init__(self, fn, version=None, do_not_specialize=None, do_not_specialize_o
         self.do_not_specialize_on_alignment = do_not_specialize_on_alignment
         self.starting_line_number = inspect.getsourcelines(fn)[1]
         self._repr = repr
-        self._fn_name = fn.__name__
+        self._fn_name = get_full_name(fn)
         self.launch_metadata = launch_metadata
 
         self.params = []
@@ -698,14 +702,15 @@ def __init__(self, fn, version=None, do_not_specialize=None, do_not_specialize_o
         # reuse docs of wrapped function
         self.__doc__ = fn.__doc__
         self.__name__ = fn.__name__
+        self.__qualname__ = fn.__qualname__
         self.__globals__ = fn.__globals__
         self.__module__ = fn.__module__
 
     @property
     def cache_key(self):
         # TODO : hash should be attribute of `self`
         if self.hash is None:
-            dependencies_finder = DependenciesFinder(name=self.__name__, globals=self.__globals__, src=self.src)
+            dependencies_finder = DependenciesFinder(name=self._fn_name, globals=self.__globals__, src=self.src)
             dependencies_finder.visit(self.parse())
             self.hash = dependencies_finder.ret + str(self.starting_line_number)
             self.used_global_vals = dict(sorted(dependencies_finder.used_global_vals.items()))
@@ -725,9 +730,9 @@ def preload(self, specialization_data):
         import triton.language as tl
         device = driver.active.get_current_device()
         deserialized_obj = json.loads(specialization_data)
-        if deserialized_obj['name'] != self.fn.__name__:
+        if deserialized_obj['name'] != self._fn_name:
             raise RuntimeError(
-                f"Specialization data is for {deserialized_obj['name']} but trying to preload for {self.fn.__name__}")
+                f"Specialization data is for {deserialized_obj['name']} but trying to preload for {self._fn_name}")
         constant_keys = map(tuple, deserialized_obj['constant_keys'])
         constant_vals = deserialized_obj['constant_vals']
         constants = {
@@ -778,7 +783,7 @@ def _unsafe_update_src(self, new_src):
         super().__setattr__('src', new_src)
 
     def __repr__(self):
-        return f"JITFunction({self.module}:{self.fn.__name__})"
+        return f"JITFunction({self.module}:{self.fn.__qualname__})"
 
 
 # -----------------------------------------------------------------------------