jinhongyii opened a new pull request, #20072:
URL: https://github.com/apache/tvm/pull/20072
## Motivation and context
Blackwell `tcgen05.mma` is issued by a single thread, while one
`Tx.gemm_async` tile call may expand into multiple dense or block-scaled MMA
instructions over its M, N, and K iterations. Persistent kernels sometimes
need to keep the surrounding state machine warp-uniform, so they compute a
one-lane predicate and apply it at the instruction level:
```python
mma_issue = T.ptx.elect_sync()
Tx.gemm_async(
C,
A,
B,
dispatch="tcgen05",
pred=mma_issue,
)
mma_bar.arrive(stage, pred=mma_issue)
```
Before this PR, the `pred` configuration was not forwarded by the
`gemm_async` tcgen05 dispatcher to every MMA it generated. In addition, the
block-scaled MMA path and `TCGen05Bar.arrive` did not support forwarding the
same predicate. Predicating only the matching commit is not sufficient: the
active lanes may still issue all internal MMA instructions.
## Changes
- Read the optional `pred` value from the `gemm_async` tile-call
configuration and forward it to every generated dense or block-scaled tcgen05
MMA.
- Add optional instruction-predicate support to
`ptx_tcgen05_mma_block_scale` and its CUDA helper generation.
- Generate a predicated PTX instruction by converting the runtime `uint32`
value to a PTX predicate and emitting `@p_issue tcgen05.mma...`.
- Add `pred` to `TCGen05Bar.arrive` and forward it to both the single-CTA
and multicast `tcgen05.commit` forms.
- Document the supported thread-scope pattern and the requirement that the
MMA sequence and matching commit use the same predicate.
The existing behavior remains unchanged when `pred` is omitted.
## Testing
- Verify predicated single-CTA and two-CTA multicast `TCGen05Bar` commits in
CUDA codegen.
- Verify that block-scaled MMA codegen emits `@p_issue tcgen05.mma...`.
- Verify that one dense or block-scaled tile call expands to four internal
MMA calls and that all calls receive the same predicate as the matching commit.
- Run changed-files pre-commit checks.
Focused result: 5 tests passed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
