Skip to content

Commit

Permalink
Extend comment for cuda schedule_reduce_impl
Browse files Browse the repository at this point in the history
  • Loading branch information
elvin-n committed Oct 20, 2022
1 parent f4f1a03 commit 32f394d
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions python/tvm/topi/cuda/reduction.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,12 +118,20 @@ def is_scheduled(stage):

def schedule_reduce_impl(outs, schedule_reduce_stage, schedule_injective_stage):
"""Schedule for inject->reduce->bcast ops.
Traverse over the stages in the schedule and schedule separate stages depending
on the position of the stage. Injecteve post-ops of reduction will be scheduled using
injection schedule, injective pre-ops of reduction will be inlined, reduction stage
will be scheduled using reduction schedule
Parameters
----------
outs: Array of Tensor
The computation graph description of reduce in the format
of an array of tensors.
schedule_reduce_stage: Function responsible for scheduling the reduction
stage
schedule_injective_stage: Function responsible for scheduling the
standalone injection stage
Returns
-------
Expand Down

0 comments on commit 32f394d

Please sign in to comment.