[Fix] Compile and dispatch a script per target

python/tilus/lang/instantiated_script.py

@@ -44,7 +44,7 @@
 from tilus.lang.script import Script
 from tilus.lang.script_dependencies import script_dependency_fingerprint
 from tilus.runtime import CompiledProgram, compiled_program_exists, load_compiled_program
-from tilus.target import get_current_target, lazy_init
+from tilus.target import Target, get_current_target, lazy_init
 from tilus.utils import benchmark_func, relative_to_with_walk_up, to_snake_case
 from tilus.utils.multiprocess import parallel_imap
 
@@ -393,7 +393,7 @@ def _make_tuple_getter(indices: Sequence[int]) -> Callable[[Sequence[Any]], tupl
 
 
 class JitInstance:
-    SPECIALIZATION_CACHE_VERSION = 3
+    SPECIALIZATION_CACHE_VERSION = 4
 
     def __init__(
         self,
@@ -671,7 +671,10 @@ def _transpile_programs(self):
         # get the cache dir for this jit instance
         concatenated_program_text = "\n".join([str(program) for program in self.transpiled_programs])
         option_text = str(self.build_options)
-        hash_string = option_text + concatenated_program_text
+        # The target is part of the key: the same program text builds a different binary and tunes to a different
+        # dispatch table on each target.
+        target_text = str(get_current_target())
+        hash_string = target_text + option_text + concatenated_program_text
         hash_key = hashlib.sha256(hash_string.encode()).hexdigest()[:8]
         jit_name_items = [str(key) for key in self.jit_key] + [hash_key]
         self.cache_dir = (
@@ -992,6 +995,14 @@ def dump_dispatch_table(self):
             f.write(tabulate.tabulate(rows, headers=headers))
 
 
+_LastDispatch = tuple[
+    Optional[Target],
+    Optional[dict[tuple[JitKey, TuningKey], tvm_ffi.Function]],
+    Optional[tuple[Any, ...]],
+    Optional[tvm_ffi.Function],
+]
+
+
 class InstantiatedScript:
     def __init__(self, script_cls: Type[Script], script_args: Sequence[Any], script_kwargs: Mapping[str, Any]):
         self.script_cls: Type[Script] = script_cls
@@ -1006,12 +1017,14 @@ def __init__(self, script_cls: Type[Script], script_args: Sequence[Any], script_
         self.kernel_params: list[int] = self.params.kernel_params
         self.tuning_params: list[int] = self.params.tuning_params
 
-        self.jit_instances: dict[JitKey, JitInstance] = {}
-        self.dispatch_table: dict[tuple[JitKey, TuningKey], tvm_ffi.Function] = {}
+        self.jit_instances: dict[tuple[Target, JitKey], JitInstance] = {}
+        # One dispatch table per target: a binary belongs to the target it was built for.
+        self.dispatch_tables: dict[Target, dict[tuple[JitKey, TuningKey], tvm_ffi.Function]] = {}
         self._dispatch_arg_getter = _make_tuple_getter(self.const_params + self.tuning_params)
         self._kernel_arg_getter = _make_tuple_getter(self.kernel_params)
-        self._last_dispatch_args: Optional[tuple[Any, ...]] = None
-        self._last_compiled_func: Optional[tvm_ffi.Function] = None
+        # The last launch as (target, dispatch table of target, dispatch args, compiled func). It is read and replaced
+        # as one tuple, so concurrent calls under different targets never see a mix of two launches.
+        self._last_dispatch: _LastDispatch = (None, None, None, None)
 
     def __call__(self, *args, **kwargs):
         if kwargs or self.with_default:
@@ -1025,32 +1038,32 @@ def __call__(self, *args, **kwargs):
                     "The number of arguments should be {}, but got {}.".format(len(self.params.param_names), len(args))
                 )
 
+        # Resolve the target and its dispatch table once, so a target change during this call (e.g., from another
+        # thread) cannot make this call record its binary under another target. Switching targets is rare, so the
+        # common path only pays for an identity check.
+        target = get_current_target()
+        last_target, table, last_dispatch_args, last_compiled_func = self._last_dispatch
+        if target is not last_target or table is None:
+            table = self.dispatch_tables.setdefault(target, {})
+            last_dispatch_args = None
+
         # Most workloads launch the same shape repeatedly. Compare the raw dispatch arguments first,
         # avoiding key normalization and hashing entirely on that common path.
         dispatch_args = self._dispatch_arg_getter(args)
-        if dispatch_args == self._last_dispatch_args:
-            compiled_func = self._last_compiled_func
-            keys = None
+        if dispatch_args == last_dispatch_args:
+            compiled_func = last_compiled_func
         else:
             keys = extract_keys(args, self.const_params, self.tuning_params)
-            compiled_func = self.dispatch_table.get(keys)
+            compiled_func = table.get(keys)
+            if compiled_func is None:
+                # slow path
+                jit_key, tuning_key = keys
+                jit_instance = self._get_jit_instance(target, jit_key)
 
-        if compiled_func is None:
-            # slow path
-            if keys is None:
-                keys = extract_keys(args, self.const_params, self.tuning_params)
-            jit_key, tuning_key = keys
-            jit_instance: Optional[JitInstance] = self.jit_instances.get(jit_key, None)
-            if jit_instance is None:
-                jit_instance = JitInstance(self.script_cls, self.params, self.build_options, self.schedules, jit_key)
-                self.jit_instances[jit_key] = jit_instance
-
-            compiled_program = jit_instance._pick_best_program(args)
-            compiled_func = compiled_program.get_launch_func()
-            self.dispatch_table[(jit_key, tuning_key)] = compiled_func
-
-        self._last_dispatch_args = dispatch_args
-        self._last_compiled_func = compiled_func
+                compiled_program = jit_instance._pick_best_program(args)
+                compiled_func = compiled_program.get_launch_func()
+                table[keys] = compiled_func
+            self._last_dispatch = (target, table, dispatch_args, compiled_func)
 
         # call the compiled function
         kernel_args = self._kernel_arg_getter(args)
@@ -1098,10 +1111,13 @@ def _jit_instance_for(self, *args: Any, **kwargs: Any) -> JitInstance:
                 )
 
         jit_key, _ = extract_keys(args, self.const_params, self.tuning_params)
-        jit_instance: Optional[JitInstance] = self.jit_instances.get(jit_key, None)
+        return self._get_jit_instance(get_current_target(), jit_key)
+
+    def _get_jit_instance(self, target: Target, jit_key: JitKey) -> JitInstance:
+        jit_instance: Optional[JitInstance] = self.jit_instances.get((target, jit_key), None)
         if jit_instance is None:
             jit_instance = JitInstance(self.script_cls, self.params, self.build_options, self.schedules, jit_key)
-            self.jit_instances[jit_key] = jit_instance
+            self.jit_instances[(target, jit_key)] = jit_instance
         return jit_instance
 
 

tests/lang/test_autotune.py

@@ -33,6 +33,7 @@
     span_space,
     tuning_metadata_matches,
 )
+from tilus.target import get_current_target, nvgpu_sm90, scope
 
 
 def test_span_space_single_and_grouped_keys():
@@ -106,15 +107,16 @@ def launch(*args):
     script.jit_instances = {}
     script._dispatch_arg_getter = _make_tuple_getter([2, 1])
     script._kernel_arg_getter = _make_tuple_getter([0, 1])
-    script._last_dispatch_args = None
-    script._last_compiled_func = None
+    script.dispatch_tables = {}
+    script._last_dispatch = (None, None, None, None)
 
     ptr = object()
     keys = extract_keys((ptr, 128, 7), script.const_params, script.tuning_params)
-    script.dispatch_table = {keys: launch}  # type: ignore[dict-item]
-    script(ptr, 128, 7)
-    script.dispatch_table.clear()
-    script(ptr, 128, 7)
+    with scope(nvgpu_sm90):
+        script.dispatch_tables[get_current_target()] = {keys: launch}  # type: ignore[dict-item]
+        script(ptr, 128, 7)
+        script.dispatch_tables[get_current_target()].clear()
+        script(ptr, 128, 7)
 
     assert calls == [(ptr, 128), (ptr, 128)]
 

tests/lang/test_compile_target.py

@@ -0,0 +1,171 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+import tilus
+import tilus.lang.instantiated_script as instantiated_script
+import tilus.target
+import torch
+from tilus import int32
+from tilus.lang.instantiated_script import InstantiatedScript
+from tilus.target import Target, get_current_target, nvgpu_sm80, nvgpu_sm90, nvgpu_sm103a, scope
+from tilus.testing import requires
+
+
+def target_marker() -> int:
+    """Encode the current target, evaluated while the script is transpiled."""
+    major, minor = get_current_target().properties.compute_capability
+    return major * 1000 + minor * 10 + (0 if get_current_target().properties.feature_suffix is None else 1)
+
+
+class TargetExample(tilus.Script):
+    def __call__(self, size: int32, out: ~int32):
+        self.attrs.blocks = 1
+        self.attrs.warps = 1
+
+        marker = target_marker()
+        g_out = self.global_view(out, dtype=int32, shape=[1])
+        with self.thread_group(0, num_threads=1):
+            self.store_global(g_out, self.register_tensor(dtype=int32, shape=[1], init=lambda i: marker), offsets=[0])
+
+
+class CompileExample(tilus.Script):
+    def __call__(self, out: ~int32):
+        self.attrs.blocks = 1
+        self.attrs.warps = 1
+
+        g_out = self.global_view(out, dtype=int32, shape=[1])
+        with self.thread_group(0, num_threads=1):
+            self.store_global(g_out, self.register_tensor(dtype=int32, shape=[1], init=lambda i: 1), offsets=[0])
+
+
+def test_jit_instance_per_target():
+    script = InstantiatedScript(CompileExample, (), {})
+
+    with scope(nvgpu_sm80):
+        first = script._jit_instance_for(None)
+    with scope(nvgpu_sm90):
+        second = script._jit_instance_for(None)
+    with scope(nvgpu_sm80):
+        again = script._jit_instance_for(None)
+
+    assert first is not second
+    assert first is again
+
+
+def test_script_cache_dir_per_target():
+    # Transpiling only, no nvcc or GPU needed. The program text is the same on both targets, the cache dir must not be.
+    script = InstantiatedScript(CompileExample, (), {})
+
+    with scope(nvgpu_sm80):
+        first = script._jit_instance_for(None)
+    with scope(nvgpu_sm90):
+        second = script._jit_instance_for(None)
+
+    assert [str(program) for program in first.transpiled_programs] == [
+        str(program) for program in second.transpiled_programs
+    ]
+    assert first.cache_dir != second.cache_dir
+
+
+def test_compile_builds_each_target(monkeypatch):
+    # Compiling needs nvcc but no GPU.
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    script = InstantiatedScript(CompileExample, (), {})
+
+    with scope(nvgpu_sm80):
+        first = script.compile(None)
+    with scope(nvgpu_sm90):
+        second = script.compile(None)
+
+    first_dir = first.compiled_programs[0].program_dir
+    second_dir = second.compiled_programs[0].program_dir
+    assert first_dir != second_dir
+    assert "sm80" in (first_dir / "options.txt").read_text()
+    assert "sm90" in (second_dir / "options.txt").read_text()
+
+
+def test_call_dispatches_on_target(monkeypatch):
+    picks = []
+    launches = []
+
+    class FakeProgram:
+        def __init__(self, target):
+            self.target = target
+
+        def get_launch_func(self):
+            return lambda *args: launches.append(self.target)
+
+    class FakeJitInstance:
+        def __init__(self, *args):
+            self.target = get_current_target()
+
+        def _pick_best_program(self, args):
+            picks.append((self.target, args[0]))
+            return FakeProgram(self.target)
+
+    monkeypatch.setattr(instantiated_script, "JitInstance", FakeJitInstance)
+    script = InstantiatedScript(TargetExample, (), {})
+
+    out = object()
+    with scope(nvgpu_sm90):
+        script(128, out)
+        script(128, out)  # the last-launch shortcut
+    with scope(nvgpu_sm103a):
+        script(128, out)  # same arguments as the shortcut, another target
+        script(256, out)
+        script(128, out)  # the dispatch table of this target
+    with scope(nvgpu_sm90):
+        script(128, out)  # back to the first target, without building again
+
+    assert launches == [nvgpu_sm90, nvgpu_sm90, nvgpu_sm103a, nvgpu_sm103a, nvgpu_sm103a, nvgpu_sm90]
+    assert picks == [(nvgpu_sm90, 128), (nvgpu_sm103a, 128), (nvgpu_sm103a, 256)]
+
+
+def compatible_target() -> Target:
+    """Find another target whose binaries run on this GPU. A binary only runs on its own major version."""
+    current = get_current_target()
+    candidates = [
+        target
+        for target in vars(tilus.target).values()
+        if isinstance(target, Target)
+        and target != current
+        and target.kind == current.kind
+        and target.properties.feature_suffix is None
+        and target.properties.compute_capability[0] == current.properties.compute_capability[0]
+        and current.supports(target)
+    ]
+    if not candidates:
+        pytest.skip("No other target can run on {}".format(current))
+    return max(candidates, key=lambda target: target.properties.compute_capability)
+
+
+@requires.nvgpu_sm80
+def test_call_runs_binary_of_target():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    script = InstantiatedScript(TargetExample, (), {})
+    out = torch.zeros(1, dtype=torch.int32, device=device)
+
+    script(128, out)
+    assert out.item() == target_marker()
+    expected = out.item()
+
+    target = compatible_target()
+    with scope(target):
+        script(128, out)
+        assert out.item() == target_marker()
+        assert out.item() != expected
+
+    script(128, out)
+    assert out.item() == expected