Skip to content

fix(passes): carry the Module's Target into the lowered TIR - #65

Open
bigSheep123 wants to merge 1 commit into
tile-ai:mainfrom
bigSheep123:fix/module-target-propagation
Open

fix(passes): carry the Module's Target into the lowered TIR#65
bigSheep123 wants to merge 1 commit into
tile-ai:mainfrom
bigSheep123:fix/module-target-propagation

Conversation

@bigSheep123

Copy link
Copy Markdown
Contributor

fix(passes): carry the Module's Target into the lowered TIR

Fixes #63.

What was wrong

HirToTirPass built every PrimFunction without target=, so the Target
declared by the owning Module was discarded and
field(default_factory=default_target) — cuda / sm_90 — took over silently.

Because group_functions_by_target keys on fn.target.name, the failure is not
just a wrong SM arch but a wrong backend: an amx Module was routed to the CUDA
emitter.

declared on Module : amx   apple_amx
landed on TIR      : cuda  sm_90        # before
  codegen groups as: ['cuda']

declared on Module : amx   apple_amx
landed on TIR      : amx   apple_amx    # after
  codegen groups as: ['amx']

The change

Resolve the target once in run() with Module.resolve_target() — it 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, for static bodies and mangled dispatch variants
  • _build_dispatch_entry, for the dispatch entry, so an entry and the variants
    it 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 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 around resolve_target() matches the idiom
ir/core/module.py::_reentered already uses for an undeclared target.

Tests

Three tests in tests/passes/test_hir_to_tir.py, which had no coverage of target
propagation before — the reason this could regress unnoticed:

  1. the declared Target reaches a static body, every mangled dispatch variant,
    and the dispatch entry;
  2. a child Module lowers against the target it inherits from its root;
  3. a Module declaring no target still lowers, on PrimFunction's default.

They assert object identity, not equality: default_target() returns a fresh
equal 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

pytest tests/ -q     644 passed (main @ e55e4b0)  ->  647 passed (+3 new)

Every .pre-commit-config.yaml hook 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.

`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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

HirToTirPass discards the Module's Target, so a Module can be compiled for the wrong backend

1 participant