Relax JIT Build UX

Development unity
#1

In this post, we would like to discuss the interface of relax build and runtime UX. The following code summarizes the current set of use-cases.

# build and directly run
ex: relax.Executable = relax.vm.build(mod, target)
vm = relax.VirtualMachine(ex, device)

# export to library and run
ex.export_library("exported.so")
loaded_mod = ex.load_library("exported.so")
vm = relax.VirtualMachine(ex, device)

The main issue we will encounter here is that direct build and run may not always work, especially when the result module contains source modules(such as cutlass BYOC). In which case export and load back is necessary. However, we can also not by default export and load back, because the loaded back library is not exportable again.

To explain the situation further, we know that modules can have a subset of the following properties:

  • dso_exportable : can be exported into object or source that then compiled into a so
  • binary_exportable: can be exported via SaveBinary, when loaded back it will become runnable
  • runnable: directly runnable and can obtain packed functions

Note-ably, some modules may have subset of these properties

  • llvm: {runnable, dso_exportable}
  • a python module backed by python function: {runnable}
  • cuda file: {dso_exportable}
  • vm bytecode: {binary_exportable, runnable}

If we encounter a module that is exportable(either dso or binary) but is not runnable, then we will need to export it and load back. If we encounter a module that is runnable, but not exportable(e.g. python native module in the same process), then we cannot simply call export library. To handle all scenarios, it is helpful to look at the two primary use-cases here:

  • U0: JIT, build and directly generate runnable code in the same process
  • U1: Export, and then load back in same or another process

We propose to update the UX to make these two use-cases explicit. Specifically, relax.Executable will no longer be directly acceptable by VirtualMachine. Instead, we introduce a jit() function to explicit generate a runnable runtime module from the executable collection.

# build.py: compile time construct
class Executable:
    def export_library(name, fcompile=None):
        # redirects to mod.export_library

    def jit(fcompile=None) -> runtime.Module:
        """Just in time compile the executable to a runtime.Module
        
        Examples
        --------
        .. code:: python
                    
                ex = relax.build(mod, target)
                relax.VirtualMachine(ex.jit("nvcc")) 	
        """
        pass

# vm.py: runtime only
class VirtualMachine:
    def __init__(self, mod: runtime.Module, device: Device):
        # avoid look at class to isolate runtime
        if (not isinstance(mod, runtime.Module) and 
                type(mod).__name__ == "Executable"):
                raise ValueError(
                        "Please explicit run mod.jit()"
                        "or export then load it back before passing to VirtualMachine")

So the new flow becomes

# build and directly run
ex: relax.Executable = relax.build(mod, target)
vm = relax.VirtualMachine(ex.jit(), device)

# export to library and run
ex.export_library("exported.so")
loaded_mod = ex.load_library("exported.so")
vm = relax.VirtualMachine(loaded_mod, device)

Base on the properties of each module in our collection, we define the behavior of jit as follows:

  • jit
    • collect all dso_exportable and binary_exportable, call export_library to put them into a dso
    • load back and re-import the runnable modules
  • export_library: require all modules to be exportable or binary_exportable

Discussion

There are several pts that can be discussed, for example, the naming choice of the relax.Executable which collects the results of the build.

1 Like
[TVM Unity] Relax Open Dev Meeting Feb 21
#2

relax.vm.build seems too verbose, and we should instead use relax.build given both bytecode and AOT mode are supported by the VM already.

#3

The name “jit” might be confusing, since people associate the term with “compilation” while here the input is already a compiled artifact (exe).

1 Like
#4

indeed terminology is hard, on the other hand, jit also precisely describes what we are doing, as we get a transformed(compiled) but not yet runnable object, and jit further connects things together to a runnable state. Other suggestions are also welcomed

#5

We discussed this topic in the community meeting today. Everyone agree that having such separation is helpful

There is a suggestion to allow the VM to still take executable, and implicitly call ex.jit() in that case. This will allow us to make things backward compatible

ex: relax.Executable = relax.build(mod, target)
# this works
vm = relax.VirtualMachine(ex.jit(), device)
# implicitly call ex.jit()
vm = relax.VirtualMachine(ex, device)

We also agree that the name “jit” can be further iterated, a few candidates were discussed and we did not come up with a better alternative, so we will stay with this name for now, while allow the common case to not having to depend on explicit calling of jit.

We will proceed with an implementation of the discussed results and continue to iterate on the UX if needed.

2 Likes
#6

The implementation now lands in the unity branch https://github.com/apache/tvm/pull/14088

1 Like