quic-sanirudh commented on code in PR #14490:
URL: https://github.com/apache/tvm/pull/14490#discussion_r1158188742
##########
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:
Sure, thanks. I'll update the PR to take a block and return output blocks in
that scope.
--
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]
