wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tvm/runtime/vm/builtin.h
+89
View File
@@ -0,0 +1,89 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/*!
* \file tvm/runtime/vm/builtin.h
* \brief Builtin runtime APIs.
*/
#ifndef TVM_RUNTIME_VM_BUILTIN_H_
#define TVM_RUNTIME_VM_BUILTIN_H_
namespace tvm {
namespace runtime {
namespace vm {
/*!
* \brief Op code used in built-in match-shape function.
*
* The function takes the following signature:
* MatchShape(input_shape, shape_heap, n, c[0], r[0], c[1], r[1], ... c[n], r[n], err_ctx)
*
* This function provides runtime shape population and checking support for match-cast.
* When a shape variable appears in the first time, we should load the shape and
* populate the variable. When a shape variable already appears, we should
* assert that it already equals an existing shape value.
*
* NOTE: It is OK to pass nullptr shape_heap if all code are AssertEqualToImm.
*/
enum class MatchShapeCode : int {
/*!
* \brief Perform an assertion that shape equals immediate.
*
* assert input_shape[i] == r[i]
*/
kAssertEqualToImm = 0,
/*!
* \brief This is the first time we see a symbolic shape variable, store to heap.
*
* shape_heap[r[i]] = input_shape[i]
*/
kStoreToHeap = 1,
/*!
* \brief skip and do not do anything.
*/
kNoOp = 2,
/*!
* \brief Peform an assertion that the shape equals a loaded value.
*
* assert input_shape[i] == shape_heap[r[i]]
*/
kAssertEqualToLoad = 3,
};
/*!
* \brief Op code used in builtin function MakeShape.
*
* MakeShape(shape_heap, n, c[0], r[0], c[1], r[1], ... c[n], r[n]).
*
* \note It is OK to pass nullptr to shape_heap if all code are UseImm.
*/
enum class MakeShapeCode : int {
/*! \brief Use the following r[i] as immediate shape value. */
kUseImm = 0,
/*!
* \brief Load shape value from the shape_heap[[r[i]].
*/
kLoadShape = 1,
};
} // namespace vm
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_VM_BUILTIN_H_
include/tvm/runtime/vm/bytecode.h
+246
View File
@@ -0,0 +1,246 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/*!
* \file tvm/runtime/vm/bytecode.h
* \brief The bytecode for the virtual machine.
*/
#ifndef TVM_RUNTIME_VM_BYTECODE_H_
#define TVM_RUNTIME_VM_BYTECODE_H_
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/error.h>
#include <iostream>
#include <vector>
namespace tvm {
namespace runtime {
namespace vm {
/*!
* \brief The storage type for the bytecode in the VM.
*/
using ExecWord = int64_t;
/*! \brief A register name. */
using RegName = ExecWord;
/*!
* \brief An alias for the integer type used ubiquitously in the VM.
*/
using Index = ExecWord;
/*!
* \brief An enumeration of Relax's opcodes.
*
* The opcode is used to implement instruction
* as a tagged union.
*/
enum class Opcode {
Call = 1U,
Ret = 2U,
Goto = 3U,
If = 4U,
};
/*! \brief A single virtual machine instruction.
*
* The representation of the instruction is as
* a tagged union.
*
* The first field represents which instruction,
* and by extension which field of the union
* is active.
*/
struct Instruction {
/*! \brief The number of bit for storing value. */
static constexpr ExecWord kKindBit = 8;
/*! \brief The number of bit for storing value. */
static constexpr ExecWord kValueBit = sizeof(ExecWord) * 8 - kKindBit;
/*! \brief The bit mask of the value part. */
static constexpr ExecWord kValueMask = (static_cast<ExecWord>(1) << kValueBit) - 1;
/*! \brief Maximum possible value, use 1 bit for sign. */
static constexpr ExecWord kValueMaxLimit = (static_cast<ExecWord>(1) << (kValueBit - 1)) - 1;
/*! \brief Minimum possible value, remove 1 slot to keep things symmetric. */
static constexpr ExecWord kValueMinLimit = -kValueMaxLimit;
/*! \brief Beginning of special register section. */
static constexpr RegName kBeginSpecialReg = static_cast<ExecWord>(1) << 54;
/*! \brief Random magic number that represents void argument, indicate null value */
static constexpr RegName kVoidRegister = kBeginSpecialReg + 0;
/*! \brief Random magic number that represents the VM context */
static constexpr RegName kVMRegister = kBeginSpecialReg + 1;
/*!
* \brief The kind of instruction's argument.
*/
enum class ArgKind : int { kRegister = 0, kImmediate = 1, kConstIdx = 2, kFuncIdx = 3 };
friend std::ostream& operator<<(std::ostream& os, const ArgKind& kind) {
switch (kind) {
case ArgKind::kRegister:
os << "kRegister";
break;
case ArgKind::kImmediate:
os << "kImmediate";
break;
case ArgKind::kConstIdx:
os << "kConstIdx";
break;
case ArgKind::kFuncIdx:
os << "kFuncIdx";
break;
default:
TVM_FFI_THROW(InternalError)
<< "Internal error: "
<< "Invalid ArgKind with integer value " << static_cast<int>(kind);
}
return os;
}
/*!
* \brief The auxiliary data structure for instruction argument.
*/
class Arg {
public:
/*! \brief Construct a void argument. */
Arg() : data_(Instruction::kVoidRegister) {}
/*!
* \brief construct Arg from data.
* \param data The data repr.
*/
static Arg FromData(ExecWord data) { return Arg(data); }
/*!
* \brief construct a register Arg.
* \param reg The register number.
* \return The constructed arg.
*/
static Arg Register(RegName reg) { return Arg(ArgKind::kRegister, reg); }
/*!
* \brief construct a ConstIdx arg.
* \param index The constant index.
* \return The constructed arg.
*/
static Arg ConstIdx(Index index) { return Arg(ArgKind::kConstIdx, index); }
/*!
* \brief construct a immediate arg.
* \param imm_value The immediate value.
* \return The constructed arg.
*/
static Arg Immediate(int64_t imm_value) { return Arg(ArgKind::kImmediate, imm_value); }
/*!
* \brief construct a FuncIdx arg.
* \param index The func index in the function table.
* \return The constructed arg.
*/
static Arg FuncIdx(Index index) { return Arg(ArgKind::kFuncIdx, index); }
/*!
* \brief Get the kind of argument.
* \return The kind of argument.
*/
ArgKind kind() const {
uint8_t kind = (data_ >> kValueBit) & 0xFF;
return Instruction::ArgKind(kind);
}
/*!
* \brief Get the value of argument.
* \return The value of argument.
* \note We store both positive and negative values by sign extension.
*/
ExecWord value() const { return ((data_ & kValueMask) << kKindBit) >> kKindBit; }
/*!
* \brief Get the raw data repr of the arg.
* \return The raw data.
*/
ExecWord data() const { return data_; }
private:
/*! \brief Construct from the data. */
explicit Arg(ExecWord data) : data_(data) {}
/*! \brief Construct from the kind and value. */
Arg(ArgKind kind, Index value) {
TVM_FFI_ICHECK_LE(value, kValueMaxLimit);
TVM_FFI_ICHECK_GE(value, kValueMinLimit);
data_ = (static_cast<ExecWord>(kind) << kValueBit) | (value & kValueMask);
}
/*! \brief The underlying stored data. */
ExecWord data_;
};
/*! \brief The instruction opcode. */
Opcode op;
union {
struct /* Call */ {
/*! \brief The destination register. */
RegName dst;
/*! \brief The index into the packed function table. */
Index func_idx;
/*! \brief The number of arguments to the packed function. */
Index num_args;
/*! \brief The arguments of the packed function. */
Arg* args;
};
struct /* Ret */ {
/*! \brief The return result. */
RegName result;
};
struct /* Goto */ {
/*! \brief The jump offset. */
Index pc_offset;
};
struct /* If */ {
/*! \brief The register containing the cond value. */
RegName cond;
/*! \brief The program counter offset for the false branch. */
Index false_offset;
};
};
/*!
* \brief Construct a Call instruction.
* \param func_idx The index of the function to call.
* \param num_args The number of arguments.
* \param args The input arguments.
* \param dst The destination register.
* \return The call instruction.
*/
static Instruction Call(Index func_idx, Index num_args, Arg* args, RegName dst);
/*!
* \brief Construct a return instruction.
* \param result The register containing the return value.
* \return The return instruction.
*/
static Instruction Ret(RegName result);
/*!
* \brief Construct a goto instruction.
* \param pc_offset The register containing the jump offset.
* \return The goto instruction.
*/
static Instruction Goto(RegName pc_offset);
/*!
* \brief Construct an If instruction.
* \param cond The register containing the cond value.
* \param false_offset The program counter offset for the false branch.
* \return The If instruction.
*/
static Instruction If(RegName cond, Index false_offset);
};
} // namespace vm
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_VM_BYTECODE_H_
include/tvm/runtime/vm/executable.h
+242
View File
@@ -0,0 +1,242 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/*!
* \file tvm/runtime/vm/executable.h
*/
#ifndef TVM_RUNTIME_VM_EXECUTABLE_H_
#define TVM_RUNTIME_VM_EXECUTABLE_H_
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/support/io.h>
#include <tvm/support/serializer.h>
#include <string>
#include <unordered_map>
#include <vector>
#include "./bytecode.h"
// Convention: this version should set to minimum TVM version it support
// NOTE: this file only changes if we change relax vm format
// for example if relax vm format do not change in 0.15, this should remain as 0.14
// if it changes in 0.16, we will change it to 0.16
#define VM_VERSION "0.14"
namespace tvm {
namespace runtime {
namespace vm {
/*!
* \brief Information entry in executable function table.
*
* Contains metadata about the compiled function, as
* well as the compiled VM instructions.
*/
struct VMFuncInfo {
/*! \brief kind of the function. */
enum class FuncKind : int {
/*! \brief system level packed function */
kPackedFunc = 0,
/*! \brief VM function. */
kVMFunc = 1,
/*! \brief VMTIR function. */
kVMTIRFunc = 2,
};
/*! \brief The kind of function. */
FuncKind kind;
/*! \brief The function's name, global symbol */
std::string name;
/*! \brief The start instruction index of the function. */
Index start_instr = 0;
/*! \brief The end instruction index of the function. */
Index end_instr = 0;
/*! \brief The number of arguments of the function. */
Index num_args = 0;
/*! \brief The register file size of the function. */
Index register_file_size = 0;
/*! \brief The function parameter names.*/
std::vector<std::string> param_names;
// defined customized loader save
void Save(support::Stream* writer) const;
bool Load(support::Stream* reader);
};
/*!
* \brief The virtual machine executable emitted by the VM compiler.
*
* The executable contains information (e.g. data in different memory regions)
* to run in a virtual machine.
*/
class TVM_RUNTIME_DLL VMExecutable : public ffi::ModuleObj {
public:
/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ffi::Module::kBinarySerializable; };
/*!
* \brief Print the detailed statistics of the given code, i.e. number of
* globals and constants, etc.
* \return The statistics represented by a string.
*/
std::string Stats() const;
/*!
* \brief Get the i-th instruction from the executable.
* \param i The index of the instruction to be fetched.
* \return The instruction.
*/
Instruction GetInstruction(Index i) const;
/*!
* \brief Set j-th byte data of i-th instruction to val.
* \param i The index of the instruction to be updated.
* \param j The index of the byte data of the instruction to be updated.
* \param val The value to be set
*/
void SetInstructionData(Index i, Index j, ExecWord val);
/*!
* \brief Print the instructions as text format.
* \return The text format of the instructions.
*/
ffi::String AsText() const;
/*!
* \brief Print the instructions as python program.
* \return The python program of the instructions, represented by a string.
*/
ffi::String AsPython() const;
/*!
* \brief Write the VMExecutable to the binary stream in serialized form.
* \return The binary bytes that save the executable to.
*/
ffi::Bytes SaveToBytes() const final;
/*!
* \brief Load VMExecutable from the binary stream in serialized form.
* \param bytes The binary bytes that load the executable from.
* \return The loaded executable, in the form of a `runtime::Module`.
*/
static ffi::Module LoadFromBytes(const ffi::Bytes& bytes);
/*!
* \brief Write the VMExecutable to the provided path as a file containing its serialized content.
* \param file_name The name of the file to write the serialized data to.
* \param format The target format of the saved file.
*/
void WriteToFile(const ffi::String& file_name, const ffi::String& format) const final;
/*! \brief Create a Relax virtual machine and load `this` as the executable. */
ffi::Module VMLoadExecutable() const;
/*! \brief Check if the VMExecutable contains a specific function. */
bool HasFunction(const ffi::String& name) const;
/*!
* \brief Load VMExecutable from the file.
* \param file_name The path of the file that load the executable from.
* \return The loaded executable, in the form of a `runtime::Module`.
*/
static ffi::Module LoadFromFile(const ffi::String& file_name);
/*! \brief The virtual machine's function table. */
std::vector<VMFuncInfo> func_table;
/*! \brief A map from globals (as strings) to their index in the function map. */
std::unordered_map<std::string, Index> func_map;
/*! \brief The global constant pool. */
std::vector<ffi::Any> constants;
/*! \brief The VDevice memory scopes */
std::unordered_map<Index, std::string> memory_scopes;
/*! \brief The offset of instruction. */
std::vector<Index> instr_offset;
/*! \brief The byte data of instruction. */
std::vector<ExecWord> instr_data;
virtual ~VMExecutable() {}
/*! \brief Module type key. */
const char* kind() const final;
/*!
* \brief Look up an exported function by name.
* \param name The function name.
* \return The function if found, otherwise std::nullopt.
*/
ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) override;
private:
/*!
* \brief Save the globals.
* \param strm The input stream.
*/
void SaveGlobalSection(support::Stream* strm) const;
/*!
* \brief Save the memory scopes.
* \param strm The output stream.
*/
void SaveMemoryScopeSection(support::Stream* strm) const;
/*!
* \brief Save the constant pool.
* \param strm The input stream.
*/
void SaveConstantSection(support::Stream* strm) const;
/*!
* \brief Save the instructions.
* \param strm The input stream.
*/
void SaveCodeSection(support::Stream* strm) const;
/*!
* \brief Save the packed functions.
* \param strm The input stream.
*/
void SavePackedFuncNames(support::Stream* strm) const;
/*!
* \brief Load the globals.
* \param strm The input stream.
*/
void LoadGlobalSection(support::Stream* strm);
/*!
* \brief Load the memory scopes.
* \param strm The input stream.
*/
void LoadMemoryScopeSection(support::Stream* strm);
/*!
* \brief Load the constant pool.
* \param strm The input stream.
*/
void LoadConstantSection(support::Stream* strm);
/*!
* \brief Load the instructions.
* \param strm The input stream.
*/
void LoadCodeSection(support::Stream* strm);
/*!
* \brief Save the packed functions.
* \param strm The input stream.
*/
void LoadPackedFuncNames(support::Stream* strm);
};
} // namespace vm
} // namespace runtime
} // namespace tvm
namespace tvm {
namespace support {
template <>
struct Serializer<::tvm::runtime::vm::VMFuncInfo> {
static constexpr bool enabled = true;
static void Write(Stream* strm, const ::tvm::runtime::vm::VMFuncInfo& data) { data.Save(strm); }
static bool Read(Stream* strm, ::tvm::runtime::vm::VMFuncInfo* data) { return data->Load(strm); }
};
} // namespace support
} // namespace tvm
#endif // TVM_RUNTIME_VM_EXECUTABLE_H_
include/tvm/runtime/vm/tensor_cache_support.h
+97
View File
@@ -0,0 +1,97 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
#ifndef TVM_RUNTIME_VM_TENSOR_CACHE_SUPPORT_H_
#define TVM_RUNTIME_VM_TENSOR_CACHE_SUPPORT_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/runtime/tensor.h>
#include <string>
#include <unordered_map>
#include <vector>
namespace tvm {
namespace runtime {
namespace vm {
/*!
* \brief Metadata for Tensor cache, which by default, is named as "tensor-cache.json".
*/
struct TensorCacheMetadata {
/*! \brief Each shard of Tensor cache, which by default, is named as "params_shard_x.bin". */
struct FileRecord {
/*! \brief Metadata of each parameter */
struct ParamRecord {
/*!
* \brief Load the parameter from raw data.
* \param device The device to load the parameter onto.
* \param raw_data The raw data stream
* \param staging_buffer The buffer to be used to avoid extra OpenCL copies. Pass in a nullptr
* in other cases
*/
TVM_RUNTIME_DLL Tensor Load(Device device, const std::string* raw_data,
ffi::Optional<Tensor>* staging_buffer = nullptr) const;
/*! \brief Name of the parameter */
std::string name;
/*! \brief Shape of the parameter */
ffi::Shape shape;
/*! \brief Data type of the parameter */
DLDataType dtype;
/*! \brief Format of the parameter */
std::string format;
/*! \brief Number of bytes */
int64_t nbytes;
/*! \brief Offset from the raw stream */
int64_t byte_offset;
};
/*! \brief Load a FileRecord into memory */
TVM_RUNTIME_DLL ffi::Array<Tensor> Load(Device device, //
const std::string& path_prefix, //
std::string* raw_data_buffer, //
ffi::Optional<Tensor>* staging_buffer = nullptr) const;
/*! \brief Relative path to the bin file */
std::string data_path;
/*! \brief Format of the file */
std::string format;
/*! \brief Size of the file */
int64_t nbytes;
/*! \brief The parameters in the file */
std::vector<ParamRecord> records;
};
/*! \brief The files in the Tensor cache */
std::vector<FileRecord> records;
/*! \brief The path to the `tensor-cache.json` file */
std::string path;
/*! \brief Load the metadata from a specific directory */
TVM_RUNTIME_DLL static TensorCacheMetadata Load(const std::string& path);
/*! \brief Load the metadata from a given JSON string */
TVM_RUNTIME_DLL static TensorCacheMetadata LoadFromStr(const std::string& json_str,
const std::string& path);
};
} // namespace vm
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_VM_TENSOR_CACHE_SUPPORT_H_
include/tvm/runtime/vm/vm.h
+226
View File
@@ -0,0 +1,226 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/*!
* \file tvm/runtime/vm/vm.h
*/
#ifndef TVM_RUNTIME_VM_VM_H_
#define TVM_RUNTIME_VM_VM_H_
#include <tvm/ffi/extra/module.h>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "../memory/memory_manager.h"
#include "./bytecode.h"
#include "./executable.h"
namespace tvm {
namespace runtime {
using memory::Allocator;
using memory::AllocatorType;
using memory::MemoryManager;
using memory::Storage;
using memory::StorageObj;
namespace vm {
/*!
* \brief Possible instrument actions.
*/
enum class VMInstrumentReturnKind : int {
/*! \brief Running as normal. */
kNoOp = 0,
/*! \brief Skip the following run, only valid in before. */
kSkipRun = 1,
};
/*!
* \brief An object representing a vm closure.
*/
class VMClosureObj : public ffi::Object {
public:
/*!
* \brief The function name. The function could be any
* function object that is compatible to the VM runtime.
*/
ffi::String func_name;
/*!
* \brief The implementation of the Closure.
* \note This function takes context pointer(VirtualMachine*)
* as the first argument. The rest of arguments follows
* the same arguments as the normal function call.
*/
ffi::Function impl;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.vm.Closure", VMClosureObj, ffi::Object);
};
/*! \brief reference to closure. */
class VMClosure : public ffi::ObjectRef {
public:
VMClosure(ffi::String func_name, ffi::Function impl);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VMClosure, ffi::ObjectRef, VMClosureObj);
/*!
* \brief Create another ffi::Function with last arguments already bound to last_args.
*
* This is a helper function to create captured closures.
* \param func The input func, can be a VMClosure or ffi::Function.
* \param last_args The arguments to bound to in the end of the function.
* \note The new function takes in arguments and append the last_args in the end.
*/
static ffi::Function BindLastArgs(ffi::Function func, std::vector<ffi::Any> last_args);
};
/*!
* \brief Represent a VM extension.
* A VM extension allows the user to extend the VM with target specific functionalities.
* The VM holds the reference of the extensions to ensure the extensions have the same lifetime
* as the VM.
*
* This is the base class for all VM extensions and should not be used directly.
*/
class VMExtensionNode : public ffi::Object {
protected:
TVM_FFI_DECLARE_OBJECT_INFO("runtime.VMExtension", VMExtensionNode, ffi::Object);
};
/*! \brief Managed reference to VM extension. */
class VMExtension : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VMExtension, ffi::ObjectRef, VMExtensionNode);
};
/*!
* \brief The virtual machine.
*
* The virtual machine contains all the current execution state,
* as well as the executable.
*
* The goal is to have a single self-contained object,
* enabling one to easily pass around VMs, execute them on
* multiple threads, or serialize them to disk or over the
* wire.
*/
class VirtualMachine : public ffi::ModuleObj {
public:
/*!
* \brief Initialize the virtual machine for a set of devices.
* \param devices The set of TVM devices.
* \param alloc_types The allocator types for each device.
*/
virtual void Init(const std::vector<Device>& devices,
const std::vector<AllocatorType>& alloc_types) = 0;
/*!
* \brief Load the executable for the virtual machine.
* \param exec The executable.
*/
virtual void LoadExecutable(ffi::ObjectPtr<VMExecutable> exec) = 0;
/*!
* \brief Get global function in the VM.
* \param func_name The name of the function.
* \return The closure
*/
virtual VMClosure GetClosure(const ffi::String& func_name) = 0;
/*!
* \brief Invoke closure or packed function using ffi::Function convention.
* \param closure_or_packedfunc A VM closure or a packed_func.
* \param args The input arguments.
* \param rv The return value.
*/
virtual void InvokeClosurePacked(const ffi::ObjectRef& closure_or_packedfunc,
ffi::PackedArgs args, ffi::Any* rv) = 0;
/*!
* \brief Set an instrumentation function.
*
* If instrument is present, the function will be called
* before/after each Call instruction.
*
* bool instrument(func, func_symbol, before_run, args...)
*
* - func: Union[VMClosure, ffi::Function], the function object.
* - func_symbol: string, the symbol of the function.
* - before_run: bool, whether it is before or after call.
* - ret_value: Only valid in after run, otherwise it is null.
* - args: the arguments being passed to call.
*
* instrument can return an int which corresponds to the action value.
* \sa VMInstrumentAction
*
* \param instrument The instrument function.
*/
virtual void SetInstrument(ffi::Function instrument) = 0;
/*!
* \brief Get or create a VM extension. Once created, the extension will be stored in the VM
* and held until the VM is destructed.
*
* \tparam T The type of the extension
* \return The extension instance
*/
template <typename T, typename = std::enable_if_t<std::is_base_of<VMExtension, T>::value>>
T GetOrCreateExtension() {
using ContainerType = typename T::ContainerType;
uint32_t key = ContainerType::RuntimeTypeIndex();
if (auto it = extensions.find(key); it != extensions.end()) {
ffi::Any value = (*it).second;
return value.cast<T>();
}
auto [it, _] = extensions.emplace(key, T::Create());
ffi::Any value = (*it).second;
return value.cast<T>();
}
/*!
* \brief Create a specific instance of VM.
* \return Created VM
*/
static ffi::ObjectPtr<VirtualMachine> Create();
/*!
* \brief Helper function for vm closure functions to get the context ptr
* \param arg The argument value.
*/
static VirtualMachine* GetContextPtr(ffi::AnyView arg) {
return static_cast<VirtualMachine*>(arg.cast<void*>());
}
~VirtualMachine() {}
//--------------------------------------------------------------------------
// The following section contains states that other builtin can depend on
//--------------------------------------------------------------------------
/*! \brief The memory allocators. */
std::vector<Allocator*> allocators;
/*! \brief Runtime physical device list. */
std::vector<Device> devices;
/*! \brief The VM extensions. Mapping from the type index of the extension to the extension
* instance. */
std::unordered_map<uint32_t, Any> extensions;
};
} // namespace vm
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_VM_VM_H_