()
| 45 | |
| 46 | |
| 47 | def initializer(): |
| 48 | @register_func("meta_schedule.builder.async_build") |
| 49 | def async_build(mod, target, _params): # pylint: disable=unused-variable, unused-argument |
| 50 | # pylint: disable=import-outside-toplevel |
| 51 | from tvm.driver import build as tvm_build |
| 52 | from tvm.tir.transform import RemoveWeightLayoutRewriteBlock |
| 53 | |
| 54 | # re-import here for local builder to register index_map_m16n8k8_matrixC |
| 55 | # pylint: disable=import-outside-toplevel, unused-import |
| 56 | from tvm.tir.tensor_intrin import cuda |
| 57 | |
| 58 | mod = RemoveWeightLayoutRewriteBlock(skip_ndarray_rewrite=True)(mod) |
| 59 | with tvm.transform.PassContext(config={"tir.use_async_copy": 1}): |
| 60 | rt_mod = tvm_build(mod, target=target) |
| 61 | return rt_mod |
| 62 | |
| 63 | |
| 64 | def multi_level_tiling_mma(out_dtype): |
nothing calls this directly
no outgoing calls
no test coverage detected