Skip to content

Commit cbf525b

Browse files
authored
[Fix][MetaSchedule] Param for rule AutoBind on Python side (apache#13454) (apache#26)
Following apache#13206, this PR brings the new parameter added to the AutoBind schedule rule to Python side.
1 parent eee21b8 commit cbf525b

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

python/tvm/meta_schedule/schedule_rule/auto_bind.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,17 +33,21 @@ class AutoBind(ScheduleRule):
3333
The maximum number of threadblock on GPU.
3434
thread_extents: Optional[List[int]]
3535
Candidates of thread axis extent.
36+
max_threads_per_block: int
37+
The maximum number of threads per block, if it is known when this schedule rule is created.
3638
"""
3739

3840
def __init__(
3941
self,
4042
max_threadblocks: int = 256,
4143
thread_extents: Optional[List[int]] = None,
44+
max_threads_per_block: int = -1,
4245
) -> None:
4346
if thread_extents is None:
4447
thread_extents = [32, 64, 128, 256, 512, 1024]
4548
self.__init_handle_by_constructor__(
4649
_ffi_api.ScheduleRuleAutoBind, # type: ignore # pylint: disable=no-member
4750
max_threadblocks,
4851
thread_extents,
52+
max_threads_per_block,
4953
)

0 commit comments

Comments
 (0)