[FRONTEND] clean-up backend/jit interface (#5512)

This is a fairly big refactor of the JIT that prepares the field for
supporting named tuples (to come in a subsequent PR):
- fixes bug in tuple specialization hashes
- simplifies launcher by always taking all function arguments (even
unused constexprs)
- enables un-annotated constant expressions (necessary for variadic
tuples containing constexprs)
- replaces attrsdescriptor by much more minimal backend interfaces
- unifies constexprs provided by users (e.g., block sizes, none,
function arguments...) and specialization (e.g., 1)

Author: ptillet

python/test/unit/language/test_core.py

@@ -4435,7 +4435,7 @@ def test_value_specialization(value: int, value_type: str, device) -> None:
 
     def repr(specialization):
         ty = specialization.signature["value1"]
-        cst = '_'.join([k for k, v in specialization.constants.items() if v == 1])
+        cst = '_'.join([k for k, v in specialization.constants.items() if isinstance(k, str) and v == 1])
         return f"kernel_{ty}_{cst}"
 
     @triton.jit(repr=repr)

python/test/unit/language/test_tuple.py

@@ -81,6 +81,7 @@ def _tuple_fn0(Ptr, cst2: tl.constexpr, tuple1):
 def _tuple_serialize(Ptr, N1, tuple1, cst1: tl.constexpr, val1, tuple2):
     tl.static_assert(N1 is None)
     tl.static_assert(tuple1[1][1] is None)
+    tl.static_assert(tuple1[1][3] == 4)
     tl.store(Ptr + 0, tl.load(tuple1[0]))
     tl.store(Ptr + 1, tuple1[1][0])
     tl.store(Ptr + 2, tl.load(tuple1[1][2]))
@@ -95,6 +96,6 @@ def test_serialize(device="cuda"):
     y0 = torch.tensor([10], dtype=torch.int32, device=device)
     z = torch.empty((10, ), dtype=torch.int32, device=device)
     # we want to check that JIT specialization propagates to tuples:
-    _tuple_serialize[(1, )](z, None, (x0, (1, None, x1)), 20, 1, (y0, ))
+    _tuple_serialize[(1, )](z, None, (x0, (1, None, x1, tl.constexpr(4))), 20, 1, (y0, ))
     ref = torch.tensor([8, 1, 12, 21, 10, 15, -1, 8, 1, 12], device=device)
     assert torch.equal(z, ref)

python/test/unit/runtime/test_bindings.py

@@ -68,7 +68,6 @@ def walk_fn(op):
         constexprs={kernel.arg_names[i]: arg
                     for i, arg in enumerate(args)
                     if not isinstance(arg, torch.Tensor)},
-        attrs=backend.get_attrs_descriptor(kernel.params, args),
     )
 
     context = triton._C.libtriton.ir.context()

python/test/unit/runtime/test_cache.py

@@ -563,7 +563,7 @@ def kernel_add(a):
 
     def cache_hook(*args, **kwargs):
         nonlocal pointer_range_32
-        pointer_range_32 = kwargs["compile"]["configs"][0].pointer_range_32
+        pointer_range_32 = [k for k, v in kwargs["compile"]["configs"][0].items() if ['tt.pointer_range', 32] in v]
 
     JITFunction.cache_hook = cache_hook
     # In warmup we assume that the pointer range is 32 bits

python/test/unit/runtime/test_subproc.py

@@ -3,14 +3,13 @@
 
 import triton
 import triton.language as tl
-from triton.backends.compiler import AttrsDescriptor
 from triton.compiler import ASTSource
 
 target = triton.runtime.driver.active.get_current_target()
 start_method = 'fork' if 'fork' in multiprocessing.get_all_start_methods() else 'spawn'
 
 
-def compile_fn(attrs):
+def compile_fn():
 
     @triton.jit
     def kernel_sub(a, b, o, N: tl.constexpr):
@@ -21,21 +20,19 @@ def kernel_sub(a, b, o, N: tl.constexpr):
         fn=kernel_sub,
         constexprs={'N': 32},
         signature={'a': "*fp32", 'b': "*fp32", 'o': "*fp32", 'N': 'constexpr'},
