[TVMScript] Ensure completed root block has no read/write

Prior to this PR, the root block of a parsed TIR TVMScript is possible
to have non-empty read/write regions, which conflicts with the design
of root blocks in TIR.

This PR updates the script completion pass and ensures that the root
block will no longer have read/write region.

Author: MasterJH5574

src/tir/ir/script/script_complete.cc

@@ -36,21 +36,18 @@ namespace tir {
 class ScriptCompleter : public StmtMutator {
  public:
   explicit ScriptCompleter(Map<Var, Buffer>* buffer_var_map) : buffer_var_map_(buffer_var_map) {}
-  /*! \brief Whether the stmt contains at least one block. */
-  bool contains_block = false;
 
  private:
   Map<Var, Buffer>* buffer_var_map_;
-  Stmt VisitStmt_(const BlockRealizeNode* op) override {
-    contains_block = true;
+  Stmt VisitStmt_(const BlockRealizeNode* op) final {
     for (const PrimExpr& value : op->iter_values) {
       CHECK(value.dtype().is_int())
           << "BlockRealize iter_value expected a IntImm, but got " << value.dtype();
     }
     return StmtMutator::VisitStmt_(op);
   }
 
-  Stmt VisitStmt_(const BlockNode* op) override {
+  Stmt VisitStmt_(const BlockNode* op) final {
     // Buffers allocated in the block can be accessed by its body.
     for (const auto& alloc_buffer : op->alloc_buffers) {
       buffer_var_map_->Set(alloc_buffer->data, alloc_buffer);
@@ -59,7 +56,12 @@ class ScriptCompleter : public StmtMutator {
       const Buffer& target_buffer = match_buffer->buffer;
       buffer_var_map_->Set(target_buffer->data, target_buffer);
     }
+
+    bool is_root_block = this->is_root_block_;
+    this->is_root_block_ = false;
     Block block = Downcast<Block>(StmtMutator::VisitStmt_(op));
+    this->is_root_block_ = is_root_block;
+
     // Remove buffers allocated inside block to detect its access region
     for (const auto& alloc_buffer : op->alloc_buffers) {
       buffer_var_map_->erase(alloc_buffer->data);
@@ -85,15 +87,19 @@ class ScriptCompleter : public StmtMutator {
           << "ValueError: Can not auto detect buffer access region from tir.Load, tir.Store or "
              "direct access by buffer data. Please annotation the access region manually";
       auto n = CopyOnWrite(block.operator->());
-      if (mask & 1) n->reads = reads;
-      if (mask & 2) n->writes = writes;
+      if (!is_root_block) {
+        if (mask & 1) n->reads = reads;
+        if (mask & 2) n->writes = writes;
+      }
       n->annotations = op->annotations;
       n->annotations.erase(attr::script_parsing_detect_access);
       return Block(n);
     } else {
       return std::move(block);
     }
   }
+
+  bool is_root_block_ = true;
 };
 
 PrimFunc ScriptComplete(PrimFunc func, const Array<Buffer>& root_allocates) {

tests/python/unittest/test_tvmscript_complete.py

@@ -153,6 +153,10 @@ def test_complete_matmul_original():
 def _check_elementwise(func):
     A, B, C = [func.buffer_map[x] for x in func.params]
 
+    root_block = func.body.block
+    assert len(root_block.reads) == 0
+    assert len(root_block.writes) == 0
+
     block1 = func.body.block.body[0].body.body.block
     assert isinstance(block1, tvm.tir.Block)
     vi, vj = [x.var for x in block1.iter_vars]