Annotations for elementary dtypes and pointers (#6152)

This introduces thorough support for annotations of elementary dtypes
and pointers. Annotated kernel arguments can under certain circumstances
skip `specialize_impl` leading to faster launch times. This is inspired
by the patch provided by @saagarjha in the comments in
#6064 but we go further by
doing the logic statically -- when constructing `dynamic_func` -- rather
than at runtime.

We also take this opportunity to do some general code simplification and
cleanup, including consistently using "u1" rather than "i1" to denote
booleans in the specialisation key.

Author: apgoucher

python/triton/runtime/jit.py

@@ -221,11 +221,29 @@ def visit_For(self, node):
 
 
 def _normalize_ty(ty) -> str:
-    if isinstance(ty, type):
-        return ty.__name__
-    elif isinstance(ty, str):
-        return ty
-    return repr(ty)
+    import triton.language.core as core
+    if isinstance(ty, str):
+        ty = ty.strip()
+        if ty.startswith("const "):
+            ty = ty.removeprefix("const")
+            ty = _normalize_ty(ty)
+            assert ty.startswith("*")
+            return "*k" + ty[1:]
+        if ty.endswith("*"):
+            return "*" + _normalize_ty(ty[:-1])
+        if ty.startswith("*"):
+            return "*" + _normalize_ty(ty[1:])
+        if ty.startswith("tl."):
+            return _normalize_ty(ty.removeprefix("tl."))
+    elif isinstance(ty, core.pointer_type):
+        return f"*{_normalize_ty(ty.element_ty)}"
+    elif isinstance(ty, core.dtype):
+        ty = ty.name
+    elif isinstance(ty, type):
+        ty = ty.__name__
+    else:
+        ty = str(ty)
+    return type_canonicalisation_dict.get(ty.replace("_t", ""), ty)
 
 
 class KernelParam:
@@ -250,13 +268,13 @@ def annotation(self):
 
     @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"
+        a = self.annotation
+        if a.startswith("*k"):
+            a = a[2:]
+        elif a.startswith("*"):
+            a = a[1:]
+        if a in set(type_canonicalisation_dict.values()):
+            return self.annotation
         return ""
 
     @cached_property
@@ -265,7 +283,9 @@ def is_constexpr(self):
 
     @cached_property
     def is_const(self):
-        return "const" in self.annotation and not self.is_constexpr
+        if self.is_constexpr:
+            return False
+        return "const" in self.annotation or self.annotation.startswith("*k")
 
     @property
     def default(self):
@@ -289,7 +309,7 @@ def specialize_impl(arg, is_const=False, specialize_value=True, align=True):
         if arg is None:
             return ("constexpr", None)
         elif isinstance(arg, bool):
-            return ("i1", None)
+            return ("u1", None)
         elif isinstance(arg, int):
             key = specialize_extra(arg, "int", align=align) if specialize_value else None
             if arg == 1 and specialize_value:
@@ -381,7 +401,15 @@ def create_function_from_signature(sig, kparams, backend):
             align = 'False' if kp.do_not_specialize_on_alignment else 'True'
             ret = f"specialize_impl({name}, {is_const}, {specialize}, {align})"
             if kp.annotation_type:
-                specialization.append(f'("{kp.annotation_type}",) + {ret}[1:]')
+                if isinstance(kp.annotation_type, str):
+                    if kp.annotation_type == "u1" or kp.annotation_type[:2] in ["fp", "bf"]:
+                        # we do not specialize non-constexpr floats and bools:
+                        specialize = False
+                if specialize:
+                    specialization.append(f'("{kp.annotation_type}",) + {ret}[1:]')
+                else:
+                    # skip runtime specialization:
+                    specialization.append(f'("{kp.annotation_type}", None)')
             else:
                 specialization.append(f"{ret}")
 
@@ -412,7 +440,12 @@ def dynamic_func({", ".join(list(map(arg, sig.parameters.items())) + ["**options
 
 
 type_canonicalisation_dict = {
-    "bool": "i1",
+    # we canonicalise all bools to be unsigned:
+    "bool": "u1",
+    "int1": "u1",
+    "uint1": "u1",
+    "i1": "u1",
+    # floating-point dtypes:
     "float8e4nv": "fp8e4nv",
     "float8e5": "fp8e5",
     "float8e4b15": "fp8e4b15",
@@ -422,14 +455,20 @@ def dynamic_func({", ".join(list(map(arg, sig.parameters.items())) + ["**options
     "float8_e5m2": "fp8e5",
     "float8e5b16": "fp8e5b16",
     "float8_e5m2fnuz": "fp8e5b16",
+    "half": "fp16",
     "float16": "fp16",
     "bfloat16": "bf16",
+    "float": "fp32",
     "float32": "fp32",
+    "double": "fp64",
     "float64": "fp64",
+    # signed integers:
     "int8": "i8",
     "int16": "i16",
+    "int": "i32",
     "int32": "i32",
     "int64": "i64",
+    # unsigned integers:
     "uint8": "u8",
     "uint16": "u16",
     "uint32": "u32",