update func attrs

Author: motinwing

python/triton/compiler/code_generator.py

@@ -1243,7 +1243,7 @@ def visit_Call(self, node):
         args = list(itertools.chain.from_iterable(x if isinstance(x, list) else [x] for x in args))
 
         # 4. Get current line number and hints
-        hint_trigger("func1", self, node, names, values)
+        flagtree_hints = hint_trigger("get_node_hints", self, node)
 
         # 5. Handle JIT function calls
         if isinstance(fn, JITFunction):
@@ -1258,7 +1258,7 @@ def visit_Call(self, node):
                 extra_kwargs['_generator'] = self
             try:
                 # Special handling for tl.load with hints
-                hint_trigger("func2", self, node, names, values)
+                hint_trigger("inject_kwargs_with_hints", fn, flagtree_hints, node.lineno, kws)
 
                 ret = fn(*args, **extra_kwargs, **kws)
                 # builtin functions return plain tuples for readability

python/triton/runtime/jit.py

@@ -706,7 +706,9 @@ def preload(self, specialization_data):
     # Our unit tests do this, for example.
     def parse(self):
         # Maps line numbers to comment hints
-        hint_trigger("maps_line_numbers_to_comment_hints", self)
+        line_flagtree_hints = hint_trigger("maps_line_numbers_to_comment_hints", self)
+        if line_flagtree_hints is None:
+            line_flagtree_hints = {}
 
         tree = ast.parse(self.src)
         assert isinstance(tree, ast.Module)

third_party/aipu/backend/aipu_hint_handler.py

@@ -5,77 +5,35 @@
 from triton.compiler.code_generator import _is_triton_value
 
 class AipuHintHandler(BaseHintHandler):
-
+    # because aipu is diff from ascend in 2 aspects
+    # 1. not backend_spec, modify triton src violently
+    # 2. modify builder, semantic, core, and so on. pollute the src, which cant be involved in hint manager
+    # for this, we just move changes in codegen & jit into hintmanager.
 
     @staticmethod
-    def func1(code_generator, node, names, values):
+    def get_node_hints(code_generator, node):
         line_num = node.lineno
-        function_def = self.jit_fn.parse()
+        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)
+        return flagtree_hints
 
-    def func2():
+    @staticmethod
+    def inject_kwargs_with_hints(fn, flagtree_hints, line_num, kws):
         if fn.__name__ == "load" and flagtree_hints is not None:
                     print(f"[FLAGTREE] tl.load at line {line_num} has annotation {flagtree_hints}")
                     if 'flagtree_hints' not in kws:
                         kws['flagtree_hints'] = ""
                     if flagtree_hints not in kws['flagtree_hints']:
                         kws['flagtree_hints'] = flagtree_hints
 
-    @staticmethod
-    def ext_CodeGenerator_visit_Assign_hint_anno(code_generator, node, names, values):
-        if not (hasattr(node, 'lineno') and hasattr(code_generator, 'jit_fn')):
-            return
-
-        if not flagtree_hints:
-            return
-
-        # 3. AIPU 特有的 Hint 处理逻辑
-        # [请对照 PR 修改此处字符串] 假设 AIPU 有一个 hint 叫 'global_memory_cache' 或者类似的
-        target_hint_key = 'aipu_specific_hint_key' # <--- 请替换为 PR 里的真实字符串
-        
-        if target_hint_key in flagtree_hints:
-            # 检查是否作用于 tl.load / tl.store
-            if (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'):
-                
-                # 例如：如果是 load 操作
-                if node.value.func.attr == 'load':
-                    for name, value in zip(names, values):
-                        if _is_triton_value(value):
-                            # 创建 AIPU 特有的 Annotation
-                            # print(f"[FLAGTREE][AIPU] Applying {target_hint_key} to {name}")
-                            hint_val = code_generator.builder.get_unit_attr()
-                            # 'aipu.hint_name' 需要与后端 LLVM IR 处理逻辑对应
-                            code_generator.builder.create_annotation(value.handle, target_hint_key, hint_val)
-
-    @staticmethod
-    def check_override_bind_sub_block(code_generator, node, bind_sub_block):
-        """
-        对应 CodeGenerator.visit_For 中决定是否开启 bind_sub_block 的逻辑
-        """
-        if not (hasattr(node, 'lineno') and hasattr(code_generator, 'jit_fn')):
-            return bind_sub_block
-
-        line_num = node.lineno
-        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)
-
-        # 检查 AIPU 是否也支持通过 hint 强制开启/关闭 sub_block 绑定
-        if flagtree_hints and 'bind_sub_block' in flagtree_hints:
-             # 如果 AIPU 后端支持此特性，则返回 True
-            return True
-        
-        return bind_sub_block
-
     @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 = self.src
+        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: