add hash key for code in libtuner

Author: meinie0826

src/flag_gems/utils/libentry.py

@@ -1,4 +1,5 @@
 import builtins
+import hashlib
 import inspect
 import math
 import os
@@ -27,6 +28,17 @@
 version = triton.__version__.split(".")
 major_version, minor_version = eval(version[0]), eval(version[1])
 
+
+def get_kernel_hash(func):
+    if hasattr(func, "fn"):
+        original_func = func.fn
+    else:
+        original_func = func
+
+    source_code = inspect.getsource(original_func)
+    return hashlib.md5(source_code.encode("utf-8")).hexdigest()[:8]
+
+
 if major_version == 2:
 
     def all_kwargs(self):
@@ -192,16 +204,23 @@ def __init__(
         self.cache = libcache[share] if share else libcache[self.__name__]
         if strategy:
             assert len(self.strategy) == len(self.keys), "Invalid number of strategies"
+        self.base_fn = fn
+        while not inspect.isfunction(self.base_fn):
+            self.base_fn = self.base_fn.fn
+        self.kernel_hash = get_kernel_hash(self.base_fn)
 
     def get_key(self, args):
         if self.strategy is None:
             key = [args[k] for k in self.keys if k in args]
-            return key
-        key = []
-        for i, k in enumerate(self.keys):
-            s = STRATEGY[self.strategy[i]]
-            v = s(args[k])
-            key.append(v)
+        else:
+            key = []
+            for i, k in enumerate(self.keys):
+                s = STRATEGY[self.strategy[i]]
+                v = s(args[k])
+                key.append(v)
+
+        key.append(f"hash_{self.kernel_hash}")
+
         return key
 
     def run(self, *args, **kwargs):

tests/test_libentry.py

@@ -8,7 +8,7 @@
 
 import flag_gems
 from flag_gems.runtime import torch_device_fn
-from flag_gems.utils import libentry
+from flag_gems.utils import libentry, libtuner
 
 
 # not_raises is copied from https://gist.github.com/oisinmulvihill/45c14271fad7794a4a52516ecb784e69
@@ -235,3 +235,57 @@ def test_threadsafety():
     for i in range(100):
         with not_raises(Exception):
             run_two_threads()
+
+
+@pytest.mark.skipif(
+    flag_gems.vendor_name == "kunlunxin",
+    reason="Test Files for Operators Not Pending Testing",
+)
+def test_hash_generation():
+    @libtuner(
+        configs=[
+            triton.Config({"TILE_N": 32}),
+            triton.Config({"TILE_N": 64}),
+            triton.Config({"TILE_N": 128}),
+            triton.Config({"TILE_N": 256}),
+            triton.Config({"TILE_N": 512}),
+            triton.Config({"TILE_N": 1024}),
+        ],
+        key=["x"],
+    )
+    @triton.jit
+    def kernel_a(x, y):
+        return x + y + 1
+
+    @libtuner(
+        configs=[
+            triton.Config({"TILE_N": 32}),
+            triton.Config({"TILE_N": 64}),
+            triton.Config({"TILE_N": 128}),
+            triton.Config({"TILE_N": 256}),
+            triton.Config({"TILE_N": 512}),
+            triton.Config({"TILE_N": 1024}),
+        ],
+        key=["x"],
+    )
+    @triton.jit
+    def kernel_b(x, y):
+        return x + y
+
+    @libtuner(
+        configs=[
+            triton.Config({"TILE_N": 32}),
+            triton.Config({"TILE_N": 64}),
+            triton.Config({"TILE_N": 128}),
+            triton.Config({"TILE_N": 256}),
+            triton.Config({"TILE_N": 512}),
+            triton.Config({"TILE_N": 1024}),
+        ],
+        key=["x"],
+    )
+    @triton.jit
+    def kernel_a_copy(x, y):
+        return x + y + 1
+
+    assert kernel_a.kernel_hash != kernel_a_copy.kernel_hash
+    assert kernel_a.kernel_hash != kernel_b.kernel_hash