[Frontend] Support binding self to JITFunction when they are methods (#6963)

📚 Stacked PRs 📚

* triton-lang/triton#6970
* ➡️ triton-lang/triton#6963

This PR makes `base_value.method` return a BoundJITMethod which keeps
`base_value` to be passed as `__self__`.

Author: Mogball

python/test/unit/language/test_frontend.py

@@ -41,11 +41,7 @@ def run_filecheck(name, module_str, check_template):
         raise ValueError(matcher.stderr.getvalue())
 
 
-def run_filecheck_test(kernel_fn):
-    assert isinstance(kernel_fn, triton.runtime.JITFunction)
-    check_template = inspect.getsource(kernel_fn.fn)
-    if check_template is None:
-        raise ValueError("kernel function must have a docstring with FileCheck template")
+def run_parser(kernel_fn):
     sigkeys = [x.name for x in kernel_fn.params]
     sigvals = [f"arg{i}" for i in range(len(sigkeys))]
     signature = {k: v for (k, v) in zip(sigkeys, sigvals)}
@@ -59,7 +55,15 @@ def run_filecheck_test(kernel_fn):
     options = stub_backend.parse_options(dict(**extra_options))
     codegen_fns = stub_backend.get_codegen_implementation(options)
     module_map = stub_backend.get_module_map()
-    mlir_module = src.make_ir(options, codegen_fns, module_map, context)
+    return src.make_ir(options, codegen_fns, module_map, context)
+
+
+def run_filecheck_test(kernel_fn):
+    assert isinstance(kernel_fn, triton.runtime.JITFunction)
+    check_template = inspect.getsource(kernel_fn.fn)
+    if check_template is None:
+        raise ValueError("kernel function must have a docstring with FileCheck template")
+    mlir_module = run_parser(kernel_fn)
 
     run_filecheck("placeholder", str(mlir_module), check_template)
 
@@ -142,6 +146,17 @@ def _flatten_ir(self, handles: List[ir.value]) -> None:
         self.first._flatten_ir(handles)
         self.second._flatten_ir(handles)
 
+    @triton.jit
+    def get_first(self):
+        return self.first
+
+    def get_second(self, _builder=None):
+        return self.second
+
+    @triton.jit
+    def unpack(self):
+        return self.get_first(), self.get_second()
+
 
 @tl.core.builtin
 def pair_value_ctor(first, second, _builder=None):
@@ -160,3 +175,19 @@ def test_assign_attribute():
     # CHECK-NEXT: call @"anchor{{.*}}"([[RANGE]], %c42_i32)
     pair.second = 42
     anchor(pair)
+
+
+@filecheck_test
+@triton.jit
+def test_jit_method():
+    # CHECK-LABEL: test_jit_method
+    # CHECK: %c11_i32 = arith.constant 11 : i32
+    # CHECK: [[RANGE:%.*]] = tt.make_range {end = 4 : i32, start = 0 : i32}
+    scalar = 11
+    # CHECK: [[V:%.*]]:2 = tt.call @"unpack{{.*}}"([[RANGE]], %c11_i32)
+    pair = pair_value_ctor(tl.arange(0, 4), scalar)
+    a, b = pair.unpack()
+    # CHECK: call @anchor{{.*}}([[V]]#0)
+    anchor(a)
+    # CHECK: call @anchor{{.*}}([[V]]#1)
+    anchor(b)

python/triton/compiler/code_generator.py

@@ -4,6 +4,7 @@
 import warnings
 import textwrap
 import itertools
+from dataclasses import dataclass
 from types import ModuleType
 from typing import Any, Callable, Dict, Optional, Tuple, Type, Union, Iterable, List
 
@@ -271,6 +272,12 @@ def make_template(ty):
         return vals
 
 
+@dataclass(frozen=True)
+class BoundJITMethod:
+    __self__: base_value
+    __func__: JITFunction
+
+
 class CodeGenerator(ast.NodeVisitor):
 
     def __init__(self, context, prototype, gscope, function_name, jit_fn: JITFunction, options, codegen_fns, module_map,
@@ -1247,6 +1254,9 @@ def visit_Call(self, node):
         kws = dict(self.visit(keyword) for keyword in node.keywords)
         args = [self.visit(arg) for arg in node.args]
         args = list(itertools.chain.from_iterable(x if isinstance(x, list) else [x] for x in args))
+        if isinstance(fn, BoundJITMethod):
+            args.insert(0, fn.__self__)
+            fn = fn.__func__
         if isinstance(fn, JITFunction):
             _check_fn_args(node, fn, args)
             return self.call_JitFunction(fn, args, kws)
@@ -1338,7 +1348,10 @@ def visit_Attribute(self, node):
         lhs = self.visit(node.value)
         if _is_triton_tensor(lhs) and node.attr == "T":
             return semantic.permute(lhs, (1, 0), builder=self.builder)
-        return getattr(lhs, node.attr)
+        attr = getattr(lhs, node.attr)
+        if _is_triton_value(lhs) and isinstance(attr, JITFunction):
+            return BoundJITMethod(lhs, attr)
+        return attr
 
     def visit_Expr(self, node):
         node.value._is_unused = True