Skip to content

Commit

Permalink
[Fix][MetaSchedule] Param for rule AutoBind on Python side (apache#13454
Browse files Browse the repository at this point in the history
) (apache#26)

Following apache#13206, this PR brings the new parameter added to the AutoBind schedule rule to Python side.
  • Loading branch information
MasterJH5574 authored Nov 22, 2022
1 parent eee21b8 commit cbf525b
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions python/tvm/meta_schedule/schedule_rule/auto_bind.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,17 +33,21 @@ class AutoBind(ScheduleRule):
The maximum number of threadblock on GPU.
thread_extents: Optional[List[int]]
Candidates of thread axis extent.
max_threads_per_block: int
The maximum number of threads per block, if it is known when this schedule rule is created.
"""

def __init__(
self,
max_threadblocks: int = 256,
thread_extents: Optional[List[int]] = None,
max_threads_per_block: int = -1,
) -> None:
if thread_extents is None:
thread_extents = [32, 64, 128, 256, 512, 1024]
self.__init_handle_by_constructor__(
_ffi_api.ScheduleRuleAutoBind, # type: ignore # pylint: disable=no-member
max_threadblocks,
thread_extents,
max_threads_per_block,
)

0 comments on commit cbf525b

Please sign in to comment.