-        attrs=attrs,
     )
     triton.compile(src=src, target=target)
 
 
 def test_compile_in_subproc() -> None:
-    config = AttrsDescriptor.from_hints({i: 16 for i in range(4)})
     mp_ctx = multiprocessing.get_context(start_method)
-    proc = mp_ctx.Process(target=compile_fn, args=(config, ))
+    proc = mp_ctx.Process(target=compile_fn)
     proc.start()
     proc.join()
     assert proc.exitcode == 0
 
 
-def compile_fn_dot(attrs):
+def compile_fn_dot():
 
     @triton.jit
     def kernel_dot(Z):
@@ -44,28 +41,27 @@ def kernel_dot(Z):
         z = tl.dot(z, z)
         tl.store(Z + offs, z)
 
-    src = ASTSource(fn=kernel_dot, signature={'Z': "*fp32"}, attrs=attrs, constexprs={})
+    src = ASTSource(fn=kernel_dot, signature={'Z': "*fp32"})
     triton.compile(src=src, target=target)
 
 
 def test_compile_in_forked_subproc(fresh_triton_cache) -> None:
-    config = AttrsDescriptor.from_hints({0: 16})
     mp_ctx = multiprocessing.get_context(start_method)
-    proc = mp_ctx.Process(target=compile_fn_dot, args=(config, ))
+    proc = mp_ctx.Process(target=compile_fn_dot)
     proc.start()
     proc.join()
     assert proc.exitcode == 0
 
 
-def compile_empty_kernel_with_gc(attrs):
+def compile_empty_kernel_with_gc():
 
     @triton.jit
     def empty_kernel():
         pass
 
     import gc
     gc.collect()
-    src = ASTSource(fn=empty_kernel, signature={}, attrs=attrs, constexprs={})
+    src = ASTSource(fn=empty_kernel, signature={})
     triton.compile(src=src, target=target)
 
 
@@ -88,13 +84,12 @@ def test_compile_in_forked_subproc_with_forced_gc(fresh_triton_cache) -> None:
     gc.disable()
 
     # stage 1.p
-    config = AttrsDescriptor.from_hints({0: 16})
-    compile_empty_kernel_with_gc(config)
+    compile_empty_kernel_with_gc()
 
     # stage 2.p
     shutil.rmtree(fresh_triton_cache)
     mp_ctx = multiprocessing.get_context(start_method)
-    proc = mp_ctx.Process(target=compile_empty_kernel_with_gc, args=(config, ))
+    proc = mp_ctx.Process(target=compile_empty_kernel_with_gc)
 
     # stage 3.c
     proc.start()

python/triton/backends/compiler.py

@@ -1,235 +1,11 @@
 import os
 import re
-import hashlib
 import subprocess
 import sysconfig
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
-from typing import Dict, List, Tuple, Union
+from typing import Dict, Union
 from types import ModuleType
