runtime JIT execution of the same kernel incurs high overhead(binder function)

[Learnmore666]

Describe the issue

as mentioned :#3503
It is very expensive to call binder() and pack call args.
bound_args, specialization, options = binder(*args, **kwargs)
When repeatedly calling a compiled kernel, JIT will still repeatedly call the function to generate key for comparison, which brings certain overhead. Can we simplify this？

triton/python/triton/runtime/jit.py

Line 282 in 64a07f8

def specialize_impl(arg, specialize_extra, is_const=False, specialize_value=True, align=True):

Environment details

Hardware independent
The latest version of Triton

[Learnmore666]

@ptillet Hi, I have read the relevant PR you submitted to JIT. Do you have any knowledge about this？

[Learnmore666]

@apgoucher Hi，will there be further improvements at present

[Jokeren]

When repeatedly calling a compiled kernel, JIT will still repeatedly call the function to generate key for comparison, which brings certain overhead. Can we simplify this？

It's not as expensive as you thought now. But if you're sure the kernel won't change under your input, you can just cache the compiled kernel and call it directly.

[apgoucher]

I've ran my profiling script and it does report a performance regression from last time (from 12us to 25us, but still much better than the 108us that we were at before my original PR). It's not easily fixable, though -- Triton's frontend expressive power has increased, as we now support tuples and functions as arguments.

You're correct that specialize_impl is the main source of the slowness. I'd probably move the import statement outside of that function and rearrange the if-else control flow so that more common types (such as tensors and integers) get matched with higher priority than more esoteric types (such as functions, tuples, and TMA descriptors); this will reduce the average runtime spent in this function.

[Learnmore666]

When repeatedly calling a compiled kernel, JIT will still repeatedly call the function to generate key for comparison, which brings certain overhead. Can we simplify this？

It's not as expensive as you thought now. But if you're sure the kernel won't change under your input, you can just cache the compiled kernel and call it directly.

There are many situations where the same kernel is repeatedly called, such as in training scenarios, dynamic shapes, or autotuning, which can help improve performance in real-world scenarios

I've ran my profiling script and it does report a performance regression from last time (from 12us to 25us, but still much better than the 108us that we were at before my original PR). It's not easily fixable, though -- Triton's frontend expressive power has increased, as we now support tuples and functions as arguments.

You're correct that specialize_impl is the main source of the slowness. I'd probably move the import statement outside of that function and rearrange the if-else control flow so that more common types (such as tensors and integers) get matched with higher priority than more esoteric types (such as functions, tuples, and TMA descriptors); this will reduce the average runtime spent in this function.

Yes,this is an heuristic principle to reduce control flow comparison and minimize overhead.

1. There are annotation_type, do_not_Specification, and do_not_Specifie_on-alignment functions here that are easy to pass explicitly to reduce the need to call specialize_impl. Is there a demo for reference here? It seems to increase the complexity of kernel programming.
2. Why not use kernel name+data type+additional keyword parameters for the unique identification of key sig? Can compute_spec_key be removed?
3. Can a python dict be used to map from Kernel (*args, **kwargs) to key(cached kernel)?
   or others? Torch inducotr runtime doesn't seem to need this key cached strategy，will it be better？

[apgoucher]

There are many situations where the same kernel is repeatedly called, such as in training scenarios

In those cases if you're bottlenecked by CPU overhead then you'd want to record a CUDA graph of the computation and then replay it; PyTorch provides functionality to make this easy for you. Then you completely bypass the CPU overhead of Python, Triton's JIT mechanism, and libcuda launch overhead -- and this no longer becomes an issue at all.

Why not use kernel name+data type+additional keyword parameters for the unique identification of key sig? Can compute_spec_key be removed?

We don't need the Triton kernel name at all, because the dict is local to that Triton kernel. But we do need to incorporate data type, alignment information, and any constexpr/specialised values when looking up the compiled specialised kernel: if we didn't do that, then you'd get a CUDA error if (for example) you first called the function on a float array starting at a multiple-of-16-byte-aligned address (where it would use LDG.128 instructions) and then tried to call the same specialisation on a float array starting at an unaligned address. So no, compute_spec_key cannot be removed.

[saagarjha]

