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 logging
 import math
@@ -31,6 +32,19 @@
 version = triton.__version__.split(".")
 major_version, minor_version = eval(version[0]), eval(version[1])
 
+
+def get_kernel_hash(func, configs):
+    if hasattr(func, "fn"):
+        original_func = func.fn
+    else:
+        original_func = func
+
+    source_code = inspect.getsource(original_func)
+    config_strs = [str(config) for config in configs]
+    combined_content = f"{source_code}{config_strs}"
+    return hashlib.md5(combined_content.encode("utf-8")).hexdigest()[:8]
+
+
 if major_version == 2:
 
     def all_kwargs(self):
@@ -199,9 +213,19 @@ def __init__(
         self.keys = key
         self.strategy = strategy
         self.share = share
-        self.cache = libcache[share] if share else libcache[self.__name__]
+        self.kernel_hash = get_kernel_hash(self.base_fn, self.configs)
+        # Use table name with hash instead of hash in key
+        self.table_name = (
+            f"{share}_{self.kernel_hash}"
+            if share
+            else f"{self.__name__}_{self.kernel_hash}"
+        )
+        self.cache = libcache[self.table_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
 
     def get_key(self, args):
         if self.strategy is None:

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