tqchen commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1239862854
##########
python/tvm/dlight/gpu/default_schedule.py:
##########
@@ -0,0 +1,89 @@
+# pylint: disable=missing-docstring
+from typing import Callable
+
+from tvm import tir
+from tvm._ffi import get_global_func
+from tvm.ir import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.target import Target
+
+
+def _default_schedule( # pylint: disable=too-many-locals
+ func: tir.PrimFunc,
+ target: Target,
+) -> tir.PrimFunc:
+ max_threads_per_block = None
+ for name in ["max_threads_per_block", "max_num_threads"]:
+ if max_threads_per_block is None:
+ max_threads_per_block = target.attrs.get(name, None)
+ if max_threads_per_block is None:
+ max_threads_per_block = 64
+
+ get_loop_iter_type = get_global_func("tir.schedule.GetLoopIterType")
Review Comment:
This can be a python function get exposed
--
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]
