fix(passes): carry the Module's Target into the lowered TIR - #65
Open
bigSheep123 wants to merge 1 commit into
Open
fix(passes): carry the Module's Target into the lowered TIR#65bigSheep123 wants to merge 1 commit into
bigSheep123 wants to merge 1 commit into
Conversation
`HirToTirPass` builds every `PrimFunction` without passing `target=`, so a
Module's declared Target is discarded and `PrimFunction`'s
`field(default_factory=default_target)` — cuda / sm_90 — silently takes over.
The visible consequence is not a wrong SM arch but a wrong *backend*:
`group_functions_by_target` keys on `fn.target.name`, so an `amx` Module is
routed to the CUDA emitter.
declared on Module : amx apple_amx
landed on TIR : cuda sm_90 # before
codegen groups as: ['cuda']
Resolve the target once in `run()` through `Module.resolve_target()`, which
walks the owner chain so a child Module selected out of a tree still gets the
root's target, and thread it into both construction sites — `_lower_function`
(static bodies and mangled dispatch variants) and `_build_dispatch_entry` (the
dispatch entry), so an entry and its callees cannot disagree on a backend.
A Module declaring no target anywhere leaves the kwarg off rather than passing
None: that keeps the choice of fallback single-sourced in `PrimFunction`, and
storing None would trip `group_functions_by_target`'s own no-resolved-Target
check. Resolving one target per Module also satisfies that function's existing
rule that a Module's cuda functions all carry identical Target facts.
`try/except ValueError` matches the idiom `ir/core/module.py::_reentered`
already uses for an undeclared target.
Tested by identity, not equality: `default_target()` returns a fresh equal
value per call, so `==` cannot tell "propagated" from "defaulted". This is also
why the bug is invisible at runtime — the fallback arch still PTX-JITs to a
correct-answering kernel — so the three tests assert propagation for a static
body, dispatch variants and the dispatch entry; inheritance by a child Module;
and that a target-less Module still lowers on the default.
pytest tests/ -q: 644 passed before, 647 after (+3 new, no regressions).
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.
fix(passes): carry the Module's Target into the lowered TIR
Fixes #63.
What was wrong
HirToTirPassbuilt everyPrimFunctionwithouttarget=, so the Targetdeclared by the owning
Modulewas discarded andfield(default_factory=default_target)— cuda / sm_90 — took over silently.Because
group_functions_by_targetkeys onfn.target.name, the failure is notjust a wrong SM arch but a wrong backend: an
amxModule was routed to the CUDAemitter.
The change
Resolve the target once in
run()withModule.resolve_target()— it walks theowner chain, so a child Module selected out of a tree still gets the root's
target — and thread it into both construction sites:
_lower_function, for static bodies and mangled dispatch variants_build_dispatch_entry, for the dispatch entry, so an entry and the variantsit dispatches to cannot end up on different backends
A Module that declares no target anywhere leaves the kwarg off rather than
passing
None. Two reasons: it keeps the choice of fallback single-sourced inPrimFunction, and storingNonewould tripgroup_functions_by_target's own"no resolved Target" check. Resolving one target per Module also satisfies that
function's existing rule that a Module's CUDA functions all carry identical
Target facts.
try/except ValueErroraroundresolve_target()matches the idiomir/core/module.py::_reenteredalready uses for an undeclared target.Tests
Three tests in
tests/passes/test_hir_to_tir.py, which had no coverage of targetpropagation before — the reason this could regress unnoticed:
and the dispatch entry;
PrimFunction's default.They assert object identity, not equality:
default_target()returns a freshequal value per call, so
==cannot distinguish "propagated" from "defaulted".Each test was confirmed to fail against the unfixed pass before being kept.
This is also why no runtime test can catch the bug — the fallback arch still
PTX-JITs to a kernel that returns correct answers.
Verification
Every
.pre-commit-config.yamlhook clean: ruff, spec-rules, spec-refs,spec-entropy, forward-references, comment-hygiene, no-machine-paths,
english-only, and clang-format over all 29 C++ files.