Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Runtime] Introduce runtime module property #14406

Merged
merged 8 commits into from
Mar 30, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 38 additions & 13 deletions include/tvm/runtime/module.h
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,35 @@
namespace tvm {
namespace runtime {

/*!
* \brief Property of runtime module
* We classify the property of runtime module into the following categories.
*/
enum ModulePropertyMask : int {
/*! \brief kBinarySerializable
* we can serialize the module to the stream of bytes. CUDA/OpenCL/JSON
* runtime are representative examples. A binary exportable module can be integrated into final
* runtime artifact by being serialized as data into the artifact, then deserialized at runtime.
* This class of modules must implement SaveToBinary, and have a matching deserializer registered
* as 'runtime.module.loadbinary_<type_key>'.
*/
kBinarySerializable = 0b001,
/*! \brief kRunnable
* we can run the module directly. LLVM/CUDA/JSON runtime, executors (e.g,
* virtual machine) runtimes are runnable. Non-runnable modules, such as CSourceModule, requires a
* few extra steps (e.g,. compilation, link) to make it runnable.
*/
kRunnable = 0b010,
/*! \brief kDSOExportable
* we can export the module as DSO. A DSO exportable module (e.g., a
* CSourceModuleNode of type_key 'c') can be incorporated into the final runtime artifact (ie
* shared library) by compilation and/or linking using the external compiler (llvm, nvcc, etc).
* DSO exportable modules must implement SaveToFile. In general, DSO exportable modules are not
* runnable unless there is a special support like JIT for `LLVMModule`.
*/
kDSOExportable = 0b100
};

class ModuleNode;
class PackedFunc;

Expand Down Expand Up @@ -193,20 +222,16 @@ class TVM_DLL ModuleNode : public Object {
const std::vector<Module>& imports() const { return imports_; }

/*!
* \brief Returns true if this module is 'DSO exportable'.
*
* A DSO exportable module (eg a CSourceModuleNode of type_key 'c') can be incorporated into the
* final runtime artifact (ie shared library) by compilation and/or linking using the external
* compiler (llvm, nvcc, etc). DSO exportable modules must implement SaveToFile.
*
* By contrast, non-DSO exportable modules (eg CUDAModuleNode of type_key 'cuda') typically must
* be incorporated into the final runtime artifact by being serialized as data into the
* artifact, then deserialized at runtime. Non-DSO exportable modules must implement SaveToBinary,
* and have a matching deserializer registered as 'runtime.module.loadbinary_<type_key>'.
*
* The default implementation returns false.
* \brief Returns bitmap of property.
* By default, none of the property is set. Derived class can override this function and set its
* own property.
*/
virtual bool IsDSOExportable() const;
virtual int GetPropertyMask() const { return 0b000; }

/*! \brief Returns true if this module is 'DSO exportable'. */
bool IsDSOExportable() const {
return (GetPropertyMask() & ModulePropertyMask::kDSOExportable) != 0;
}

/*!
* \brief Returns true if this module has a definition for a function of \p name. If
Expand Down
3 changes: 3 additions & 0 deletions include/tvm/runtime/vm/executable.h
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,9 @@ class TVM_DLL Executable : public ModuleNode {
*/
PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ModulePropertyMask::kBinarySerializable; };

/*!
* \brief Write the Executable to the binary stream in serialized form.
*
Expand Down
3 changes: 3 additions & 0 deletions include/tvm/runtime/vm/vm.h
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,9 @@ class TVM_DLL VirtualMachine : public runtime::ModuleNode {
*/
virtual void LoadExecutable(const ObjectPtr<Executable>& exec);

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ModulePropertyMask::kRunnable; }

protected:
/*! \brief Push a call frame on to the call stack. */
void PushFrame(Index arg_count, Index ret_pc, const VMFunction& vm_func);
Expand Down
48 changes: 47 additions & 1 deletion python/tvm/runtime/module.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,14 @@ def __str__(self):
)


class ModulePropertyMask(object):
"""Runtime Module Property Mask."""

BINARY_SERIALIZABLE = 0b001
RUNNABLE = 0b010
DSO_EXPORTABLE = 0b100


class Module(object):
"""Runtime Module."""

Expand Down Expand Up @@ -239,6 +247,40 @@ def imported_modules(self):
nmod = _ffi_api.ModuleImportsSize(self)
return [_ffi_api.ModuleGetImport(self, i) for i in range(nmod)]

def get_property_mask(self):
"""Get the runtime module property mask. The mapping is stated in ModulePropertyMask.

Returns
-------
mask : int
Bitmask of runtime module property
"""
return _ffi_api.ModuleGetPropertyMask(self)

@property
def is_binary_serializable(self):
"""Returns true if module is 'binary serializable', ie can be serialzed into binary
stream and loaded back to the runtime module.

Returns
-------
b : Bool
True if the module is binary serializable.
"""
return (self.get_property_mask() & ModulePropertyMask.BINARY_SERIALIZABLE) != 0

@property
def is_runnable(self):
"""Returns true if module is 'runnable'. ie can be executed without any extra
compilation/linking steps.

Returns
-------
b : Bool
True if the module is runnable.
"""
return (self.get_property_mask() & ModulePropertyMask.RUNNABLE) != 0

