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/06/22 23:17:20 UTC
[GitHub] [tvm] yzh119 opened a new pull request, #15144: [TensorIR][Schedule] New schedule primitive `hide_buffer_access`
yzh119 opened a new pull request, #15144:
URL: https://github.com/apache/tvm/pull/15144
# Motivation
Currently, our `tensorize` schedule primitives rely on buffer read/write regions in the given block to perform pattern matching. However, for workloads such as block sparse operators, the read/write regions include some indices arraies that may fail `tensorize` primitive.
In SparseTIR we introduce a new schedule primitive called `hide_buffer_access` which allows us to hide certain buffer regions in a block so that the read/write buffer regions would be recognized by the `tensorize` primitive to further utilize tensor acceleration units.
This PR upstreams this schedule primitive to TensorIR mainline.
# The schedule primitive interface
```python
def hide_buffer_access(self, block: BlockRV, buf_type: str, buf_index_array: List[int]) -> None:
"""Hide some buffer access in a given block.
Parameters
----------
block : BlockRV
The block where we hide read access.
buf_type : str
The buffer type: "read"/"write".
buf_index_array : List[int]
The array of buffer indices we hide access.
"""
pass
```
# Example
```python
@T.prim_func
def indirect_mem_access(a: T.handle, idx_a: T.handle, b: T.handle, idx_b: T.handle) -> None:
A = T.match_buffer(a, [128], dtype="float32")
IA = T.match_buffer(idx_a, [10], dtype="int32")
B = T.match_buffer(b, [128], dtype="float32")
IB = T.match_buffer(idx_b, [10], dtype="int32")
for i in range(10):
with T.block("B"):
vi = T.axis.spatial(10, i)
T.reads(A[IA[vi]], IA[vi])
T.writes(B[IB[vi]], IB[vi])
B[IB[vi]] = A[IA[vi]]
```
After we perform hiding buffer access to `IA[vi]` via:
```python
sch = tir.Schedule(indirect_mem_access, debug_mask="all")
block_b = sch.get_block("B")
sch.hide_buffer_access(block_b, "write", [1])
```
the desired transformed IR would be:
```python
@T.prim_func
def indirect_mem_access_hide_ia(a: T.handle, idx_a: T.handle, b: T.handle, idx_b: T.handle) -> None:
A = T.match_buffer(a, [128], dtype="float32")
IA = T.match_buffer(idx_a, [10], dtype="int32")
B = T.match_buffer(b, [128], dtype="float32")
IB = T.match_buffer(idx_b, [10], dtype="int32")
for i in range(10):
with T.block("B"):
vi = T.axis.spatial(10, i)
T.reads(A[IA[vi]])
T.writes(B[IB[vi]], IB[vi])
B[IB[vi]] = A[IA[vi]]
```
The existing passes/schedules would not be influenced by this PR.
cc @junrushao @MasterJH5574 @masahi
--
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 #15144: [TensorIR][Schedule] New schedule primitive `unsafe_hide_buffer_access`
Posted by "quic-sanirudh (via GitHub)" <gi...@apache.org>.
quic-sanirudh commented on PR #15144:
URL: https://github.com/apache/tvm/pull/15144#issuecomment-1605253437
> Hi @quic-sanirudh, thank you for your suggestions, I have marked the schedule as unsafe and added some docstrings explaining it.
Thanks for taking my suggestion. Looks good to me now
--
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 #15144: [TensorIR][Schedule] New schedule primitive `hide_buffer_access`
Posted by "tvm-bot (via GitHub)" <gi...@apache.org>.
tvm-bot commented on PR #15144:
URL: https://github.com/apache/tvm/pull/15144#issuecomment-1603427562
<!---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`, `schedule` <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] yzh119 merged pull request #15144: [TensorIR][Schedule] New schedule primitive `unsafe_hide_buffer_access`
Posted by "yzh119 (via GitHub)" <gi...@apache.org>.
yzh119 merged PR #15144:
URL: https://github.com/apache/tvm/pull/15144
--
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 #15144: [TensorIR][Schedule] New schedule primitive `hide_buffer_access`
Posted by "yzh119 (via GitHub)" <gi...@apache.org>.
yzh119 commented on code in PR #15144:
URL: https://github.com/apache/tvm/pull/15144#discussion_r1240391281
##########
src/tir/schedule/primitive/hide_buffer_access.cc:
##########
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include "../../transforms/ir_utils.h"
+#include "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+/******** Error Classes ********/
+
+namespace {
+class BufTypeError : public ScheduleError {
+ public:
+ explicit BufTypeError(IRModule mod, const String& buf_type)
+ : mod_(std::move(mod)), buf_type_(buf_type) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: Invalid buffer type for hide_buffer_access schedule.";
+ }
+
+ String DetailRenderTemplate() const final {
+ return "The buffer type for hide_buffer_access schedule should either be 'read'"
+ " or 'write', got " +
+ buf_type_ + " instead.";
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ private:
+ IRModule mod_;
+ String buf_type_;
+};
+
+class InvalidIndexError : public ScheduleError {
+ public:
+ explicit InvalidIndexError(IRModule mod, int num_access_regions, int buf_idx)
+ : mod_(std::move(mod)), num_access_regions_(num_access_regions), buf_idx_(buf_idx) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: Invalid buffer index array for hide_buffer_access schedule.";
+ }
+
+ String DetailRenderTemplate() const final {
+ return "The buffer index array for hide_buffer_access schedule should be a list of integers"
+ " between 0 and " +
+ std::to_string(num_access_regions_ - 1) + ", got " + std::to_string(buf_idx_) +
+ " instead.";
+ }
+
+ IRModule mod() const final { return mod_; }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ private:
+ IRModule mod_;
+ int num_access_regions_;
+ int buf_idx_;
+};
+
+} // namespace
+
+/******** Implementation ********/
+
+void HideBufferAccess(ScheduleState self, const StmtSRef& block_sref, const String& buf_type,
+ const Array<IntImm>& buf_index_array) {
+ /*!
+ * Check:
+ * - validity of buf_index_array
+ * - validity of buf_type
+ */
+ const BlockNode* block = TVM_SREF_TO_BLOCK(block_sref);
+ int num_access_regions = 0;
+ if (buf_type == "read") {
+ num_access_regions = block->reads.size();
+ } else if (buf_type == "write") {
+ num_access_regions = block->writes.size();
+ } else {
+ throw BufTypeError(self->mod, buf_type);
+ }
+
+ std::set<int> buf_indices;
+ for (const IntImm& buf_idx : buf_index_array) {
+ int buf_idx_val = buf_idx->value;
+ if (buf_idx_val >= 0 && buf_idx_val < num_access_regions) {
+ buf_indices.insert(buf_idx_val);
+ } else {
+ throw InvalidIndexError(self->mod, num_access_regions, buf_idx_val);
+ }
+ }
+
+ /* Step 0: Collect new buffer access regions. */
+
+ Array<BufferRegion> reads, writes;
+
+ if (buf_type == "read") {
+ for (size_t i = 0; i < block->reads.size(); ++i) {
+ if (!buf_indices.count(i)) {
+ reads.push_back(block->reads[i]);
+ }
+ }
+ writes = block->writes;
+ } else if (buf_type == "write") {
+ for (size_t i = 0; i < block->writes.size(); ++i) {
+ if (!buf_indices.count(i)) {
+ writes.push_back(block->writes[i]);
+ }
+ }
+ reads = block->reads;
+ } else {
+ CHECK(false) << "Unrecognized buffer type " << buf_type << ", only support read/write";
+ }
+
+ /* Step 1: Replace old block with the new block */
+
+ auto n = make_object<BlockNode>(*block);
+ n->reads = reads;
+ n->writes = writes;
+ Block new_block = Block(n);
+ Map<Block, Block> blk_map;
+ blk_map.Set(GetRef<Block>(block), new_block);
+ self->Replace(block_sref, new_block, blk_map);
Review Comment:
Yes you are right, this schedule primitive should be used as last steps (right before tensorize) during scheduling.
It makes sense to me to mark it as `unsafe`.
--
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 a diff in pull request #15144: [TensorIR][Schedule] New schedule primitive `hide_buffer_access`
Posted by "quic-sanirudh (via GitHub)" <gi...@apache.org>.
quic-sanirudh commented on code in PR #15144:
URL: https://github.com/apache/tvm/pull/15144#discussion_r1239867418
##########
src/tir/schedule/primitive/hide_buffer_access.cc:
##########
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include "../../transforms/ir_utils.h"
+#include "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+/******** Error Classes ********/
+
+namespace {
+class BufTypeError : public ScheduleError {
+ public:
+ explicit BufTypeError(IRModule mod, const String& buf_type)
+ : mod_(std::move(mod)), buf_type_(buf_type) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: Invalid buffer type for hide_buffer_access schedule.";
+ }
+
+ String DetailRenderTemplate() const final {
+ return "The buffer type for hide_buffer_access schedule should either be 'read'"
+ " or 'write', got " +
+ buf_type_ + " instead.";
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ private:
+ IRModule mod_;
+ String buf_type_;
+};
+
+class InvalidIndexError : public ScheduleError {
+ public:
+ explicit InvalidIndexError(IRModule mod, int num_access_regions, int buf_idx)
+ : mod_(std::move(mod)), num_access_regions_(num_access_regions), buf_idx_(buf_idx) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: Invalid buffer index array for hide_buffer_access schedule.";
+ }
+
+ String DetailRenderTemplate() const final {
+ return "The buffer index array for hide_buffer_access schedule should be a list of integers"
+ " between 0 and " +
+ std::to_string(num_access_regions_ - 1) + ", got " + std::to_string(buf_idx_) +
+ " instead.";
+ }
+
+ IRModule mod() const final { return mod_; }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ private:
+ IRModule mod_;
+ int num_access_regions_;
+ int buf_idx_;
+};
+
+} // namespace
+
+/******** Implementation ********/
+
+void HideBufferAccess(ScheduleState self, const StmtSRef& block_sref, const String& buf_type,
+ const Array<IntImm>& buf_index_array) {
+ /*!
+ * Check:
+ * - validity of buf_index_array
+ * - validity of buf_type
+ */
+ const BlockNode* block = TVM_SREF_TO_BLOCK(block_sref);
+ int num_access_regions = 0;
+ if (buf_type == "read") {
+ num_access_regions = block->reads.size();
+ } else if (buf_type == "write") {
+ num_access_regions = block->writes.size();
+ } else {
+ throw BufTypeError(self->mod, buf_type);
+ }
+
+ std::set<int> buf_indices;
+ for (const IntImm& buf_idx : buf_index_array) {
+ int buf_idx_val = buf_idx->value;
+ if (buf_idx_val >= 0 && buf_idx_val < num_access_regions) {
+ buf_indices.insert(buf_idx_val);
+ } else {
+ throw InvalidIndexError(self->mod, num_access_regions, buf_idx_val);
+ }
+ }
+
+ /* Step 0: Collect new buffer access regions. */
+
+ Array<BufferRegion> reads, writes;
+
+ if (buf_type == "read") {
+ for (size_t i = 0; i < block->reads.size(); ++i) {
+ if (!buf_indices.count(i)) {
+ reads.push_back(block->reads[i]);
+ }
+ }
+ writes = block->writes;
+ } else if (buf_type == "write") {
+ for (size_t i = 0; i < block->writes.size(); ++i) {
+ if (!buf_indices.count(i)) {
+ writes.push_back(block->writes[i]);
+ }
+ }
+ reads = block->reads;
+ } else {
+ CHECK(false) << "Unrecognized buffer type " << buf_type << ", only support read/write";
+ }
+
+ /* Step 1: Replace old block with the new block */
+
+ auto n = make_object<BlockNode>(*block);
+ n->reads = reads;
+ n->writes = writes;
+ Block new_block = Block(n);
+ Map<Block, Block> blk_map;
+ blk_map.Set(GetRef<Block>(block), new_block);
+ self->Replace(block_sref, new_block, blk_map);
Review Comment:
Considering that we use the `Replace` function to update the `reads/writes` for the new block, the dependence analysis would also be updated now, but the new updated dependences might not be correct as there could be missing dependences which in turn could cause correctness issues if I understand correctly.
Could we rename the primitive to `unsafe_hide_buffer_access` similar to the existing `unsafe_set_dtype` to notify users of the potential to get incorrect output?
--
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 #15144: [TensorIR][Schedule] New schedule primitive `hide_buffer_access`
Posted by "yzh119 (via GitHub)" <gi...@apache.org>.
yzh119 commented on PR #15144:
URL: https://github.com/apache/tvm/pull/15144#issuecomment-1605195313
Hi @quic-sanirudh, thank you for your suggestions, I have marked the schedule as unsafe and added some docstrings explaining 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