-from .._utils import find_paths_if
-
-# Table that associates strings to AttrsDescriptor (sub)classes.
-# In this way we can dynamically select the correct class
-# constructor
-_descriptor_table = {}
-
-
-def register_descriptor(cls):
-    """
-    Register a descriptor into the descriptor table
-    """
-    _descriptor_table[cls.__name__] = cls
-    return cls
-
-
-@register_descriptor
-class AttrsDescriptor:
-    """
-    This class handles compile-time properties for specific function parameters.
-
-    Different backends can add more properties to the common ones. The class
-    contains two fields:
-
-    `arg_properties`: a dictionary containing the different compile-time properties for different
-        parameters. I.e., the dictionary is a map from property names to parameter indices
-        {
-        "prop0": (0, 2, 3)
-        "prop1": (0, 4, 5)
-        }
-        Different backends might need different properties on those paraemters to enable
-        specific optimizations. The common compile time properties contained in this class
-        are :
-        - "tt.divisibility", i.e., is the given parameter divisible by 16
-        - "tt.equal_to_1", i.e., is the given parameter an integer constant 1
-
-    `property_values`: a dictionary containing the value of the different compile-time properties, like:
-        {
-            "prop0": val0
-            "prop1": val1
-        }
-
-    `constant_properties`: a set containing the properties that can be used to determine if a parameter is constant
-
-    """
-    __slots__ = ('divisibility_16', 'equal_to_1', 'equal_to_none', 'arg_properties', 'property_values',
-                 'constant_properties')
-
-    def __init__(self, params=None, values=None):
-        """
-        Initialize the compile-time properties
-
-        We can initialize the AttrsDescriptor class by passing the list of params
-        of the function and their `values`. The function will try to apply the properties
-        to the values and save the parameters in the `arg_properties` list. If we don't pass
-        either the `params` or the `values` we should initialize the class via an alternative method
-        (see `from_dict` or `from_hints`)
-        """
-        # Default initialization
-        self.arg_properties = {}
-        self.property_values = {}
-        self.equal_to_none = {}
-        self.constant_properties = set()
-
-        self._add_common_properties(params, values)
-        self._add_backend_properties(params, values)
-        self._init_slots()
-
-    def _add_common_properties(self, params, values):
-        """ Add common compile-time properties """
-        self.property_values["tt.divisibility"] = 16
-        self.property_values["tt.equal_to"] = 1
-        self.constant_properties.add("tt.equal_to")
-
-        if (params is None) or (values is None):
-            return
-
-        # Compile properties deduction
-        assert (len(params) == len(values))
-
-        # Divisibility property
-        divisibility_16 = []
-        for param, arg in zip(params, values):
-            if param.do_not_specialize or \
-               param.do_not_specialize_on_alignment:
-                continue
-            paths = find_paths_if(arg, lambda path, val: AttrsDescriptor.is_divisible_by_16(val))
-            divisibility_16 += [(param.num, ) + x for x in paths]
-        self.arg_properties["tt.divisibility"] = divisibility_16
-
-        # Equal to 1 property
-        equal_to_1 = []
-        for param, arg in zip(params, values):
-            if param.do_not_specialize:
-                continue
-            paths = find_paths_if(arg, lambda path, val: AttrsDescriptor.is_equal_to_1(val))
-            equal_to_1 += [(param.num, ) + x for x in paths]
-        self.arg_properties["tt.equal_to"] = equal_to_1
-
-        # Equal to None property
-        equal_to_none = []
-        for param, arg in zip(params, values):
-            paths = find_paths_if(arg, lambda path, val: val is None)
-            equal_to_none += [(param.num, ) + x for x in paths]
-        self.equal_to_none = equal_to_none
-
-    def _add_backend_properties(self, params=None, values=None):
-        """ This method is for different subclasses to implement their own compile-time properties """
-        pass
-
-    def _init_slots(self):
-        """ Initialize the slots of this class """
-        for name, val in self.arg_properties.items():
-            setattr(self, name.removeprefix('tt.') + '_' + str(self.property_values[name]), val)
-
-    def get_fn_attrs(self) -> Dict:
-        """
-        Get the function attributes as a dictionary.
-
-        The returned dictionary will look like :
-            {
-            "arg0" : [(prop_name00, val00), (prop_name01, val01), ...)]}
-            "arg1" : [(prop_name10, val10), (prop_name11, val11), ...)]}
-            }
-        """
-        attrs = {}
-        for prop_name, arg_set in self.arg_properties.items():
-            prop_val = self.property_values[prop_name]
-            for arg in arg_set:
-                attrs[arg] = attrs.get(arg, []) + [(prop_name, prop_val)]
-        return attrs
-
-    def get_constants(self) -> Dict:
-        """ Return a mapping of constant parameters to their values """
-        constants = {}
-        for prop_name in self.constant_properties:
-            for p in self.arg_properties.get(prop_name, []):
-                constants[p] = self.property_values[prop_name]
-        for v in self.equal_to_none:
-            constants[v] = None
-        return constants
-
-    def filter_out_constants(self):
-        """ Return the same object, without properties marked as constants"""
-        import copy
-        c = copy.deepcopy(self)
-        for prop_name in c.constant_properties:
-            c.arg_properties.pop(prop_name, None)
-            c.property_values.pop(prop_name, None)
-        c.constant_properties = {}
-        return c
-
-    def hash(self):
-        values = [sorted(self.arg_properties.values())]
-        values += [sorted(self.property_values.values())]
-        values += [sorted(self.constant_properties)]
-        key = str(values)
-        return hashlib.sha256(key.encode("utf-8")).hexdigest()
-
-    def to_dict(self):
-        """
-        Store the fields of this class in a serializable dictionary
-        """
-        # We need to only store the `arg_properties` field. To initialize the
-        # other fields we relay on the class type. We store it as a string in
-        # the dictionary so that we can use it to invoke the appropriate
-        # (sub)class constructor in the `from_dict` method.
-        return {"arg_properties": self.arg_properties, "cls": type(self).__name__}
-
-    @staticmethod
-    def from_dict(data):
-        """
-        Create the object from a serializable dictionary
-        """
-        attrs_descriptor = _descriptor_table[data["cls"]]()
-        for prop_name, param_ids in data["arg_properties"].items():
-            attrs_descriptor.arg_properties[prop_name] = list(map(tuple, param_ids))
-        attrs_descriptor._init_slots()
-        return attrs_descriptor
-
-    @classmethod
-    def from_hints(cls, hints: List[Tuple[int, int]]):
-        """
-        Create the class from a set of hints that are passed in.
-
-        Instead of deducing the properties from a list of paramaters and values,
-        the user can pass in a list of `hints=[(param_index, val)]` and if `val`
-        matches one of the values of the properties (e.g., `prop_val[prop0]`),
-        then we insert `param_index` into the correct list (e.g., in
-        `arg_properties[prop0]`)
-        """
-        attrs_descriptor = cls()
-        for prop_name, prop_val in attrs_descriptor.property_values.items():
-            attrs_descriptor.arg_properties[prop_name] = [i for i, h in hints.items() if h == prop_val]
-        attrs_descriptor._init_slots()
-        return attrs_descriptor
-
-    @staticmethod
-    def is_divisible_by_16(x):
-        """ Return if the argument is a multiple of 16"""
-        if hasattr(x, "data_ptr"):
-            return x.data_ptr() % 16 == 0
-        elif isinstance(x, int):
-            return x % 16 == 0
-        if x is None:
-            return True
-        return False
-
-    @staticmethod
-    def is_equal_to_1(x):
-        """ Return if the argument is a constant 1"""
-        return True if isinstance(x, int) and not isinstance(x, bool) and x == 1 else False
-
-    @staticmethod
-    def get_property_key(val, align):
-        if align and AttrsDescriptor.is_divisible_by_16(val):
-            return "D"
-        if AttrsDescriptor.is_equal_to_1(val):
-            return "1"
-        return "N"
-
-    def __repr__(self):
-        return f"AttrsDescriptor.from_dict({self.to_dict()!r})"
 
 
 @dataclass(frozen=True)
@@ -308,16 +84,21 @@ def get_module_map(self) -> Dict[str, ModuleType]:
         """
         raise NotImplementedError
 
-    def get_attrs_descriptor(self, params, args):
-        """
-        Return an attribute descriptor: given a set of parameters and arguments
-        the descriptor stores a set of compile time properties that can improve code
-        generation. Different backends might benefit from different properties
-        """
-        return AttrsDescriptor(params, args)
+    @staticmethod
+    def parse_attr(desc):
+        assert isinstance(desc, str)
+        ret = []
+        if "D" in desc:
+            ret += [["tt.divisibility", 16]]
+        return ret
 
-    def compute_spec_key(self, arg, align):
+    @staticmethod
+    def get_arg_specialization(arg, ty, **kwargs):
         """
-        Return the ascii key for a given argument with a given set of properties
+        Return a string unique to each possible specialization of the argument
         """
-        return AttrsDescriptor.get_property_key(arg, align)
+        if ty == "int" and arg % 16 == 0 and kwargs.get("align", False):
+            return "D"
+        if ty == "tensor" and arg.data_ptr() % 16 == 0 and kwargs.get("align", False):
+            return "D"
+        return ""