You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "masahi (via GitHub)" <gi...@apache.org> on 2023/05/04 20:39:21 UTC
[GitHub] [tvm] masahi opened a new pull request, #14770: [Vulkan] Add VK_NV_cooperative_matrix support
masahi opened a new pull request, #14770:
URL: https://github.com/apache/tvm/pull/14770
See https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/NV/SPV_NV_cooperative_matrix.html
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187808268
##########
python/tvm/tir/op.py:
##########
@@ -3037,6 +3037,148 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr):
return call_intrin("int32", "tir.TVMBackendFreeWorkspace", device_type, device_id, ptr)
+# Intrinsics for the VK_NV_cooperative_matrix Vulkan extension.
+
+
+def cooperative_matrix_load_NV(buffer_mat, offset, src, rows, cols, stride, column_major):
Review Comment:
I see, I wanted to define those intrinsics to be as similar as possible to the spec, rather than trying to fit an existing TVM concept to it.
That said, I deviate from the spec when it makes sense - for example the load intrinsic in the spec is a pure one that returns a matrix as value, but TIR doesn't have a matrix as a first-class value. Also this definition is very similar to how we define CUDA WMMA intrinsics, so it would be familiar to developers and porting CUDA schedules becomes trivial.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1190398473
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Looks like my claim "Cooperative matrices cannot be loaded from / stored to a buffer" was based on a wrong assumption. I'd still argue that my approach is simpler and cleaner than the load / store based approach, even if there is no other benefit. A con of my approach is that it differs from how we treat vector types - I don't have a good answer for the question "why cooperative matrices should get special handling" other than somewhat subjective opinions (for example, I don't like how we support CUDA WMMA today, which is based on load / store into an array of fragments). One thing that's not clear to me: I believe the spec doesn't say how many bytes each matrix takes up, so I am not sure it is possible to allocate a buffer of and address them. Maybe we don't need to worry about such issues since memory is typed in SPIRV and size / offset are in terms of the number of elements.
I'm happy to change the implementation to be based on conventional load / store based approach if people prefer that way. I'd love to get @Lunderberg's opinion on this point (see Mei's PR https://github.com/apache/tvm/pull/14817 that implements the same extension in the alternative approach very similar to how our cuda codegen supports WMMA).
I should also remove the term "unroll" in the comment since nothing is really unrolled by codegen. For example here is a good 4k matmul TIR https://gist.github.com/masahi/dc7173d39f4b376884f49019cde24826 and the generated spv https://gist.github.com/masahi/3fd481c8037b1ed55876c89305ec02ef. The two representations mostly map one to one (no more unrolling done on spv).
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1190398473
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Looks like my claim "Cooperative matrices cannot be loaded from / stored to a buffer" was based on a wrong assumption. I'd still argue that my approach is simpler and cleaner than the load / store based approach, even if there is no other benefit. A con of my approach is that it differs from how we treat vector types - I don't have a good answer for the question "why cooperative matrices should get special handling" other than somewhat subjective opinions. One thing that's not clear to me: I believe the spec doesn't say how many bytes each matrix takes up, so I am not sure it is possible to allocate a buffer of and address them. Maybe we don't need to worry about such issues since memory is typed in SPIRV and size / offset are in terms of the number of elements.
I'm happy to change the implementation to be based on conventional load / store based approach if people prefer that way. I'd love to get @Lunderberg's opinion on this point.
I should also remove the term "unroll" in the comment since nothing is really unrolled by codegen. For example here is a good 4k matmul TIR https://gist.github.com/masahi/dc7173d39f4b376884f49019cde24826 and the generated spv https://gist.github.com/masahi/3fd481c8037b1ed55876c89305ec02ef. The two representations mostly map one to one (no more unrolling done on spv).
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1194650947
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Being able to reuse existing schedule and auto tensorization for WMMA is a big advantage. I'm wondering if we can reuse wmma intrinsics since we can still use fragment index to identify matrix (if we want to scalarize it)
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187816725
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
+ // memory location corresponding to the scalar c. But when C, A, and B are cooperative matrices,
+ // C += A * B cannot be implemented in the same way, since they cannot be load from or stored in
+ // a buffer.
+
+ // We leverage SPIRV's SSA semantics to implement multiply-add update on C matrices.
+ // The original buffer is subdivided into fixed-size matrices, and multiply-add on individual
+ // matrix is unrolled. For example, if the buffer for the C matrix is of shape 64x64, we
+ // materialze 16 16x16 matrices and generate 16 multiply-add instructions. Each matrix is
+ // represented by a phi value, and each iteration of the reduction loop reads the phi values for
+ // the C matrices and update them.
+
+ // A map from a buffer with kCooperativeMatrixNV storage scope (whose size can be arbitrary)
+ // to cooperative matrices "C" (whose sizes are fixed, e.g. 16x16).
+ // Each matrix is identified by an element offset into the buffer. This encoding lets us bridge
+ // the gap between a typical TIR matmul schedule (which uses
+ // `cache_read(..., "cooperative_matrix_nv")` from arbitrary-sized shared memory) and
+ // the fixed-size matrices as expected by the extension.
+ std::unordered_map<const VarNode*, std::unordered_set<int>> accum_matrices;
+
+ if (op->kind == ForKind::kSerial) {
+ // If this is a serial loop, record which C matrices are updated in this loop.
+ tir::PostOrderVisit(op->body, [&accum_matrices](const ObjectRef& obj) {
Review Comment:
hmm I haven't thought about such possibilities... this code only collects the matrices that are accumulated in the reduction loop ("C" matrix). I haven't put a deep thought into it, but I don't see how cache_read / write would pose a problem. The C matrices should be allocated outside of the loop that it is accumulated.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on PR #14770:
URL: https://github.com/apache/tvm/pull/14770#issuecomment-1536702664
@Lunderberg can you help review this PR?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1194831846
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
I'm open to adopting Mei's approach and actually leaning toward it. I'll experiment with the WMMA-based coop matrix support and make sure that I can get the same perf on my matmul schedule. The PR #14817 is a bit messy right now, so I hope @mei-ye can make it easier to review.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi closed pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi closed pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
URL: https://github.com/apache/tvm/pull/14770
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187808268
##########
python/tvm/tir/op.py:
##########
@@ -3037,6 +3037,148 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr):
return call_intrin("int32", "tir.TVMBackendFreeWorkspace", device_type, device_id, ptr)
+# Intrinsics for the VK_NV_cooperative_matrix Vulkan extension.
+
+
+def cooperative_matrix_load_NV(buffer_mat, offset, src, rows, cols, stride, column_major):
Review Comment:
I see, I wanted to define those intrinsics to be as similar as possible to the spec, rather than trying to fit an existing TVM concept to it.
That said, I deviate from the spec when it makes sense - for example the load intrinsic in the spec is a pure one that returns a matrix as value, but TIR doesn't have a matrix as a first-class value. Also this definition is very similar to how we define CUDA WMMA intrinsics, so it would be familiar to developers and porting CUDA schedules become trivial.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1190398473
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Looks like my claim "Cooperative matrices cannot be loaded from / stored to a buffer" was based on a wrong assumption. I'd still argue that my approach is simpler and cleaner than the load / store based approach, even if there is no other benefit. A con of my approach is that it differs from how we treat vector types - I don't have a good answer for the question "why cooperative matrices should get special handling" other than somewhat subjective opinions (for example, I don't like how we support CUDA WMMA today, which is based load / store into an array of fragments). One thing that's not clear to me: I believe the spec doesn't say how many bytes each matrix takes up, so I am not sure it is possible to allocate a buffer of and address them. Maybe we don't need to worry about such issues since memory is typed in SPIRV and size / offset are in terms of the number of elements.
I'm happy to change the implementation to be based on conventional load / store based approach if people prefer that way. I'd love to get @Lunderberg's opinion on this point.
I should also remove the term "unroll" in the comment since nothing is really unrolled by codegen. For example here is a good 4k matmul TIR https://gist.github.com/masahi/dc7173d39f4b376884f49019cde24826 and the generated spv https://gist.github.com/masahi/3fd481c8037b1ed55876c89305ec02ef. The two representations mostly map one to one (no more unrolling done on spv).
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1186473675
##########
tests/python/unittest/test_target_codegen_vulkan.py:
##########
@@ -600,5 +602,387 @@ def func(A: T.Buffer((N, 2), "int32")):
np.testing.assert_array_equal(a[:, 1], (np.arange(N) - offset) % divisor)
+@T.prim_func
Review Comment:
Can this `@T.prim_func` definition be moved inside the test that uses it? Otherwise, any failure to generate the PrimFunc prevents the entire test suite from being run.
##########
python/tvm/tir/op.py:
##########
@@ -3037,6 +3037,148 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr):
return call_intrin("int32", "tir.TVMBackendFreeWorkspace", device_type, device_id, ptr)
+# Intrinsics for the VK_NV_cooperative_matrix Vulkan extension.
+
+
+def cooperative_matrix_load_NV(buffer_mat, offset, src, rows, cols, stride, column_major):
Review Comment:
The parameters for this intrinsic look very similar to those of [BufferRegion](https://github.com/apache/tvm/blob/main/include/tvm/tir/stmt.h#L1099), but I can't think of any way to use it in this context. If we could, that would be a nice way to name as many parameters as possible.
##########
python/tvm/tir/op.py:
##########
@@ -3037,6 +3037,148 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr):
return call_intrin("int32", "tir.TVMBackendFreeWorkspace", device_type, device_id, ptr)
+# Intrinsics for the VK_NV_cooperative_matrix Vulkan extension.
+
+
+def cooperative_matrix_load_NV(buffer_mat, offset, src, rows, cols, stride, column_major):
+ """The intrinsic corresponding to the OpCooperativeMatrixLoadNV instruction.
+
+ Parameters
+ ----------
+ buffer_mat: Var
+ The destination buffer with "cooperative_matrix_nv" scope.
+
+ offset: IntImm
+ The element offset for the matrix to be loaded in the destination buffer.
+
+ src: Expr
+ The source pointer expression.
+
+ rows: int
+ The number of rows in the matrix.
+
+ cols: int
+ The number of columns in the matrix.
+
+ stride: Expr
Review Comment:
Can we specify that this is stride is the number of elements? Otherwise, there's potential confusion with a stride specified in number of bytes.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187816725
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
+ // memory location corresponding to the scalar c. But when C, A, and B are cooperative matrices,
+ // C += A * B cannot be implemented in the same way, since they cannot be load from or stored in
+ // a buffer.
+
+ // We leverage SPIRV's SSA semantics to implement multiply-add update on C matrices.
+ // The original buffer is subdivided into fixed-size matrices, and multiply-add on individual
+ // matrix is unrolled. For example, if the buffer for the C matrix is of shape 64x64, we
+ // materialze 16 16x16 matrices and generate 16 multiply-add instructions. Each matrix is
+ // represented by a phi value, and each iteration of the reduction loop reads the phi values for
+ // the C matrices and update them.
+
+ // A map from a buffer with kCooperativeMatrixNV storage scope (whose size can be arbitrary)
+ // to cooperative matrices "C" (whose sizes are fixed, e.g. 16x16).
+ // Each matrix is identified by an element offset into the buffer. This encoding lets us bridge
+ // the gap between a typical TIR matmul schedule (which uses
+ // `cache_read(..., "cooperative_matrix_nv")` from arbitrary-sized shared memory) and
+ // the fixed-size matrices as expected by the extension.
+ std::unordered_map<const VarNode*, std::unordered_set<int>> accum_matrices;
+
+ if (op->kind == ForKind::kSerial) {
+ // If this is a serial loop, record which C matrices are updated in this loop.
+ tir::PostOrderVisit(op->body, [&accum_matrices](const ObjectRef& obj) {
Review Comment:
hmm I haven't thought about such possibilities... this code only collects the matrices that are accumulated in the reduction loop ("C" matrix). I haven't put a deep thought into it, but I don't see how cache_read / write would pose a problem. The C matrices should be allocated outside of the loop where they are accumulated.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1188800164
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
+ // memory location corresponding to the scalar c. But when C, A, and B are cooperative matrices,
+ // C += A * B cannot be implemented in the same way, since they cannot be load from or stored in
+ // a buffer.
+
+ // We leverage SPIRV's SSA semantics to implement multiply-add update on C matrices.
+ // The original buffer is subdivided into fixed-size matrices, and multiply-add on individual
+ // matrix is unrolled. For example, if the buffer for the C matrix is of shape 64x64, we
+ // materialze 16 16x16 matrices and generate 16 multiply-add instructions. Each matrix is
+ // represented by a phi value, and each iteration of the reduction loop reads the phi values for
+ // the C matrices and update them.
+
+ // A map from a buffer with kCooperativeMatrixNV storage scope (whose size can be arbitrary)
+ // to cooperative matrices "C" (whose sizes are fixed, e.g. 16x16).
+ // Each matrix is identified by an element offset into the buffer. This encoding lets us bridge
+ // the gap between a typical TIR matmul schedule (which uses
+ // `cache_read(..., "cooperative_matrix_nv")` from arbitrary-sized shared memory) and
+ // the fixed-size matrices as expected by the extension.
+ std::unordered_map<const VarNode*, std::unordered_set<int>> accum_matrices;
+
+ if (op->kind == ForKind::kSerial) {
+ // If this is a serial loop, record which C matrices are updated in this loop.
+ tir::PostOrderVisit(op->body, [&accum_matrices](const ObjectRef& obj) {
Review Comment:
The issue I'm imagining would occur when the reduction loop occurs within another loop. For example, suppose I took the matrix multiplication unit test, but added a batch dimension (rough sketch below). The cooperative matrix loads wouldn't depend on the batch number for their shapes, and the declaration/allocation of the cooperative matrices could occur within the body of `batch_i`. However, the loop over `batch_i` shouldn't have an `OpPhi` associated with the cooperative matrices, because there is no value of the matrix prior to entering the loop over `batch_i`.
```python
@T.prim_func
def main(
X: T.Buffer((4, 16, 32), "float16"),
W: T.Buffer((4, 32, 16), "float16"),
compute: T.Buffer((4, 16, 16), "float32"),
):
for batch_i in range(4):
X_shared = T.decl_buffer((16, 32), "float16", scope="shared")
W_shared = T.decl_buffer((32, 16), "float16", scope="shared")
X_shared_cooperative_matrix_nv = T.decl_buffer(
(16, 32), "float16", scope="cooperative_matrix_nv"
)
W_shared_cooperative_matrix_nv = T.decl_buffer(
(32, 16), "float16", scope="cooperative_matrix_nv"
)
compute_cooperative_matrix_nv = T.decl_buffer(
(16, 16), scope="cooperative_matrix_nv"
)
i_0_j_0_fused = T.launch_thread("blockIdx.x", 1)
with T.launch_thread("threadIdx.x", 32) as tx:
T.cooperative_matrix_fill_NV(
compute_cooperative_matrix_nv.data, 0, 16, 16, T.float32(0)
)
# Continues with remainder of matrix multiplication
for k_0 in range(2):
...
```
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187821900
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
+ // memory location corresponding to the scalar c. But when C, A, and B are cooperative matrices,
+ // C += A * B cannot be implemented in the same way, since they cannot be load from or stored in
+ // a buffer.
+
+ // We leverage SPIRV's SSA semantics to implement multiply-add update on C matrices.
+ // The original buffer is subdivided into fixed-size matrices, and multiply-add on individual
+ // matrix is unrolled. For example, if the buffer for the C matrix is of shape 64x64, we
+ // materialze 16 16x16 matrices and generate 16 multiply-add instructions. Each matrix is
+ // represented by a phi value, and each iteration of the reduction loop reads the phi values for
+ // the C matrices and update them.
+
+ // A map from a buffer with kCooperativeMatrixNV storage scope (whose size can be arbitrary)
+ // to cooperative matrices "C" (whose sizes are fixed, e.g. 16x16).
+ // Each matrix is identified by an element offset into the buffer. This encoding lets us bridge
+ // the gap between a typical TIR matmul schedule (which uses
+ // `cache_read(..., "cooperative_matrix_nv")` from arbitrary-sized shared memory) and
+ // the fixed-size matrices as expected by the extension.
+ std::unordered_map<const VarNode*, std::unordered_set<int>> accum_matrices;
+
+ if (op->kind == ForKind::kSerial) {
+ // If this is a serial loop, record which C matrices are updated in this loop.
+ tir::PostOrderVisit(op->body, [&accum_matrices](const ObjectRef& obj) {
Review Comment:
Or, are you imagining a scenario where the matrix to be accumulated is initialized from `cooperative_matrix_load`, rather than zero initialized inside the kernel before the reduction loop?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187465440
##########
python/tvm/tir/op.py:
##########
@@ -3037,6 +3037,148 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr):
return call_intrin("int32", "tir.TVMBackendFreeWorkspace", device_type, device_id, ptr)
+# Intrinsics for the VK_NV_cooperative_matrix Vulkan extension.
+
+
+def cooperative_matrix_load_NV(buffer_mat, offset, src, rows, cols, stride, column_major):
Review Comment:
Mostly that the parameters used for the intrinsic are specifying a source and destination region for a copy to be performed. There are some restrictions to it, such as requiring that the source and destination buffers are both 2-d, and that one of the two values in `source_region->buffer->strides` is one. I was thinking that it might be a cleaner representation to write the intrinsic as `cooperative_matrix_load_NV(dest_region, source_region)`, if possible.
Looking at it again this morning, though, writing this intrinsic in terms of BufferRegion would only work if the buffers are maintained in their non-flattened state all the way through to codegen, as determining the rows/cols/stride/column_major parameters would require a 2-d buffer. With that restriction, along with needing to compute some of the strides during codegen, writing the intrinsic in terms of BufferRegion probably isn't a good idea.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mei-ye commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "mei-ye (via GitHub)" <gi...@apache.org>.
mei-ye commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1190357098
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
I assume the intention here is to scalarize each 16x16 fragment and represent it using a variable. This transformation does not bring performance gain compared to the approach of using a vector of 16x16 fragments and load/store from it since shader compiler can easily does the same scalarization. Also, if you want to unroll the loop, the address computation in each unrolled iteration is not CSEed. So you will see a larger code size with all these non-CSEed address computations.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1190398473
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Looks like my claim "Cooperative matrices cannot be loaded from / stored to a buffer" was based on a wrong assumption. I'd still argue that my approach is simpler and cleaner than the load / store based approach, even if there is no other benefit. A con of my approach is that it differs from how we treat vector types - I don't have a good answer for the question "why cooperative matrices should get special handling" other than somewhat subjective opinions (for example, I don't like how we support CUDA WMMA today, which is based on load / store into an array of fragments). One thing that's not clear to me: I believe the spec doesn't say how many bytes each matrix takes up, so I am not sure it is possible to allocate a buffer of and address them. Maybe we don't need to worry about such issues since memory is typed in SPIRV and size / offset are in terms of the number of elements.
I'm happy to change the implementation to be based on conventional load / store based approach if people prefer that way. I'd love to get @Lunderberg's opinion on this point.
I should also remove the term "unroll" in the comment since nothing is really unrolled by codegen. For example here is a good 4k matmul TIR https://gist.github.com/masahi/dc7173d39f4b376884f49019cde24826 and the generated spv https://gist.github.com/masahi/3fd481c8037b1ed55876c89305ec02ef. The two representations mostly map one to one (no more unrolling done on spv).
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187816725
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
+ // memory location corresponding to the scalar c. But when C, A, and B are cooperative matrices,
+ // C += A * B cannot be implemented in the same way, since they cannot be load from or stored in
+ // a buffer.
+
+ // We leverage SPIRV's SSA semantics to implement multiply-add update on C matrices.
+ // The original buffer is subdivided into fixed-size matrices, and multiply-add on individual
+ // matrix is unrolled. For example, if the buffer for the C matrix is of shape 64x64, we
+ // materialze 16 16x16 matrices and generate 16 multiply-add instructions. Each matrix is
+ // represented by a phi value, and each iteration of the reduction loop reads the phi values for
+ // the C matrices and update them.
+
+ // A map from a buffer with kCooperativeMatrixNV storage scope (whose size can be arbitrary)
+ // to cooperative matrices "C" (whose sizes are fixed, e.g. 16x16).
+ // Each matrix is identified by an element offset into the buffer. This encoding lets us bridge
+ // the gap between a typical TIR matmul schedule (which uses
+ // `cache_read(..., "cooperative_matrix_nv")` from arbitrary-sized shared memory) and
+ // the fixed-size matrices as expected by the extension.
+ std::unordered_map<const VarNode*, std::unordered_set<int>> accum_matrices;
+
+ if (op->kind == ForKind::kSerial) {
+ // If this is a serial loop, record which C matrices are updated in this loop.
+ tir::PostOrderVisit(op->body, [&accum_matrices](const ObjectRef& obj) {
Review Comment:
hmm I haven't thought about such possibilities... this code only collects the matrices that are accumulated in the reduction loop ("C" matrix). I haven't put a deep thought into it, but I don't see how cache_read / write would pose a problem. The C matrices should be allocated outside of the loop where it is accumulated.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1195361585
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Looking through, I like Mei's approach as well. It avoids making usage restrictions (e.g. buffers must be accessed through coop matrix primitives or load/store, but not both), and maintains the mutability semantics of both spirv/tir by avoiding long-lived coop matrix values in the generated spir-v.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1194831846
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
I'm open to adopting Mei's approach and actually leaning toward it. I'll experiment with the WMMA based coop matrix support and make sure that I can get the same perf on my matmul schedule. The PR #14817 is a bit messy right now, so I hope @mei-ye can make it easier to review.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187500787
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -672,6 +804,9 @@ void CodeGenSPIRV::VisitStmt_(const AllocateNode* op) {
size_t num_bytes = op->dtype.bytes() * op->dtype.lanes() * static_cast<uint32_t>(constant_size);
shared_memory_bytes_used_ += num_bytes;
+ } else if (storage_scope.rank == runtime::StorageRank::kCooperativeMatrixNV) {
Review Comment:
To avoid the over-collection issue, we could track the currently-active buffers with `kCooperativeMatrixNV` here. If we track which allocations are currently open, then the `post_order_visit` in the `VisitStmt_(ForNode*)` could limit its collection to buffers that that are currently in scope.
```c++
active_cooperative_matrices_.insert(op->buffer_var);
this->VisitStmt(op->body);
active_cooperative_matrices_.erase(op->buffer_var);
```
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
If I understand correctly, the special handling for cooperative matrices is required because (1) the cooperative matrices have buffer-type semantics in TIR and (2) the cooperative matrices have value-type semantics in SPIR-V. Due to (1), the cooperative matrices may be mutated, unlike a variable in a LetStmt. Due to (2), the mutation requires a `OpPhi` in order to join the values from the two predecessor branches, unlike in a BufferStore/BufferLoad. Is that understanding correct?
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
+ // memory location corresponding to the scalar c. But when C, A, and B are cooperative matrices,
+ // C += A * B cannot be implemented in the same way, since they cannot be load from or stored in
+ // a buffer.
+
+ // We leverage SPIRV's SSA semantics to implement multiply-add update on C matrices.
+ // The original buffer is subdivided into fixed-size matrices, and multiply-add on individual
+ // matrix is unrolled. For example, if the buffer for the C matrix is of shape 64x64, we
+ // materialze 16 16x16 matrices and generate 16 multiply-add instructions. Each matrix is
+ // represented by a phi value, and each iteration of the reduction loop reads the phi values for
+ // the C matrices and update them.
+
+ // A map from a buffer with kCooperativeMatrixNV storage scope (whose size can be arbitrary)
+ // to cooperative matrices "C" (whose sizes are fixed, e.g. 16x16).
+ // Each matrix is identified by an element offset into the buffer. This encoding lets us bridge
+ // the gap between a typical TIR matmul schedule (which uses
+ // `cache_read(..., "cooperative_matrix_nv")` from arbitrary-sized shared memory) and
+ // the fixed-size matrices as expected by the extension.
+ std::unordered_map<const VarNode*, std::unordered_set<int>> accum_matrices;
+
+ if (op->kind == ForKind::kSerial) {
+ // If this is a serial loop, record which C matrices are updated in this loop.
+ tir::PostOrderVisit(op->body, [&accum_matrices](const ObjectRef& obj) {
Review Comment:
For deeply-nested loops, does this cause a performance issue, as it repeats the visit of the loop body once for every nested loop?
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
+ // memory location corresponding to the scalar c. But when C, A, and B are cooperative matrices,
+ // C += A * B cannot be implemented in the same way, since they cannot be load from or stored in
+ // a buffer.
+
+ // We leverage SPIRV's SSA semantics to implement multiply-add update on C matrices.
+ // The original buffer is subdivided into fixed-size matrices, and multiply-add on individual
+ // matrix is unrolled. For example, if the buffer for the C matrix is of shape 64x64, we
+ // materialze 16 16x16 matrices and generate 16 multiply-add instructions. Each matrix is
+ // represented by a phi value, and each iteration of the reduction loop reads the phi values for
+ // the C matrices and update them.
+
+ // A map from a buffer with kCooperativeMatrixNV storage scope (whose size can be arbitrary)
+ // to cooperative matrices "C" (whose sizes are fixed, e.g. 16x16).
+ // Each matrix is identified by an element offset into the buffer. This encoding lets us bridge
+ // the gap between a typical TIR matmul schedule (which uses
+ // `cache_read(..., "cooperative_matrix_nv")` from arbitrary-sized shared memory) and
+ // the fixed-size matrices as expected by the extension.
+ std::unordered_map<const VarNode*, std::unordered_set<int>> accum_matrices;
+
+ if (op->kind == ForKind::kSerial) {
+ // If this is a serial loop, record which C matrices are updated in this loop.
+ tir::PostOrderVisit(op->body, [&accum_matrices](const ObjectRef& obj) {
Review Comment:
Would this collection loop collect too many buffers in the case of a cache_read/write that occurs inside a loop? If a cooperative buffer's `AllocateNode` occurs within the body of a loop, then it cannot have a pre-loop value, nor can its value be passed to the next loop iteration.
I think this case would cause an error to occur in the `auto mat = builder_->GetCooperativeMatrix` that occurs below. The error would be suppressed if an earlier pass has hoisted the allocation to function-scope, but I don't think the codegen should require that to be the case.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1190398473
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Looks like my claim "Cooperative matrices cannot be loaded from / stored to a buffer" was based on a wrong assumption. I'd still argue that my approach is simpler and cleaner than the load / store based approach, even if there is no other benefit. A con of my approach is that it differs from how we treat vector types - I don't have a good answer for the question "why cooperative matrices should get special handling" other than somewhat subjective opinions (for example, I don't like how we support CUDA WMMA today, which is based on load / store into an array of fragments). One thing that's not clear to me: I believe the spec doesn't say how many bytes each matrix takes up, so I am not sure it is possible to allocate a buffer of and address them. Maybe we don't need to worry about such issues since memory is typed in SPIRV and size / offset are in terms of the number of elements.
I'm happy to change the implementation to be based on a conventional load / store based approach if people prefer that way. I'd love to get @Lunderberg's opinion on this point (see Mei's PR https://github.com/apache/tvm/pull/14817 that implements the same extension in the alternative approach very similar to how our cuda codegen supports WMMA).
I should also remove the term "unroll" in the comment since nothing is really unrolled by codegen. For example here is a good 4k matmul TIR https://gist.github.com/masahi/dc7173d39f4b376884f49019cde24826 and the generated spv https://gist.github.com/masahi/3fd481c8037b1ed55876c89305ec02ef. The two representations mostly map one to one (no more unrolling done on spv).
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1190428194
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
The Mei's approach in https://github.com/apache/tvm/pull/14817, being based on the WMMA intrinsics originally developed for CUDA, inherits both pros and cons of CUDA WMMA support in TVM. In particular, it allows the same schedule to be compiled for CUDA and VK, and since CUDA WMMA is supported by our auto tensorization infra, auto tensorization might also work out of the box for VK. That's a big advantage. cc @vinx13 @tqchen
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1190398473
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Looks like my claim "Cooperative matrices cannot be loaded from / stored to a buffer" was based on a wrong assumption. I'd still argue that my approach is simpler and cleaner than the load / store based approach, even if there is no other benefit. A con of my approach is that it differs from how we treat vector types - I don't have a good answer for the question "why cooperative matrices should get special handling" other than somewhat subjective opinions (for example, I don't like how we support CUDA WMMA today, which is based on load / store into an array of fragments). One thing that's not clear to me: I believe the spec doesn't say how many bytes each matrix takes up, so I am not sure it is possible to allocate a buffer of and address them. Maybe we don't need to worry about such issues since memory is typed in SPIRV and size / offset are in terms of the number of elements.
I'm happy to change the implementation to be based on a conventional load / store based approach if people prefer that way. I'd love to get @Lunderberg's opinion on this point (see Mei's PR https://github.com/apache/tvm/pull/14817 that implements the same extension in the alternative approach very similar to how our cuda codegen supports WMMA).
I should also remove the term "unroll" in the comment since nothing is really unrolled by codegen. For example here is a good 4k matmul TIR https://gist.github.com/masahi/dc7173d39f4b376884f49019cde24826 and the generated spv https://gist.github.com/masahi/3fd481c8037b1ed55876c89305ec02ef. The two representations mostly map one to one (no more unrolling done on spv). But it does assume that unrolling has been done in a schedule so that elem offset becomes a constant. I admit this is another con of my approach.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1186658661
##########
python/tvm/tir/op.py:
##########
@@ -3037,6 +3037,148 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr):
return call_intrin("int32", "tir.TVMBackendFreeWorkspace", device_type, device_id, ptr)
+# Intrinsics for the VK_NV_cooperative_matrix Vulkan extension.
+
+
+def cooperative_matrix_load_NV(buffer_mat, offset, src, rows, cols, stride, column_major):
Review Comment:
What do you mean? I don't see a similarity.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1195598468
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Ok thanks for the discussion, I'm closing my PR in favor of #14817 cc @mei-ye
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tvm-bot commented on pull request #14770: [Vulkan] Add VK_NV_cooperative_matrix support
Posted by "tvm-bot (via GitHub)" <gi...@apache.org>.
tvm-bot commented on PR #14770:
URL: https://github.com/apache/tvm/pull/14770#issuecomment-1535382959
<!---bot-comment-->
Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from [Reviewers](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @-ing them in a comment.
<!--bot-comment-ccs-start-->
* No users to tag found in teams: `vulkan` <sub>See [#10317](https://github.com/apache/tvm/issues/10317) for details</sub><!--bot-comment-ccs-end-->
<sub>Generated by [tvm-bot](https://github.com/apache/tvm/blob/main/ci/README.md#github-actions)</sub>
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187794220
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Exactly right. I really liked your explanation, so I'll update the comment here based on it.
From TIR perspective, cooperative matrices (and also the corresponding concept in CUDA WMMA) are nothing but a normal buffer with a specialized scope. In particular it can have any shape. On the other hand, in SPIRV they have a fixed shape and has value semantics.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187794220
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
Exactly right. I really liked your explanation, so I'll update the comment here based on it.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187465440
##########
python/tvm/tir/op.py:
##########
@@ -3037,6 +3037,148 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr):
return call_intrin("int32", "tir.TVMBackendFreeWorkspace", device_type, device_id, ptr)
+# Intrinsics for the VK_NV_cooperative_matrix Vulkan extension.
+
+
+def cooperative_matrix_load_NV(buffer_mat, offset, src, rows, cols, stride, column_major):
Review Comment:
Mostly that the parameters used for the intrinsic are specifying a source and destination region for a copy to be performed. There are some restrictions to it, such as requiring that the source and destination buffers are both 2-d, and that one of the two values in `source_region->buffer->strides` is one. I was thinking that it might be a cleaner representation to write the intrinsic as `cooperative_matrix_load_NV(dest_region, source_region)`, if possible.
Looking at it again this morning, though, writing this intrinsic in terms of BufferRegion would only work if the buffers are maintained in their non-flattened state all the way through to codegen, as determining the rows/cols/stride/column_major parameters would require a 2-d buffer. That restriction, along with needing to compute some of the strides during codegen, is making me like the original implementation more, with the intrinsic directly matching the SPIR-V opcode.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1190428194
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
Review Comment:
The Mei's approach in https://github.com/apache/tvm/pull/14817, being based on the WMMA intrinsics originally developed for CUDA, inherits both pros and cons of CUDA WMMA support in TVM. In particular, it allows the same schedule to be compiled for CUDA and VK, and since CUDA WMMA is supported by our auto tensorization infra, auto tensorization might also work out of the box for VK. That's a big advantage. cc @vinx13
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14770: [Vulkan] Add `VK_NV_cooperative_matrix` support
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14770:
URL: https://github.com/apache/tvm/pull/14770#discussion_r1187802759
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -561,6 +630,55 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// Loop head
builder_->StartLabel(head_label);
+
+ // In a normal matmul, the update semantics c += a * b is implemented by load and store to the
+ // memory location corresponding to the scalar c. But when C, A, and B are cooperative matrices,
+ // C += A * B cannot be implemented in the same way, since they cannot be load from or stored in
+ // a buffer.
+
+ // We leverage SPIRV's SSA semantics to implement multiply-add update on C matrices.
+ // The original buffer is subdivided into fixed-size matrices, and multiply-add on individual
+ // matrix is unrolled. For example, if the buffer for the C matrix is of shape 64x64, we
+ // materialze 16 16x16 matrices and generate 16 multiply-add instructions. Each matrix is
+ // represented by a phi value, and each iteration of the reduction loop reads the phi values for
+ // the C matrices and update them.
+
+ // A map from a buffer with kCooperativeMatrixNV storage scope (whose size can be arbitrary)
+ // to cooperative matrices "C" (whose sizes are fixed, e.g. 16x16).
+ // Each matrix is identified by an element offset into the buffer. This encoding lets us bridge
+ // the gap between a typical TIR matmul schedule (which uses
+ // `cache_read(..., "cooperative_matrix_nv")` from arbitrary-sized shared memory) and
+ // the fixed-size matrices as expected by the extension.
+ std::unordered_map<const VarNode*, std::unordered_set<int>> accum_matrices;
+
+ if (op->kind == ForKind::kSerial) {
+ // If this is a serial loop, record which C matrices are updated in this loop.
+ tir::PostOrderVisit(op->body, [&accum_matrices](const ObjectRef& obj) {
Review Comment:
Good point, although in a gpu kernel having such a deep loop nest is rare, I'll polish this.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org