@ziheng Sorry for my late reply. The code of tvm.build in build_module.py on master branch is attached below.
After we define the TIR PrimFunc with the script, we should be able to put it inside an IRModule and using tvm.build as normal.
if isinstance(inputs, schedule.Schedule):
if args is None:
raise ValueError("args must be given for build from schedule")
input_mod = lower(inputs, args, name=name, binds=binds)
elif isinstance(inputs, (list, tuple, container.Array)):
merged_mod = tvm.IRModule({})
for x in inputs:
merged_mod.update(x)
input_mod = merged_mod
elif isinstance(inputs, tvm.IRModule):
input_mod = inputs
elif not isinstance(inputs, (dict, container.Map)):
raise ValueError("inputs must be Schedule, IRModule or dict of target to IRModule")