Hzfengsy commented on code in PR #14161:
URL: https://github.com/apache/tvm/pull/14161#discussion_r1122570607
##########
include/tvm/tir/schedule/schedule.h:
##########
@@ -405,6 +405,36 @@ class ScheduleNode : public runtime::Object {
virtual BlockRV CacheWrite(const BlockRV& block_rv, int write_buffer_index,
const String& storage_scope,
const Array<BlockRV> consumer_blocks = {}) = 0;
+ /*!
+ * \brief Create a block that reads a buffer region into a read cache. It
requires:
+ * 1) There is at most one block who writes the buffer in the scope.
+ * 2) The scope block have stage-pipeline property.
+ * Compared to cache read, the indices to access allocated cache buffer is
customized by user.
+ * \param block_rv The consumer block of the target buffer.
+ * \param read_buffer_index The index of the buffer in block's read region.
+ * \param storage_scope The target storage scope.
+ * \param index_map User defined indices to access allocated cache buffer,
maps from block iter
+ * vars. \param consumer_blocks An optional list of consumers to read from
cache directly. \return
+ * The cache stage block.
+ */
Review Comment:
```suggestion
* \param read_buffer_index The index of the buffer in block's read region.
* \param storage_scope The target storage scope.
* \param index_map User defined indices to access allocated cache buffer,
maps from block iter
* vars.
* \param consumer_blocks An optional list of consumers to read from cache
directly.
* \return The cache stage block.
*/
```
##########
include/tvm/tir/schedule/schedule.h:
##########
@@ -405,6 +405,36 @@ class ScheduleNode : public runtime::Object {
virtual BlockRV CacheWrite(const BlockRV& block_rv, int write_buffer_index,
const String& storage_scope,
const Array<BlockRV> consumer_blocks = {}) = 0;
+ /*!
+ * \brief Create a block that reads a buffer region into a read cache. It
requires:
+ * 1) There is at most one block who writes the buffer in the scope.
+ * 2) The scope block have stage-pipeline property.
+ * Compared to cache read, the indices to access allocated cache buffer is
customized by user.
+ * \param block_rv The consumer block of the target buffer.
+ * \param read_buffer_index The index of the buffer in block's read region.
+ * \param storage_scope The target storage scope.
+ * \param index_map User defined indices to access allocated cache buffer,
maps from block iter
+ * vars. \param consumer_blocks An optional list of consumers to read from
cache directly. \return
+ * The cache stage block.
+ */
+ virtual BlockRV ReindexCacheRead(const BlockRV& block_rv, int
read_buffer_index,
+ const String& storage_scope, const
IndexMap& index_map,
+ Array<BlockRV> consumer_blocks) = 0;
+ /*!
+ * \brief Create a block that writes a buffer region into a write cache. It
requires:
+ * 1) There is only one block who writes the target buffer.
+ * 2) The scope block have stage-pipeline property.
+ * Compared to cache write, the indices to access allocated cache buffer is
customized by user.
+ * \param block_rv The producer of the buffer
+ * \param write_buffer_index The index of the buffer in block's write region
+ * \param storage_scope The target storage scope
+ * \param index_map User defined indices to access allocated cache buffer,
maps from block iter
+ * vars. \param consumer_blocks An optional list of consumers to read from
cache directly. \return
Review Comment:
dito
##########
src/tir/schedule/primitive/cache_read_write.cc:
##########
@@ -94,6 +94,132 @@ Optional<BufferRegion> GetBufferRegionFromBuffer(const
Array<BufferRegion>& buff
return res;
}
+struct ReindexCacheStageInfo : CacheStageInfo {
+ /* Indices used to access the allocated cache buffer. */
+ Array<PrimExpr> indices;
+ /* Touched loop variable related information. */
+ Array<Var> loop_vars;
+ Array<Range> loop_ranges;
+ /* Touched block variable related information. */
+ Array<IterVar> block_iter_vars;
+ Array<PrimExpr> block_iter_values;
+};
+
+/* \brief The schedule error that accessed buffer region is not a single point
for
+ * reindex_cache_read/write. */
Review Comment:
```suggestion
/*
* \brief The schedule error that accessed buffer region is not a single
point for
* reindex_cache_read/write.
*/
```
Please use multi-line style here
##########
src/tir/schedule/primitive.h:
##########
@@ -269,6 +269,39 @@ TVM_DLL StmtSRef CacheRead(ScheduleState self, const
StmtSRef& block_sref, int r
TVM_DLL StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref,
int write_buffer_index,
const String& storage_scope,
const Array<StmtSRef> consumer_blocks = {});
+/*!
+ * \brief Create a block that reads a buffer region into a read cache. It
requires:
+ * 1) There is at most one block who writes the buffer in the scope.
+ * 2) The scope block have stage-pipeline property.
+ * Compared to cache read, the indices to access allocated cache buffer is
customized by user.
+ * \param self The state of the schedule
+ * \param block_sref The consumer block of the target buffer.
+ * \param read_buffer_index The index of the buffer in block's read region.
+ * \param storage_scope The target storage scope.
+ * \param index_map User defined indices to access allocated cache buffer,
maps from block iter
+ * vars. \param consumer_blocks Array of blocks that consume the cache.
\return The cache stage
+ * block.
+ */
+TVM_DLL StmtSRef ReindexCacheRead(ScheduleState self, const StmtSRef&
block_sref,
+ int read_buffer_index, const String&
storage_scope,
+ const IndexMap& index_map, Array<StmtSRef>
consumer_blocks = {});
+/*!
+ * \brief Create a block that writes a buffer region into a write cache. It
requires:
+ * 1) There is only one block that writes the target buffer.
+ * 2) The scope block have stage-pipeline property.
+ * Compared to cache write, the indices to access allocated cache buffer is
customized by user.
+ * \param self The state of the schedule
+ * \param block_sref The producer of the buffer
+ * \param write_buffer_index The index of the buffer in block's write region
+ * \param storage_scope The target storage scope
+ * \param index_map User defined indices to access allocated cache buffer,
maps from block iter
+ * vars. \param consumer_blocks Array of blocks that consume the cache.
\return The cache stage
Review Comment:
dito
##########
src/tir/schedule/primitive.h:
##########
@@ -269,6 +269,39 @@ TVM_DLL StmtSRef CacheRead(ScheduleState self, const
StmtSRef& block_sref, int r
TVM_DLL StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref,
int write_buffer_index,
const String& storage_scope,
const Array<StmtSRef> consumer_blocks = {});
+/*!
+ * \brief Create a block that reads a buffer region into a read cache. It
requires:
+ * 1) There is at most one block who writes the buffer in the scope.
+ * 2) The scope block have stage-pipeline property.
+ * Compared to cache read, the indices to access allocated cache buffer is
customized by user.
+ * \param self The state of the schedule
+ * \param block_sref The consumer block of the target buffer.
+ * \param read_buffer_index The index of the buffer in block's read region.
+ * \param storage_scope The target storage scope.
+ * \param index_map User defined indices to access allocated cache buffer,
maps from block iter
+ * vars. \param consumer_blocks Array of blocks that consume the cache.
\return The cache stage
Review Comment:
dito
##########
src/tir/schedule/primitive/cache_read_write.cc:
##########
@@ -1297,6 +1603,317 @@ StmtSRef CacheWrite(ScheduleState self, const StmtSRef&
block_sref, int write_bu
return result_block_sref;
}
+/*! \brief A visitor that collects variables appeared in expressions, stored
in `touched` field.*/
+class VarCollector : public ExprVisitor {
Review Comment:
Is it the same as
https://github.com/yzh119/tvm/blob/7bca3c2ca9fb144244f1a6b5968f5911d47a956b/include/tvm/tir/analysis.h#L102
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
