Skip to content

Commit

Permalink
[MetaSchedule] Parallel-Vectorize-Unroll & Random-Compute-Location (a…
Browse files Browse the repository at this point in the history
…pache#516)

* parallel vectorize unroll & random compute location

* rebased
  • Loading branch information
spectrometerHBH authored and zxybazh committed Jan 7, 2022
1 parent 49915c2 commit 535984c
Showing 1 changed file with 25 additions and 0 deletions.
25 changes: 25 additions & 0 deletions python/tvm/tir/schedule/schedule.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,31 @@ def sample_compute_location(
decision,
)

def sample_compute_location(
self,
block: BlockRV,
decision: Optional[int] = None,
) -> LoopRV:
"""Sample a compute-at location on a BlockRV so that its producer can compute at that loop
Parameters
----------
block : BlockRV
The consumer block to be computed at
decision : Optional[int]
The sampling decision
Returns
-------
result : LoopRV
The sampled loop to be computed at
"""
return _ffi_api.ScheduleSampleComputeLocation( # pylint: disable=no-member
self,
block,
decision,
)

########## Schedule: Get blocks & loops ##########
@type_checked
def get_block(
Expand Down

0 comments on commit 535984c

Please sign in to comment.