Repository navigation
[Fix] Compile and dispatch a script per target - #173
Merged
Merged
Conversation
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 <dezhen.lu@student.uni-tuebingen.de>
…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 <noreply@anthropic.com>
Member
|
/ok to test fab36d2 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
tilus.target.scopehas no effect on a script that has already been compiled or launched. Under another target,compile()and__call__return the first target's binary. This PR gives each target its own JIT instance and launch function.Why
Script(...)returns a cachedInstantiatedScript, which keeps its compiled state in three places. None of them includes the target:jit_instances, keyed byjit_keydispatch_table, keyed by(jit_key, tuning_key)__call__, which compares the raw call argumentsSo after compiling for
sm_90a,compile()underscope(nvgpu_sm103a)returns thesm_90abuild. This matters for the documented use ofcompile()in CI, and forrequires.*, which compiles underscope(...)in the same process.Adding the target only to the
jit_instanceskey is not enough for__call__: the shortcut and the dispatch table would still return the old function.How
jit_instancesis keyed by(target, jit_key).dispatch_tablesholds one table per target, anddispatch_tablestays the table of the current target.__call__compares the current target with the last one by identity. When it changes,__call__switches the table and drops the shortcut.The tiers below that check are unchanged and do not hash the target, because
Targetis a nested frozen dataclass and hashing it on every launch would be wasted.tests/lang/test_autotune.pybuilds the script by hand, so it now sets the last target and runs underscope.Testing
Host tests ran on the Grace CPU of a GH200: GCC 13.3, CUDA 13.2, PyTorch 2.14.1+cu130, tilus
mainat f23018a.tests/lang/test_compile_target.pyis new:maintest_jit_instance_per_targettest_compile_builds_each_target(sm_80, sm_90)test_call_dispatches_on_target: sm_90 → sm_103a → sm_90, stubbed JIT instancetest_call_runs_binary_of_target: the kernel writes a marker of the target it was transpiled for; default target, then a compatible base target, then default againcompile()underscopein one process with this PR, reading the architecture from the built library withcuobjdump -lelf:sm_90asm_90asm_103asm_103asm_90asm_90a(same program dir as 1)sm_90sm_90sm_103sm_103Cost of
__call__with a no-op launch, median ns per call over 400k calls, two runs, pinned to one Grace core:mainThat is about 30 ns for the target check on the shortcut path. The second run of the alternating case on this PR looks noisy, since its minimum was 725 ns. Reading the module global directly instead of calling
get_current_target()was no faster (medians 358 to 424 ns against 349 to 355 ns), so the PR keeps the call. Real launches take far longer than this.GPU tests ran on GH200 (
sm_90, CUDA 13.2). Not run on B300.The rest of the CPU-runnable suite (
tests/lang,tests/ir,tests/drivers,tests/transformsand the top-level tests) gives the same result onmainand here: the same 5 failed because they need a GPU. Passed: 192 onmain, 196 here with the 4 new tests, of which 3 run without a GPU and the GPU test compiles only.ruff checkandruff format --check(0.11.0 frompyproject.toml, 0.12.3 from pre-commit),mypy(1.15.0) and the copyright check pass on the changed files.Not in this PR
The script-level cache directory is keyed by program text and build options, not by target. When two targets produce identical programs, they share
dispatch_table.jsonand theprograms/symlinks. The binaries are still separate. The saved metadata already includes the target, so a table from another target is ignored and the kernel is re-tuned rather than misdispatched. The only cost is repeated tuning when one cache alternates between targets.