jcf94 commented on a change in pull request #6107:
URL: https://github.com/apache/incubator-tvm/pull/6107#discussion_r458569182
##########
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:
Updated.
----------------------------------------------------------------
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]
