junrushao commented on code in PR #16101:
URL: https://github.com/apache/tvm/pull/16101#discussion_r1388798899
##########
python/tvm/dlight/gpu/reduction.py:
##########
@@ -198,12 +198,20 @@ def _sch_inner_spatial(
sch: tir.Schedule,
_: Target,
block: tir.schedule.BlockRV,
+ block_info: BlockInfo,
unroll_spatial_factor: Optional[int],
epilogue_info: Optional[BlockInfo],
):
# pylint: disable=invalid-name
s, r, _ = sch.get_loops(block)
len_tx, len_ty = 16, 16
+ s_factor = [i.dom for i in block_info.iters if i.kind == "S"][-1]
+ # get perfect spatial factor, spatial factor should be divide the
innermost spatial loop so
+ # that the block after r_factor and be reversed compute at the
original scope
+ while len_tx > 1:
+ if s_factor % len_tx == 0:
+ break
+ len_tx -= 1
Review Comment:
Would be nice if we have a util method `_get_perfect_tile_factor` instead of
adding an extra `block_info` parameter
--
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]
