How to save module with external codegen as .so or .dll

#1

I am running a model with external codegen. After calling relay.build, I call update_lib. However, when I call lib.export_library() or lib.save() on the library returned from the function, I get an error: Module[library] does not support SaveToFile.

What is the recommended way to save this to a single file? Or is it recommended to save the original lib separately, and call update_lib every time we load it?

I would like to do this in C++ as well, and am currently working on Windows.

cc @zhiics @comaniac

Here is the update_lib code:

def update_lib(lib):
    cur_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
    tvm_home = os.getenv("TVM_HOME", None)

    contrib_path = os.path.join(tvm_home, "src", "runtime", "contrib")

    kwargs = {}
    kwargs["options"] = ["-O3", "-std=c++11", "-I" + contrib_path]

    tmp_path = util.tempdir()
    lib_name = 'lib.so'
    lib_path = tmp_path.relpath(lib_name)
    lib.export_library(lib_path, fcompile=False, **kwargs)
    lib = tvm.runtime.load_module(lib_path)

    lib.export_library("test.dll") # this causes the error

    return lib
#2

Ping, has anyone successfully done this?

#3

This should not happen at least from my experience…

#4

I’ve successfully done something like this, but I’m not sure how similar your external codegen is. If you’ve implemented a custom runtime module, you may also need to implement a custom SaveToFile function. Could you give a bit more detail on how you get the initial ‘lib’?

#5

My custom codegen is almost the same as the DNNL example - it’s a CSourceModule. The original lib comes from:

MergeComposite(...)
Annotate(...)
PartitionGraph(...)
graph, lib, params = relay.build(...)

lib = update_lib(lib) # implementation is the same as the tutorial
#6

A weird thing to me is this error message. library is the type key of the host module LibraryModule, which should be a host module. It seems like export_library tries to call SaveToFile of the host module but failed because it is not implemented.

#7

Would you be able to share some example code on how you’ve exported it?

#8

@comaniac do you have any example code here? I just want to make sure I’m not doing anything silly.

#9

I don’t have other examples except for the unittest (the one you posted). The reason we exported the module to lib.so and loaded it back is because the compilation for generated kernels actually happens when exporting the module. It means if we skip this process and directly run the module, the kernel code you generated will not be compiled. We should fix this behavior in the future, but for your question I think it should be fine to have a single library file (lib.so). I am not sure why lib.export_library("test.dll") causes the error tho, have you tried lib.export_library("test.so") to make sure it’s not the target file format issue?

#10

Yeah, I tried .so and it had the same issue.

To be clear, after I call export_library and compile the generated code, update_lib then calls tvm.runtime.load_module. When I call export_library on the module returned from that function, I see the error.

Is that workflow correct? @mbaret what did your workflow look like?

Btw, when I print lib, it is of type “library”, not “llvm”.

Update: when I call lib.imported_modules[0].export_library, it seems like it tries to recompile.

#11

@lhutton1 - maybe you can help here since you’ve also been looking in this area.

#12

I haven’t run into this issue but the flow for ACL is slightly different, we cross compile the library and load it on a remote device. I’m probably missing something obvious here but is there a reason for calling export_library twice?