HirToTirPass discards the Module's Target, so a Module can be compiled for the wrong backend
Affects: main @ e55e4b0 (reproduced on the current tip)
Summary
HirToTirPass constructs every PrimFunction without passing target=, so the
Target declared by the owning Module never reaches the lowered TIR.
PrimFunction.target's field(default_factory=default_target) (cuda / sm_90 /
h200) silently takes over.
The consequence is not limited to a wrong SM arch. group_functions_by_target
keys on fn.target.name, so a Module declaring a non-CUDA backend is handed to
the CUDA emitter.
Reproduction
Only the Module's target= varies; the function, params and body are identical.
from tilefoundry.codegen.registry import group_functions_by_target
from tilefoundry.ir.core import Var
from tilefoundry.ir.core.module import Module
from tilefoundry.ir.hir.function import Function as HirFunction
from tilefoundry.ir.types import DType, TensorType
from tilefoundry.passes.transforms import HirToTirPass
from tilefoundry.target.amx.target import AmxTarget
from tilefoundry.target.cuda.target import CudaTarget
for declared in (CudaTarget("nvidia.h200_sxm"), AmxTarget()):
ty = TensorType(shape=(8,), dtype=DType.f32, layout=None, storage="gmem")
x = Var(type=ty, name="x")
fn = HirFunction.build(name="k", params=(x,), body=x, return_type=ty)
module = Module(name="m", functions=(fn,), entry="k", target=declared)
lowered = HirToTirPass().run(module)
pf = lowered.functions[0]
print(f"declared on Module : {module.resolve_target().name:5s} "
f"{module.resolve_target().architecture.name}")
print(f"landed on TIR : {pf.target.name:5s} {pf.target.architecture.name}")
print(f" propagated? : {pf.target is declared}")
print(f" codegen groups as: {list(group_functions_by_target(lowered))}\n")
Actual
declared on Module : cuda sm_90
landed on TIR : cuda sm_90
propagated? : False # a fresh default, not the declared value
codegen groups as: ['cuda']
declared on Module : amx apple_amx
landed on TIR : cuda sm_90 # backend silently swapped
propagated? : False
codegen groups as: ['cuda'] # routed to the CUDA emitter
Expected
The declared Target reaches the lowered PrimFunction, and an amx Module
groups as ['amx'].
Notes on detecting this
Both CUDA lines print sm_90 because the shipped hardware corpus has only
nvidia.sm90; the amx case is what makes the loss visible. For the CUDA case
the loss cannot be observed with == at all — default_target() returns a
fresh equal value on each call, so only object identity distinguishes
"propagated" from "defaulted".
That is also why this has no runtime symptom: codegen builds for the fallback
arch and the driver PTX-JITs the result, so the kernel still computes correct
answers. There is currently no test asserting target propagation, which is what
allowed it to go unnoticed.
Root cause
Two sites in passes/transforms/hir_to_tir.py build PrimFunction with no
target=:
_lower_function — one per static body and per mangled dispatch variant
_build_dispatch_entry — the dispatch entry
Since 45bd840 ("give execution context and hardware facts one owner each",
#32) the Target is owned by Module (with owner-chain inheritance via
Module.resolve_target()), and hir_to_tir never consults it.
Suggested fix
Resolve once in HirToTirPass.run via module.resolve_target() and thread it
into both sites. A Module declaring no target anywhere should leave the kwarg
off rather than pass None, so PrimFunction's own default stays
single-sourced and group_functions_by_target's no-resolved-Target check is not
tripped. Resolving one target per Module also satisfies that function's existing
requirement that a Module's CUDA functions carry identical Target facts.
PR follows.
HirToTirPass discards the Module's Target, so a Module can be compiled for the wrong backend
Affects:
main@e55e4b0(reproduced on the current tip)Summary
HirToTirPassconstructs everyPrimFunctionwithout passingtarget=, so theTarget declared by the owning
Modulenever reaches the lowered TIR.PrimFunction.target'sfield(default_factory=default_target)(cuda / sm_90 /h200) silently takes over.
The consequence is not limited to a wrong SM arch.
group_functions_by_targetkeys on
fn.target.name, so a Module declaring a non-CUDA backend is handed tothe CUDA emitter.
Reproduction
Only the Module's
target=varies; the function, params and body are identical.Actual
Expected
The declared Target reaches the lowered
PrimFunction, and anamxModulegroups as
['amx'].Notes on detecting this
Both CUDA lines print
sm_90because the shipped hardware corpus has onlynvidia.sm90; theamxcase is what makes the loss visible. For the CUDA casethe loss cannot be observed with
==at all —default_target()returns afresh equal value on each call, so only object identity distinguishes
"propagated" from "defaulted".
That is also why this has no runtime symptom: codegen builds for the fallback
arch and the driver PTX-JITs the result, so the kernel still computes correct
answers. There is currently no test asserting target propagation, which is what
allowed it to go unnoticed.
Root cause
Two sites in
passes/transforms/hir_to_tir.pybuildPrimFunctionwith notarget=:_lower_function— one per static body and per mangled dispatch variant_build_dispatch_entry— the dispatch entrySince
45bd840("give execution context and hardware facts one owner each",#32) the Target is owned by
Module(with owner-chain inheritance viaModule.resolve_target()), andhir_to_tirnever consults it.Suggested fix
Resolve once in
HirToTirPass.runviamodule.resolve_target()and thread itinto both sites. A Module declaring no target anywhere should leave the kwarg
off rather than pass
None, soPrimFunction's own default stayssingle-sourced and
group_functions_by_target's no-resolved-Target check is nottripped. Resolving one target per Module also satisfies that function's existing
requirement that a Module's CUDA functions carry identical Target facts.
PR follows.