Support JITFunction in preload (#8794)

Author: ThomasRaoux

python/test/unit/runtime/test_cache.py

@@ -825,3 +825,62 @@ def kernel(out_ptr, FUNC: tl.constexpr) -> None:
 Compiling with fn_a
 Compiling with fn_a after modification
 """)
+
+
+def test_preload_higher_order_kernels(device, fresh_triton_cache) -> None:
+
+    @triton.jit
+    def fn_a():
+        return 17
+
+    @triton.jit
+    def fn_b():
+        return 31
+
+    @triton.jit
+    def kernel(out_ptr, FUNC: tl.constexpr) -> None:
+        val = FUNC()
+        tl.store(out_ptr, val)
+
+    device = getattr(torch, device).current_device()
+
+    # get the serialized specialization data
+    specialization_data = None
+
+    def cache_hook(*args, **kwargs):
+        nonlocal specialization_data
+        specialization_data = kwargs["compile"]["specialization_data"]
+
+    triton.knobs.runtime.jit_cache_hook = cache_hook
+    output = torch.empty((), device=device, dtype=torch.int32)
+    compiled_kernel = kernel[(1, )](output, fn_a)
+    assert output.item() == 17
+    hash = compiled_kernel.hash
+    assert specialization_data is not None
+
+    # clear the cache
+    shutil.rmtree(fresh_triton_cache)
+    kernel.device_caches[device][0].clear()
+
+    # preload the kernel
+    kernel_preload = kernel.preload(specialization_data)
+    assert kernel_preload.hash == hash
+    assert len(kernel.device_caches[device][0]) == 1
+
+    # we should hit the cache and not compile anything
+    counter = 0
+
+    def inc_counter(*args, **kwargs):
+        nonlocal counter
+        counter += 1
+
+    triton.knobs.runtime.jit_cache_hook = inc_counter
+    final_kernel = kernel[(1, )](output, fn_a)
+    assert counter == 0
+    assert len(kernel.device_caches[device][0]) == 1
+    assert final_kernel.hash == hash
+
+    # different function should compile and not hit the cache
+    kernel[(1, )](output, fn_b)
+    assert counter == 1
+    assert output.item() == 31

python/triton/runtime/jit.py

@@ -373,8 +373,9 @@ def __getitem__(self, grid) -> T:
 
 def serialize_specialization_data(name, signature, constants, attrs, options, key):
     constants = {
-        key: str(value) if value.__class__.__name__ == "dtype" else
-        {"constexpr": value.value} if value.__class__.__name__ == "constexpr" else value
+        key: str(value) if value.__class__.__name__ == "dtype" else {"constexpr": value.value}
+        if value.__class__.__name__ == "constexpr" else {"jit_function": f"{value.module}:{value.fn.__qualname__}"}
+        if value.__class__.__name__ == "JITFunction" else value
         for key, value in constants.items()
     }
 
@@ -560,6 +561,9 @@ def _get_src(self):
     src = property(fget=_get_src, fset=_set_src)
 
 
+_triton_jit_function_registry = {}
+
+
 @dataclass
 class JitFunctionInfo:
     module: ModuleType
@@ -771,6 +775,8 @@ def __init__(self, fn, version=None, do_not_specialize=None, do_not_specialize_o
         self.do_not_specialize_on_alignment = do_not_specialize_on_alignment
         self._repr = repr
         self.launch_metadata = launch_metadata
+        # Register for simple deserialization of JITFunction constants
+        _triton_jit_function_registry[f"{self.module}:{self.fn.__qualname__}"] = self
 
         self.params = []
         for i, param in enumerate(self.signature.parameters.values()):
@@ -805,12 +811,21 @@ def preload(self, specialization_data):
                 f"Specialization data is for {deserialized_obj['name']} but trying to preload for {self._fn_name}")
         constant_keys = map(tuple, deserialized_obj['constant_keys'])
         constant_vals = deserialized_obj['constant_vals']
-        constexprs = {
-            key:
-            tl.dtype(value) if tl.dtype.is_dtype(value) else
-            tl.constexpr(value['constexpr']) if isinstance(value, dict) and 'constexpr' in value else value
-            for key, value in zip(constant_keys, constant_vals)
-        }
+
+        def _decode_constant(value):
+            if tl.dtype.is_dtype(value):
+                return tl.dtype(value)
+            if isinstance(value, dict):
+                if 'constexpr' in value:
+                    return tl.constexpr(value['constexpr'])
+                if 'jit_function' in value:
+                    jf_key = value['jit_function']
+                    if jf_key in _triton_jit_function_registry:
+                        return _triton_jit_function_registry[jf_key]
+                    raise RuntimeError(f"Unable to resolve JITFunction {jf_key} for preload")
+            return value
+
+        constexprs = {key: _decode_constant(value) for key, value in zip(constant_keys, constant_vals)}
         attrs_keys = map(tuple, deserialized_obj['attrs_keys'])
         attrs_vals = deserialized_obj['attrs_vals']
         attrs = dict(zip(attrs_keys, attrs_vals))