merrymercy commented on a change in pull request #6107:
URL: https://github.com/apache/incubator-tvm/pull/6107#discussion_r459825492
##########
File path: python/tvm/auto_scheduler/compute_dag.py
##########
@@ -126,11 +126,16 @@ def infer_bound_from_state(self, state):
Returns
-------
- state : State
+ updated_state : State
The State with complete bound information.
"""
state_obj = state if isinstance(state, StateObject) else
state.state_object
- return State(_ffi_api.ComputeDAGInferBoundFromState(self, state_obj),
self)
+ updated_state = State(_ffi_api.ComputeDAGInferBoundFromState(self,
state_obj), self)
+ # Copy the stage_id_map from the original state
Review comment:
```suggestion
# Copy the stage_id_map from the original state to make the old
indices still valid.
```
##########
File path: src/auto_scheduler/loop_state.h
##########
@@ -225,6 +238,16 @@ class StateNode : public Object {
* operation.
*/
AttachMap attach_map;
+ /*!
+ * \brief The up-to-date ComputeDAG of this state, used for some steps that
may change the
+ * stage structure of the ComputeDAG (e.g. CacheReadStep/CacheWriteStep).
This will alway be kept
+ * up-to-date, while the original ComputeDAG may not be up-to-date.
+ * The default value is an empty NullOpt, means no modification to the
original DAG.
+ * Typical usage for this is when acquiring information from ComputeDAG
(e.g. check for its
+ * AccessAnalyzer), use the `current_compute_dag` first, if it's Null, use
the original
+ * ComputeDAG.
+ */
Review comment:
```suggestion
* \brief The up-to-date ComputeDAG of this state. If the value is null,
it means no steps modify the DAG.
* The dag is the same as the compute_dag in the `SearchTask`. Otherwise,
it means some steps
* (e.g., CacheReadStep/CacheWriteStep) modify the dag, then the stored
value is the new and up-to-date
* DAG for this state.
*/
```
##########
File path: src/auto_scheduler/transform_step.h
##########
@@ -194,20 +201,25 @@ void StepApplyToState(const Step& step, State* state,
const ComputeDAG& dag);
/*!
* \brief Apply the step to tvm.schedule.
* \param step The step to be applied to tvm.schedule.
- * \param stages A pointer to a `te::Stage` Array.
- * \param stage_to_axes A pointer to a StageToAxesMap.
+ * \param stages A mutable pointer to a `te::Stage` Array.
Review comment:
This kind of doc only repeats the type of the parameters and gives no
information
##########
File path: src/auto_scheduler/transform_step.cc
##########
@@ -923,5 +958,275 @@ String
ComputeRootStepNode::PrintAsPythonAPI(Array<te::Stage>* stages,
return ss.str();
}
+/********** Primitives adding new stages **********/
+
+/*!
+ * \brief Common part for steps that add new stages(e.g. CacheReadStep,
CacheWriteStep,
+ * RfactorStep). This will filter out all steps that can change the number of
stages in a
Review comment:
```suggestion
* RfactorStep). This returns all steps that can change the number of stages
in a
```
"Filter out" is the antonym of "return"
----------------------------------------------------------------
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:
[email protected]
