Hzfengsy commented on code in PR #14490:
URL: https://github.com/apache/tvm/pull/14490#discussion_r1158133904
##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -540,6 +540,23 @@ def get_consumers(self, block: Union[BlockRV, str]) ->
List[BlockRV]:
block = self._normalize_block_arg(block)
return list(_ffi_api.ScheduleGetConsumers(self, block)) # type:
ignore # pylint: disable=no-member
+ @type_checked
+ def get_output_blocks(
+ self,
+ func_name: Optional[str] = None,
+ ) -> List[BlockRV]:
+ """Get the list of output blocks
+ An output block is a block which has atleast one buffer being written
+ to, but is not allocated within the PrimFunc
+
+ Returns
+ -------
+ output_blocks : List[BlockRV]
+ A list of all blocks that write to some output buffer
+
+ """
+ return list(_ffi_api.ScheduleGetOutputBlocks(self, func_name)) #
type: ignore # pylint: disable=no-member
Review Comment:
Thanks for your update.
Taking no arguments is a good choice for the root block. And also "take a
block and return output blocks in the scope of that block." makes more sense to
me. That helps if we have nested blocks.
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
