[JIT] Compile group of functions/modules

[edgarriba]

🚀 Feature

I would like to be able to jit compile a group of functions or modules in order to serialize the obtained library in a single file so that later can be loaded in C++.

Motivation

The main motivation for this feature is to port Kornia to C++ avoiding the massive effort to re-implement the whole thing from scratch in C++. We expect to do at some point in the future but as a first iteration I think we could bootstrap the process of compiling from python to C++ and producing something functional out of the box.

In addition, this could somehow being extended to port to other frameworks by serializing into ONNX format.

Any thoughts about how to proceed with this will be very welcomed :)

cc @suo

[edgarriba]

/cc @suo @fmassa @soumith

[zdevito]

If you put your functions, modules into a top-level module, and annotate them with export, then all of them will get saved:

class Top(torch.nn.Module):
   def __init__(self_:
      # mark the functions in these modules you want to save with @torch.jit.export
      self.submodule1 = ...
      self.submodule2 = ...
  @torch.jit.export
  def func1():
    ...
   @torch.jit.export
   def func2():
      ...

If this is not sufficient for what you want to save, can you elaborate more on your use case?

[edgarriba]

Sounds like It could work. Then from the C++ side I guess I will have to provide an interface for each of the exported functions, right? Which brings me to question of how the saved binary with functions works - Is it loaded at running time or can be cross compiled and linked to my C++ like a normal library ?

[suo]

It will be loaded at runtime.