Skip to content

feat(cuda): declare CUDA MMA atoms and async copy instructions - #187

Merged
zhen8838 merged 19 commits into
tile-ai:mainfrom
zhen8838:feat/atomsched-stage2b-cuda-instructions
Sep 26, 2026
Merged

zhen8838 merged 19 commits into
tile-ai:mainfrom
zhen8838:feat/atomsched-stage2b-cuda-instructions

Conversation

@zhen8838

@zhen8838 zhen8838 commented Sep 26, 2026 •

Copy link
Copy Markdown
Collaborator

Why

  • Continue the AtomSched migration by replacing hard-coded CUDA instruction descriptions with declarative atoms, operand patterns, and framework-enforced contracts.
  • Preserve authored schedule programs as canonical evidence while moving reusable layout logic to its owning layers.

What

  • Add declarative SM80 MMA and SM90 WGMMA atoms, T.tiled_mma verification, and fragment-only MMA runtime dispatch.
  • Add async-copy, bulk-copy, tensor-copy, and ldmatrix declarations with supported codegen/runtime paths.
  • Import and round-trip 21 authored CUDA schedules without changing their layouts.
  • Isolate CuTe swizzle algebra, add profile-driven coalescing, align pattern structures with IR, and contain mesh suffix replacement.
  • Replace embedded layout flags and specialized pattern classes with named predicates for forwarding, injectivity, whole vectors, TMA boxes/maps, and plain arrangements.
  • Drive effect-Op operand and between checks from declarations, while native verifiers add rules declarations cannot express.
  • Chain captures across operands in ParamDef order.

Contract

  • Remove legacy T.mma, MmaOpSpec, and MmaAtom descriptors in favor of T.tiled_mma and declared atoms; rename T.tma_copy to T.copy_async_bulk.
  • Restrict ops::mma to atom-shaped lane fragments; the Tile tier is removed.
  • Rename ShardLayoutPattern.arrangement to layout, make mesh a MeshPattern, remove TensorPattern.rank, and expose LayoutPattern.from_layout.
  • LayoutPattern now accepts optional structure plus named predicates; Forward and Injective own per-mode behavior.
  • Add coalesce(layout, trg_profile); move apply and NotProjectable to ir/types/layout.py.
  • make_mesh suffix replacement preserves per-level starts, re-encodes sliced offsets, and rejects runs outside the parent scope.
  • tir.md now requires declaration checks for every effect Op; native verifiers supplement them, and an Op with neither declarations nor a verifier is invalid.
  • core-ir.md now specifies ordered input-pattern matching against one shared capture environment, allowing later operands to require earlier captures.

Risk

  • Removing MmaOpSpec also removes the only remaining coverage of section-key disambiguation.
  • Broadcast-to-split runtime copy repeats local values; MmaTile uses a fragment view for B, while A uses T.ldmatrix.
  • Mesh containment supports selections reducible to one continuous run; other grids are rejected.
  • T.copy_async_tensor has no emitter, and WGMMA atoms have no T.tiled_mma emitter.
  • The unreachable ShardLayoutPattern extra-Broadcast-prefix path was removed; instrumentation saw 372 calls with no extra prefix.
  • physical_frames_match now compares per-level frames. Multi-level behavior differs, but all 124 observed MMA checks used one thread level.
  • Framework-driven between diagnostics omit the former copy_async_tensor-specific lead while retaining field-level constraint details.

@zhen8838 zhen8838 changed the title feat(cuda): add declarative CUDA instruction scheduling feat(cuda): declare CUDA MMA atoms and async copy instructions Sep 26, 2026
@zhen8838
zhen8838 merged commit 612dcc8 into tile-ai:main Sep 26, 2026
1 check passed
@zhen8838
zhen8838 deleted the feat/atomsched-stage2b-cuda-instructions branch September 26, 2026 12:03
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.

1 participant