We have the same issues and found that skipping specialize_impl was a significant win when we could (for example, we have many parameters that don't improve performance much by being specialized, like tensors of only a handful of elements). Unfortunately, annotations are currently broken, so we had to patch that up to make it work (see for example #4690). Here is the patch we are using:

diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py
index 749f4870d..2d86760dc 100644
--- a/python/triton/runtime/jit.py
+++ b/python/triton/runtime/jit.py
@@ -221,10 +221,23 @@ class DependenciesFinder(ast.NodeVisitor):
 
 
 def _normalize_ty(ty) -> str:
+    import triton.language.core as core
+
     if isinstance(ty, type):
         return ty.__name__
+    elif isinstance(ty, core.dtype):
+        return type_canonicalisation_dict[ty.name]
+    elif isinstance(ty, core.pointer_type):
+        return f"*{_normalize_ty(ty.element_ty)}"
     elif isinstance(ty, str):
+        # Support legacy users who annotate their parameters with "tl.*" strings
+        if ty.startswith("tl."):
+            return _normalize_ty(getattr(core, ty.removeprefix("tl.")))
         return ty
+    elif isinstance(ty, tuple):
+        return type(ty)(*map(_normalize_ty, ty)) if hasattr(ty, "_fields") else tuple(map(_normalize_ty, ty))
+    elif type is None:
+        return "constexpr"
     return repr(ty)
 
 
@@ -245,27 +258,16 @@ class KernelParam:
     @cached_property
     def annotation(self):
         if not self._param.annotation or self._param.annotation == inspect.Parameter.empty:
-            return ""
+            return None
         return _normalize_ty(self._param.annotation)
 
-    @cached_property
-    def annotation_type(self):
-        annotation = self.annotation
-        for ty1, ty2 in [("uint", 'u'), ("int", 'i')]:
-            width = annotation[annotation.find(ty1) + len(ty1):]
-            if width and ty1 in annotation:
-                return f"{ty2}{width}"
-        if annotation == "bool":
-            return "u1"
-        return ""
-
     @cached_property
     def is_constexpr(self):
-        return "constexpr" in self.annotation
+        return self.annotation == "constexpr"
 
     @cached_property
     def is_const(self):
-        return "const" in self.annotation and not self.is_constexpr
+        return self.annotation == "const"
 
     @property
     def default(self):
@@ -279,15 +281,24 @@ class KernelParam:
 dtype2str = {}
 
 
-def specialize_impl(arg, specialize_extra, is_const=False, specialize_value=True, align=True):
-    from ..language import constexpr
+def specialize_impl(arg, annotation, specialize_extra, is_const=False, specialize_value=True, align=True, constexpr_cache=[], types_cache={}):
+    if not constexpr_cache:
+        from ..language import constexpr
+        constexpr_cache.append(constexpr)
+    else:
+        constexpr = constexpr_cache[0]
+
     if arg is None:
         return ("constexpr", None)
     elif isinstance(arg, JITFunction):
         return ("constexpr", arg.cache_key)
     elif isinstance(arg, constexpr):
         return ("constexpr", arg)
-    elif isinstance(arg, bool):
+
+    if not specialize_value and annotation:
+        return (annotation, None)
+
+    if isinstance(arg, bool):
         return ("i1", None)
     elif isinstance(arg, int):
         key = specialize_extra(arg, "int", align=align) if specialize_value else None
@@ -304,7 +315,7 @@ def specialize_impl(arg, specialize_extra, is_const=False, specialize_value=True
     elif hasattr(arg, "tma_desc_cpu_ptr"):
         return ("nvTmaDesc", None)
     elif isinstance(arg, tuple):
-        spec = [specialize_impl(x, specialize_extra) for x in arg]
+        spec = [specialize_impl(x, None, specialize_extra) for x in arg]
         make_tuple = lambda vals: type(arg)(*vals) if hasattr(arg, "_fields") else tuple(vals)
         tys = make_tuple([x[0] for x in spec])
         keys = make_tuple([x[1] for x in spec])
@@ -321,7 +332,7 @@ def specialize_impl(arg, specialize_extra, is_const=False, specialize_value=True
 
 
 def mangle_type(arg, specialize=False):
-    return specialize_impl(arg, lambda _, **kwargs: None, specialize_value=specialize)[0]
+    return specialize_impl(arg, None, lambda _, **kwargs: None, specialize_value=specialize)[0]
 
 
 class KernelInterface(Generic[T]):
@@ -359,6 +370,8 @@ def create_function_from_signature(sig, kparams, backend):
     assert len(sig.parameters) == len(kparams)
     # Create the function argument list and the dict entries for the return statement
     specialization = []
+    # Names of named tuples in annotations so they can be constructed later
+    extra_types = {}
     # signature
     for name, kp in zip(sig.parameters.keys(), kparams):
         if kp.is_constexpr:
@@ -367,11 +380,24 @@ def create_function_from_signature(sig, kparams, backend):
             is_const = 'True' if kp.is_const else 'False'
             specialize = 'False' if kp.do_not_specialize else 'True'
             align = 'False' if kp.do_not_specialize_on_alignment else 'True'
-            ret = f"specialize_impl({name}, specialize_extra, {is_const}, {specialize}, {align})"
-            if kp.annotation_type:
-                specialization.append(f'("{kp.annotation_type}",) + {ret}[1:]')
+            if isinstance(kp.annotation, tuple):
+                def find_named_tuples(annotation):
+                    if isinstance(annotation, tuple):
+                        if hasattr(annotation, "_fields"):
+                            extra_types[type(annotation).__name__] = type(annotation)
+
+                        for element in annotation:
+                            find_named_tuples(element)
+
+                find_named_tuples(kp.annotation)
+                annotation = kp.annotation
+            elif kp.annotation is not None:
+                annotation = f'"{kp.annotation}"'
             else:
-                specialization.append(f"{ret}")
+                annotation = "None"
+
+            ret = f"specialize_impl({name}, {annotation}, specialize_extra, {is_const}, {specialize}, {align})"
+            specialization.append(f"{ret}")
 
     # compute argument string for a given parameter
     arg = lambda x: x[0] if x[1].default is inspect.Parameter.empty else f"{x[0]}=default_{x[0]}"
@@ -393,6 +419,9 @@ def dynamic_func({", ".join(list(map(arg, sig.parameters.items())) + ["**options
     func_namespace["specialize_impl"] = specialize_impl
     func_namespace["specialize_extra"] = backend.get_arg_specialization
 
+    for ty in extra_types:
+        func_namespace[ty] = extra_types[ty]
+
     # Execute the function string in func_namespace to create the function
     exec(func_body, func_namespace)
 

The idea here is that we can short-circuit specialize_impl if the user specifies the type and indicates it does not need value specialization. Most of the changes are there to make it so that annotations can actually capture complex types. While I don't think this is really ready for a PR it might be useful to reference as part of this discussion.

[Learnmore666]

dea here is that we can short-circuit specialize_impl if the user specifies the type and indicates it does not need value specialization. Most of the changes are there to make it so that annotations can actually capture complex types.

great！
I have merged these modifications in the latest Triton version，but there have been some issues here：
example #4690
@triton.jit def unary_pointwise_kernel(X: tl.pointer_type, O: tl.pointer_type, n: int, BLOCK_N: tl.constexpr):
the issue still exists because the pointer_type always enters branch:
isinstance(ty, type): in the _normalize_ty(ty)
"pointer_type" is generated in key
[('pointer_type', None), ('pointer_type', None), ('int', None), ('constexpr', 1024)]
but what we need is their element types for codegen to create pointer_type. Here, " *type " should be generated, but it does not satisfy:
isinstance(ty, core.pointer_type):

What version are you using on your end? Can it pass normally?

[saagarjha]

Hmm, I feel like I checked and it wasn't a type, but it seems like it is. You should be able to fix it by moving it be the first thing checked?

[apgoucher]

I think that the issue is that you're using a bare tl.pointer_type annotation (which is useless for specialisation anyway as it matters what the underlying dtype is) rather than annotating it with an instance of tl.pointer_type (e.g. int32*, float*, bfloat16*, etc.). tl.pointer_type (the type!) is a type, whereas individual instantiations are not.

The solution is to just never use bare tl.pointer_type annotations, and the correct behaviour is to throw an error (or at least a warning). Admittedly KeyError: "ier_type" is probably not the most informative error message; sorry about that!

[apgoucher]

@saagarjha this doesn't need to be on the hot codepath:

    if not constexpr_cache:
        from ..language import constexpr
        constexpr_cache.append(constexpr)
    else:
        constexpr = constexpr_cache[0]

If there's some reason why we can't import constexpr directly inside jit.py, then we should have a higher-order function like so:

specialize_impl_cache = []
def create_specialize_impl():
    if len(specialize_impl_cache) > 0:
        return specialize_impl_cache[-1]
    from .language import constexpr
    def specialize_impl(...):
        ...
    specialize_impl_cache.append(specialize_impl)
    return specialize_impl

which can be executed at runtime whenever we need specialize_impl -- such as when we inject it into func_namespace. Importantly this doesn't at all happen on the hot codepath.

[Learnmore666]

Hmm, I feel like I checked and it wasn't a type, but it seems like it is. You should be able to fix it by moving it be the first thing checked?

Of course, I have tried this approach,thanks

I think that the issue is that you're using a bare tl.pointer_type annotation (which is useless for specialisation anyway as it matters what the underlying dtype is) rather than annotating it with an instance of tl.pointer_type (e.g. int32*, float*, bfloat16*, etc.). tl.pointer_type (the type!) is a type, whereas individual instantiations are not.

The solution is to just never use bare tl.pointer_type annotations, and the correct behaviour is to throw an error (or at least a warning). Admittedly KeyError: "ier_type" is probably not the most informative error message; sorry about that!

I agree with your point of view.
I tried to use a string representation of similar data type abbreviations in specialize_impl, such as "* fp32" for annotation, and it was approved.

[apgoucher]

However, I still believe that this is not an elegant way for front-end interfaces because it requires a simplified unconventional writing style. We need a more standardized form.

To clarify, the thing that you should write is tl.pointer_type(tl.int32) (for example); I was just writing int32* in my comment for brevity 😄

[apgoucher]

then we should have a higher-order function like so:

I've just tried my idea locally (on top of Triton main) and my benchmark runtime already improves from 25.7us to 18.4us (without any need for kernel annotations or suchlike). Imports on the hot codepath are expensive! PR incoming soon...