Skip to content

feat(ir): represent a scheduler-produced thread execution domain #135

Description

@zhen8838

Context

The atom schedule search (external prototype, docs/plans/atom-schedule-search/) decides the
thread-level execution structure of an op: which warps issue a TMA, which warpgroup runs a
WGMMA, and whether those domains overlap. The CTA-level domain is given by the input program;
the thread level is the search's output.

Reading that input works today. An HIR matmul authored inside a CTA Mesh carries the scope:

with Mesh(("cta",), layout=(128,), names=("m",)) as cta:
    xs = tf.reshard(x, (4096 @ cta.m, 5120), "gmem")
    return tf.matmul(xs, w)
# get_metadata(call, ExecutionDomainMetadata).at("cta") -> Mesh(Topology('cta',128), ...)

Writing the decision back has no home.

Observation

  • src/tilefoundry/ir/core/metadata.py:18 states ExecutionDomainMetadata is "a fact about
    where it was written". A scheduler-chosen thread domain was never authored, so recording
    it there contradicts the documented contract rather than extending it.
  • src/tilefoundry/analysis/check.py:264 consumes it to separate execution placement from
    result layout, and rejects an occurrence with no domain at the selected level. The same check
    is what a scheduled program must eventually satisfy at thread level.
  • src/tilefoundry/ir/tir/stmts.py:67 MeshScope(mesh, binding, body) is the only construct
    carrying a concrete Mesh as a scope, but it is TIR-only and is built during lowering
    (src/tilefoundry/passes/transforms/hir_to_tir.py:1641), i.e. after the point where the
    schedule decision is made.
  • Result layout cannot substitute: the metadata docstring says so directly — a value may be laid
    out across threads while the work producing it ran on one CTA.

Request

Decide where a scheduler-produced thread-level execution domain lives, and how it reaches
_execution_placement() so a scheduled program can be checked the same way an authored one is.

Options worth weighing: a distinct derived metadata separate from the authored stack; extending
ExecutionDomainMetadata with provenance (authored vs derived); or a scheduling-stage scope
construct that is neither the TIR MeshScope nor a parser artifact.

Related to #113, which asks whether the authored representation is more state than necessary.
This issue is the other direction: the authored fact alone does not cover a decision the
scheduler makes. Whichever way #113 resolves, the derived case still needs an answer.

No change is requested in this repository right now — the prototype is external and records the
gap in its own backport list. Filing so the contract question is not lost.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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