Hzfengsy commented on code in PR #14490:
URL: https://github.com/apache/tvm/pull/14490#discussion_r1158111988
##########
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:
One thing a bit tricky is that the output block is based on the parent
scope/block instead of PrimFunc. Would be great to use root block as input if
you'd like to get the output blocks of the prim function
--
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]
