vinx13 commented on code in PR #11225:
URL: https://github.com/apache/tvm/pull/11225#discussion_r867038052
##########
src/tir/schedule/transform.h:
##########
@@ -66,6 +72,122 @@ Array<BufferRegion> ReplaceBuffer(Array<BufferRegion>
regions, const Buffer& sou
Array<MatchBufferRegion> ReplaceBuffer(Array<MatchBufferRegion> match_buffers,
const Buffer& source,
const Buffer& target);
+/*!
+ * \brief A helper mutator which recursively replaces the old buffer with the
new buffer and
+ * collects the block sref reuse information for the following replacement.
+ *
+ * If the buffer to be replaced in used as the source in `match_buffers`,
depending the specific
+ * use cases, the target buffers in `match_buffers` may also need to be
mutated. In this
+ * case, this class should be subclassed to explicitly handle `match_buffers`.
+ */
+class ReplaceBufferMutator : public StmtExprMutator {
Review Comment:
`transform.h` is for transformational utilities, it is included in
`../schedule/utils.h`.
##########
src/tir/schedule/transform.h:
##########
@@ -66,6 +72,122 @@ Array<BufferRegion> ReplaceBuffer(Array<BufferRegion>
regions, const Buffer& sou
Array<MatchBufferRegion> ReplaceBuffer(Array<MatchBufferRegion> match_buffers,
const Buffer& source,
const Buffer& target);
+/*!
+ * \brief A helper mutator which recursively replaces the old buffer with the
new buffer and
+ * collects the block sref reuse information for the following replacement.
+ *
+ * If the buffer to be replaced in used as the source in `match_buffers`,
depending the specific
+ * use cases, the target buffers in `match_buffers` may also need to be
mutated. In this
+ * case, this class should be subclassed to explicitly handle `match_buffers`.
+ */
+class ReplaceBufferMutator : public StmtExprMutator {
+ public:
+ ReplaceBufferMutator(const Buffer& old_buffer, Buffer new_buffer,
+ Map<Block, Block>* block_sref_reuse)
+ : block_sref_reuse_(block_sref_reuse) {
+ buffer_var_map_[old_buffer->data.get()] = std::move(new_buffer);
+ }
Review Comment:
sounds good, updated
--
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-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
