[GitHub] [incubator-tvm] jcf94 commented on a change in pull request #6107: [Ansor][AutoTVM v2.0] Phase 1: Add cache_read/cache_write steps

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

jcf94 commented on a change in pull request #6107:
URL: https://github.com/apache/incubator-tvm/pull/6107#discussion_r458555837



##########
File path: python/tvm/auto_scheduler/loop_state.py
##########
@@ -371,6 +433,23 @@ def _update_stage_id_map(self):
         for index, stage in enumerate(self.stages):
             self.stage_id_map[stage.op] = index
 
+    def _insert_new_stage(self, new_stage_id):
+        added_op = self.stages[new_stage_id].op
+
+        # Add a new stage will change all ops. But we still want to use the old ops to index stages,
+        # So we keep updating them and do not remove the old ops.
+
+        # Update stage_id_map for old ops, so we can still use the old ops to index stages.
+        for key, value in self.stage_id_map.items():
+            if value >= new_stage_id:
+                self.stage_id_map[key] = value + 1
+        self.stage_id_map[added_op] = new_stage_id
+
+        # Update stage_id_map for new ops
+        self._update_stage_id_map()

Review comment:
       After cache_read/cache_write, actually all the stages behind the new added stage will be update, their ops are different.
   This is to keep the original Tensor/ops works and add those new ops to the id map.




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

For queries about this service, please contact Infrastructure at:
users@infra.apache.org