support mm and autotune

Signed-off-by: jinjieliu <jinjie.liu@usc.edu>

python/triton_tvm_ffi/jit.py

@@ -2,24 +2,38 @@ from __future__ import annotations
 
 from functools import cached_property
 import inspect
-from typing import Any, Callable, Dict, Final, List, Optional, Tuple, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Final,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import torch
 from triton.compiler import CompiledKernel
-from triton.runtime import JITFunction
+from triton.runtime import Autotuner, JITFunction
 import tvm_ffi
 
 from .utils import type_canonicalize
 
 
 class TVMFFIJITFunction(object):
-    def __init__(self, fn: JITFunction, *args, **kwargs) -> None:
+    def __init__(self, fn: Union[Autotuner, JITFunction], *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
-        self.fn: Final[JITFunction] = fn
-        self.signature: Optional[List[str]] = None
+        self.fn: Final[Union[Autotuner, JITFunction]] = fn
+        self.signature: List[str] = [*inspect.signature(self.basefn).parameters.keys()]
+        self.best_config: Optional[Dict[str, Any]] = None
         self.ctypes: Optional[List[Optional[str]]] = None
         self.kernel: Optional[bytes] = None
         self.num_warps: Optional[int] = None
+        self.shmem: int = 0
 
         @tvm_ffi.register_global_func(self.fullname)
         def _(
@@ -29,22 +43,23 @@ class TVMFFIJITFunction(object):
             _: int,
             num_warps: Optional[int],
             num_stages: Optional[int],
-            *args,
-            **kwargs,
+            args: Sequence[Any],
+            kwargs: Mapping[str, Any],
         ):
-            args: List[Any] = map(self.canonicalize, args)
+            args: Iterator[Any] = map(self.canonicalize, args)
             kwargs: Dict[str, Any] = {
-                k: self.canonicalize(v) for k, v in kwargs.items()
-            }
+                k: v for k, v in zip(self.signature, args) if v is not None
+            } | {k: self.canonicalize(v) for k, v in kwargs.items()}
             if num_warps is not None:
                 kwargs["num_warps"] = num_warps
             if num_stages is not None:
                 kwargs["num_stages"] = num_stages
             kernel: CompiledKernel = self.fn[grid](*args, **kwargs)
-            self.num_warps, _, _ = kernel.packed_metadata
-            self.signature = [*inspect.signature(self.fn.fn).parameters.keys()]
+            self.num_warps, _, self.shmem = kernel.packed_metadata
             self.ctypes = [type_canonicalize(v) for v in kernel.src.signature.values()]
             self.kernel = kernel.kernel
+            if isinstance(self.fn, Autotuner):
+                self.best_config = self.fn.best_config.all_kwargs()
             return kernel
 
     def __getitem__(
@@ -55,6 +70,10 @@ class TVMFFIJITFunction(object):
     ):
         return self.fn[grid]
 
+    @cached_property
+    def basefn(self) -> Callable:
+        return self.jitfn.fn
+
     @property
     def cache_hash(self) -> int:
         return self.ctypes_hash ^ self.kernel_hash
@@ -63,21 +82,35 @@ class TVMFFIJITFunction(object):
     def ctypes_hash(self) -> int:
         return hash(tuple(self.ctypes) if self.ctypes is not None else None)
 
-    @property
-    def kernel_hash(self) -> int:
-        return hash(self.kernel)
-
     @cached_property
     def fnname(self) -> str:
-        return self.fn.fn.__name__
+        return self.basefn.__name__
 
     @cached_property
     def fullname(self) -> str:
         return f"triton.{self.name}"
 
+    @cached_property
+    def jitfn(self) -> JITFunction:
+        fn: Union[Autotuner, JITFunction] = self.fn
+        while not isinstance(fn, JITFunction):
+            fn = fn.fn
+        return fn
+
+    @property
+    def kernel_hash(self) -> int:
+        return hash(self.kernel)
+
+    @property
+    def kernel_cstr(self) -> Optional[str]:
+        if self.kernel is not None:
+            return "".join(f"\\x{byte:02x}" for byte in self.kernel)
+        else:
+            return None
+
     @cached_property
     def name(self) -> str:
-        return f"{self.fnname}_{hash(self.fn.fn)}"
+        return f"{self.fnname}_{hash(self.basefn)}"
 
     @staticmethod
     def canonicalize(val: Any) -> Any: