update hintmanager, wrap additional code into hintmanager, back no-hint-related handler func into spec, update import, change jit implement into hintmanager, simplify trigger call

Author: starrryz

python/triton/compiler/code_generator.py

@@ -15,6 +15,7 @@
 from ..runtime import JITFunction
 from .errors import (CompilationError, CompileTimeAssertionFailure, UnsupportedLanguageConstruct)
 from types import ModuleType
+from .hintmanager import hint_trigger
 
 
 def mangle_ty(ty):
@@ -247,11 +248,6 @@ def __init__(self, context, prototype, gscope, attributes, constants, function_n
         # special handling.
         self.visiting_arg_default_value = False
 
-        # adding unified hint manager init
-        from .hint_manager import HintManager
-        from .hint_manager import hint_get_flagtree_backend
-        self.hint_manager = HintManager(hint_get_flagtree_backend())
-
     builtin_namespace: Dict[str, Any] = {_.__name__: _ for _ in (len, list, range, float, int, isinstance, getattr)}
     builtin_namespace.update((
         ('print', language.core.device_print),
@@ -522,7 +518,7 @@ def visit_Assign(self, node):
             self.set_value(name, value)
 
         # switch into hintmanager
-        self.hint_manager.handler.trigger("ext_CodeGenerator_visit_Assign_hint_anno", self, node, names, values)
+        hint_trigger("ext_CodeGenerator_visit_Assign_hint_anno", self, node, names, values)
 
     def visit_AugAssign(self, node):
         name = node.target.id
@@ -927,9 +923,10 @@ def visit_For(self, node):
         # flagtree backend specialization: add more ForOp attributes
         for_op_ext_attrs = (False, False, False, False)
 
+        # flagtree backend specialization
+        from triton.runtime.driver import spec
         bind_sub_block = None
-        ext_it_class_support = [language.range] # why?
-        ext_it_class_support += self.hint_manager.handler.trigger("visit_For_ext_support")
+        ext_it_class_support = spec("visit_For_ext_support")
         ext_it_class_support = [] if ext_it_class_support is None else ext_it_class_support
         if IteratorClass in [language.range] + ext_it_class_support:
             iterator = IteratorClass(*iter_args, **iter_kwargs)
@@ -942,7 +939,10 @@ def visit_For(self, node):
             num_stages = iterator.num_stages
             loop_unroll_factor = iterator.loop_unroll_factor
 
-            new_bind_sub_block = self.hint_manager.handler.trigger("set_bind_sub_block_when_parallel", IteratorClass, iterator, bind_sub_block)
+            # flagtree backend specialization
+            for_op_ext_attrs = spec("for_op_ext_attrs", iterator)
+            # flagtree backend specialization
+            new_bind_sub_block = spec("set_bind_sub_block_when_parallel", IteratorClass, iterator, bind_sub_block)
             if new_bind_sub_block is not None:
                 bind_sub_block = new_bind_sub_block
         elif IteratorClass is range:
@@ -955,7 +955,7 @@ def visit_For(self, node):
         else:
             raise RuntimeError('Only `range` and `static_range` iterators are currently supported')
 
-        new_bind_sub_block = self.hint_manager.handler.trigger("check_override_bind_sub_block", self, node, bind_sub_block)
+        new_bind_sub_block = hint_trigger("check_override_bind_sub_block", self, node, bind_sub_block)
         if new_bind_sub_block is not None:
             bind_sub_block = new_bind_sub_block
 
@@ -1027,7 +1027,7 @@ def visit_For(self, node):
             spec("for_op_set_ext_attrs", for_op, self.builder, for_op_ext_attrs)
             # flagtree backend specialization
             if bind_sub_block:
-                self.hint_manager.handler.trigger("forop_setattr_for_bind_sub_block", self, for_op, bind_sub_block)
+                hint_trigger("forop_setattr_for_bind_sub_block", self, for_op, bind_sub_block)
 
             self.scf_stack.append(node)
             self.builder.set_insertion_point_to_start(for_op.get_body(0))
@@ -1112,8 +1112,9 @@ def call_JitFunction(self, fn: JITFunction, args, kwargs):
             except Exception as e:
                 # Wrap the error in the callee with the location of the call.
 
-
-                if self.hint_manager.handler.trigger("need_repr_in_CodeGenerator_CompilationError"):
+                # flagtree backend specialization
+                from triton.runtime.driver import spec
+                if spec('need_repr_in_CodeGenerator_CompilationError'):
                     raise CompilationError(self.jit_fn.src, self.cur_node, repr(e)) from e
                 raise CompilationError(self.jit_fn.src, self.cur_node, None) from e
 
@@ -1160,7 +1161,9 @@ def visit_Call(self, node):
                 # preserve the traceback of the original error, which may e.g.
                 # be in core.py.
 
-                if self.hint_manager.handler.trigger("need_repr_in_CodeGenerator_CompilationError"):
+                #flagtree backend specialization
+                from triton.runtime.driver import spec
+                if spec('need_repr_in_CodeGenerator_CompilationError'):
                     raise CompilationError(self.jit_fn.src, node, repr(e)) from e
                 raise CompilationError(self.jit_fn.src, node, None) from e
 

python/triton/compiler/hint_manager.py

@@ -152,8 +152,16 @@ def hint_get_flagtree_backend() -> str:
         return canonical_backend
     else:
         # version and backend mismatch
-        logging.warning(
+        msg = (
             f"[Flagtree] Hint ignored: Detected backend '{canonical_backend}' but current Triton version "
             f"'{current_triton_version}' matches no supported versions {supported_versions}."
         )
-    return ""
+        print(msg, file=sys.stderr)
+    return ""
+# lazy load after first call hint trigger
+_global_hint_manager = None
+
+def hint_trigger(hook_name, *args, **kwargs):
+    if _global_hint_manager is None:
+        _global_hint_manager = HintManager(hint_get_flagtree_backend())
+    return _global_hint_manager.handler.trigger(hook_name, *args, **kwargs)

python/triton/runtime/jit.py

@@ -11,6 +11,7 @@
 from typing import Callable, Generic, Iterable, Optional, TypeVar, Union, overload, Dict, Any, Tuple
 from ..runtime.driver import driver
 from types import ModuleType
+from ..compiler.hintmanager import hint_trigger
 
 TRITON_MODULE = __name__[:-len(".runtime.jit")]
 
@@ -826,26 +827,16 @@ def get_flagtree_backend():
     # the user might want to monkey-patch self.src dynamically.
     # Our unit tests do this, for example.
     def parse(self):
-        # remove flagtree backend specialization, because the implementation of 2 method is totally same
-        line_flagtree_hints = {}
-        code_str = self.src
-        g = tokenize.generate_tokens(StringIO(code_str).readline)
-        for tok_type, tok_text, start, end, _ in g:
-            if tok_type == tokenize.COMMENT:
-                comment = tok_text.replace(" ", "").strip()
-                if comment.startswith('#@hint:'):
-                    flagtree_hints = comment[len('#@hint:'):].strip()
-                    # Record the line number of the comment
-                    line_num = start[0]
-                    line_flagtree_hints[line_num] = flagtree_hints
+        # after removing flagtree backend specialization, hiding the implementation into hintmanager
+        line_flagtree_hints = hint_trigger("maps_line_numbers_to_comment_hints", self)
 
         tree = ast.parse(self.src)
         assert isinstance(tree, ast.Module)
         assert len(tree.body) == 1
         assert isinstance(tree.body[0], ast.FunctionDef)
 
         # Attach the line number to comment mapping to the function definition node
-        tree.body[0].line_flagtree_hints = line_flagtree_hints
+        hint_trigger('attach_line_number_to_comment_mapping', tree, line_flagtree_hints)
 
         return tree
 

third_party/ascend/backend/ascend_hint_handler.py

@@ -6,44 +6,36 @@
 
 class AscendHintHandler(BaseHintHandler):
 
-   def ext_CodeGenerator_visit_Assign_hint_anno(self, code_generator, node, names, values):
-    import ast
-    from triton.compiler.code_generator import _is_triton_value
-    # flagtree: After normal processing, check if we need to add hint annotation
-    if hasattr(node, 'lineno') and hasattr(code_generator, 'jit_fn'):
-        line_num = node.lineno
-        # TODO: reparse needed in case we need to deal with complex cases, will be redesigned later
-        function_def = code_generator.jit_fn.parse()
-        line_flagtree_hints = getattr(function_def.body[0], 'line_flagtree_hints', {})
-        flagtree_hints = line_flagtree_hints.get(line_num)
-
-        # Check if this is a tl.load call with dot_pad_only_k hint
-        if (flagtree_hints and 'dot_pad_only_k' in flagtree_hints and
-            isinstance(node.value, ast.Call) and
-            isinstance(node.value.func, ast.Attribute) and
-            isinstance(node.value.func.value, ast.Name) and
-            node.value.func.value.id == 'tl' and
-            node.value.func.attr == 'load'):
-
-            # Add hint annotation to the loaded tensor(s)
-            for name, value in zip(names, values):
-                if _is_triton_value(value):
-                    # print(f"[FLAGTREE] Creating hint annotation for tensor: {flagtree_hints}")
-                    # Create hint annotation
-                    hint_val = code_generator.builder.get_unit_attr()
-                    code_generator.builder.create_annotation(value.handle, 'dot_pad_only_k', hint_val)
+    @staticmethod
+    def ext_CodeGenerator_visit_Assign_hint_anno(code_generator, node, names, values):
+        import ast
+        from triton.compiler.code_generator import _is_triton_value
+        # flagtree: After normal processing, check if we need to add hint annotation
+        if hasattr(node, 'lineno') and hasattr(code_generator, 'jit_fn'):
+            line_num = node.lineno
+            # TODO: reparse needed in case we need to deal with complex cases, will be redesigned later
+            function_def = code_generator.jit_fn.parse()
+            line_flagtree_hints = getattr(function_def.body[0], 'line_flagtree_hints', {})
+            flagtree_hints = line_flagtree_hints.get(line_num)
 
-    def visit_For_ext_support(self):
-        import triton.language as language
-        return [language.parallel]
+            # Check if this is a tl.load call with dot_pad_only_k hint
+            if (flagtree_hints and 'dot_pad_only_k' in flagtree_hints and
+                isinstance(node.value, ast.Call) and
+                isinstance(node.value.func, ast.Attribute) and
+                isinstance(node.value.func.value, ast.Name) and
+                node.value.func.value.id == 'tl' and
+                node.value.func.attr == 'load'):
 
-    def set_bind_sub_block_when_parallel(self, IteratorClass, iterator, bind_sub_block):
-        import triton.language as language
-        if (IteratorClass is language.parallel):
-            return iterator.bind_sub_block
-        return bind_sub_block
+                # Add hint annotation to the loaded tensor(s)
+                for name, value in zip(names, values):
+                    if _is_triton_value(value):
+                        # print(f"[FLAGTREE] Creating hint annotation for tensor: {flagtree_hints}")
+                        # Create hint annotation
+                        hint_val = code_generator.builder.get_unit_attr()
+                        code_generator.builder.create_annotation(value.handle, 'dot_pad_only_k', hint_val)
 
-    def check_override_bind_sub_block(self, code_generator, node, bind_sub_block):
+    @staticmethod
+    def check_override_bind_sub_block(code_generator, node, bind_sub_block):
         # flagtree: After normal processing, check if we need to override bind_sub_block
         if hasattr(node, 'lineno') and hasattr(code_generator, 'jit_fn'):
             line_num = node.lineno
@@ -58,8 +50,33 @@ def check_override_bind_sub_block(self, code_generator, node, bind_sub_block):
                 # print(f"[FLAGTREE] Found bind_sub_block hint at line {line_num}")
         return bind_sub_block
 
-    def forop_setattr_for_bind_sub_block(self, code_generator, for_op, bind_sub_block):
+    @staticmethod
+    def forop_setattr_for_bind_sub_block(code_generator, for_op, bind_sub_block):
         for_op.set_attr("bind_sub_block", code_generator.builder.get_bool_attr(bind_sub_block))
 
-    def need_repr_in_CodeGenerator_CompilationError(self):
-        return True
+
+    @staticmethod
+    def maps_line_numbers_to_comment_hints(jit_fn):
+    import tokenize
+    from io import StringIO
+    # Maps line numbers to comment hints
+    line_flagtree_hints = {}
+    code_str = jit_fn.src
+    g = tokenize.generate_tokens(StringIO(code_str).readline)
+    for tok_type, tok_text, start, end, _ in g:
+        if tok_type == tokenize.COMMENT:
+            comment = tok_text.replace(" ", "").strip()
+            if comment.startswith('#@hint:'):
+                flagtree_hints = comment[len('#@hint:'):].strip()
+                # Record the line number of the comment
+                line_num = start[0]
+                line_flagtree_hints[line_num] = flagtree_hints
+
+                # print(f"[FLAGTREE] Parsed hint at line {line_num}: {flagtree_hints}")
+
+    return line_flagtree_hints
+
+    @staticmethod
+    def attach_line_number_to_comment_mapping(tree, line_flagtree_hints):
+    # Attach the line number to comment mapping to the function definition node
+    tree.body[0].line_flagtree_hints = line_flagtree_hints