You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/08/09 20:32:20 UTC
[GitHub] [tvm] vinx13 opened a new pull request #8693: [TensorIR][M2a] Storage Align
vinx13 opened a new pull request #8693:
URL: https://github.com/apache/tvm/pull/8693
This PR is part of the TensorIR upstreaming effort (#7527), which adds the one schedule primitive `storage_align`.
Co-authored-by: Siyuan Feng <Hz...@sjtu.edu.cn>
Co-authored-by: Ruihang Lai <la...@qq.com>
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r685605104
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,313 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment " << buffer_->name << " has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
Review comment:
Yeah that's correct
--
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 change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
vinx13 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r687999574
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,315 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment of, " << buffer_->name << ", has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ static int CheckAndUpdate(const IRModule& mod, const Buffer& buffer, int axis) {
+ int ndim = static_cast<int>(buffer->shape.size());
+ if (axis < -ndim || axis >= ndim) {
+ throw StorageAlignAxisOutOfRangeError(mod, buffer, axis);
+ }
+ // If axis is negative, convert it to a non-negative one.
+ if (axis < 0) {
+ axis += ndim;
+ }
+ return axis;
+ }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+ int axis_;
+};
+
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ // Climb up along the sref tree, and find the block where `buffer` is in alloc_buffers or
+ // match_buffers.
+ const StmtSRefNode* defining_site_sref = block_sref.get();
+ while (defining_site_sref != nullptr) {
+ const auto* block = defining_site_sref->StmtAs<BlockNode>();
+ // If this sref is not a block sref, skip it.
+ if (block == nullptr) {
+ defining_site_sref = defining_site_sref->parent;
+ continue;
+ }
+ // Try to find the buffer in `allloc_buffers`
+ for (const Buffer& alloc_buffer : block->alloc_buffers) {
+ if (buffer.same_as(alloc_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), true};
+ }
+ }
+ // We do not allow the buffer being defined in `match_buffer`.
+ for (const MatchBufferRegion match_buffer : block->match_buffers) {
+ if (buffer.same_as(match_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), false};
+ }
+ }
+ defining_site_sref = defining_site_sref->parent;
+ }
+ // If we cannot find the defining site block, it means that the buffer must be in the function's
+ // buffer_map, which isn't an intermediate buffer.
+ return {NullOpt, false};
+}
+
+class NonAllocatedBufferError : public ScheduleError {
+ public:
+ explicit NonAllocatedBufferError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input buffer is not allocated by a block. This means the buffer is "
+ " either a function parameter or defined in `match_buffer` of a block.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input buffer " << buffer_->name
+ << " is not allocated by a block. This means the buffer is either a function parameter or "
+ "defined in `match_buffer` of a block.";
+ return os.str();
+ }
+
+ static void CheckBufferAllocated(const IRModule& mod, const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ Optional<StmtSRef> defining_site_sref;
+ bool is_alloc;
+ std::tie(defining_site_sref, is_alloc) = GetBufferDefiningSite(block_sref, buffer);
+ if (!defining_site_sref || !is_alloc) {
+ throw NonAllocatedBufferError(mod, buffer);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+};
+
+class StorageAlignInvalidFactorError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidFactorError(IRModule mod, int factor)
+ : mod_(std::move(mod)), factor_(factor) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `factor` of storage_align is expected to be a positive "
+ "number.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input `factor` of storage_align is expected to be a positive number. However, the "
+ "input `factor` is "
+ << factor_ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ static void Check(const IRModule& mod, int factor) {
+ if (factor <= 0) {
+ throw StorageAlignInvalidFactorError(mod, factor);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ int factor_;
+};
+
+class StorageAlignInvalidAnnotationError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidAnnotationError(IRModule mod, Block block)
+ : mod_(std::move(mod)), block_(std::move(block)) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The block annotation for storage align is expected to be an array of "
+ "3-integer-tuples (axis, factor, offset).";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The block annotation for storage align is expected to be an array of 3-integer-tuples "
+ "(axis, factor, offset). However, the block annotation with key "
+ << attr::buffer_dim_align << " of the block {0} is "
+ << block_->annotations.at(attr::buffer_dim_align) << ", which is unexpected.";
+ return os.str();
+ }
+
+ static Array<Array<Array<Integer>>> CheckAndGetAnnotation(const IRModule& mod,
+ const Block& block) {
+ // Get existing annotation value.
+ auto it = block->annotations.find(attr::buffer_dim_align);
+ if (it != block->annotations.end()) {
+ if (!IsValidAnnotation(block, (*it).second)) {
+ throw StorageAlignInvalidAnnotationError(mod, block);
+ }
+ return Downcast<Array<Array<Array<Integer>>>>((*it).second);
+ }
+
+ // Create new annotation value
+ Array<Array<Array<Integer>>> storage_align_annotation;
+ storage_align_annotation.resize(block->writes.size());
Review comment:
We requires the storage_align annotation array to be the same size as `block->writes` (if it exists), each element is used for each write buffer. Alternatively, we can also store buffer index in the annotation. What do you think?
--
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] junrushao1994 merged pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 merged pull request #8693:
URL: https://github.com/apache/tvm/pull/8693
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r686489336
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,315 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment of, " << buffer_->name << ", has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ static int CheckAndUpdate(const IRModule& mod, const Buffer& buffer, int axis) {
+ int ndim = static_cast<int>(buffer->shape.size());
+ if (axis < -ndim || axis >= ndim) {
+ throw StorageAlignAxisOutOfRangeError(mod, buffer, axis);
+ }
+ // If axis is negative, convert it to a non-negative one.
+ if (axis < 0) {
+ axis += ndim;
+ }
+ return axis;
+ }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+ int axis_;
+};
+
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ // Climb up along the sref tree, and find the block where `buffer` is in alloc_buffers or
+ // match_buffers.
+ const StmtSRefNode* defining_site_sref = block_sref.get();
+ while (defining_site_sref != nullptr) {
+ const auto* block = defining_site_sref->StmtAs<BlockNode>();
+ // If this sref is not a block sref, skip it.
+ if (block == nullptr) {
+ defining_site_sref = defining_site_sref->parent;
+ continue;
+ }
+ // Try to find the buffer in `allloc_buffers`
+ for (const Buffer& alloc_buffer : block->alloc_buffers) {
+ if (buffer.same_as(alloc_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), true};
+ }
+ }
+ // We do not allow the buffer being defined in `match_buffer`.
+ for (const MatchBufferRegion match_buffer : block->match_buffers) {
+ if (buffer.same_as(match_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), false};
+ }
+ }
+ defining_site_sref = defining_site_sref->parent;
+ }
+ // If we cannot find the defining site block, it means that the buffer must be in the function's
+ // buffer_map, which isn't an intermediate buffer.
+ return {NullOpt, false};
+}
+
+class NonAllocatedBufferError : public ScheduleError {
+ public:
+ explicit NonAllocatedBufferError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input buffer is not allocated by a block. This means the buffer is "
+ " either a function parameter or defined in `match_buffer` of a block.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input buffer " << buffer_->name
+ << " is not allocated by a block. This means the buffer is either a function parameter or "
+ "defined in `match_buffer` of a block.";
+ return os.str();
+ }
+
+ static void CheckBufferAllocated(const IRModule& mod, const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ Optional<StmtSRef> defining_site_sref;
+ bool is_alloc;
+ std::tie(defining_site_sref, is_alloc) = GetBufferDefiningSite(block_sref, buffer);
+ if (!defining_site_sref || !is_alloc) {
+ throw NonAllocatedBufferError(mod, buffer);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+};
+
+class StorageAlignInvalidFactorError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidFactorError(IRModule mod, int factor)
+ : mod_(std::move(mod)), factor_(factor) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `factor` of storage_align is expected to be a positive "
+ "number.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input `factor` of storage_align is expected to be a positive number. However, the "
+ "input `factor` is "
+ << factor_ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ static void Check(const IRModule& mod, int factor) {
+ if (factor <= 0) {
+ throw StorageAlignInvalidFactorError(mod, factor);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ int factor_;
+};
+
+class StorageAlignInvalidAnnotationError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidAnnotationError(IRModule mod, Block block)
+ : mod_(std::move(mod)), block_(std::move(block)) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The block annotation for storage align is expected to be an array of "
+ "3-integer-tuples (axis, factor, offset).";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The block annotation for storage align is expected to be an array of 3-integer-tuples "
+ "(axis, factor, offset). However, the block annotation with key "
+ << attr::buffer_dim_align << " of the block {0} is "
+ << block_->annotations.at(attr::buffer_dim_align) << ", which is unexpected.";
+ return os.str();
+ }
+
+ static Array<Array<Array<Integer>>> CheckAndGetAnnotation(const IRModule& mod,
+ const Block& block) {
+ // Get existing annotation value.
+ auto it = block->annotations.find(attr::buffer_dim_align);
+ if (it != block->annotations.end()) {
+ if (!IsValidAnnotation(block, (*it).second)) {
+ throw StorageAlignInvalidAnnotationError(mod, block);
+ }
+ return Downcast<Array<Array<Array<Integer>>>>((*it).second);
Review comment:
The Array of Array of Array of Integers looks a little bit confusing to me 😢
What about we do this:
- using StorageAlignTuple = Array<Integer>; // (buffer_idx, axis, factor, offset)
- using StorageAlignAnnotation = Array<StorageAlignTuple>; // unordered array of `StorageAlignTuple`
--
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] junrushao1994 commented on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896533659
I dont have strong opinion though, just quick thoughts when reading the error checking code path. @tqchen What do you think?
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r686491273
##########
File path: src/tir/schedule/analysis.h
##########
@@ -224,6 +237,17 @@ std::vector<TypedPackedFunc<CommReducer(DataType)>> GetReducerGetters();
bool FromIdentityCombiner(const PrimExpr& identity, const BufferStore& combiner,
CommReducer* result_reducer, PrimExpr* lhs, PrimExpr* rhs);
+/******** Annotation ********/
+
+/*!
+ * \brief Create a new block with the given annotation added
+ * \param block The block with original annotation
+ * \param attr_key The annotation key to be added
+ * \param attr_value The annotation value to be added
+ * \return A new block with the given annotation as its last annotation
+ */
+Block WithAnnotation(const BlockNode* block, const String& attr_key, const ObjectRef& attr_value);
Review comment:
Move this to a new file `transform.h` given it is not analysis
--
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] MasterJH5574 commented on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
MasterJH5574 commented on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896919488
Yes, this is also my concern. But on the other hand, it's also not straightforward for users to know the `write_buffer_index` of a buffer in a block. So the two ways are similar from this angle.
To clarify, most of the time a block only writes one buffer, and most of the time a block allocate zero or one buffer... Hard to choose...
--
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] junrushao1994 commented on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896501084
@vinx13 Hey I am curious on the semantics of `StorageAlign`: why it works on an output buffer, instead of the buffer allocation side? I imagine that it could simplify the logic a bit (don't need to go up straight to find the alloc side)
--
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] tkonolige commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
tkonolige commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r686141549
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,313 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment " << buffer_->name << " has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
Review comment:
Could we add where buffers are declared to locations of interest? Then we can point to them. But it might not be helpful if it always point to the function inputs.
--
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] MasterJH5574 removed a comment on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
MasterJH5574 removed a comment on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896446136
Thanks @jinhongyii @junrushao1994! I'm gonna merge it and send it to mainline.
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r686485786
##########
File path: src/tir/transforms/compact_buffer_region.cc
##########
@@ -303,28 +303,82 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
support::Arena arena_;
};
+/*! \brief Collect storage alignment information from block annotations. */
+class StorageAlignCollector : public StmtVisitor {
+ public:
+ static std::unordered_map<Buffer, Array<Array<Integer>>, ObjectPtrHash, ObjectPtrEqual> Collect(
+ const PrimFunc& f) {
+ StorageAlignCollector collector;
+ collector(f->body);
+ return std::move(collector.storage_align_);
+ }
+
+ private:
+ void VisitStmt_(const BlockNode* op) final {
+ auto it = op->annotations.find(attr::buffer_dim_align);
+ if (it != op->annotations.end()) {
+ const auto& storage_align = Downcast<Array<Array<Array<Integer>>>>((*it).second);
Review comment:
We don't need the const ref here
```suggestion
auto storage_align = Downcast<Array<Array<Array<Integer>>>>((*it).second);
```
--
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 pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
vinx13 commented on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896516267
> @vinx13 Hey I am curious on the semantics of `StorageAlign`: why it works on an output buffer, instead of the buffer allocation side? I imagine that it could simplify the logic a bit (don't need to go up straight to find the alloc side)
The semantic is to annotate producer block (multiple producers are allowed, the annotations will be unified during lowering). IMO the easiest way to specify the buffer in the API is block + write_buffer_index as `write_buffer_index` is more obvious than `alloc_buffer_index`, what do you think?
--
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] junrushao1994 commented on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896305858
Will do the review today!
--
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] junrushao1994 commented on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-897841794
Okay, seems that we reach a conclusion to use `write_buffer_index`.
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r686489336
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,315 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment of, " << buffer_->name << ", has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ static int CheckAndUpdate(const IRModule& mod, const Buffer& buffer, int axis) {
+ int ndim = static_cast<int>(buffer->shape.size());
+ if (axis < -ndim || axis >= ndim) {
+ throw StorageAlignAxisOutOfRangeError(mod, buffer, axis);
+ }
+ // If axis is negative, convert it to a non-negative one.
+ if (axis < 0) {
+ axis += ndim;
+ }
+ return axis;
+ }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+ int axis_;
+};
+
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ // Climb up along the sref tree, and find the block where `buffer` is in alloc_buffers or
+ // match_buffers.
+ const StmtSRefNode* defining_site_sref = block_sref.get();
+ while (defining_site_sref != nullptr) {
+ const auto* block = defining_site_sref->StmtAs<BlockNode>();
+ // If this sref is not a block sref, skip it.
+ if (block == nullptr) {
+ defining_site_sref = defining_site_sref->parent;
+ continue;
+ }
+ // Try to find the buffer in `allloc_buffers`
+ for (const Buffer& alloc_buffer : block->alloc_buffers) {
+ if (buffer.same_as(alloc_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), true};
+ }
+ }
+ // We do not allow the buffer being defined in `match_buffer`.
+ for (const MatchBufferRegion match_buffer : block->match_buffers) {
+ if (buffer.same_as(match_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), false};
+ }
+ }
+ defining_site_sref = defining_site_sref->parent;
+ }
+ // If we cannot find the defining site block, it means that the buffer must be in the function's
+ // buffer_map, which isn't an intermediate buffer.
+ return {NullOpt, false};
+}
+
+class NonAllocatedBufferError : public ScheduleError {
+ public:
+ explicit NonAllocatedBufferError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input buffer is not allocated by a block. This means the buffer is "
+ " either a function parameter or defined in `match_buffer` of a block.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input buffer " << buffer_->name
+ << " is not allocated by a block. This means the buffer is either a function parameter or "
+ "defined in `match_buffer` of a block.";
+ return os.str();
+ }
+
+ static void CheckBufferAllocated(const IRModule& mod, const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ Optional<StmtSRef> defining_site_sref;
+ bool is_alloc;
+ std::tie(defining_site_sref, is_alloc) = GetBufferDefiningSite(block_sref, buffer);
+ if (!defining_site_sref || !is_alloc) {
+ throw NonAllocatedBufferError(mod, buffer);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+};
+
+class StorageAlignInvalidFactorError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidFactorError(IRModule mod, int factor)
+ : mod_(std::move(mod)), factor_(factor) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `factor` of storage_align is expected to be a positive "
+ "number.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input `factor` of storage_align is expected to be a positive number. However, the "
+ "input `factor` is "
+ << factor_ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ static void Check(const IRModule& mod, int factor) {
+ if (factor <= 0) {
+ throw StorageAlignInvalidFactorError(mod, factor);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ int factor_;
+};
+
+class StorageAlignInvalidAnnotationError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidAnnotationError(IRModule mod, Block block)
+ : mod_(std::move(mod)), block_(std::move(block)) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The block annotation for storage align is expected to be an array of "
+ "3-integer-tuples (axis, factor, offset).";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The block annotation for storage align is expected to be an array of 3-integer-tuples "
+ "(axis, factor, offset). However, the block annotation with key "
+ << attr::buffer_dim_align << " of the block {0} is "
+ << block_->annotations.at(attr::buffer_dim_align) << ", which is unexpected.";
+ return os.str();
+ }
+
+ static Array<Array<Array<Integer>>> CheckAndGetAnnotation(const IRModule& mod,
+ const Block& block) {
+ // Get existing annotation value.
+ auto it = block->annotations.find(attr::buffer_dim_align);
+ if (it != block->annotations.end()) {
+ if (!IsValidAnnotation(block, (*it).second)) {
+ throw StorageAlignInvalidAnnotationError(mod, block);
+ }
+ return Downcast<Array<Array<Array<Integer>>>>((*it).second);
Review comment:
The Array of Array of Array of Integers looks a little bit confusing to me 😢
What about we do this:
- `using StorageAlignTuple = Array\<Integer>;` // (buffer_idx, axis, factor, offset)
- `using StorageAlignAnnotation = Array\<StorageAlignTuple>;` // unordered array of `StorageAlignTuple`
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,315 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment of, " << buffer_->name << ", has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ static int CheckAndUpdate(const IRModule& mod, const Buffer& buffer, int axis) {
+ int ndim = static_cast<int>(buffer->shape.size());
+ if (axis < -ndim || axis >= ndim) {
+ throw StorageAlignAxisOutOfRangeError(mod, buffer, axis);
+ }
+ // If axis is negative, convert it to a non-negative one.
+ if (axis < 0) {
+ axis += ndim;
+ }
+ return axis;
+ }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+ int axis_;
+};
+
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ // Climb up along the sref tree, and find the block where `buffer` is in alloc_buffers or
+ // match_buffers.
+ const StmtSRefNode* defining_site_sref = block_sref.get();
+ while (defining_site_sref != nullptr) {
+ const auto* block = defining_site_sref->StmtAs<BlockNode>();
+ // If this sref is not a block sref, skip it.
+ if (block == nullptr) {
+ defining_site_sref = defining_site_sref->parent;
+ continue;
+ }
+ // Try to find the buffer in `allloc_buffers`
+ for (const Buffer& alloc_buffer : block->alloc_buffers) {
+ if (buffer.same_as(alloc_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), true};
+ }
+ }
+ // We do not allow the buffer being defined in `match_buffer`.
+ for (const MatchBufferRegion match_buffer : block->match_buffers) {
+ if (buffer.same_as(match_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), false};
+ }
+ }
+ defining_site_sref = defining_site_sref->parent;
+ }
+ // If we cannot find the defining site block, it means that the buffer must be in the function's
+ // buffer_map, which isn't an intermediate buffer.
+ return {NullOpt, false};
+}
+
+class NonAllocatedBufferError : public ScheduleError {
+ public:
+ explicit NonAllocatedBufferError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input buffer is not allocated by a block. This means the buffer is "
+ " either a function parameter or defined in `match_buffer` of a block.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input buffer " << buffer_->name
+ << " is not allocated by a block. This means the buffer is either a function parameter or "
+ "defined in `match_buffer` of a block.";
+ return os.str();
+ }
+
+ static void CheckBufferAllocated(const IRModule& mod, const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ Optional<StmtSRef> defining_site_sref;
+ bool is_alloc;
+ std::tie(defining_site_sref, is_alloc) = GetBufferDefiningSite(block_sref, buffer);
+ if (!defining_site_sref || !is_alloc) {
+ throw NonAllocatedBufferError(mod, buffer);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+};
+
+class StorageAlignInvalidFactorError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidFactorError(IRModule mod, int factor)
+ : mod_(std::move(mod)), factor_(factor) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `factor` of storage_align is expected to be a positive "
+ "number.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input `factor` of storage_align is expected to be a positive number. However, the "
+ "input `factor` is "
+ << factor_ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ static void Check(const IRModule& mod, int factor) {
+ if (factor <= 0) {
+ throw StorageAlignInvalidFactorError(mod, factor);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ int factor_;
+};
+
+class StorageAlignInvalidAnnotationError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidAnnotationError(IRModule mod, Block block)
+ : mod_(std::move(mod)), block_(std::move(block)) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The block annotation for storage align is expected to be an array of "
+ "3-integer-tuples (axis, factor, offset).";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The block annotation for storage align is expected to be an array of 3-integer-tuples "
+ "(axis, factor, offset). However, the block annotation with key "
+ << attr::buffer_dim_align << " of the block {0} is "
+ << block_->annotations.at(attr::buffer_dim_align) << ", which is unexpected.";
+ return os.str();
+ }
+
+ static Array<Array<Array<Integer>>> CheckAndGetAnnotation(const IRModule& mod,
+ const Block& block) {
+ // Get existing annotation value.
+ auto it = block->annotations.find(attr::buffer_dim_align);
+ if (it != block->annotations.end()) {
+ if (!IsValidAnnotation(block, (*it).second)) {
+ throw StorageAlignInvalidAnnotationError(mod, block);
+ }
+ return Downcast<Array<Array<Array<Integer>>>>((*it).second);
Review comment:
The Array of Array of Array of Integers looks a little bit confusing to me 😢
What about we do this:
- `using StorageAlignTuple = Array<Integer>;` // (buffer_idx, axis, factor, offset)
- `using StorageAlignAnnotation = Array<StorageAlignTuple>;` // unordered array of `StorageAlignTuple`
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r686490760
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,315 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment of, " << buffer_->name << ", has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ static int CheckAndUpdate(const IRModule& mod, const Buffer& buffer, int axis) {
+ int ndim = static_cast<int>(buffer->shape.size());
+ if (axis < -ndim || axis >= ndim) {
+ throw StorageAlignAxisOutOfRangeError(mod, buffer, axis);
+ }
+ // If axis is negative, convert it to a non-negative one.
+ if (axis < 0) {
+ axis += ndim;
+ }
+ return axis;
+ }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+ int axis_;
+};
+
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ // Climb up along the sref tree, and find the block where `buffer` is in alloc_buffers or
+ // match_buffers.
+ const StmtSRefNode* defining_site_sref = block_sref.get();
+ while (defining_site_sref != nullptr) {
+ const auto* block = defining_site_sref->StmtAs<BlockNode>();
+ // If this sref is not a block sref, skip it.
+ if (block == nullptr) {
+ defining_site_sref = defining_site_sref->parent;
+ continue;
+ }
+ // Try to find the buffer in `allloc_buffers`
+ for (const Buffer& alloc_buffer : block->alloc_buffers) {
+ if (buffer.same_as(alloc_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), true};
+ }
+ }
+ // We do not allow the buffer being defined in `match_buffer`.
+ for (const MatchBufferRegion match_buffer : block->match_buffers) {
+ if (buffer.same_as(match_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), false};
+ }
+ }
+ defining_site_sref = defining_site_sref->parent;
+ }
+ // If we cannot find the defining site block, it means that the buffer must be in the function's
+ // buffer_map, which isn't an intermediate buffer.
+ return {NullOpt, false};
+}
+
+class NonAllocatedBufferError : public ScheduleError {
+ public:
+ explicit NonAllocatedBufferError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input buffer is not allocated by a block. This means the buffer is "
+ " either a function parameter or defined in `match_buffer` of a block.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input buffer " << buffer_->name
+ << " is not allocated by a block. This means the buffer is either a function parameter or "
+ "defined in `match_buffer` of a block.";
+ return os.str();
+ }
+
+ static void CheckBufferAllocated(const IRModule& mod, const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ Optional<StmtSRef> defining_site_sref;
+ bool is_alloc;
+ std::tie(defining_site_sref, is_alloc) = GetBufferDefiningSite(block_sref, buffer);
+ if (!defining_site_sref || !is_alloc) {
+ throw NonAllocatedBufferError(mod, buffer);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+};
+
+class StorageAlignInvalidFactorError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidFactorError(IRModule mod, int factor)
+ : mod_(std::move(mod)), factor_(factor) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `factor` of storage_align is expected to be a positive "
+ "number.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input `factor` of storage_align is expected to be a positive number. However, the "
+ "input `factor` is "
+ << factor_ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ static void Check(const IRModule& mod, int factor) {
+ if (factor <= 0) {
+ throw StorageAlignInvalidFactorError(mod, factor);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ int factor_;
+};
+
+class StorageAlignInvalidAnnotationError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidAnnotationError(IRModule mod, Block block)
+ : mod_(std::move(mod)), block_(std::move(block)) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The block annotation for storage align is expected to be an array of "
+ "3-integer-tuples (axis, factor, offset).";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The block annotation for storage align is expected to be an array of 3-integer-tuples "
+ "(axis, factor, offset). However, the block annotation with key "
+ << attr::buffer_dim_align << " of the block {0} is "
+ << block_->annotations.at(attr::buffer_dim_align) << ", which is unexpected.";
+ return os.str();
+ }
+
+ static Array<Array<Array<Integer>>> CheckAndGetAnnotation(const IRModule& mod,
+ const Block& block) {
+ // Get existing annotation value.
+ auto it = block->annotations.find(attr::buffer_dim_align);
+ if (it != block->annotations.end()) {
+ if (!IsValidAnnotation(block, (*it).second)) {
+ throw StorageAlignInvalidAnnotationError(mod, block);
+ }
+ return Downcast<Array<Array<Array<Integer>>>>((*it).second);
+ }
+
+ // Create new annotation value
+ Array<Array<Array<Integer>>> storage_align_annotation;
+ storage_align_annotation.resize(block->writes.size());
+ return storage_align_annotation;
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {block_}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ static bool IsValidAnnotation(const Block& block, const ObjectRef& anno_value) {
+ if (!anno_value->IsInstance<ArrayNode>()) {
+ return false;
+ }
+ const auto& buffer_annotations = Downcast<Array<ObjectRef>>(anno_value);
+ if (buffer_annotations.size() != block->writes.size()) {
+ return false;
+ }
+ for (const ObjectRef& buffer_annotation : buffer_annotations) {
+ if (!buffer_annotation->IsInstance<ArrayNode>()) {
+ return false;
+ }
+ const auto& dim_annotations = Downcast<Array<ObjectRef>>(buffer_annotation);
+ for (const ObjectRef& dim_annotation : dim_annotations) {
+ if (!dim_annotation->IsInstance<ArrayNode>()) {
+ return false;
+ }
+ const auto& dim_anno = Downcast<Array<ObjectRef>>(dim_annotation);
+ // Check if the annotations are consist of 3-tuples.
+ if (dim_anno.size() != 3) {
+ return false;
+ }
+ for (const ObjectRef& dim_anno_element : dim_anno) {
+ if (!dim_anno_element->IsInstance<IntImmNode>()) {
+ return false;
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ IRModule mod_;
+ Block block_;
+};
+
+void StorageAlign(ScheduleState self, const StmtSRef& block_sref, int buffer_index, int axis,
+ int factor, int offset) {
+ const BlockNode* block_ptr = TVM_SREF_TO_BLOCK(block_ptr, block_sref);
+ Buffer buffer = GetNthWriteBuffer(self, GetRef<Block>(block_ptr), buffer_index);
+ StorageAlignInvalidFactorError::Check(self->mod, factor);
+ axis = StorageAlignAxisOutOfRangeError::CheckAndUpdate(self->mod, buffer, axis);
+ NonAllocatedBufferError::CheckBufferAllocated(self->mod, block_sref, buffer);
+
+ // Step 1: Get existing or create new annotation value.
+ auto storage_align_annotation = StorageAlignInvalidAnnotationError::CheckAndGetAnnotation(
Review comment:
do not use `auto` in this 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] MasterJH5574 edited a comment on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
MasterJH5574 edited a comment on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896746839
I wrote the origin code that "go up straight to find the allocation site", which is first used in primitive `SetScope`. `SetScope` also specify the target buffer by one of its producer block and a `write_buffer_index`. Probably Wuwei just reused the part of code, so it's my fault 🤐:sweat:.
From my POV, the interface was designed like this mainly because `CacheRead` and `CacheWrite` specify the target buffer by a block and a index. Therefore I inherited the interface in `CacheRead` and `CacheWrite`.
At that time I thought that "specifying a buffer directly on the allocation side is impossible", but after a second thought today... Of course we can do that.
So ATM I think specifying the target buffer from the allocation side by `alloc_buffer_index` is much more convenient. Would like to hear your opinions :-)
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r686373290
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,314 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment of, " << buffer_->name << ", has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ static int CheckAndUpdate(const IRModule& mod, const Buffer& buffer, int axis) {
+ int ndim = static_cast<int>(buffer->shape.size());
+ if (axis < -ndim || axis >= ndim) {
+ throw StorageAlignAxisOutOfRangeError(mod, buffer, axis);
+ }
+ // If axis is negative, convert it to a non-negative one.
+ if (axis < 0) {
+ axis += ndim;
+ }
+ return axis;
+ }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+ int axis_;
+};
+
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<StmtSRef, bool> GetBufferDefiningSite(const StmtSRef& block_sref, const Buffer& buffer) {
Review comment:
Use `Optional` to explicitly indicate the case that the buffer is defined in `PrimFuncNode::buffer_map`
```suggestion
std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref, const Buffer& buffer) {
```
--
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] tkonolige commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
tkonolige commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r685575038
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,313 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment " << buffer_->name << " has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
Review comment:
Does the TIR error handling not allow you to point to the buffer?
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,313 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment " << buffer_->name << " has " << ndim
Review comment:
```suggestion
os << "The buffer to set storage alignment of, " << buffer_->name << ", has " << ndim
```
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r686490641
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,315 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment of, " << buffer_->name << ", has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ static int CheckAndUpdate(const IRModule& mod, const Buffer& buffer, int axis) {
+ int ndim = static_cast<int>(buffer->shape.size());
+ if (axis < -ndim || axis >= ndim) {
+ throw StorageAlignAxisOutOfRangeError(mod, buffer, axis);
+ }
+ // If axis is negative, convert it to a non-negative one.
+ if (axis < 0) {
+ axis += ndim;
+ }
+ return axis;
+ }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+ int axis_;
+};
+
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ // Climb up along the sref tree, and find the block where `buffer` is in alloc_buffers or
+ // match_buffers.
+ const StmtSRefNode* defining_site_sref = block_sref.get();
+ while (defining_site_sref != nullptr) {
+ const auto* block = defining_site_sref->StmtAs<BlockNode>();
+ // If this sref is not a block sref, skip it.
+ if (block == nullptr) {
+ defining_site_sref = defining_site_sref->parent;
+ continue;
+ }
+ // Try to find the buffer in `allloc_buffers`
+ for (const Buffer& alloc_buffer : block->alloc_buffers) {
+ if (buffer.same_as(alloc_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), true};
+ }
+ }
+ // We do not allow the buffer being defined in `match_buffer`.
+ for (const MatchBufferRegion match_buffer : block->match_buffers) {
+ if (buffer.same_as(match_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), false};
+ }
+ }
+ defining_site_sref = defining_site_sref->parent;
+ }
+ // If we cannot find the defining site block, it means that the buffer must be in the function's
+ // buffer_map, which isn't an intermediate buffer.
+ return {NullOpt, false};
+}
+
+class NonAllocatedBufferError : public ScheduleError {
+ public:
+ explicit NonAllocatedBufferError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input buffer is not allocated by a block. This means the buffer is "
+ " either a function parameter or defined in `match_buffer` of a block.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input buffer " << buffer_->name
+ << " is not allocated by a block. This means the buffer is either a function parameter or "
+ "defined in `match_buffer` of a block.";
+ return os.str();
+ }
+
+ static void CheckBufferAllocated(const IRModule& mod, const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ Optional<StmtSRef> defining_site_sref;
+ bool is_alloc;
+ std::tie(defining_site_sref, is_alloc) = GetBufferDefiningSite(block_sref, buffer);
+ if (!defining_site_sref || !is_alloc) {
+ throw NonAllocatedBufferError(mod, buffer);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+};
+
+class StorageAlignInvalidFactorError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidFactorError(IRModule mod, int factor)
+ : mod_(std::move(mod)), factor_(factor) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `factor` of storage_align is expected to be a positive "
+ "number.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input `factor` of storage_align is expected to be a positive number. However, the "
+ "input `factor` is "
+ << factor_ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ static void Check(const IRModule& mod, int factor) {
+ if (factor <= 0) {
+ throw StorageAlignInvalidFactorError(mod, factor);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ int factor_;
+};
+
+class StorageAlignInvalidAnnotationError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidAnnotationError(IRModule mod, Block block)
+ : mod_(std::move(mod)), block_(std::move(block)) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The block annotation for storage align is expected to be an array of "
+ "3-integer-tuples (axis, factor, offset).";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The block annotation for storage align is expected to be an array of 3-integer-tuples "
+ "(axis, factor, offset). However, the block annotation with key "
+ << attr::buffer_dim_align << " of the block {0} is "
+ << block_->annotations.at(attr::buffer_dim_align) << ", which is unexpected.";
+ return os.str();
+ }
+
+ static Array<Array<Array<Integer>>> CheckAndGetAnnotation(const IRModule& mod,
+ const Block& block) {
+ // Get existing annotation value.
+ auto it = block->annotations.find(attr::buffer_dim_align);
+ if (it != block->annotations.end()) {
+ if (!IsValidAnnotation(block, (*it).second)) {
+ throw StorageAlignInvalidAnnotationError(mod, block);
+ }
+ return Downcast<Array<Array<Array<Integer>>>>((*it).second);
+ }
+
+ // Create new annotation value
+ Array<Array<Array<Integer>>> storage_align_annotation;
+ storage_align_annotation.resize(block->writes.size());
Review comment:
The problem of doing this is that it could be sort of confusing when we dump the TVM script, so let's keep the array as flat as possible but keep the necessary checks
--
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] MasterJH5574 edited a comment on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
MasterJH5574 edited a comment on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896446136
Thanks @jinhongyii @junrushao1994! I'm gonna merge it and send it to mainline.
--
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] junrushao1994 commented on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-895714416
Thanks @MasterJH5574 @vinx13! Yes, we should add the primitives accordingly in TracedSchedule
--
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 change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
vinx13 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r685581513
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,313 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment " << buffer_->name << " has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
Review comment:
Good point, @junrushao1994 do we have convention what should be put in location of interests?
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r688010398
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,315 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment of, " << buffer_->name << ", has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+
+ static int CheckAndUpdate(const IRModule& mod, const Buffer& buffer, int axis) {
+ int ndim = static_cast<int>(buffer->shape.size());
+ if (axis < -ndim || axis >= ndim) {
+ throw StorageAlignAxisOutOfRangeError(mod, buffer, axis);
+ }
+ // If axis is negative, convert it to a non-negative one.
+ if (axis < 0) {
+ axis += ndim;
+ }
+ return axis;
+ }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+ int axis_;
+};
+
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ // Climb up along the sref tree, and find the block where `buffer` is in alloc_buffers or
+ // match_buffers.
+ const StmtSRefNode* defining_site_sref = block_sref.get();
+ while (defining_site_sref != nullptr) {
+ const auto* block = defining_site_sref->StmtAs<BlockNode>();
+ // If this sref is not a block sref, skip it.
+ if (block == nullptr) {
+ defining_site_sref = defining_site_sref->parent;
+ continue;
+ }
+ // Try to find the buffer in `allloc_buffers`
+ for (const Buffer& alloc_buffer : block->alloc_buffers) {
+ if (buffer.same_as(alloc_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), true};
+ }
+ }
+ // We do not allow the buffer being defined in `match_buffer`.
+ for (const MatchBufferRegion match_buffer : block->match_buffers) {
+ if (buffer.same_as(match_buffer)) {
+ return {GetRef<StmtSRef>(defining_site_sref), false};
+ }
+ }
+ defining_site_sref = defining_site_sref->parent;
+ }
+ // If we cannot find the defining site block, it means that the buffer must be in the function's
+ // buffer_map, which isn't an intermediate buffer.
+ return {NullOpt, false};
+}
+
+class NonAllocatedBufferError : public ScheduleError {
+ public:
+ explicit NonAllocatedBufferError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input buffer is not allocated by a block. This means the buffer is "
+ " either a function parameter or defined in `match_buffer` of a block.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input buffer " << buffer_->name
+ << " is not allocated by a block. This means the buffer is either a function parameter or "
+ "defined in `match_buffer` of a block.";
+ return os.str();
+ }
+
+ static void CheckBufferAllocated(const IRModule& mod, const StmtSRef& block_sref,
+ const Buffer& buffer) {
+ Optional<StmtSRef> defining_site_sref;
+ bool is_alloc;
+ std::tie(defining_site_sref, is_alloc) = GetBufferDefiningSite(block_sref, buffer);
+ if (!defining_site_sref || !is_alloc) {
+ throw NonAllocatedBufferError(mod, buffer);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ Buffer buffer_;
+};
+
+class StorageAlignInvalidFactorError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidFactorError(IRModule mod, int factor)
+ : mod_(std::move(mod)), factor_(factor) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `factor` of storage_align is expected to be a positive "
+ "number.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The input `factor` of storage_align is expected to be a positive number. However, the "
+ "input `factor` is "
+ << factor_ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ static void Check(const IRModule& mod, int factor) {
+ if (factor <= 0) {
+ throw StorageAlignInvalidFactorError(mod, factor);
+ }
+ }
+
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
+ IRModule mod() const final { return mod_; }
+
+ private:
+ IRModule mod_;
+ int factor_;
+};
+
+class StorageAlignInvalidAnnotationError : public ScheduleError {
+ public:
+ explicit StorageAlignInvalidAnnotationError(IRModule mod, Block block)
+ : mod_(std::move(mod)), block_(std::move(block)) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The block annotation for storage align is expected to be an array of "
+ "3-integer-tuples (axis, factor, offset).";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ os << "The block annotation for storage align is expected to be an array of 3-integer-tuples "
+ "(axis, factor, offset). However, the block annotation with key "
+ << attr::buffer_dim_align << " of the block {0} is "
+ << block_->annotations.at(attr::buffer_dim_align) << ", which is unexpected.";
+ return os.str();
+ }
+
+ static Array<Array<Array<Integer>>> CheckAndGetAnnotation(const IRModule& mod,
+ const Block& block) {
+ // Get existing annotation value.
+ auto it = block->annotations.find(attr::buffer_dim_align);
+ if (it != block->annotations.end()) {
+ if (!IsValidAnnotation(block, (*it).second)) {
+ throw StorageAlignInvalidAnnotationError(mod, block);
+ }
+ return Downcast<Array<Array<Array<Integer>>>>((*it).second);
+ }
+
+ // Create new annotation value
+ Array<Array<Array<Integer>>> storage_align_annotation;
+ storage_align_annotation.resize(block->writes.size());
Review comment:
Yes, I was proposing the lazily store buffer_index as well: https://github.com/apache/tvm/pull/8693#discussion_r686489336
--
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] MasterJH5574 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
MasterJH5574 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r685601750
##########
File path: src/tir/schedule/primitive/block_annotate.cc
##########
@@ -0,0 +1,313 @@
+/*
+ * 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 "../utils.h"
+
+namespace tvm {
+namespace tir {
+
+class StorageAlignAxisOutOfRangeError : public ScheduleError {
+ public:
+ explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis)
+ : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {}
+
+ String FastErrorString() const final {
+ return "ScheduleError: The input `axis` is out of range. It is required to be in range "
+ "[-ndim, ndim) where `ndim` is the number of dimensions of the buffer to set "
+ "storage alignment.";
+ }
+
+ String DetailRenderTemplate() const final {
+ std::ostringstream os;
+ int ndim = static_cast<int>(buffer_->shape.size());
+ os << "The buffer to set storage alignment " << buffer_->name << " has " << ndim
+ << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim
+ << ") for storage_align. However, the input `axis` is " << axis_
+ << ", which is out of the expected range.";
+ return os.str();
+ }
+
+ IRModule mod() const final { return mod_; }
+ Array<ObjectRef> LocationsOfInterest() const final { return {}; }
Review comment:
Yes, we only put "components of IR", that means, blocks and loops, in the location of interests.
--
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 pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
vinx13 commented on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896914345
Using alloc_buffer_index is possible. My only concern is that the index is not obvious to users when calling schedule primitives without converting the module to tir script and looking at 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] MasterJH5574 commented on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
MasterJH5574 commented on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896446136
Thanks @jinhongyii @junrushao1994! I'm gonna to merge it and send it to mainline.
--
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] MasterJH5574 commented on pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
MasterJH5574 commented on pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#issuecomment-896746839
I wrote the origin code that "go up straight to find the allocation site", which is first used in primitive `SetScope`. `SetScope` also specify the target buffer by one of its producer block and a `write_buffer_index`.
From my POV, the interface was designed like this mainly because `CacheRead` and `CacheWrite` specify the target buffer by a block and a index. Therefore I inherited the interface in `CacheRead` and `CacheWrite`.
At that time I thought that "specifying a buffer directly on the allocation side is impossible", but after a second thought today... Of course we can do that.
So ATM I think specifying the target buffer from the allocation side by `alloc_buffer_index` is much more convenient. Would like to hear your opinions :-)
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r686482634
##########
File path: python/tvm/tir/schedule/schedule.py
##########
@@ -710,6 +710,77 @@ def after_rfactor(a: ty.handle, b: ty.handle) -> None:
"""
return _ffi_api.ScheduleRFactor(self, loop, factor_axis) # type: ignore # pylint: disable=no-member
+ ######## Schedule: Block annotatoin ########
+
+ def storage_align( # pylint: disable=too-many-arguments
+ self, block: BlockRV, buffer_index: int, axis: int, factor: int, offset: int
+ ) -> None:
+ """Set alignment requirement for specific dimension such that
+ stride[axis] == k * factor + offset for some k.
Review comment:
Would you like to clarify what `k` is here?
--
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] junrushao1994 commented on a change in pull request #8693: [TensorIR][M2a] Storage Align
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #8693:
URL: https://github.com/apache/tvm/pull/8693#discussion_r686483010
##########
File path: python/tvm/tir/schedule/schedule.py
##########
@@ -710,6 +710,77 @@ def after_rfactor(a: ty.handle, b: ty.handle) -> None:
"""
return _ffi_api.ScheduleRFactor(self, loop, factor_axis) # type: ignore # pylint: disable=no-member
+ ######## Schedule: Block annotatoin ########
+
+ def storage_align( # pylint: disable=too-many-arguments
+ self, block: BlockRV, buffer_index: int, axis: int, factor: int, offset: int
+ ) -> None:
+ """Set alignment requirement for specific dimension such that
+ stride[axis] == k * factor + offset for some k.
Review comment:
There are some TE docs that we could probably add here: https://tvm.apache.org/docs/api/python/te.html#tvm.te.Stage.storage_align
--
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