Skip to content

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

Description

@bigSheep123

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions