From 7da03a7bc08b69b0507185e7b3fa0f3a81b27cce Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sat, 7 Nov 2026 12:52:53 +0100 Subject: [PATCH 1/2] [Fix] Compile and dispatch a script per target InstantiatedScript cached its JIT instances and launch functions by call arguments only. After compiling or running under one target, a call under scope(other_target) reused the first binary, so compiling for sm_103a after sm_90a returned the sm_90a build. Key the JIT instances by target as well and keep one dispatch table per target. The last-launch shortcut stays valid until the target changes, so a repeated launch only adds an identity check of the current target. Signed-off-by: 0z5a --- python/tilus/lang/instantiated_script.py | 27 ++-- tests/lang/test_autotune.py | 9 +- tests/lang/test_compile_target.py | 156 +++++++++++++++++++++++ 3 files changed, 181 insertions(+), 11 deletions(-) create mode 100644 tests/lang/test_compile_target.py diff --git a/python/tilus/lang/instantiated_script.py b/python/tilus/lang/instantiated_script.py index ad184935..80f6d38e 100644 --- a/python/tilus/lang/instantiated_script.py +++ b/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 @@ -1006,8 +1006,11 @@ 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.jit_instances: dict[tuple[Target, JitKey], JitInstance] = {} + # One table per target, dispatch_table is the one of _last_target. + self.dispatch_tables: dict[Target, dict[tuple[JitKey, TuningKey], tvm_ffi.Function]] = {} self.dispatch_table: dict[tuple[JitKey, TuningKey], tvm_ffi.Function] = {} + self._last_target: Optional[Target] = None 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 @@ -1025,6 +1028,14 @@ def __call__(self, *args, **kwargs): "The number of arguments should be {}, but got {}.".format(len(self.params.param_names), len(args)) ) + # A binary belongs to the target it was built for. Switching targets is rare, so the launches below only pay + # for this identity check. + target = get_current_target() + if target is not self._last_target: + self.dispatch_table = self.dispatch_tables.setdefault(target, {}) + self._last_target = target + self._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) @@ -1040,10 +1051,7 @@ def __call__(self, *args, **kwargs): 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 + jit_instance = self._get_jit_instance(target, jit_key) compiled_program = jit_instance._pick_best_program(args) compiled_func = compiled_program.get_launch_func() @@ -1098,10 +1106,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 diff --git a/tests/lang/test_autotune.py b/tests/lang/test_autotune.py index 15f6fae3..34ae3559 100644 --- a/tests/lang/test_autotune.py +++ b/tests/lang/test_autotune.py @@ -33,6 +33,7 @@ span_space, tuning_metadata_matches, ) +from tilus.target import nvgpu_sm90, scope def test_span_space_single_and_grouped_keys(): @@ -108,13 +109,15 @@ def launch(*args): script._kernel_arg_getter = _make_tuple_getter([0, 1]) script._last_dispatch_args = None script._last_compiled_func = None + script._last_target = nvgpu_sm90 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(ptr, 128, 7) + script.dispatch_table.clear() + script(ptr, 128, 7) assert calls == [(ptr, 128), (ptr, 128)] diff --git a/tests/lang/test_compile_target.py b/tests/lang/test_compile_target.py new file mode 100644 index 00000000..07516f66 --- /dev/null +++ b/tests/lang/test_compile_target.py @@ -0,0 +1,156 @@ +# 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_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 From fab36d2b14a0827d82e73d643f9c44aafea3dae3 Mon Sep 17 00:00:00 2001 From: Yaoyao Ding Date: Mon, 9 Nov 2026 12:23:31 -0500 Subject: [PATCH 2/2] Address review: pin the dispatch table per call and key script cache dirs by target - Resolve the target and its dispatch table once at the start of InstantiatedScript.__call__, and keep the last launch as a single (target, table, args, func) tuple, so a target change during a call cannot record a binary under another target or mix fast-path state. - Include the target in the scripts/ cache-dir hash, so two targets with the same program text no longer share programs/ links and dispatch_table.json. Bump SPECIALIZATION_CACHE_VERSION so manifests that point at the old shared dirs are not reused. - Test that the script cache dir differs per target, and update the hand-built script in test_autotune to the new attributes. Co-Authored-By: Claude Opus 5.5 --- python/tilus/lang/instantiated_script.py | 67 +++++++++++++----------- tests/lang/test_autotune.py | 11 ++-- tests/lang/test_compile_target.py | 15 ++++++ 3 files changed, 56 insertions(+), 37 deletions(-) diff --git a/python/tilus/lang/instantiated_script.py b/python/tilus/lang/instantiated_script.py index 80f6d38e..4206936a 100644 --- a/python/tilus/lang/instantiated_script.py +++ b/python/tilus/lang/instantiated_script.py @@ -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 @@ -1007,14 +1018,13 @@ def __init__(self, script_cls: Type[Script], script_args: Sequence[Any], script_ self.tuning_params: list[int] = self.params.tuning_params self.jit_instances: dict[tuple[Target, JitKey], JitInstance] = {} - # One table per target, dispatch_table is the one of _last_target. + # 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_table: dict[tuple[JitKey, TuningKey], tvm_ffi.Function] = {} - self._last_target: Optional[Target] = None 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: @@ -1028,37 +1038,32 @@ def __call__(self, *args, **kwargs): "The number of arguments should be {}, but got {}.".format(len(self.params.param_names), len(args)) ) - # A binary belongs to the target it was built for. Switching targets is rare, so the launches below only pay - # for this identity check. + # 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() - if target is not self._last_target: - self.dispatch_table = self.dispatch_tables.setdefault(target, {}) - self._last_target = target - self._last_dispatch_args = None + 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) - - 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 = self._get_jit_instance(target, jit_key) - - 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_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) + + 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) diff --git a/tests/lang/test_autotune.py b/tests/lang/test_autotune.py index 34ae3559..b4b9ff01 100644 --- a/tests/lang/test_autotune.py +++ b/tests/lang/test_autotune.py @@ -33,7 +33,7 @@ span_space, tuning_metadata_matches, ) -from tilus.target import nvgpu_sm90, scope +from tilus.target import get_current_target, nvgpu_sm90, scope def test_span_space_single_and_grouped_keys(): @@ -107,16 +107,15 @@ 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._last_target = nvgpu_sm90 + 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] with scope(nvgpu_sm90): + script.dispatch_tables[get_current_target()] = {keys: launch} # type: ignore[dict-item] script(ptr, 128, 7) - script.dispatch_table.clear() + script.dispatch_tables[get_current_target()].clear() script(ptr, 128, 7) assert calls == [(ptr, 128), (ptr, 128)] diff --git a/tests/lang/test_compile_target.py b/tests/lang/test_compile_target.py index 07516f66..3e372a80 100644 --- a/tests/lang/test_compile_target.py +++ b/tests/lang/test_compile_target.py @@ -64,6 +64,21 @@ def test_jit_instance_per_target(): 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)