You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/07/22 04:06:06 UTC
[GitHub] [incubator-tvm] comaniac commented on a change in pull request #6107: [Ansor][AutoTVM v2.0] Phase 1: Add cache_read/cache_write steps
comaniac commented on a change in pull request #6107:
URL: https://github.com/apache/incubator-tvm/pull/6107#discussion_r458507744
##########
File path: python/tvm/auto_scheduler/loop_state.py
##########
@@ -351,6 +351,68 @@ def compute_root(self, stage):
self.state_object = _ffi_api.StateComputeRoot(self.state_object,
self._resolve_stage_id(stage))
+ def cache_read(self, stage, scope_name, reader_stages):
+ """ Schedule primitive corresponds to te.schedule.cache_read.
+
+ Parameters
+ ----------
+ stage : Union[int, Operation, Tensor]
+ The Stage to be cache read, which can be specified by the integer index, Operation,
+ or output tensor of the stage.
+ scope_name : str
+ The scope name to be set for the new added read stage.
+ reader_stages : List[Union[int, Operation, Tensor]]
+ The reader stages. Each of the list can be specified by the integer index, Operation,
+ or output tensor of the stage.
+
+ Returns
+ -------
+ new_stage_op : Operator
+ The Operator of the new added stage.
+
+ Notes
+ -----
+ Cache read step will add an extra stage to the original ComputeDAG.
+ """
+ if isinstance(reader_stages, list):
+ reader_stage_ids = [self._resolve_stage_id(id) for id in reader_stages]
+ else:
+ raise ValueError("reader_stages must be a list of the integer index, Operation, " + \
+ "or output tensor of the stage")
Review comment:
Since `reader_stages` must be a list, looks like we don't need this check.
##########
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:
It seems to me that this call will override the mapped stage ID of all ops from stages but don't fully understand. Could you elaborate more about `_insert_new_stage`?
##########
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.
Review comment:
- This description is better to just be the docstring of this function.
- The function name is misleading. This function does not insert a new stage. Instead, it updates the stage ID map by taking a new stage ID. The name like `_update_stage_id_map` might be better.
- The return type of this function is improper, since it does nothing with `added_op`. Better to not return anything.
- I feel this function should not be a standalone function. This is tightly coupled CacheRead/CacheWrite implementation. In this case, it is improper to explicitly call this function.
##########
File path: src/auto_scheduler/loop_state.h
##########
@@ -195,6 +197,17 @@ class AttachMap : public ObjectRef {
void UpdateIters(const std::vector<IterKey>& original_iters,
const std::vector<IterKey>& new_iters);
+ /*!
+ * \brief Traverse through `stage_to_attach_iter` and `iter_to_attached_stages` map, add offset
+ * to stage indexes that are larger than the start_id. Used for steps that inserts net stages to
Review comment:
You mean "and update offset"?
##########
File path: src/auto_scheduler/compute_dag.h
##########
@@ -114,6 +114,15 @@ class ComputeDAG : public ObjectRef {
*/
State InferBound(const State& state) const;
+ /*!
+ * \brief Some steps may change the structure of ComputeDAG(e.g. CacheRead/CacheWrite Step), this
+ * is to replay the transform steps and get the up-to-date ComputeDAG.
Review comment:
```suggestion
* \brief Since some steps may change the ComputeDAG (e.g. CacheRead/CacheWrite),
* the initial ComputeDAG may not be up-to-date. This function replays the given transform steps
* and return an up-to-date ComputeDAG.
```
##########
File path: python/tvm/auto_scheduler/loop_state.py
##########
@@ -351,6 +351,68 @@ def compute_root(self, stage):
self.state_object = _ffi_api.StateComputeRoot(self.state_object,
self._resolve_stage_id(stage))
+ def cache_read(self, stage, scope_name, reader_stages):
+ """ Schedule primitive corresponds to te.schedule.cache_read.
+
+ Parameters
+ ----------
+ stage : Union[int, Operation, Tensor]
+ The Stage to be cache read, which can be specified by the integer index, Operation,
+ or output tensor of the stage.
+ scope_name : str
+ The scope name to be set for the new added read stage.
+ reader_stages : List[Union[int, Operation, Tensor]]
+ The reader stages. Each of the list can be specified by the integer index, Operation,
+ or output tensor of the stage.
+
+ Returns
+ -------
+ new_stage_op : Operator
+ The Operator of the new added stage.
+
+ Notes
+ -----
+ Cache read step will add an extra stage to the original ComputeDAG.
Review comment:
Maybe explain how stage ID will change. Ditto to cache write.
##########
File path: python/tvm/auto_scheduler/loop_state.py
##########
@@ -351,6 +351,68 @@ def compute_root(self, stage):
self.state_object = _ffi_api.StateComputeRoot(self.state_object,
self._resolve_stage_id(stage))
+ def cache_read(self, stage, scope_name, reader_stages):
+ """ Schedule primitive corresponds to te.schedule.cache_read.
+
+ Parameters
+ ----------
+ stage : Union[int, Operation, Tensor]
+ The Stage to be cache read, which can be specified by the integer index, Operation,
+ or output tensor of the stage.
+ scope_name : str
+ The scope name to be set for the new added read stage.
+ reader_stages : List[Union[int, Operation, Tensor]]
+ The reader stages. Each of the list can be specified by the integer index, Operation,
+ or output tensor of the stage.
+
+ Returns
+ -------
+ new_stage_op : Operator
+ The Operator of the new added stage.
+
+ Notes
+ -----
+ Cache read step will add an extra stage to the original ComputeDAG.
+ """
+ if isinstance(reader_stages, list):
+ reader_stage_ids = [self._resolve_stage_id(id) for id in reader_stages]
+ else:
+ raise ValueError("reader_stages must be a list of the integer index, Operation, " + \
+ "or output tensor of the stage")
+
+ self.state_object, new_stage_id = _ffi_api.StateCacheRead(self.state_object,
+ self._resolve_stage_id(stage),
+ scope_name, reader_stage_ids,
+ self.compute_dag)
+ return self._insert_new_stage(int(new_stage_id))
+
+ def cache_write(self, stage, scope_name):
+ """ Schedule primitive corresponds to te.schedule.cache_write.
+
+ Parameters
+ ----------
+ stage : Union[int, Operation, Tensor]
+ The Stage to be cache write, which can be specified by the integer index, Operation,
+ or output tensor of the stage.
+ scope_name : str
+ The scope name to be set for the new added write stage.
+
+ Returns
+ -------
+ new_stage_op : Operator
+ The Operator of the new added stage.
+
+ Notes
+ -----
+ Cache write step will add an extra stage to the original ComputeDAG, a up-to-date
+ ComputeDAG is stored in State's `current_compute_dag`.
+ This step will cache write all output tensors of the target stage.
Review comment:
```suggestion
Cache write step will add an extra stage to the original ComputeDAG. The up-to-date
ComputeDAG is stored in `self.current_compute_dag`.
This step will cache write all output tensors of the target stage.
```
##########
File path: src/auto_scheduler/loop_state.h
##########
@@ -195,6 +197,17 @@ class AttachMap : public ObjectRef {
void UpdateIters(const std::vector<IterKey>& original_iters,
const std::vector<IterKey>& new_iters);
+ /*!
+ * \brief Traverse through `stage_to_attach_iter` and `iter_to_attached_stages` map, add offset
+ * to stage indexes that are larger than the start_id. Used for steps that inserts net stages to
+ * ComputeDAG(e.g. CacheRead/CacheWrite step).
+ * \param start_id The index threshold, stage indexes in AttachMap which are larger than this
+ * will be applied the extra offset.
+ * \param offset The index offset to be added to the stage index.
+ * \return The updated AttachMap after applying stage index offset.
+ */
+ AttachMap ApplyStageIdOfffset(int start_id, int offset) const;
Review comment:
s/ApplyStageIdOfffset/ApplyStageIdOffset/
----------------------------------------------------------------
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