Skip to content

fix(codegen): emit effect-form Cast, Fill through kernel-param views, and MeshCoord/Dim-arith slice coordinates - #240

Merged
zhen8838 merged 5 commits into
tile-ai:mainfrom
sablin39:fix/tir-cuda-emitter-fixes
Oct 19, 2026
Merged

zhen8838 merged 5 commits into
tile-ai:mainfrom
sablin39:fix/tir-cuda-emitter-fixes

Conversation

@sablin39

Copy link
Copy Markdown
Contributor

What

Four CUDA TIR emission gaps that fire on scheduled programs finalized
against the composed cta×thread mesh (the shape every multi-CTA
schedule finalize produces). One commit, src/tilefoundry/codegen/cuda/tir
only, +45/−10.

  1. arith.py — effect-form Cast emitter. Cast in effect form had
    no registered EMIT handler, so any scheduled program whose TIR carries
    one raised "no codegen handler". The store into the destination
    converts, so the elementwise functor is the identity — the same shape
    as the existing UnaryKind.CAST path, dynamic-extent clip included.
  2. fill.py — Fill on a kernel parameter. Fill named the raw
    pointer for a kernel-param destination while tensor_view establishes
    the <name>_tensor view for params; the emitted kernel referenced an
    undeclared name and did not compile. Now routes through
    ctx.is_kernel_param, matching tensor_view's naming.
  3. tensor_view.py — slice pointer offsets through _coord_ref.
    PtrOf(Slice(...)) offsets were rendered with render_scalar_expr,
    which cannot resolve mesh-relative Dim arithmetic; slice coordinates on
    Dim expressions failed or mis-rendered. _coord_ref resolves them
    against the mesh binding. Also teaches _COORD_OPERATORS
    DimFloorDiv/DimMod and renders MeshCoord call targets as
    cute::get<axis>(cute::flatten(mesh_coords(...))).
  4. tensor_view.py — named refusal for unknown coordinate calls.
    An unsupported Call target in a coordinate previously fell through to
    a generic leaf path; it now raises a NotImplementedError naming the
    target type.

May close part of #230 (sliced pointers / nested layouts) — the slice
coordinate half overlaps; happy to rebase if that lands first.

Validation

  • Upstream test suite at aaa8054 with this commit: 1802 passed,
    11 skipped
    (pytest tests -n 4). The single failure in that run,
    tests/cli/test_cli.py::test_target_list_expressions_execute_and_show_accepts_their_identities,
    is an artifact of the development environment (a locally registered
    custom target in the CLI registry); it passes with a pristine registry
    and is unrelated to this change.
  • The four gaps were hit and the fixes first exercised end-to-end while
    driving a scheduled RMSNorm through scaffold/compile() on sm_90.

Not in scope

  • CopyAsyncTensor (TMA), WGMMA (TiledMma non-SM80 atoms), and Where
    emitters remain absent as documented in
    docs/spec/codegen.md
    and
    docs/spec/tir.md
    — the schedule tutorial's optimal GEMM still cannot compile. Deliberate
    raises today; unchanged here.
  • A separate runtime bug: the reduce dispatch mis-counts
    warps_per_group for a composed cta×thread mesh (single is_thread
    per mesh in
    reduce_common_impl.h,
    so thread axes inside a cta-scoped composed mesh count their full
    extent as warps → OOB smem reads on an RMSNorm repro). Not touched by
    this PR; will file separately with the repro.

… and MeshCoord/Dim-arith slice coordinates

Four CUDA TIR emission gaps hit by every scheduled program that finalizes
against the composed cta x thread mesh:

- arith: register an EMIT handler for tir Cast (effect form). The store into
  the destination converts, so the elementwise functor is the identity, as
  for UnaryKind.CAST. Previously any effect-form Cast raised 'no codegen
  handler'.
- fill: write through the <name>_tensor view when the destination is a
  kernel parameter, matching the naming tensor_view establishes. Previously
  Fill on a kernel param named the raw pointer and the emitted kernel did
  not compile.
- tensor_view: price slice pointer offsets with _coord_ref instead of
  render_scalar_expr so Dim arithmetic resolves against mesh coordinates;
  teach _COORD_OPERATORS DimFloorDiv/DimMod; render MeshCoord call targets
  as cute::get<axis>(cute::flatten(mesh_coords(...))); refuse unknown Call
  targets with a named error instead of falling through.

Validated against upstream's test suite at aaa8054 (tests/codegen: 27
passed).

@zhen8838 zhen8838 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I’m sorry that I haven‘t had time to do much work on CodeGen and Runtime recently. I really want to further improve the runtime of tilefoundry to the extent that it can be used as a distributed runtime lib. So thank you very much for your contribution to the tilefoundry project. Help us take a step closer to that direction.

sablin39 and others added 4 commits October 17, 2026 00:02
Register effect-form Cast on the Unary emitter as UnaryKind.CAST instead of a
copied body, which also clears the D205/D209 docstring lint that failed CI.
Fill names its destination through tensor_view's _tensor_ref, and the
MeshCoord rendering is one helper shared by the coordinate and scalar
renderers. A CUDA test narrows a mesh-coordinate window with T.cast and fills
a kernel parameter; it fails on main for each of the three fixes.
@zhen8838
zhen8838 merged commit a5501c9 into tile-ai:main Oct 19, 2026
1 check passed
sablin39 added a commit to sablin39/tilelang-cuda-skills that referenced this pull request Oct 20, 2026
…iagnostics fixed

Upstream main moved aaa8054 -> 25fbbaa (6 commits). Notables:
- Our four emitter fixes (slice coords, MeshCoord, Cast, Fill on kernel
  params) merged upstream as a5501c9 (PR tile-ai/TileFoundry#240).
- I-3 fixed via #241 (closes #237): schedule finalize errors are now
  actionable -- they name the site, list instruction candidates, and
  point at tf.schedule. Skill guidance updated to the new error shape.
- #244 adds AMX (Intel CPU) support; harmless for this GPU suite.

Still open and reproducing byte-identically at the new pin: #239
(I-1/I-2 predicted-ns billing artifacts -- trust-table numbers
unchanged), TMA/WGMMA/Where emitter gaps, warps_per_group
reduce-dispatch bug, I-5 example bit-rot.

Verification at the new pin: probe harness 12/12 PASS (two finalize
fingerprints updated to the actionable error text); pin ritual green
(tutorial extraction + analyze ideal-ns=13, scan twin check PASS both
outputs, tests/ops/tir/test_copy 1 passed); rms_norm_quant TileLang
twin check PASS (equal on fp8, allclose on f32 scale); trust-table
spot checks byte-identical (scan fused ideal-ns=125063, GEMM scheduled
ideal-ns=143256 bound-by=compute).
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.

2 participants