@property
def is_dso_exportable(self):
"""Returns true if module is 'DSO exportable', ie can be included in result of
Expand All @@ -249,7 +291,7 @@ def is_dso_exportable(self):
b : Bool
True if the module is DSO exportable.
"""
return _ffi_api.ModuleIsDSOExportable(self)
return (self.get_property_mask() & ModulePropertyMask.DSO_EXPORTABLE) != 0

def save(self, file_name, fmt=""):
"""Save the module to file.
Expand Down Expand Up @@ -383,6 +425,10 @@ def _collect_from_import_tree(self, filter_func):
stack.append(self)
while stack:
module = stack.pop()
assert (
module.is_dso_exportable or module.is_binary_serializable
), f"Module {module.type_key} should be either dso exportable or binary serializable."

if filter_func(module):
dso_modules.append(module)
for m in module.imported_modules:
Expand Down
3 changes: 3 additions & 0 deletions src/relay/backend/aot_executor_codegen.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1358,6 +1358,9 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {

const char* type_key() const final { return "RelayGraphRuntimeCodegenModule"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return runtime::ModulePropertyMask::kRunnable; }

private:
void init(void* mod, const Array<Target>& targets) {
codegen_ =
Expand Down
3 changes: 3 additions & 0 deletions src/relay/backend/build_module.cc
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,9 @@ class RelayBuildModule : public runtime::ModuleNode {
*/
const char* type_key() const final { return "RelayBuildModule"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return runtime::ModulePropertyMask::kRunnable; }

/*!
* \brief Build relay IRModule for graph executor
*
Expand Down
3 changes: 2 additions & 1 deletion src/relay/backend/contrib/ethosu/source_module.cc
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,8 @@ class EthosUModuleNode : public ModuleNode {
return Module(n);
}

bool IsDSOExportable() const final { return true; }
/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const { return ModulePropertyMask::kDSOExportable; }

bool ImplementsFunction(const String& name, bool query_imports) final {
return std::find_if(compilation_artifacts_.begin(), compilation_artifacts_.end(),
Expand Down
3 changes: 3 additions & 0 deletions src/relay/backend/graph_executor_codegen.cc
Original file line number Diff line number Diff line change
Expand Up @@ -687,6 +687,9 @@ class GraphExecutorCodegenModule : public runtime::ModuleNode {

const char* type_key() const final { return "RelayGraphExecutorCodegenModule"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return runtime::ModulePropertyMask::kRunnable; }

private:
std::shared_ptr<GraphExecutorCodegen> codegen_;
LoweredOutput output_;
Expand Down
3 changes: 3 additions & 0 deletions src/relay/backend/vm/compiler.h
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,9 @@ class VMCompiler : public runtime::ModuleNode {

const char* type_key() const final { return "VMCompiler"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ModulePropertyMask::kRunnable; }

/*!
* \brief Set the parameters
*
Expand Down
3 changes: 3 additions & 0 deletions src/relay/printer/model_library_format_printer.cc
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,9 @@ class ModelLibraryFormatPrinter : public ::tvm::runtime::ModuleNode {

const char* type_key() const final { return "model_library_format_printer"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return runtime::ModulePropertyMask::kRunnable; }

std::string Print(const ObjectRef& node) {
std::ostringstream oss;
oss << node;
Expand Down
3 changes: 3 additions & 0 deletions src/runtime/aot_executor/aot_executor.h
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,9 @@ class TVM_DLL AotExecutor : public ModuleNode {
*/
const char* type_key() const final { return "AotExecutor"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ModulePropertyMask::kRunnable; }

void Run();

/*!
Expand Down
3 changes: 3 additions & 0 deletions src/runtime/aot_executor/aot_executor_factory.h
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,9 @@ class TVM_DLL AotExecutorFactory : public runtime::ModuleNode {
*/
const char* type_key() const final { return "AotExecutorFactory"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ModulePropertyMask::kBinarySerializable; }

/*!
* \brief Save the module to binary stream.
* \param stream The binary stream to save to.
Expand Down
3 changes: 3 additions & 0 deletions src/runtime/const_loader_module.cc
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,9 @@ class ConstLoaderModuleNode : public ModuleNode {

const char* type_key() const final { return "const_loader"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ModulePropertyMask::kBinarySerializable; };

/*!
* \brief Get the list of constants that is required by the given module.
* \param symbol The symbol that is being queried.
Expand Down
5 changes: 5 additions & 0 deletions src/runtime/contrib/coreml/coreml_runtime.h
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,11 @@ class CoreMLRuntime : public ModuleNode {
*/
virtual PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self);

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final {
return ModulePropertyMask::kBinarySerializable | ModulePropertyMask::kRunnable;
}

/*!
* \brief Serialize the content of the mlmodelc directory and save it to
* binary stream.
Expand Down
5 changes: 5 additions & 0 deletions src/runtime/contrib/ethosn/ethosn_runtime.h
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,11 @@ class EthosnModule : public ModuleNode {

const char* type_key() const override { return "ethos-n"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final {
return ModulePropertyMask::kBinarySerializable | ModulePropertyMask::kRunnable;
};

private:
/*! \brief A map between ext_symbols (function names) and ordered compiled networks. */
std::map<std::string, OrderedCompiledNetwork> network_map_;
Expand Down
5 changes: 5 additions & 0 deletions src/runtime/contrib/json/json_runtime.h
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,11 @@ class JSONRuntimeBase : public ModuleNode {

const char* type_key() const override { return "json"; } // May be overridden

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const {
return ModulePropertyMask::kBinarySerializable | ModulePropertyMask::kRunnable;
}

/*! \brief Initialize a specific json runtime. */
virtual void Init(const Array<NDArray>& consts) = 0;

Expand Down
4 changes: 4 additions & 0 deletions src/runtime/contrib/libtorch/libtorch_runtime.cc
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,10 @@ class TorchModuleNode : public ModuleNode {
: symbol_name_(symbol_name), module_(module) {}

const char* type_key() const { return "torch"; }
/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final {
return ModulePropertyMask::kBinarySerializable | ModulePropertyMask::kRunnable;
}

/*!
* \brief Get a packed function.
Expand Down
3 changes: 3 additions & 0 deletions src/runtime/contrib/onnx/onnx_module.cc
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,9 @@ class ONNXSourceModuleNode : public runtime::ModuleNode {
: code_(code), symbol_(symbol), const_vars_(const_vars) {}
const char* type_key() const { return "onnx"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ModulePropertyMask::kRunnable; };

PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final {
if (name == "get_symbol") {
return PackedFunc(
Expand Down
5 changes: 5 additions & 0 deletions src/runtime/contrib/tensorrt/tensorrt_runtime.cc
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,11 @@ class TensorRTRuntime : public JSONRuntimeBase {
*/
const char* type_key() const final { return "tensorrt"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final {
return ModulePropertyMask::kBinarySerializable | ModulePropertyMask::kRunnable;
}

/*!
* \brief Initialize runtime. Create TensorRT layer from JSON
* representation.
Expand Down
3 changes: 3 additions & 0 deletions src/runtime/contrib/tflite/tflite_runtime.h
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,9 @@ class TFLiteRuntime : public ModuleNode {
*/
const char* type_key() const { return "TFLiteRuntime"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ModulePropertyMask::kRunnable; };

/*!
* \brief Invoke the internal tflite interpreter and run the whole model in
* dependency order.
Expand Down
5 changes: 5 additions & 0 deletions src/runtime/contrib/vitis_ai/vitis_ai_runtime.h
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,11 @@ class VitisAIRuntime : public ModuleNode {
*/
const char* type_key() const { return "VitisAIRuntime"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final {
return ModulePropertyMask::kBinarySerializable | ModulePropertyMask::kRunnable;
};

/*!
* \brief Serialize the content of the pyxir directory and save it to
* binary stream.
Expand Down
5 changes: 5 additions & 0 deletions src/runtime/cuda/cuda_module.cc
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,11 @@ class CUDAModuleNode : public runtime::ModuleNode {

const char* type_key() const final { return "cuda"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final {
return ModulePropertyMask::kBinarySerializable | ModulePropertyMask::kRunnable;
}

PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;

void SaveToFile(const std::string& file_name, const std::string& format) final {
Expand Down
3 changes: 3 additions & 0 deletions src/runtime/graph_executor/graph_executor.h
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,9 @@ class TVM_DLL GraphExecutor : public ModuleNode {
const char* type_key() const final { return "GraphExecutor"; }
void Run();

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ModulePropertyMask::kRunnable; }

/*!
* \brief Initialize the graph executor with graph and device.
* \param graph_json The execution graph.
Expand Down
3 changes: 3 additions & 0 deletions src/runtime/graph_executor/graph_executor_factory.h
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,9 @@ class TVM_DLL GraphExecutorFactory : public runtime::ModuleNode {
*/
const char* type_key() const final { return "GraphExecutorFactory"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ModulePropertyMask::kBinarySerializable; }

/*!
* \brief Save the module to binary stream.
* \param stream The binary stream to save to.
Expand Down
4 changes: 4 additions & 0 deletions src/runtime/hexagon/hexagon_module.h
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,10 @@ class HexagonModuleNode : public runtime::ModuleNode {
PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) override;
std::string GetSource(const std::string& format) override;
const char* type_key() const final { return "hexagon"; }
/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const {
return ModulePropertyMask::kBinarySerializable | ModulePropertyMask::kRunnable;
}
void SaveToFile(const std::string& file_name, const std::string& format) override;
void SaveToBinary(dmlc::Stream* stream) override;

Expand Down
5 changes: 5 additions & 0 deletions src/runtime/library_module.cc
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,11 @@ class LibraryModuleNode final : public ModuleNode {

const char* type_key() const final { return "library"; }

/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final {
return ModulePropertyMask::kBinarySerializable | ModulePropertyMask::kRunnable;
};

PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final {
TVMBackendPackedCFunc faddr;
if (name == runtime::symbol::tvm_module_main) {
Expand Down
Loading