Fix AutoScheduler for anaconda python

[dlexplorer]

In case of non cpython flavour of python, the task passed to measure process should be serialized using pickle approach. The task includes workload which is a list of Tensors. This list should be serialized and deserialized as an atomic object.

If list is not serialized as the atomic object, each tensor will be serialized/deserialized independently. And during deserialization of output tensor, it will recreate input tensors one more time instead of reusing existing from the list because it will not know that input tensors are created during deserialization of list elements. As a result, the function SchedulePostProcToPrimFunc will fail on assert in GetOrAllocBuffer dues to missmatching of input tensors.

[comaniac]

LGTM

[comaniac]

@dlexplorer looks like the CI failed at the RPC runner. Please check if your changes could pass the unit tests locally first (not in the anaconda).

[dlexplorer]

@dlexplorer looks like the CI failed at the RPC runner. Please check if your changes could pass the unit tests locally first (not in the anaconda).

@comaniac I see some inconsistency in workload content. It can store

• workload key->list of tensors (the way of tasks generated through auto_scheduler.extract_tasks)
• function name-> function
  If we have the first case, the list of tensors can be serialized successfully using SaveJSON
  If we have second case, exactly what happens in test, the function matmul_auto_scheduler_test is registered, then it canno tbe serialized using SaveJSON.

And I suspect it cannot be serialized without some processing.
On the other hand, the process of measure assume creation of second process and passing the task for measure in the second process. The object should be serializable. I believe eveything work with cpython and functino by the same reason as with workload id/list of tensor. content of the task is not serialized, due to implementation of cpython process just forks and all objects are on the same place as in original process.

I can add workaround - to verify the type of the content and if it is a list of tensors, apply SaveJSON, in other case do not call. And in deserialize I can also verify value and call LoadJSON only if it has string type.

But it is a design issue if we can have situation when we have object which should be passed to another process and which cannot be serialized.

Can we process/modify function somehow to be serialized?
Is function usecase happen only in tests? or it is assumed that it will be used in real life as well?

[comaniac]

@comaniac I see some inconsistency in workload content. It can store

• workload key->list of tensors (the way of tasks generated through auto_scheduler.extract_tasks)
• function name-> function
  If we have the first case, the list of tensors can be serialized successfully using SaveJSON
  If we have second case, exactly what happens in test, the function matmul_auto_scheduler_test is registered, then it canno tbe serialized using SaveJSON.

Yes this is the current design: auto_scheduler supports workload registration from both compute function or DAG.

And I suspect it cannot be serialized without some processing.

The current implementation doesn't serialize the function, because the receiver must have TVM deployed so the function is already there. We only need to look at the registry for the function by its name. This is even simpler because the only thing you need to serialize and pass around is the function name and arguments.

On the other hand, the process of measure assume creation of second process and passing the task for measure in the second process. The object should be serializable. I believe eveything work with cpython and functino by the same reason as with workload id/list of tensor. content of the task is not serialized, due to implementation of cpython process just forks and all objects are on the same place as in original process.

I can add workaround - to verify the type of the content and if it is a list of tensors, apply SaveJSON, in other case do not call. And in deserialize I can also verify value and call LoadJSON only if it has string type.

Based on my above explanation, the only thing we need to do is just ignoring functions in the registry.

cc @merrymercy

[jcf94]

Is is easy to have a UT test for this case? Since we're continue to bring new features to the Autoscheduler, I'm just afriad this will still be broken in the future.

[dlexplorer]

Is is easy to have a UT test for this case?

It definitely make sense. Will figure out which scenario to select which better reproduce issue on cpython

[dlexplorer]

Re-run task_python_integration.sh on my machine (Ubuntu 18.04/Python 3.6.9) twice. All tests passed. Is test_resize flaky test?

[dlexplorer]

Added test which fails on main branch with anaconda and passed on this branch.
But test passes for cpython even on main branch. Will try to develop test on serialization of workload using pickle even with cpython.

[dlexplorer]

Modified test to repeat serialize/deserialize issue with cpython.
feature with test is ready

[jcf94]

@comaniac @dlexplorer Thanks!