Custom datatypes in AutoScheduler?

@vinx13

I modified reset_global_scope() in env.py into the following:

import tvm
def reset_global_scope(global_scope):
    """Reset global autotvm state. This is needed to initialize PopenPool workers."""
    global GLOBAL_SCOPE
    GLOBAL_SCOPE.deep_copy(global_scope)
    AutotvmGlobalScope.current = global_scope

     # Defition of the custom datatype and a few operators
    tvm.target.datatype.register("cmpl", 150)
    tvm.target.datatype.register_op(
        tvm.target.datatype.create_lower_func({(32, 64): "Float32ToComplex64"}),
        "Cast",
        "llvm",
        "float",
        "cmpl",
    )
    tvm.target.datatype.register_op(
        tvm.target.datatype.create_lower_func({64: "Complex64Add"}),
        "Add",
        "llvm",
        "cmpl",
    )
    tvm.target.datatype.register_op(
        tvm.target.datatype.create_lower_func({64: "Complex64Sub"}),
        "Sub",
        "llvm",
        "cmpl",
    )
    tvm.target.datatype.register_op(
        tvm.target.datatype.create_lower_func({64: "Complex64Mul"}),
        "Mul",
        "llvm",
        "cmpl",
    )
    tvm.target.datatype.register_op(
        tvm.target.datatype.create_lower_func({64: "Complex64Div"}),
        "Div",
        "llvm",
        "cmpl",
    )
    tvm.target.datatype.register_op(
        tvm.target.datatype.lower_call_pure_extern,
        "Call",
        "llvm",
        "cmpl",
        intrinsic_name="tir.call_pure_extern",
    )
    tvm.target.datatype.register_op(
        tvm.target.datatype.create_lower_func({(64, 32): "Complex64ToFloat32"}),
        "Cast",
        "llvm",
        "cmpl",
        "float",
    )

But I still get the same error message:

Traceback (most recent call last):
  File "/usr/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/victor/hpc/tvm-build/tvmgpucopy2/python/tvm/exec/popen_worker.py", line 105, in <module>
    main()
  File "/home/victor/hpc/tvm-build/tvmgpucopy2/python/tvm/exec/popen_worker.py", line 77, in main
    fn, args, kwargs, timeout = cloudpickle.loads(reader.read(bytes_size))
  File "/home/victor/hpc/tvm-build/tvmgpucopy2/python/tvm/runtime/object.py", line 93, in __setstate__
    self.__init_handle_by_constructor__(_ffi_node_api.LoadJSON, handle)
  File "/home/victor/hpc/tvm-build/tvmgpucopy2/python/tvm/_ffi/_ctypes/object.py", line 136, in __init_handle_by_constructor__
    handle = __init_by_constructor__(fconstructor, args)
  File "/home/victor/hpc/tvm-build/tvmgpucopy2/python/tvm/_ffi/_ctypes/packed_func.py", line 260, in __init_handle_by_constructor__
    raise get_last_ffi_error()
tvm._ffi.base.TVMError: Traceback (most recent call last):
  9: TVMFuncCall
  8: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::runtime::ObjectRef (std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)>::AssignTypedLambda<tvm::runtime::ObjectRef (*)(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)>(tvm::runtime::ObjectRef (*)(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >), std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  7: tvm::LoadJSON(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)
  6: tvm::JSONAttrSetter::Set(tvm::runtime::ObjectPtr<tvm::runtime::Object>*, tvm::JSONNode*)
  5: tvm::JSONAttrSetter::Visit(char const*, tvm::runtime::DataType*)
  4: tvm::runtime::String2DLDataType(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)
  3: tvm::runtime::ParseCustomDatatype(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, char const**)
  2: tvm::runtime::GetCustomTypeCode(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)
  1: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::datatype::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#2}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  0: tvm::datatype::Registry::GetTypeCode(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)
  File "/home/victor/hpc/tvm-build/tvmgpucopy2/src/target/datatype/registry.cc", line 59
TVMError: 
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------
  Check failed: (name_to_code_.find(type_name) != name_to_code_.end()) is false: Type name cmpl not registered. Elements in type registration list: 0

Note that I extended the error message a bit by also printing name_to_code_.size() in Registry::GetTypeCode in the file registery.cc. And indeed the length of that std::unordered_map is equal to 0. So indeed it seems as if type registrations are not properly done.

Do you have any suggestions?