[usmp] Also remap VarNode to USMP-allocated buffer #12880
Conversation
* Fixes cases where a USMP buffer is used with tensorization
intrinsics. In these cases, the buffer data var is referenced
directly in the call_extern args.
Before this patch, ConvertPoolAllocationsToOffsets would generate TIR
like the following:
let dense_let: Pointer(global int32) = @tir.address_of(global_workspace_37_buffer_var[69952], dtype=handle)
for (k.outer: int32, 0, 64) {
@tir.call_extern("gemm_1x1x1_update_UKVNAEBL", ..., dense, ...)
}
T_multiply[ax1] = @tir.q_multiply_shift(((dense: Buffer(dense_let,
int32, [10], [], align=32)[ax1], ...)
This caused CodegenSourceBase to later fail with this error:
"src/target/source/codegen_source_base.cc", line 67
Check failed: (it != var_idmap_.end()) is false: Find undefined
Variable dense
After this patch, "dense" in the call_extern is changed to read "dense_let."
|
NOTE: sending this through CI to see what I broke, then will add tests. |
|
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 by @-ing them in a comment.
Generated by tvm-bot |
github-actions
bot
requested review from
Mousius,
leandron,
lhutton1,
manupak and
mehrdadh
|
also cc @d-smirnov |
|
Hmmmm... I suppose this scenario want to pass a buffer itself to the call (as opposed to a BufferLoad) ? |
|
@manupak I think it's ok to pass the buffer's edit: I suppose I should test the call_packed case, though... |
|
Thanks for the fix -- yea I double checked whether buffer info extraction captures this scenario as well -- it seems so. If we can get the tests and LOG(INFO)s removed -- I think we can get this in... |
areusch
force-pushed
the
remap-varnode
branch
from
28ded98 to
9c538e0
Compare
|
@manupak great, deleted the LOG lines and added a test. please take a look! |
|
I think we need a commit message that is bit plain text and describes the changes. |
|
to clarify a bit more -- I just meant to say, according to https://github.com/apache/tvm/blob/main/docs/contribute/pull_request.rst, we might want to put a textual description in here in the body as the test in the PR shows what is being fixed. Also I think we can remove the first bullet point too (just the bullet and replace it with a text describing the change). I m happy merge this by taking that body in -- it felt odd to have a code example as a commit in the tree. |
|
gotcha, updated the message. |
Before this patch, ConvertPoolAllocationsToOffsets would generate TIR
like the following:
let dense_let: Pointer(global int32) = @tir.address_of(global_workspace_37_buffer_var[69952], dtype=handle)
for (k.outer: int32, 0, 64) {
@tir.call_extern("gemm_1x1x1_update_UKVNAEBL", ..., dense, ...)
}
T_multiply[ax1] = @tir.q_multiply_shift(((dense: Buffer(dense_let,
int32, [10], [], align=32)[ax1], ...)
This caused CodegenSourceBase to later fail with this error:
"src/target/source/codegen_source_base.cc", line 67
Check failed: (it != var_idmap_.end()) is false: Find undefined
Variable dense
After this patch, "dense" in the call_extern is changed to read "dense_let."
Before this patch, ConvertPoolAllocationsToOffsets would generate TIR
like the following:
let dense_let: Pointer(global int32) = @tir.address_of(global_workspace_37_buffer_var[69952], dtype=handle)
for (k.outer: int32, 0, 64) {
@tir.call_extern("gemm_1x1x1_update_UKVNAEBL", ..., dense, ...)
}
T_multiply[ax1] = @tir.q_multiply_shift(((dense: Buffer(dense_let,
int32, [10], [], align=32)[ax1], ...)
This caused CodegenSourceBase to later fail with this error:
"src/target/source/codegen_source_base.cc", line 67
Check failed: (it != var_idmap_.end()) is false: Find undefined
Variable dense
After this patch, "dense" in the call_extern is changed to read "dense_let."
This PR fixes cases where a USMP buffer is used with tensorization intrinsics. In these cases, it is common to
write schedules that pass Buffers to extern functions using
access_ptr(). When transforming tir.allocate intooffsets into global_workspace buffers, USMP inserts a let expression that replaces the old Buffer var with a new
var named "<buffer_var>_let." It then attempts to replace all references to the old var with the new one. When
a Buffer was passed via
access_ptr, USMP would would previously ignore such instances, producing this errorlater on in codegen (dense being the name of the Buffer in use):
This patch fixes the problem by replacing the old VarNode with the new one.
cc @lhutton1 @Mousius @leandron @manupak @mehrdadh