[Cherry-pick] [IMP] add test for launch with options (#7970) (#524)

Summary:
Cherry-picked from upstream OAI repository.

Original Commit: 0be0607
Original Author: Maximilian Stadler
Original Date: 2025-08-28 22:40:47 +0200

Original commit message:
```
[IMP] add test for launch with `options` (#7970)

follow-up to #7961
```

This PR was automatically cherry-picked from the upstream triton-lang/triton repository.

Pull Request resolved: #524

Reviewed By: minjang

Differential Revision: D85686941

Pulled By: agron911

fbshipit-source-id: 7d50ded4187a70bb858a60b5151b9b89ab163461

Author: stadlmax

python/test/unit/runtime/test_launch.py

@@ -1,18 +1,13 @@
 import gc
-# import importlib
-# import os
-# import sys
-# import tempfile
-# import textwrap
-# import time
 import tracemalloc
+import pytest
+import pathlib
+import os
 
 import torch
-
 import triton
 import triton.language as tl
-
-# from typing import Tuple
+from triton._internal_testing import is_cuda, is_hip
 
 
 def test_metadata() -> None:
@@ -149,3 +144,63 @@ def kernel(x):
     triton.knobs.runtime.kernel_load_end_hook.remove(hook_end0)
     triton.knobs.runtime.kernel_load_start_hook.remove(hook_start1)
     triton.knobs.runtime.kernel_load_end_hook.remove(hook_end1)
+
+
+@pytest.mark.parametrize("options", [
+    {"num_warps": 1},
+    {"enable_fp_fusion": False},
+    {"extern_libs": {}},
+])
+def test_launch_with_options(options) -> None:
+    if "extern_libs" in options:
+        # copied from tutorials/07-extern-functions.py
+        current_dir = pathlib.Path(os.path.dirname(os.path.abspath(__file__)))
+        if is_cuda():
+            libdir = current_dir.parent.parent.parent.parent / 'third_party/nvidia/backend/lib'
+            options["extern_libs"] = {"libdevice": str(libdir / 'libdevice.10.bc')}
+        elif is_hip():
+            libdir = current_dir.parent.parent.parent.parent / 'third_party/amd/backend/lib'
+            options["extern_libs"] = {"ocml": str(libdir / 'ocml.bc'), "ockl": str(libdir / 'ockl.bc')}
+
+    compile_info = {}
+    counter = 0
+
+    def compile_info_hook(key, repr, fn, compile, is_manual_warmup, already_compiled):
+        nonlocal compile_info
+        compile_info = compile
+
+    def cache_hook(*args, **kwargs):
+        nonlocal counter
+        counter += 1
+
+    @triton.jit
+    def kernel(x):
+        pass
+
+    triton.knobs.runtime.jit_post_compile_hook = compile_info_hook
+    triton.knobs.runtime.jit_cache_hook = cache_hook
+
+    # run first without options
+    kernel[(1, 1, 1)](6)
+    assert counter == 1
+
+    # run with options, should lead to new compilation
+    kernel[(1, 1, 1)](6, **options)
+    assert counter == 2
+
+    # run a second time for testing kernel-cache look-up
+    kernel[(1, 1, 1)](6, **options)
+    assert counter == 2
+
+    # check the options are passed on to compile_info correctly
+    option_key, option_val = next(iter(options.items()))
+    if option_key == "extern_libs":
+        # HIPOptions overwrite the extern_libs option, so we skip the test
+        # passing and specializing options still is tested
+        if not is_hip():
+            assert compile_info[option_key] == tuple(option_val.items())
+    else:
+        assert compile_info[option_key] == option_val
+
+    triton.knobs.runtime.jit_post_compile_hook = None
+    triton.knobs.runtime.jit_cache_hook = None