import os import re import hashlib import subprocess from abc import ABCMeta, abstractmethod, abstractclassmethod from dataclasses import dataclass from typing import Dict, List, Tuple, Union from types import ModuleType # 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', '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.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 self.arg_properties["tt.divisibility"] = [ param.num for param, arg in zip(params, values) if AttrsDescriptor.is_divisible_by_16(arg) and not param.do_not_specialize and not param.do_not_specialize_on_alignment ] # Equal to 1 property self.arg_properties["tt.equal_to"] = [ param.num for param, arg in zip(params, values) if AttrsDescriptor.is_equal_to_1(arg) and not param.do_not_specialize ] 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] 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] = 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) class GPUTarget(object): # Target backend, e.g., cuda, hip backend: str # Target architecture, e.g., 90 (for cuda compute capability), gfx940 (for hip) arch: Union[int, str] warp_size: int class BaseBackend(metaclass=ABCMeta): def __init__(self, target: GPUTarget) -> None: self.target = target assert self.supports_target(target) @staticmethod def _path_to_binary(binary: str): base_dir = os.path.join(os.path.dirname(__file__), os.pardir) paths = [ os.environ.get(f"TRITON_{binary.upper()}_PATH", ""), os.path.join(base_dir, "third_party", "cuda", "bin", binary), ] for p in paths: bin = p.split(" ")[0] if os.path.exists(bin) and os.path.isfile(bin): result = subprocess.check_output([bin, "--version"], stderr=subprocess.STDOUT) if result is not None: version = re.search(r".*release (\d+\.\d+).*", result.decode("utf-8"), flags=re.MULTILINE) if version is not None: return p, version.group(1) raise RuntimeError(f"Cannot find {binary}") @abstractclassmethod def supports_target(target: GPUTarget): raise NotImplementedError @abstractmethod def hash(self) -> str: """Returns a unique identifier for this backend""" raise NotImplementedError @abstractmethod def parse_options(self, options: dict) -> object: """ Converts an `options` dictionary into an arbitrary object and returns it. This function may contain target-specific heuristics and check the legality of the provided options """ raise NotImplementedError @abstractmethod def add_stages(self, stages: dict, options: object) -> None: """ Populates `stages` dictionary with entries of the form: ir_name [str] => Function[(src: str, metadata: dict) -> str|bytes] The value of each entry may populate a `metadata` dictionary. Stages will be run sequentially (in inseriton order) and can communicate using `metadata`. All stages are expected to return a `str` object, except for the last stage which returns a `bytes` object for execution by the launcher. """ raise NotImplementedError @abstractmethod def load_dialects(self, context): """ Load additional MLIR dialects into the provided `context` """ raise NotImplementedError @abstractmethod def get_module_map(self) -> Dict[str, ModuleType]: """ Return a map of interface modules to their device-specific implementations """ 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) def compute_spec_key(self, arg, align): """ Return the ascii key for a given argument with a given set of properties """ return AttrsDescriptor.get_property_key(arg, align)