Skip to content

fix(tma): select CTA-local G2S on SM120/SM121 - #1249

Open
lucifer1004 wants to merge 1 commit into
NVlabs:mainfrom
lucifer1004:fix/sm120-g2s-address-space-1238
Open

lucifer1004 wants to merge 1 commit into
NVlabs:mainfrom
lucifer1004:fix/sm120-g2s-address-space-1238

Conversation

@lucifer1004

@lucifer1004 lucifer1004 commented Sep 6, 2026

Copy link
Copy Markdown
Contributor

What this changes

Select CTA-local TMA global-to-shared loads on SM120/SM121, including their a and f targets. For local destinations, this avoids the compatibility call that conflicts with register reallocation. Partially addresses #1238.

Shape of the change

  • Select local inline PTX for non-multicast loads of ranks 1–5.
  • Retain the cluster route for other targets and multicast operations.
  • Reject explicitly cluster-addressed destinations on the local route.
  • Fix the target hint before lowering so automatic and explicit target selection agree.
SM120 local load -> direct TMA instruction -> register reallocation retained

Validation

  • Target/rank controls and native, real NVVM and ptxas checks passed.
  • Valid local GPU examples passed memory and synchronization checks.
  • SM121 hardware and peer-CTA operations were not exercised.

Remaining work

The existing API permits cluster-shared memory. A GPU target alone does not prove that the destination and barrier belong to the executing CTA (thread block), and a helper call can carry these addresses through a generic pointer.

Prove locality for both operands or expose an explicit local API while preserving the cluster contract. The successful local controls do not resolve this API gap.

Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
@nihalpasham nihalpasham added codegen Device code-generation pipeline (Rust MIR to IR to PTX) intrinsics Device intrinsics and libdevice math mappings needs-changes Review found changes required before this PR can land safety Memory safety, soundness, or undefined behavior labels Sep 15, 2026

@nihalpasham nihalpasham 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.

Reviewed 5620feea: changes requested.

  • The local path removes the compatibility call and retains register reallocation; rank/target controls and local GPU memory/synchronization checks pass.
  • The existing API permits cluster-shared destinations. A GPU target alone does not prove that both destination and barrier belong to the executing CTA (thread block).
same GPU target:
  local destination + barrier -> CTA form is valid
  peer-CTA destination        -> preserve the cluster contract
  • Establish locality or introduce an explicit local API. Checking only explicitly cluster-typed pointers misses generic pointers across helper calls. Cover both operands and preserve existing cluster behavior.
  • Rebase and sign the revision.

#1238 remains partial.

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

Labels

codegen Device code-generation pipeline (Rust MIR to IR to PTX) intrinsics Device intrinsics and libdevice math mappings needs-changes Review found changes required before this PR can land safety Memory safety, soundness, or undefined behavior

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants