MetaSchedule Builder with Relay Build

Hi, I can’t wrap my head around how to use the MetaSchedule LocalBuilder with relay.build. I tried to copy the implementation for microTVM and so far ended up with this:

def _my_f_build(
	mod: IRModule, target: Target, _params: Optional[Dict[str, NDArray]]
) -> AOTExecutorFactoryModule:

	mod = RemoveWeightLayoutRewriteBlock(skip_ndarray_rewrite=True)(mod)
	with pass_context, ScheduleFnDatabase(schedule_func):
		module = relay_build(mod, params=_params, target=target, executor=executor, runtime=runtime)
	return module
	
	
def _my_f_export(module: AOTExecutorFactoryModule, ) -> str:
	tmpdir = tvm.contrib.utils.tempdir()
	model_library_format_tar_path = tvm.micro.export_model_library_format(module, tmpdir / "model.tar")
	with tempfile.NamedTemporaryFile() as tar_tmp_file:

		template_project_path = pathlib.Path(tvm.micro.get_microtvm_template_projects("torrent"))
		project_options = {"project_type": project_type}

		generated_project_dir = pathlib.Path(tempfile.mkdtemp(), "tvm_tmp_mod.generated_project")
		generated_project = tvm.micro.generate_project(
			template_project_path, module, generated_project_dir, project_options
		)
		generated_project.build()

I don’t see how this is supposed to interact with the ScheduleSpace function that I need to pass to the meta_schedule.tune_relay call because currently in both cases the scheduling is applied. Once during the generation of the search space when traces are generated and once here during the build process, and I don’t think that is the correct way to go about it.

@junrushao

For a bit more context: I am doing the scheduling in TIR, which is why in the pass_context I have set relay.backend.use_meta_schedule to true. That is also why I need to apply the database in the f_build function.