[GitHub] [tvm] wrongtest-intellif commented on a diff in pull request #12939: [TIR][Schedule] Add cache_buffer primitive to cache opaque buffer

[GitBox <gi...@apache.org>]

wrongtest-intellif commented on code in PR #12939:
URL: https://github.com/apache/tvm/pull/12939#discussion_r984430980


##########
src/tir/schedule/primitive/cache_read_write.cc:
##########
@@ -1146,6 +1250,102 @@ StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref, int write_bu
   return result_block_sref;
 }
 
+Array<StmtSRef> CacheBuffer(ScheduleState self, const StmtSRef& block_sref, int write_buffer_index,
+                            const String& storage_scope) {
+  /*!
+   * Do cache read then cache write
+   */
+
+  // Check 0. Check the input storage scope.
+  CheckStorageScope(self, storage_scope);
+
+  // Check 1. Check index, get the target buffer and the parent scope
+  const BlockNode* block = TVM_SREF_TO_BLOCK(block_sref);
+  Buffer buffer =
+      GetNthAccessBuffer(self, GetRef<Block>(block), write_buffer_index, BufferIndexType::kRead);
+  StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false);
+
+  // Check 3. Check required region cover for cache_read
+  CheckRegionCover(self, scope_sref);
+
+  Array<StmtSRef> results_block_sref;
+  Buffer new_buffer = WithScope(buffer, storage_scope);
+
+  // Do cache read
+  // Cache read step 0. Create CacheStageInfo
+  CacheStageInfo info;
+  info.read_buffer = buffer;
+  // Create the corresponding buffer to be written for cache_read
+  info.write_buffer = new_buffer;
+  // Create the corresponding buffer allocation
+  info.alloc = info.write_buffer;
+  // Indicate which buffers should consume the cache.
+  info.consumer_blocks.push_back(block_sref);
+
+  // Cache read step 1. Update cache stage info for cache_read.
+  BufferRegion cache_region{nullptr};
+  Optional<StmtSRef> _write_block_sref = GetOnlyWriteBlock(self, scope_sref, buffer);
+
+  StmtSRef write_block_sref = _write_block_sref.value();
+  const BlockNode* write_block = TVM_SREF_TO_BLOCK(write_block_sref);
+  // Find the producing region
+  BufferRegion region = GetBufferRegionFromBuffer(write_block->writes, buffer).value();
+  StmtSRef parent_sref = GetRef<StmtSRef>(write_block_sref->parent);
+
+  // Detect insert position
+  CacheBufferLocDetector::Detect(self, write_block_sref, scope_sref, &info);
+  cache_region = RelaxBufferRegion(self, region, write_block_sref, parent_sref, info.loc_sref);
+
+  // Cache read step 2. Making new cache stage block and rewrite readers.
+  Block cache_read_stage = MakeCacheStage(/*cache_region=*/cache_region, /*info=*/&info,
+                                          /*storage_scope=*/storage_scope);
+  Stmt new_scope = CacheReadRewriter::Rewrite(/*scope_sref=*/scope_sref, /*info=*/&info);
+
+  // Cache read step 3. Replacing and updating flags for cache read.
+  self->Replace(scope_sref, new_scope, info.block_reuse);
+  StmtSRef result_block_sref = self->stmt2ref.at(cache_read_stage.get());
+  BlockInfo& block_info_read = self->block_info[result_block_sref];
+  block_info_read.affine_binding = CalculateAffineFlag(self, result_block_sref);
+  block_info_read.region_cover = true;
+  block_info_read.scope->stage_pipeline = true;

Review Comment:
   why it is state_pipeline?



##########
src/tir/schedule/primitive/cache_read_write.cc:
##########
@@ -563,8 +649,17 @@ class CacheReadRewriter : public StmtExprMutator {
     if (block == scope_sref_->stmt) {
       // If so, put buffer allocation on the parent scope
       ObjectPtr<BlockNode> n = make_object<BlockNode>(*stmt.as<BlockNode>());
-      n->alloc_buffers.push_back(info_->alloc);
-      stmt = Block(n);
+      bool alloc_buffer_exists = false;
+      for (const Buffer& it : n->alloc_buffers) {
+        if (it.same_as(info_->alloc)) {
+          alloc_buffer_exists = true;
+        }
+      }
+      // In cache_buffer case, alloc_buffer may be already exits.

Review Comment:
   If alloc_buffer must exists in `cache_buffer` case, maybe we can change info_->alloc to Optional<Buffer> and leave it undefined? 



-- 
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