You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "yzh119 (via GitHub)" <gi...@apache.org> on 2023/03/16 09:59:05 UTC
[GitHub] [tvm] yzh119 opened a new pull request, #14316: [TensorIR] New schedule primitive `set_dtype`
yzh119 opened a new pull request, #14316:
URL: https://github.com/apache/tvm/pull/14316
# Motivation
Currently, we miss a schedule primitive to change the data type of allocated buffer (e.g. via `cache_read`/`cache_write`), and thus we cannot perform type conversion while loading data from global to shared memory.
This PR adds a new schedule primitive `set_dtype` that follows the interface of `set_scope` and allows users to customize the allocated buffers' data type.
# Example
Before running `set_dtype`:
```python
@T.prim_func
def main(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
# with T.block("root"):
B = T.alloc_buffer((128, 128), "float16")
for i, j in T.grid(128, 128):
with T.block("B"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(A[vi, vj])
T.writes(B[vi, vj])
B_subregion0 = T.match_buffer(B[vi, vj], (), "float16", offset_factor=1)
B_subregion0[()] = T.Cast("float16", A[vi, vj] * T.float32(2))
for i, j in T.grid(128, 128):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(B[vi, vj])
T.writes(C[vi, vj])
B_subregion1 = T.match_buffer(B[vi, vj], (), "float16", offset_factor=1)
C[vi, vj] = T.Cast("float32", B_subregion1[()]) + T.float32(1)
```
then we perform the `set_dtype` schedule:
```python
sch = tir.Schedule(before_set_dtype)
sch.set_dtype("B", buffer_index=0, dtype="float16")
print(sch.mod["main"].script())
```
we get transformed code:
```python
@T.prim_func
def after_set_dtype(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
B = T.alloc_buffer((128, 128), dtype="float16")
for i, j in T.grid(128, 128):
with T.block("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = T.cast(A[vi, vj] * 2.0, "float16")
for i, j in T.grid(128, 128):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j]
C[vi, vj] = T.cast(B[vi, vj], "float32") + 1.0
```
where data type conversions are inserted automatically.
# Other Usage
Using the combination of `cache_read` + `set_dtype` can help us load data from the memory hierarchy while converting data to the desired type.
cc @Hzfengsy @vinx13 @junrushao
--
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] yzh119 commented on pull request #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "yzh119 (via GitHub)" <gi...@apache.org>.
yzh119 commented on PR #14316:
URL: https://github.com/apache/tvm/pull/14316#issuecomment-1472182957
Adding a `danger` flag sounds good to me, any ideas on how to make proper abstractions for `danger`?
--
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] junrushao commented on pull request #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on PR #14316:
URL: https://github.com/apache/tvm/pull/14316#issuecomment-1475282481
@yzh119 any updates
--
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] Hzfengsy commented on pull request #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "Hzfengsy (via GitHub)" <gi...@apache.org>.
Hzfengsy commented on PR #14316:
URL: https://github.com/apache/tvm/pull/14316#issuecomment-1472168710
Thanks @quic-sanirudh for the suggestion. I agree that it's a really useful primitive/tool for optimizations. It makes sense if we introduce a `danger` mode or flag :-)
--
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] yzh119 commented on pull request #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "yzh119 (via GitHub)" <gi...@apache.org>.
yzh119 commented on PR #14316:
URL: https://github.com/apache/tvm/pull/14316#issuecomment-1475603285
`unsafe_set_dtype` sounds good to me, I have updated code accordingly.
--
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] quic-sanirudh commented on pull request #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "quic-sanirudh (via GitHub)" <gi...@apache.org>.
quic-sanirudh commented on PR #14316:
URL: https://github.com/apache/tvm/pull/14316#issuecomment-1472146828
> It's a bit tricky since it would change the behavior of the PrimFunc (aka get different result before and after the schedule)
There seems to be some interest in introducing schedule primitives like these that might cause different output. This certainly seems like a useful primitive to have, but since we don't normally want schedule primitives to cause the output to change, primitives like these might be hard to introduce.
Would it make sense to have some discussion on perhaps introducing a different scheduling class that allows these kinds of primitives, or explicitly setting a "danger" flag in the Schedule class before using these primitives. That way, we could still allow primitives like these and make sure that a user knows the risk of using a primitive that could cause the output to change.
--
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] yzh119 commented on a diff in pull request #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "yzh119 (via GitHub)" <gi...@apache.org>.
yzh119 commented on code in PR #14316:
URL: https://github.com/apache/tvm/pull/14316#discussion_r1142962077
##########
include/tvm/tir/schedule/schedule.h:
##########
@@ -584,13 +584,21 @@ class ScheduleNode : public runtime::Object {
virtual void StorageAlign(const BlockRV& block_rv, int buffer_index, int axis, int factor,
int offset) = 0;
/*!
- * \brief Set the storage scope of a buffer, where the buffer is specified by the a block and a
+ * \brief Set the storage scope of a buffer, where the buffer is specified by a block and a
* write-index
* \param block_rv The producer block of the buffer
* \param buffer_index The index of the buffer in block's write region
* \param storage_scope The storage scope to be set
*/
virtual void SetScope(const BlockRV& block_rv, int buffer_index, const String& storage_scope) = 0;
+ /*!
+ * \brief Set the data type of a buffer, where the buffer is specified by a block and a
+ * write-index
+ * \param block_rv The producer block of the buffer
+ * \param buffer_index the index of the buffer in block's write region
+ * \param dtype The data type to be set
+ */
+ virtual void SetDType(const BlockRV& block_rv, int buffer_index, const String& dtype) = 0;
Review Comment:
Fixed :)
--
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 #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "tvm-bot (via GitHub)" <gi...@apache.org>.
tvm-bot commented on PR #14316:
URL: https://github.com/apache/tvm/pull/14316#issuecomment-1471642090
<!---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: `tensorir` <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] junrushao commented on a diff in pull request #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on code in PR #14316:
URL: https://github.com/apache/tvm/pull/14316#discussion_r1141615769
##########
include/tvm/tir/schedule/schedule.h:
##########
@@ -584,13 +584,21 @@ class ScheduleNode : public runtime::Object {
virtual void StorageAlign(const BlockRV& block_rv, int buffer_index, int axis, int factor,
int offset) = 0;
/*!
- * \brief Set the storage scope of a buffer, where the buffer is specified by the a block and a
+ * \brief Set the storage scope of a buffer, where the buffer is specified by a block and a
* write-index
* \param block_rv The producer block of the buffer
* \param buffer_index The index of the buffer in block's write region
* \param storage_scope The storage scope to be set
*/
virtual void SetScope(const BlockRV& block_rv, int buffer_index, const String& storage_scope) = 0;
+ /*!
+ * \brief Set the data type of a buffer, where the buffer is specified by a block and a
+ * write-index
+ * \param block_rv The producer block of the buffer
+ * \param buffer_index the index of the buffer in block's write region
+ * \param dtype The data type to be set
+ */
+ virtual void SetDType(const BlockRV& block_rv, int buffer_index, const String& dtype) = 0;
Review Comment:
looks like we missed a couple of places lol
--
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 #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on PR #14316:
URL: https://github.com/apache/tvm/pull/14316#issuecomment-1476217562
oops, @yzh119 please also check the c++ apis @junrushao mentioned :)
--
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 #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on PR #14316:
URL: https://github.com/apache/tvm/pull/14316#issuecomment-1473977633
how about `unsafe_set_dtype`, with document to users that this operation can change the model behavior.
--
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] junrushao merged pull request #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao merged PR #14316:
URL: https://github.com/apache/tvm/pull/14316
--
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] junrushao commented on pull request #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on PR #14316:
URL: https://github.com/apache/tvm/pull/14316#issuecomment-1478874076
I don't know why but now it requires committer approval to run github actions for existing contributors. I'm merging this in after the actions got green
--
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] Hzfengsy commented on pull request #14316: [TensorIR] New schedule primitive `set_dtype`
Posted by "Hzfengsy (via GitHub)" <gi...@apache.org>.
Hzfengsy commented on PR #14316:
URL: https://github.com/apache/tvm/pull/14316#issuecomment-1471836797
It's a bit tricky since it would change the behavior of the PrimFunc (aka get different result before and after the schedule)
--
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