[PTX] Support mma.sp to use Sparse Tensor Cores and refactor mma codegen #10339
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
yzh119
requested review from
junrushao,
vinx13,
tqchen,
kparzysz-quic,
ZihengJiang,
masahi,
Hzfengsy,
comaniac,
jroesch,
areusch,
yzhliu,
merrymercy and
icemelon
as code owners
|
@vinx13 one thing I'm not sure is if we need offset for metadata (which stores the indices information and is always 32bit). |
|
If metadata is stored in a larger buffer, and any some elements of it are passed to mma.sp each time, offset is needed. #9727 introduced some breaking changes in the semantic of |
yzh119
changed the title
|
@vinx13 , refactor of MMA codegen is also finished. I deleted some unittests such as p.s. I can add these tests back if necessary. |
|
okay, it turns out they mean elements inside the two multiplicands must have the same data type, but the two multiplicands can have different data types. cutlass also support u8s8s32 mma's. |
vinx13
approved these changes
Mar 8, 2022
ziqiangxu8457
pushed a commit
to ziqiangxu8457/tvm
that referenced
this pull request
Mar 9, 2022
…gen (apache#10339) * init * upd * upd * lint * lint again * upd * add m16n8k32 testcase * format * use make_tuple instead of initializer list * add metadata offset * upd * docstring and sanity * add u8s8s32 back * improvement * compatible apache#9727
junrushao
pushed a commit
that referenced
this pull request
Apr 3, 2022
…y to warp memory (#10855) We already have PTX mma and mma.sp builtin support in #9909 and #10339 . However, we have not supported corresponding data movement builtins for these mma instructions, so the data movement would not be as fast as wmma. This PR brings the `ldmatrix` builtin, which is a native PTX warp-level instruction (https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix), and we can use it to load several (1/2/4) 8x8 matrices from shared memory to warp memory.
pfk-beta
pushed a commit
to pfk-beta/tvm
that referenced
this pull request
Apr 11, 2022
…gen (apache#10339) * init * upd * upd * lint * lint again * upd * add m16n8k32 testcase * format * use make_tuple instead of initializer list * add metadata offset * upd * docstring and sanity * add u8s8s32 back * improvement * compatible apache#9727
pfk-beta
pushed a commit
to pfk-beta/tvm
that referenced
this pull request
Apr 11, 2022
…y to warp memory (apache#10855) We already have PTX mma and mma.sp builtin support in apache#9909 and apache#10339 . However, we have not supported corresponding data movement builtins for these mma instructions, so the data movement would not be as fast as wmma. This PR brings the `ldmatrix` builtin, which is a native PTX warp-level instruction (https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix), and we can use it to load several (1/2/4) 8x8 matrices from shared memory to warp memory.
mehrdadh
pushed a commit
to mehrdadh/tvm
that referenced
this pull request
Apr 11, 2022
…y to warp memory (apache#10855) We already have PTX mma and mma.sp builtin support in apache#9909 and apache#10339 . However, we have not supported corresponding data movement builtins for these mma instructions, so the data movement would not be as fast as wmma. This PR brings the `ldmatrix` builtin, which is a native PTX warp-level instruction (https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix), and we can use it to load several (1/2/4) 8x8 matrices from shared memory to warp memory.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Sparse Tensor Cores(STC) was first introduced in the Ampere architecture of NVIDIA GPUs (see whitepaper). Up to now, developers can only activate STC by calling wrapped APIs such as
cuSparseLtor writing PTX assembly code.Following #9909, this PR support
ptx_mma_spintrinsic so as to expose the interface of STC at the TIR level.This PR also refactors the
ptx_mma.ccto use template-based codegen.cc @vinx13 @junrushao1994 @Hzfengsy