TypeError: tir.Schedule is not registered via TVM_REGISTER_NODE_TYPE when using multiprocessing

Troubleshooting
#1

I have been playing around with different tuning approaches and wanted to use multiprocessing.

I am trying to have multiple processes work on a list of schedules. I implemented the following, and it works fine when using a ThreadPoolExecutor instead of a ProcessPoolExecutor. It takes a Schedule from a list of Schedules, sends it to a tuning function, and returns the tuned Schedule.

The code I am using:

with ProcessPoolExecutor(max_workers=self.context.num_threads) as executor:
    futures = [executor.submit(self._tune_schedule, tune_schedules[id], id)
               for id in range(num_sch_to_tuner)]

    for future in as_completed(futures):
        tuned_schedule, id = future.result()
        tune_schedules[id] = tuned_schedule

The error I get:

Process SpawnProcess-1:
Traceback (most recent call last):
  File "/Users/felix/anaconda3/envs/tvm-llvm/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/Users/felix/anaconda3/envs/tvm-llvm/lib/python3.10/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/Users/felix/anaconda3/envs/tvm-llvm/lib/python3.10/concurrent/futures/process.py", line 240, in _process_worker
    call_item = call_queue.get(block=True)
  File "/Users/felix/anaconda3/envs/tvm-llvm/lib/python3.10/multiprocessing/queues.py", line 122, in get
    return _ForkingPickler.loads(res)
  File "/Users/felix/Documents/Projects/ACS-Project/tvm/python/tvm/runtime/object.py", line 101, in __setstate__
    self.__init_handle_by_constructor__(_ffi_node_api.LoadJSON, handle)
  File "/Users/felix/Documents/Projects/ACS-Project/tvm/python/tvm/_ffi/_ctypes/object.py", line 145, in __init_handle_by_constructor__
    handle = __init_by_constructor__(fconstructor, args)
  File "/Users/felix/Documents/Projects/ACS-Project/tvm/python/tvm/_ffi/_ctypes/packed_func.py", line 262, in __init_handle_by_constructor__
    raise_last_ffi_error()
  File "/Users/felix/Documents/Projects/ACS-Project/tvm/python/tvm/_ffi/base.py", line 481, in raise_last_ffi_error
    raise py_err
TypeError: Traceback (most recent call last):
  File "/Users/felix/Documents/Projects/ACS-Project/tvm/src/node/reflection.cc", line 158
TypeError: tir.Schedule is not registered via TVM_REGISTER_NODE_TYPE

I have seen that the typical suggestion for this problem is to rebuild TVM, which I have done, but it did not resolve the issue. I am likely doing something wrong in creating the schedule copy for the new process. I would be thankful if someone could help me figure out the proper approach to this.

#2

I’d encountered the same issue and the solution I applied is using sch.mod as parameters for subprocesses, instead of tir.Schedule

#3

Thank you for the advice; that fixed it for me. I have been trying to pass the IRModule and Trace to the process to reconstruct the Schedule there. However, accessing the trace in the subprocess leads to a segmentation fault. Have you run into this issue as well?

#4

Serializing the trace as JSON resolves it.