quic-sanirudh commented on code in PR #14490:
URL: https://github.com/apache/tvm/pull/14490#discussion_r1158121564
##########
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 the suggestion. I thought about that, but since the root block
does not have to be passed in, I thought of not taking any arguments, but I
guess I can modify it to take a block and return output blocks in the scope of
that block.
Do you think that makes more sense? We can then call it with the root block
if we need all output blocks for the PrimFunc
--
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]
