You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2023/02/20 15:53:52 UTC
[tvm] 01/46: [Unity] Relax VM (#13878)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity-staging
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 6bb0c0594301ca1c542485be04bd6de9e4b4da32
Author: Yuchen Jin <yu...@cs.washington.edu>
AuthorDate: Wed Feb 1 08:52:16 2023 -0500
[Unity] Relax VM (#13878)
This PR implements a flexible register-based VM to execute relax programs with dynamic shape and control flow. Design: https://github.com/tlc-pack/relax/wiki/Relax-VM-Design.
Co-Authored-by: Ziheng Jiang <zi...@apache.org>
Co-Authored-by: Ruihang Lai <ru...@cs.cmu.edu>
Co-Authored-by: Sunghyun Park <49...@users.noreply.github.com>
Co-Authored-by: Junru Shao <ju...@gmail.com>
Co-Authored-by: Prakalp Srivastava <pr...@octoml.ai>
Co-Authored-by: Yong Wu <yo...@gmail.com>
Co-Authored-by: Steven S. Lyubomirsky <sl...@octoml.ai>
Co-Authored-by: Tianqi Chen <ti...@gmail.com>
Co-Authored-by: Hongyi Jin <32...@qq.com>
---
CMakeLists.txt | 2 +
include/tvm/relax/exec_builder.h | 181 ++++++
include/tvm/runtime/relax_vm/builtin.h | 89 +++
include/tvm/runtime/relax_vm/bytecode.h | 223 +++++++
include/tvm/runtime/relax_vm/executable.h | 209 +++++++
include/tvm/runtime/relax_vm/memory_manager.h | 142 +++++
include/tvm/runtime/relax_vm/vm.h | 148 +++++
python/tvm/relax/__init__.py | 24 +
python/tvm/relax/_ffi_api.py | 20 +
python/tvm/relax/exec_builder.py | 147 +++++
python/tvm/relax/testing/vm.py | 85 +++
python/tvm/relax/vm.py | 609 +++++++++++++++++++
src/relax/backend/vm/exec_builder.cc | 399 +++++++++++++
src/runtime/relax_vm/builtin.cc | 445 ++++++++++++++
src/runtime/relax_vm/bytecode.cc | 68 +++
src/runtime/relax_vm/executable.cc | 576 ++++++++++++++++++
src/runtime/relax_vm/memory_manager.cc | 181 ++++++
src/runtime/relax_vm/naive_allocator.h | 65 +++
src/runtime/relax_vm/pooled_allocator.h | 111 ++++
src/runtime/relax_vm/vm.cc | 811 ++++++++++++++++++++++++++
tests/python/relax/test_vm_execbuilder.py | 262 +++++++++
21 files changed, 4797 insertions(+)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 818e8b50ad..ed2afc3920 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -289,6 +289,7 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
src/driver/*.cc
src/support/*.cc
src/script/*.cc
+ src/relax/backend/vm/*.cc
)
tvm_file_glob(GLOB CODEGEN_SRCS
@@ -335,6 +336,7 @@ tvm_file_glob(GLOB RUNTIME_SRCS
src/runtime/*.cc
src/runtime/vm/*.cc
src/runtime/minrpc/*.cc
+ src/runtime/relax_vm/*.cc
)
if(BUILD_FOR_HEXAGON)
diff --git a/include/tvm/relax/exec_builder.h b/include/tvm/relax/exec_builder.h
new file mode 100644
index 0000000000..03e58392c2
--- /dev/null
+++ b/include/tvm/relax/exec_builder.h
@@ -0,0 +1,181 @@
+/*
+ * 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/relax/exec_builder.h
+ */
+#ifndef TVM_RELAX_EXEC_BUILDER_H_
+#define TVM_RELAX_EXEC_BUILDER_H_
+
+#include <tvm/ir/expr.h>
+#include <tvm/node/reflection.h>
+#include <tvm/node/repr_printer.h>
+#include <tvm/runtime/object.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/runtime/relax_vm/bytecode.h>
+#include <tvm/runtime/relax_vm/executable.h>
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace tvm {
+namespace relax {
+
+namespace vm = tvm::runtime::relax_vm;
+
+class ExecBuilder;
+
+/*!
+ * \brief A builder provides api to build VM executable with instructions.
+ */
+class ExecBuilderNode : public Object {
+ public:
+ /*!
+ * \brief Declare a function, it is OK to have multiple declarations.
+ * \param func The function name.
+ * \param kind The kind of the function.
+ */
+ void DeclareFunction(const std::string& func, vm::VMFuncInfo::FuncKind kind);
+ /*!
+ * \brief To annotate the start of a vm function.
+ * \param func The function name.
+ * \param num_inputs The number of inputs.
+ * \param param_names The function parameter names.
+ * \param kind The kind of the function.
+ * \param init_register_size Initial setting of register file size.
+ */
+ void EmitFunction(const std::string& func, int64_t num_inputs,
+ Optional<Array<String>> param_names,
+ vm::VMFuncInfo::FuncKind kind = vm::VMFuncInfo::FuncKind::kVMFunc,
+ int64_t init_register_size = 0);
+ /*!
+ * \brief Annotate the end of a vm function.
+ * \param func The function name.
+ */
+ void EndFunction(const std::string& func);
+ /*!
+ * \brief Emit a call instruction for a packed function.
+ * \param func The packed function name.
+ * \param args The arguments of the function.
+ * \param ret The return register.
+ */
+ void EmitCall(const std::string& func, std::vector<vm::Instruction::Arg> args, vm::RegName ret);
+ /*!
+ * \brief Emit a call instruction with func as argument.
+ * \param func The packed function index.
+ * \param args The arguments of the function.
+ * \param ret The return register.
+ */
+ void EmitCall(vm::Instruction::Arg func, std::vector<vm::Instruction::Arg> args, vm::RegName ret);
+ /*!
+ * \brief Emit a ret instruction.
+ * \param result The return result.
+ * \note result must be a register.
+ */
+ void EmitRet(vm::Instruction::Arg result);
+ /*!
+ * \brief Emit a goto instruction.
+ * \param pc_offset The program counter offset as the jump offset.
+ */
+ void EmitGoto(vm::Index pc_offset);
+ /*!
+ * \brief Emit an If instruction.
+ * \param cond The register containing the cond value.
+ * \param false_offset The program counter offset for the false branch.
+ * \note result must be a register.
+ */
+ void EmitIf(vm::Instruction::Arg cond, vm::Index false_offset);
+ /*!
+ * \brief Get function index by its name.
+ * \param name The name of the function.
+ * \return The argument corresponding to the function index.
+ */
+ vm::Instruction::Arg GetFunction(const std::string& name);
+ /*!
+ * \brief Convert a constant value something that exec builder can understand.
+ *
+ * This function may update the constant pool to include the obj value.
+ *
+ * \param value The input constant value
+ * \return An Arg that represents the result of constant argument.
+ */
+ template <typename T>
+ vm::Instruction::Arg ConvertConstant(T value) {
+ TVMRetValue rv;
+ rv = value;
+ return ConvertConstant_(rv);
+ }
+ /*!
+ * \brief Raw access to underlying executable build in progress.
+ */
+ vm::Executable* exec() const;
+ /*!
+ * \brief Finalize the build, run formalize and get the final result.
+ * \note This function should not be called during construction.
+ */
+ ObjectPtr<vm::Executable> Get();
+ /*!
+ * \brief Create an ExecBuilder.
+ * \return The ExecBuilder.
+ */
+ TVM_DLL static ExecBuilder Create();
+
+ void VisitAttrs(AttrVisitor* v) {}
+
+ static constexpr const uint32_t _type_index = TypeIndex::kDynamic;
+ static constexpr const char* _type_key = "relax.ExecBuilder";
+ TVM_DECLARE_FINAL_OBJECT_INFO(ExecBuilderNode, Object);
+
+ private:
+ /*!
+ * \brief Convert a constant value something that exec builder can understand.
+ *
+ * This function may update the constant pool to include the obj value.
+ *
+ * \param obj The constant value to be emitted
+ * \return An Arg that represents the result of constant argument.
+ */
+ vm::Instruction::Arg ConvertConstant_(TVMRetValue obj);
+
+ /*!
+ * \brief A helper function to check if an executable is legal by checking if registers are used
+ * properly
+ */
+ void CheckExecutable();
+ /*!
+ * \brief Formalize the executable.
+ */
+ void Formalize();
+
+ /*! \brief The mutable internal executable. */
+ ObjectPtr<vm::Executable> exec_; // mutable
+ /*! \brief internal dedup map when creating index for a new constant */
+ std::unordered_map<ObjectRef, vm::Index, StructuralHash, StructuralEqual> const_dedup_map_;
+};
+
+class ExecBuilder : public ObjectRef {
+ public:
+ TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ExecBuilder, ObjectRef, ExecBuilderNode);
+};
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_EXEC_BUILDER_H_
diff --git a/include/tvm/runtime/relax_vm/builtin.h b/include/tvm/runtime/relax_vm/builtin.h
new file mode 100644
index 0000000000..b994e44ae8
--- /dev/null
+++ b/include/tvm/runtime/relax_vm/builtin.h
@@ -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/relax_vm/builtin.h
+ * \brief Builtin runtime APIs.
+ */
+#ifndef TVM_RUNTIME_RELAX_VM_BUILTIN_H_
+#define TVM_RUNTIME_RELAX_VM_BUILTIN_H_
+
+namespace tvm {
+namespace runtime {
+namespace relax_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 relax_vm
+} // namespace runtime
+} // namespace tvm
+#endif // TVM_RUNTIME_RELAX_VM_BUILTIN_H_
diff --git a/include/tvm/runtime/relax_vm/bytecode.h b/include/tvm/runtime/relax_vm/bytecode.h
new file mode 100644
index 0000000000..91d1823258
--- /dev/null
+++ b/include/tvm/runtime/relax_vm/bytecode.h
@@ -0,0 +1,223 @@
+/*
+ * 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/relax_vm/bytecode.h
+ * \brief The bytecode for the virtual machine.
+ */
+#ifndef TVM_RUNTIME_RELAX_VM_BYTECODE_H_
+#define TVM_RUNTIME_RELAX_VM_BYTECODE_H_
+
+#include <tvm/runtime/data_type.h>
+#include <tvm/runtime/logging.h>
+
+#include <iostream>
+#include <vector>
+
+namespace tvm {
+namespace runtime {
+namespace relax_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 Begining 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 };
+ /*!
+ * \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) {
+ ICHECK_LE(value, kValueMaxLimit);
+ 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 relax_vm
+} // namespace runtime
+} // namespace tvm
+
+#endif // TVM_RUNTIME_RELAX_VM_BYTECODE_H_
diff --git a/include/tvm/runtime/relax_vm/executable.h b/include/tvm/runtime/relax_vm/executable.h
new file mode 100644
index 0000000000..316a03a3b8
--- /dev/null
+++ b/include/tvm/runtime/relax_vm/executable.h
@@ -0,0 +1,209 @@
+/*
+ * 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/relax_vm/executable.h
+ */
+#ifndef TVM_RUNTIME_RELAX_VM_EXECUTABLE_H_
+#define TVM_RUNTIME_RELAX_VM_EXECUTABLE_H_
+
+#include <tvm/runtime/container/closure.h>
+#include <tvm/runtime/object.h>
+#include <tvm/runtime/registry.h>
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "./bytecode.h"
+
+namespace tvm {
+namespace runtime {
+namespace relax_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(dmlc::Stream* writer) const;
+ bool Load(dmlc::Stream* reader);
+};
+
+/*!
+ * \brief The executable emitted by the VM compiler.
+ *
+ * The executable contains information (e.g. data in different memory regions)
+ * to run in a virtual machine.
+ */
+class Executable : public runtime::ModuleNode {
+ public:
+ /*!
+ * \brief Get a PackedFunc from the executable module.
+ * \param name the name of the function.
+ * \param sptr_to_self The shared_ptr that points to this module node.
+ * \return PackedFunc or nullptr when it is not available.
+ */
+ PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
+ /*!
+ * \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.
+ */
+ String AsText() const;
+ /*!
+ * \brief Print the instructions as python program.
+ * \return The python program of the instructions, represented by a string.
+ */
+ String AsPython() const;
+ /*!
+ * \brief Write the Executable to the binary stream in serialized form.
+ * \param stream The binary stream to save the executable to.
+ */
+ void SaveToBinary(dmlc::Stream* stream) final;
+ /*!
+ * \brief Load Executable from the binary stream in serialized form.
+ * \param stream The binary stream that load the executable from.
+ * \return The loaded executable, in the form of a `runtime::Module`.
+ */
+ static Module LoadFromBinary(void* stream);
+ /*!
+ * \brief Write the Executable 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 SaveToFile(const std::string& file_name, const std::string& format) final;
+ /*!
+ * \brief Load Executable 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 Module LoadFromFile(const std::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<TVMRetValue> constants;
+ /*! \brief The offset of instruction. */
+ std::vector<Index> instr_offset;
+ /*! \brief The byte data of instruction. */
+ std::vector<ExecWord> instr_data;
+
+ virtual ~Executable() {}
+
+ const char* type_key() const final { return "relax.Executable"; }
+
+ private:
+ /*!
+ * \brief Save the globals.
+ * \param strm The input stream.
+ */
+ void SaveGlobalSection(dmlc::Stream* strm);
+ /*!
+ * \brief Save the constant pool.
+ * \param strm The input stream.
+ */
+ void SaveConstantSection(dmlc::Stream* strm);
+ /*!
+ * \brief Save the instructions.
+ * \param strm The input stream.
+ */
+ void SaveCodeSection(dmlc::Stream* strm);
+ /*!
+ * \brief Save the packed functions.
+ * \param strm The input stream.
+ */
+ void SavePackedFuncNames(dmlc::Stream* strm);
+ /*!
+ * \brief Load the globals.
+ * \param strm The input stream.
+ */
+ void LoadGlobalSection(dmlc::Stream* strm);
+ /*!
+ * \brief Load the constant pool.
+ * \param strm The input stream.
+ */
+ void LoadConstantSection(dmlc::Stream* strm);
+ /*!
+ * \brief Load the instructions.
+ * \param strm The input stream.
+ */
+ void LoadCodeSection(dmlc::Stream* strm);
+ /*!
+ * \brief Save the packed functions.
+ * \param strm The input stream.
+ */
+ void LoadPackedFuncNames(dmlc::Stream* strm);
+};
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
+
+namespace dmlc {
+DMLC_DECLARE_TRAITS(has_saveload, ::tvm::runtime::relax_vm::VMFuncInfo, true);
+} // namespace dmlc
+#endif // TVM_RUNTIME_RELAX_VM_EXECUTABLE_H_
diff --git a/include/tvm/runtime/relax_vm/memory_manager.h b/include/tvm/runtime/relax_vm/memory_manager.h
new file mode 100644
index 0000000000..e5ae8cfcfb
--- /dev/null
+++ b/include/tvm/runtime/relax_vm/memory_manager.h
@@ -0,0 +1,142 @@
+/*
+ * 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/relax_vm/memory_manager.h
+ * \brief Abstract device memory management API
+ */
+#ifndef TVM_RUNTIME_RELAX_VM_MEMORY_MANAGER_H_
+#define TVM_RUNTIME_RELAX_VM_MEMORY_MANAGER_H_
+
+#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/ndarray.h>
+
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+#include <vector>
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+struct Buffer {
+ /*! \brief The pointer to the allocated block of memory. */
+ void* data{nullptr};
+ /*! \brief The size of the block. */
+ size_t size{0};
+ /*! \brief The device of the allocated buffers. */
+ Device device;
+};
+
+enum AllocatorType {
+ kNaive = 1,
+ kPooled,
+};
+
+class Allocator {
+ public:
+ explicit Allocator(AllocatorType type) : type_(type) {}
+ virtual ~Allocator() = default;
+ /*! \brief Allocate an empty NDArray using from the allocator.
+ * \param shape The shape of the NDArray.
+ * \param dtype The datatype of the NDArray.
+ * \param dev The device where the array is allocated.
+ * \return The empty NDArray.
+ */
+ runtime::NDArray Empty(std::vector<int64_t> shape, DLDataType dtype, Device dev);
+ /*! \brief Return the allocator type. */
+ inline AllocatorType type() const { return type_; }
+ /*! \brief Allocate a buffer given a size, alignment and type.
+ * \param nbytes The size of the buffer.
+ * \param alignment The alignment of the buffer.
+ * \param type_hint A type hint to the allocator.
+ * \return A sized allocation in the form of a buffer.
+ */
+ virtual Buffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) = 0;
+ /*! \brief Free a buffer allocated by the allocator.
+ * \param buffer The buffer to free.
+ */
+ virtual void Free(const Buffer& buffer) = 0;
+
+ private:
+ AllocatorType type_;
+};
+
+class MemoryManager {
+ public:
+ static MemoryManager* Global();
+ /*!
+ * \brief Get or create an allocator given the device and allocator type.
+ * \param dev The TVM device
+ * \param type The allocator type
+ * \return The memory allocator.
+ */
+ static Allocator* GetOrCreateAllocator(Device dev, AllocatorType type);
+ /*!
+ * \brief Get an allocator given the device.
+ * \param dev The TVM device
+ * \return The memory allocator.
+ */
+ static Allocator* GetAllocator(Device dev);
+
+ private:
+ MemoryManager() {}
+
+ private:
+ std::mutex mutex_;
+ std::unordered_map<Device, std::unique_ptr<Allocator>> allocators_;
+};
+
+/*! \brief An object representing a storage allocation. */
+class StorageObj : public Object {
+ public:
+ /*! \brief The index into the VM function table. */
+ Buffer buffer;
+
+ /*! \brief Allocate an NDArray from a given piece of storage. */
+ runtime::NDArray AllocNDArray(uint64_t offset, ShapeTuple shape, DLDataType dtype);
+
+ /*! \brief The deleter for an NDArray when allocated from underlying storage. */
+ static void Deleter(Object* ptr);
+
+ ~StorageObj() {
+ auto alloc = MemoryManager::Global()->GetAllocator(buffer.device);
+ alloc->Free(buffer);
+ }
+
+ static constexpr const uint32_t _type_index = runtime::TypeIndex::kDynamic;
+ static constexpr const char* _type_key = "relax.Storage";
+ TVM_DECLARE_FINAL_OBJECT_INFO(StorageObj, Object);
+};
+
+/*! \brief reference to storage. */
+class Storage : public ObjectRef {
+ public:
+ explicit Storage(Buffer buffer);
+
+ TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Storage, ObjectRef, StorageObj);
+};
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
+
+#endif // TVM_RUNTIME_RELAX_VM_MEMORY_MANAGER_H_
diff --git a/include/tvm/runtime/relax_vm/vm.h b/include/tvm/runtime/relax_vm/vm.h
new file mode 100644
index 0000000000..cfe3880904
--- /dev/null
+++ b/include/tvm/runtime/relax_vm/vm.h
@@ -0,0 +1,148 @@
+/*
+ * 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/relax_vm/vm.h
+ */
+#ifndef TVM_RUNTIME_RELAX_VM_VM_H_
+#define TVM_RUNTIME_RELAX_VM_VM_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "./bytecode.h"
+#include "./executable.h"
+#include "./memory_manager.h"
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+/*!
+ * \brief An object representing a vm closure.
+ */
+class VMClosureObj : public ClosureObj {
+ public:
+ /*!
+ * \brief The function name. The function could be any
+ * function object that is compatible to the VM runtime.
+ */
+ 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.
+ */
+ PackedFunc impl;
+
+ static constexpr const uint32_t _type_index = TypeIndex::kDynamic;
+ static constexpr const char* _type_key = "relax.vm.Closure";
+ TVM_DECLARE_FINAL_OBJECT_INFO(VMClosureObj, ClosureObj);
+};
+
+/*! \brief reference to closure. */
+class VMClosure : public Closure {
+ public:
+ VMClosure(String func_name, PackedFunc impl);
+ TVM_DEFINE_OBJECT_REF_METHODS(VMClosure, Closure, VMClosureObj);
+
+ /*!
+ * \brief Create another PackedFunc 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 PackedFunc.
+ * \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 PackedFunc BindLastArgs(PackedFunc func, std::vector<TVMRetValue> last_args);
+};
+
+/*!
+ * \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 runtime::ModuleNode {
+ 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(ObjectPtr<Executable> exec) = 0;
+ /*!
+ * \brief Get global function in the VM.
+ * \param func_name The name of the function.
+ * \return The closure
+ */
+ virtual VMClosure GetClosure(const String& func_name) = 0;
+ /*!
+ * \brief Invoke closure or packed function using PackedFunc 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 ObjectRef& closure_or_packedfunc, TVMArgs args,
+ TVMRetValue* rv) = 0;
+ /*!
+ * \brief Create a specific instance of VM.
+ * \return Created VM
+ */
+ static ObjectPtr<VirtualMachine> Create();
+ /*!
+ * \brief Helper function for vm closure functions to get the context ptr
+ * \param arg The argument value.
+ */
+ static VirtualMachine* GetContextPtr(TVMArgValue arg) {
+ return static_cast<VirtualMachine*>(arg.operator void*());
+ }
+
+ ~VirtualMachine() {}
+
+ const char* type_key() const final { return "relax.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;
+};
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
+
+#endif // TVM_RUNTIME_RELAX_VM_VM_H_
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
new file mode 100644
index 0000000000..c070fa4791
--- /dev/null
+++ b/python/tvm/relax/__init__.py
@@ -0,0 +1,24 @@
+# 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.
+# pylint: disable=invalid-name, wrong-import-position
+"""The Relax IR namespace containing the IR, type, operator, builder, vm, etc."""
+from . import exec_builder
+from . import vm
+
+# VM
+from .exec_builder import ExecBuilder
+from .vm import VirtualMachine
diff --git a/python/tvm/relax/_ffi_api.py b/python/tvm/relax/_ffi_api.py
new file mode 100644
index 0000000000..a127e1c813
--- /dev/null
+++ b/python/tvm/relax/_ffi_api.py
@@ -0,0 +1,20 @@
+# 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.
+"""FFI API for Relax."""
+import tvm._ffi
+
+tvm._ffi._init_api("relax", __name__)
diff --git a/python/tvm/relax/exec_builder.py b/python/tvm/relax/exec_builder.py
new file mode 100644
index 0000000000..1e28c967d1
--- /dev/null
+++ b/python/tvm/relax/exec_builder.py
@@ -0,0 +1,147 @@
+# 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.
+# pylint: disable=invalid-name
+"""A builder to build Relax VM executable."""
+from enum import IntEnum
+from typing import Optional, Union, List
+import tvm
+from tvm.runtime import Object
+from tvm.runtime.container import ShapeTuple
+from .vm import Executable
+from . import _ffi_api
+
+
+class SpecialReg(IntEnum):
+ """Magic numbers that represent special registers in vm."""
+
+ VOID_ARG = (1 << 54) + 0
+ VM_STATE = (1 << 54) + 1
+
+
+class VMFuncKind(IntEnum):
+ """VM function kind code."""
+
+ PACKED_FUNC = 0
+ VM_FUNC = 1
+
+
+class VMFuncScope(object):
+ """An object corresponds to each VM function, working as a context manager."""
+
+ stack: List["VMFuncScope"] = []
+
+ def __init__(self, exit_callback):
+ self.exit_callback = exit_callback
+
+ def __enter__(self):
+ VMFuncScope.stack.append(self)
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ VMFuncScope.stack.pop()
+ self.exit_callback()
+
+
+@tvm._ffi.register_object("relax.ExecBuilder")
+class ExecBuilder(Object):
+ """A builder to emit instructions and build executable for the virtual machine."""
+
+ def __init__(self) -> None:
+ self.__init_handle_by_constructor__(_ffi_api.ExecBuilderCreate) # type: ignore
+
+ def r(self, idx: int) -> int:
+ """set instruction's argument as a register."""
+ return _ffi_api.ExecBuilderR(self, idx) # type: ignore
+
+ def imm(self, value: int) -> int:
+ """set instruction's argument as an immediate."""
+ return _ffi_api.ExecBuilderImm(self, value) # type: ignore
+
+ def c(self, idx: int) -> int:
+ """set instruction's argument as a constant."""
+ return _ffi_api.ExecBuilderC(self, idx) # type: ignore
+
+ def f(self, name: str) -> int:
+ """set instruction's argument as a function."""
+ return _ffi_api.ExecBuilderF(self, name) # type: ignore
+
+ def void_arg(self) -> int:
+ return self.r(SpecialReg.VOID_ARG)
+
+ def vm_state(self) -> int:
+ return self.r(SpecialReg.VM_STATE)
+
+ def declare_function(self, func_name: str, kind: VMFuncKind = VMFuncKind.PACKED_FUNC) -> None:
+ """Declare a function"""
+ _ffi_api.ExecBuilderDecalreFunction(self, func_name, kind) # type: ignore
+
+ def function(
+ self, func_name: str, num_inputs: Optional[int] = 0, param_names: List[str] = None
+ ) -> VMFuncScope:
+ """annotate a VM function."""
+ _ffi_api.ExecBuilderEmitFunction(self, func_name, num_inputs, param_names) # type: ignore
+ return VMFuncScope(lambda: _ffi_api.ExecBuilderEndFunction(self, func_name)) # type: ignore
+
+ def _check_scope(self) -> None:
+ if len(VMFuncScope.stack) == 0:
+ raise ValueError("emit should happen in a function scope")
+
+ def convert_constant(self, const: object) -> int:
+ return _ffi_api.ExecBuilderConvertConstant(self, const) # type: ignore
+
+ def emit_call(
+ self,
+ name: str,
+ args: Optional[List[Union[tvm.nd.NDArray, tvm.DataType]]] = None,
+ dst: int = None,
+ ) -> None:
+ """emit a call instruction which calls a packed function."""
+ self._check_scope()
+ if dst is None:
+ dst = SpecialReg.VOID_ARG
+ args_ = []
+ if args is not None:
+ for arg in args:
+ if isinstance(arg, tuple):
+ shape_tuple = ShapeTuple(arg)
+ new_arg = self.convert_constant(shape_tuple)
+ args_.append(new_arg)
+ elif isinstance(arg, (tvm.nd.NDArray, tvm.DataType, ShapeTuple)):
+ new_arg = self.convert_constant(arg)
+ args_.append(new_arg)
+ else:
+ args_.append(arg)
+ _ffi_api.ExecBuilderEmitCall(self, name, args_, dst) # type: ignore
+
+ def emit_ret(self, result: int) -> None:
+ """emit a return instruction"""
+ self._check_scope()
+ _ffi_api.ExecBuilderEmitRet(self, result) # type: ignore
+
+ def emit_goto(self, pc_offset):
+ """emit a goto instruction"""
+ self._check_scope()
+ _ffi_api.ExecBuilderEmitGoto(self, pc_offset) # type: ignore
+
+ def emit_if(self, cond, false_offset):
+ """emit an if instruction"""
+ self._check_scope()
+ _ffi_api.ExecBuilderEmitIf(self, cond, false_offset) # type: ignore
+
+ def get(self) -> Executable:
+ """return the executable"""
+ return Executable(_ffi_api.ExecBuilderGet(self)) # type: ignore
diff --git a/python/tvm/relax/testing/vm.py b/python/tvm/relax/testing/vm.py
new file mode 100644
index 0000000000..79da54be10
--- /dev/null
+++ b/python/tvm/relax/testing/vm.py
@@ -0,0 +1,85 @@
+# 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.
+# pylint: disable=invalid-name
+"""Testing utilities for relax VM"""
+from typing import Any, List
+import numpy as np # type: ignore
+
+import tvm
+from tvm import relax
+from tvm.runtime.object import Object
+
+
+@tvm.register_func("test.vm.move")
+def move(src):
+ return src
+
+
+@tvm.register_func("test.vm.add")
+def add(a, b):
+ ret = a.numpy() + b.numpy()
+ return tvm.nd.array(ret)
+
+
+@tvm.register_func("test.vm.mul")
+def mul(a, b):
+ ret = a.numpy() * b.numpy()
+ return tvm.nd.array(ret)
+
+
+@tvm.register_func("test.vm.equal_zero")
+def equal_zero(a):
+ ret = np.all((a.numpy() == 0))
+ return tvm.nd.array(ret)
+
+
+@tvm.register_func("test.vm.subtract_one")
+def subtract_one(a):
+ ret = np.subtract(a.numpy(), 1)
+ return tvm.nd.array(ret)
+
+
+@tvm.register_func("test.vm.identity")
+def identity_packed(a, b):
+ b[:] = tvm.nd.array(a.numpy())
+
+
+@tvm.register_func("test.vm.tile")
+def tile_packed(a, b):
+ b[:] = tvm.nd.array(np.tile(a.numpy(), (1, 2)))
+
+
+@tvm.register_func("test.vm.add_scalar")
+def add_scalar(a, b):
+ return a + b
+
+
+@tvm.register_func("test.vm.get_device_id")
+def get_device_id(device):
+ return device.device_id
+
+
+def check_saved_func(vm: relax.VirtualMachine, func_name: str, *inputs: List[Any]) -> Object:
+ # uses save_function to create a closure with the given inputs
+ # and ensure the result is the same
+ # (assumes the functions return tensors and that they're idempotent)
+ saved_name = f"{func_name}_saved"
+ vm.save_function(func_name, saved_name, *inputs)
+ res1 = vm[func_name](*inputs)
+ res2 = vm[saved_name]()
+ tvm.testing.assert_allclose(res1.numpy(), res2.numpy(), rtol=1e-7, atol=1e-7)
+ return res1
diff --git a/python/tvm/relax/vm.py b/python/tvm/relax/vm.py
new file mode 100644
index 0000000000..ba16dfb079
--- /dev/null
+++ b/python/tvm/relax/vm.py
@@ -0,0 +1,609 @@
+# 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.
+# pylint: disable=invalid-name, redefined-builtin, no-else-return
+"""The Relax virtual machine"""
+from typing import Callable, List, Optional, Union, Dict, Tuple, Any
+import numpy as np # type: ignore
+
+from tvm._ffi import base as _base
+import tvm
+from tvm import relax
+from tvm.ir.module import IRModule
+from tvm.runtime import Device, Module, PackedFunc, container
+from tvm.runtime.object import Object
+from tvm.tir.function import PrimFunc
+from . import _ffi_api
+from ..rpc.base import RPC_SESS_MASK
+
+
+class Executable(object):
+ """The executable object emitted by the VM compiler or the ExecBuilder."""
+
+ def __init__(self, mod: Module):
+ self.mod = mod
+ self._stats = self.mod["stats"]
+ self._as_text = self.mod["as_text"]
+ self._as_python = self.mod["as_python"]
+
+ def stats(self) -> str:
+ """print the detailed statistics of the executable."""
+ return self._stats()
+
+ def as_text(self) -> str:
+ """print the instructions as text format."""
+ return self._as_text()
+
+ def as_python(self) -> str:
+ """print the instructions as python program."""
+ return self._as_python()
+
+
+class VirtualMachine(object):
+ """Relax VM runtime."""
+
+ NAIVE_ALLOCATOR = 1
+ POOLED_ALLOCATOR = 2
+
+ def __init__(
+ self,
+ exec: Union[Executable, Module],
+ device: Union[Device, List[Device]],
+ memory_cfg: Optional[Union[str, Dict[Device, str]]] = None,
+ ) -> None:
+ """
+ Construct a VirtualMachine wrapper object.
+
+ Parameters
+ ----------
+ exec: Union[Executable, Module]
+ The VM executable or Runtime Module
+
+ device : Union[Device, List[Device]]
+ The device to deploy the module.
+
+ memory_cfg : Optional[Union[str, Dict[Device, str]]]
+ Config the type of memory allocator. The allocator type can be ["naive",
+ "pooled"]. If memory_cfg is None, all devices will use pooled allocator
+ by default. If memory_cfg is string, all devices will use the specified
+ allocator type. If memory_cfg is a dict, each device uses the allocator
+ type specified in the dict, or pooled allocator if not specified in the
+ dict.
+ """
+ self.module = (
+ exec.mod["vm_load_executable"]()
+ if isinstance(exec, Executable)
+ else exec["vm_load_executable"]()
+ )
+ self._invoke_closure = self.module["invoke_closure"]
+ self._save_function = self.module["save_function"]
+ self._set_input = self.module["set_input"]
+ self._invoke_stateful = self.module["invoke_stateful"]
+ self._get_output = self.module["get_output"]
+ self._get_output_arity = self.module["get_output_arity"]
+ self._get_function_arity = self.module["get_function_arity"]
+ self._get_function_param_name = self.module["get_function_param_name"]
+ self._setup_device(device, memory_cfg)
+
+ def _setup_device(self, dev: Device, memory_cfg: Union[str, Dict[Device, str]]) -> None:
+ """init devices and allocators."""
+ devs = dev
+ if not isinstance(dev, (list, tuple)):
+ if not isinstance(dev, tvm.runtime.Device):
+ raise TypeError(
+ "dev is expected to be Device or \
+ List[Device]"
+ )
+ devs = [dev]
+
+ if any(dev.device_type % RPC_SESS_MASK == tvm.cpu().device_type for dev in devs[:-1]):
+ raise RuntimeError(
+ "CPU host is required to be the last element of the device list if provided."
+ )
+
+ # CPU is required for executing shape functions
+ if devs[-1].device_type % RPC_SESS_MASK != tvm.cpu().device_type:
+ devs.append(tvm.cpu())
+
+ default_alloc_type = VirtualMachine.POOLED_ALLOCATOR
+ if memory_cfg is None:
+ memory_cfg = {}
+ elif isinstance(memory_cfg, str):
+ assert memory_cfg in ["naive", "pooled"]
+ if memory_cfg == "naive":
+ default_alloc_type = VirtualMachine.NAIVE_ALLOCATOR
+ memory_cfg = {}
+ elif not isinstance(memory_cfg, dict):
+ raise TypeError(
+ "memory_cfg is expected be string or dictionary, "
+ + "but received {}".format(type(memory_cfg))
+ )
+ init_args = []
+ for device in devs:
+ init_args.append(device.device_type % RPC_SESS_MASK)
+ init_args.append(device.device_id)
+ alloc_type = memory_cfg[device] if device in memory_cfg else default_alloc_type
+ init_args.append(alloc_type)
+ self.module["vm_initialization"](*init_args)
+
+ def __getitem__(self, key: str) -> PackedFunc:
+ return self.module[key]
+
+ def invoke_closure(self, closure: Object, *args: Any) -> Object:
+ """Invoke a closure.
+
+ Parameters
+ ----------
+ closure : Object
+ The VMClosure Object.
+
+ args : list[tvm.runtime.NDArray] or list[np.ndarray]
+ The arguments to the closure.
+
+ Returns
+ -------
+ result : Object
+ The output.
+ """
+ return self._invoke_closure(closure, *args)
+
+ def save_function(
+ self,
+ func_name: str,
+ saved_name: str,
+ *args: List[Any],
+ include_return: bool = True,
+ **kwargs: Dict[str, Any],
+ ) -> None:
+ """
+ Convenience function. Takes a function from the module and saves
+ a `PackedFunc` that, when called, will invoke the function with the given arguments.
+ The `PackedFunc` can be accessed from the module using `saved_name`.
+ This is included to facilitate timing trials:
+ Invoking the returned `PackedFunc` will have less overhead from dictionary lookups
+ than normally running through the VM.
+
+ If the saved name is taken, it can be overridden, though it cannot override
+ the name of a function defined in the Relax source.
+
+ This is really creating a closure, but the function has a different name
+ to avoid confusion with `invoke_closure` (they are not meant to be used together).
+
+ Parameters
+ ----------
+ func_name : str
+ The function that should be packaged up.
+
+ saved_name : str
+ The name that the resulting closure should be saved under.
+
+ include_return : bool
+ Whether the saved PackedFunc should return its output.
+ If timing over RPC, it may not be desirable to send output
+ between machines.
+
+ args : List[Any]
+ The arguments to package up with the function.
+
+ kwargs : Dict[str, Any]
+ Any named arguments to package up with the function
+ """
+ cargs: List[Any] = []
+ if kwargs:
+ args = self._convert_func_named_args(func_name, args, **kwargs)
+ for arg in args:
+ self._convert(arg, cargs)
+ self._save_function(func_name, saved_name, int(include_return), *cargs)
+
+ def _convert(self, arg: Any, cargs: List) -> None:
+ """helper function to convert arguments to vm function."""
+
+ def _gettype(arg):
+ if isinstance(arg, np.float16):
+ return "float16"
+ elif isinstance(arg, (_base.integer_types, bool)):
+ return "int32"
+ else:
+ return "float32"
+
+ if isinstance(arg, Object):
+ cargs.append(arg)
+ elif isinstance(arg, np.ndarray):
+ nd_arr = tvm.nd.array(arg, device=tvm.cpu(0))
+ cargs.append(nd_arr)
+ elif isinstance(arg, tvm.runtime.NDArray):
+ cargs.append(arg)
+ elif isinstance(arg, (tuple, list)):
+ field_args: List[Any] = []
+ for field in arg:
+ self._convert(field, field_args)
+ cargs.append(container.tuple_object(field_args))
+ elif isinstance(arg, (_base.numeric_types, bool)):
+ dtype = _gettype(arg)
+ value = tvm.nd.array(np.array(arg, dtype=dtype), device=tvm.cpu(0))
+ cargs.append(value)
+ elif isinstance(arg, str):
+ cargs.append(arg)
+ else:
+ raise TypeError("Unsupported type: %s" % (type(arg)))
+
+ def _convert_func_named_args(self, func_name: str, args: Any, **kwargs: Any) -> Any:
+ """
+ Takes named function parameters and returns a list of those needed,
+ in the order they should appear
+ """
+ # kwargs can be a super set of the required function parameters.
+ # We only find the ones that are needed.
+ func_arity = self._get_function_arity(func_name)
+ func_params = [self._get_function_param_name(func_name, i) for i in range(func_arity)]
+ new_args = [None] * len(func_params)
+ cnt = 0
+ for k in kwargs:
+ if k in func_params:
+ idx = func_params.index(k)
+ new_args[idx] = kwargs[k]
+ cnt += 1
+ else:
+ print(f'Warning: Keyword argument "{k}" is unused in {func_name}')
+ assert len(args) + cnt == len(func_params)
+ idx = 0
+ for i, arg in enumerate(new_args):
+ if arg is None:
+ new_args[i] = args[idx]
+ idx += 1
+ return new_args
+
+ def set_input(self, func_name: str, *args: Any, **kwargs: Any) -> None:
+ """Set the inputs to a function.
+ This interface works when using VM over RPC by internally converting NDArray in
+ the arguments to DLTensor, which is supported in RPC where remote could only
+ have a minimal C runtime.
+
+ Note: If `set_input` is used, the function *must* be called using `invoke_stateful`
+ and the results must be obtained using `get_outputs`.
+
+ Parameters
+ ----------
+ func_name : str
+ The name of the function.
+ args: List[tvm.runtime.NDArray] or List[np.ndarray]
+ The arguments to the function.
+ kwargs: dict of str to tvm.runtime.NDArray or np.ndarray
+ Named arguments to the function.
+ """
+ cargs: List[Any] = []
+
+ if kwargs:
+ args = self._convert_func_named_args(func_name, args, **kwargs)
+
+ for arg in args:
+ self._convert(arg, cargs)
+
+ self._set_input(func_name, *cargs)
+
+ def invoke_stateful(self, func_name: str) -> None:
+ """
+ Call the named function from the VM module using the arguments set using `set_input`.
+ It is an error to call `invoke_stateful` without using `set_input` first
+ (even if it's to set 0 inputs); conversely, if `set_input` has been called,
+ it is an error to call the function without using `invoke_stateful`.
+
+ The results of the call can be obtained by calling `get_outputs`.
+
+ Parameters
+ ----------
+ func_name: str
+ The name of the function to call.
+ """
+ self._invoke_stateful(func_name)
+
+ def get_outputs(self, func_name: str) -> Union[tvm.Object, Tuple[Any]]:
+ """
+ Get the value output by the function by the given name
+ after a call of `invoke_stateful`.
+
+ It is an error to call this function without first calling `invoke_stateful`.
+
+ Parameters
+ ----------
+ func_name: str
+ The name of the function whose output should be fetched.
+
+ Returns
+ -------
+ ret: Union[tvm.Object, Tuple[Any]]
+ The result of the earlier call to the function via `invoke_stateful`.
+ If the result is a tuple, it returns a list of the fields.
+ The fields are potentially also tuples, so these can be arbitrily nested.
+ """
+ # to deal with potentially nested tuples, we need to query for arity recursively
+ def get_output_rec(func_name, *idx):
+ arity = self._get_output_arity(func_name, *idx)
+ if arity == -1:
+ return self._get_output(func_name, *idx)
+ # otherwise we need to specify more indices
+ idx_list = list(idx)
+ return tuple(get_output_rec(func_name, *(idx_list + [i])) for i in range(arity))
+
+ return get_output_rec(func_name)
+
+ def time_evaluator(
+ self,
+ func_name,
+ dev,
+ number=10,
+ repeat=1,
+ min_repeat_ms=0,
+ cooldown_interval_ms=0,
+ repeats_to_cooldown=1,
+ f_preproc="",
+ ) -> Callable[..., tvm.runtime.module.BenchmarkResult]:
+ """
+ Returns an evaluator that times a function in the module.
+ This follows the same convention as time_evaluator in tvm.runtime.module.
+ This can be used in combination with save_function() so that the
+ timings avoid extra dictionary lookups.
+
+ Parameters
+ ----------
+ func_name: str
+ The name of the function in the module.
+
+ dev: Device
+ The device we should run this function on.
+
+ number: int
+ The number of times to run this function for taking average.
+ We call these runs as one `repeat` of measurement.
+
+ repeat: int, optional
+ The number of times to repeat the measurement.
+ In total, the function will be invoked (1 + number x repeat) times,
+ where the first one is warm up and will be discarded.
+ The returned result contains `repeat` costs,
+ each of which is an average of `number` costs.
+
+ min_repeat_ms: int, optional
+ The minimum duration of one `repeat` in milliseconds.
+ By default, one `repeat` contains `number` runs. If this parameter is set,
+ the parameters `number` will be dynamically adjusted to meet the
+ minimum duration requirement of one `repeat`.
+ i.e., When the run time of one `repeat` falls below this time, the `number` parameter
+ will be automatically increased.
+
+ cooldown_interval_ms: int, optional
+ The cooldown interval in milliseconds between the number of repeats defined by
+ `repeats_to_cooldown`.
+
+ repeats_to_cooldown: int, optional
+ The number of repeats before the cooldown is activated.
+
+ f_preproc: str, optional
+ The preprocess function name we want to execute before executing the time evaluator.
+
+ Note
+ ----
+ The function will be invoked (1 + number x repeat) times,
+ with the first call discarded in case there is lazy initialization.
+
+ Example
+ -------
+ Normal use with a VM function (may not work over RPC if the function returns a tuple):
+
+ .. code-block:: python
+
+ target = tvm.target.Target("llvm", host="llvm")
+ ex = relax.vm.build(TestTimeEvaluator, target)
+ vm = relax.VirtualMachine(mod, tvm.cpu())
+ timing_res = vm.time_evaluator("func_name", tvm.cpu())(arg0, arg1, ..., argn)
+
+ Use with the stateful API:
+
+ .. code-block:: python
+
+ target = tvm.target.Target("llvm", host="llvm")
+ ex = relax.vm.build(TestTimeEvaluator, target)
+ vm = relax.VirtualMachine(mod, tvm.cpu())
+ vm.set_input("func_name", arg0, arg1, ..., argn)
+ timing_res = vm.time_evaluator("invoke_stateful", tvm.cpu())("func_name")
+
+ With saved closures via `save_function` (this results in
+ fewer dictionary lookups in the timed portion):
+
+ .. code-block:: python
+
+ target = tvm.target.Target("llvm", host="llvm")
+ ex = relax.vm.build(TestTimeEvaluator, target)
+ vm = relax.VirtualMachine(mod, tvm.cpu())
+ vm.save_function("func_name", "func_name_saved", arg0, arg1, ..., argn)
+ timing_res = vm.time_evaluator("func_name_saved", tvm.cpu())()
+
+ Returns
+ -------
+ ftimer : function
+ The function that takes same argument as func and returns a BenchmarkResult.
+ The ProfileResult reports `repeat` time costs in seconds.
+
+ """
+ return self.module.time_evaluator(
+ func_name,
+ dev,
+ number=number,
+ repeat=repeat,
+ min_repeat_ms=min_repeat_ms,
+ cooldown_interval_ms=cooldown_interval_ms,
+ repeats_to_cooldown=repeats_to_cooldown,
+ f_preproc=f_preproc,
+ )
+
+
+def _vmcodegen(
+ builder: "relax.ExecBuilder",
+ mod: tvm.IRModule,
+ exec_mode: str = "bytecode",
+) -> tvm.IRModule:
+ """Running VM codegen.
+
+ Parameters
+ ----------
+ builder: relax.ExecBuilder
+ ExecBuilder to collect the vm executable.
+
+ mod: IRModule
+ The input IRModule to be built.
+
+ exec_mode: {"bytecode", "compiled"}
+ The execution mode.
+
+ Return
+ ------
+ leftover: IRModule
+ Left over IRModule that may contain extra functions.
+ """
+
+ if exec_mode == "bytecode":
+ return _ffi_api.VMCodeGen(builder, mod) # type:ignore
+ if exec_mode == "compiled":
+ return _ffi_api.VMTIRCodeGen(builder, mod) # type: ignore
+ raise ValueError("Unknown exec_mode %s" % exec_mode)
+
+
+def _vmlink(
+ builder: "relax.ExecBuilder",
+ target: Union[str, tvm.target.Target],
+ tir_mod: Optional[tvm.IRModule] = None,
+ ext_libs: List[tvm.runtime.Module] = None,
+ params: Optional[Dict[str, list]] = None,
+):
+ """
+ Internal codegen function to make executable.
+
+ This function is only used for unit-testing purpoes.
+
+ Use build instead.
+
+ Parameters
+ ----------
+ builder: relax.ExecBuilder
+ Builder used to collect executables.
+
+ target : Union[str, tvm.target.Target]
+ A build target which can have optional host side compilation target.
+
+ tir_mod: IRModule
+ The input TIR IRModule to be linked together.
+
+ ext_libs: List[tvm.runtime.Module]
+ List of compiled external modules.
+
+ params: Optional[Dict[str, list]]
+ Extra parameter mappings.
+
+ Returns
+ -------
+ ex: tvm.relax.vm.Executable
+ An executable that can be loaded by virtual machine.
+ """
+ if isinstance(target, str):
+ target = tvm.target.Target(target)
+ if params is None:
+ params = {}
+ if ext_libs is None:
+ ext_libs = []
+ lib = None
+ if tir_mod is not None:
+ lib = tvm.build(tir_mod, target=target)
+ return Executable(_ffi_api.VMLink(builder, target, lib, ext_libs, params)) # type: ignore
+
+
+def build(
+ mod: tvm.IRModule,
+ target: Union[str, tvm.target.Target],
+ params: Optional[Dict[str, list]] = None,
+ exec_mode: str = "bytecode",
+) -> Executable:
+ """
+ Build an IRModule to VM executable.
+
+ Parameters
+ ----------
+ mod: IRModule
+ The input IRModule to be built.
+
+ target : Union[str, tvm.target.Target]
+ A build target which can have optional host side compilation target.
+
+ When TVM compiles device specific program such as CUDA,
+ we also need host(CPU) side code to interact with the driver
+ to setup the dimensions and parameters correctly.
+ host is used to specify the host side codegen target.
+ By default, llvm is used if it is enabled,
+ otherwise a stackvm interpreter is used.
+
+ params: Optional[Dict[str, list]]
+ Parameters for the input IRModule that will be bound.
+
+ exec_mode: {"bytecode", "compiled"}
+ The execution mode.
+
+ Returns
+ -------
+ ex: tvm.relax.vm.Executable
+ An executable that can be loaded by virtual machine.
+
+ Example
+ -------
+
+ .. code-block:: python
+ class InputModule:
+ @R.function
+ def foo(x: Tensor((3, 4), "float32"), y: Tensor((3, 4), "float32")):
+ z = R.add(x, y)
+ return z
+
+ mod = InputModule
+ target = tvm.target.Target("llvm", host="llvm")
+ ex = relax.vm.build(mod, target)
+ """
+ if isinstance(target, str):
+ target = tvm.target.Target(target)
+
+ passes = [relax.transform.ToNonDataflow()]
+ passes.append(relax.transform.CallTIRRewrite())
+ passes.append(relax.transform.VMBuiltinLower())
+ passes.append(relax.transform.VMShapeLower())
+ passes.append(relax.transform.AttachGlobalSymbol())
+ seq = tvm.transform.Sequential(passes)
+ new_mod = seq(mod)
+
+ # Extract external runtime modules if exist.
+ ext_libs = []
+ if mod.attrs and "external_mods" in mod.attrs:
+ ext_libs = mod.attrs["external_mods"]
+
+ # builder collects the executable
+ builder = relax.ExecBuilder()
+ leftover_mod = _vmcodegen(builder, new_mod, exec_mode=exec_mode)
+ tir_mod = _filter_tir(leftover_mod)
+ return _vmlink(builder, target, tir_mod, ext_libs, params)
+
+
+def _filter_tir(mod: tvm.IRModule) -> tvm.IRModule:
+ tir_mod = IRModule({})
+ for gv in mod.get_global_vars():
+ if isinstance(mod[gv], PrimFunc):
+ tir_mod[gv] = mod[gv]
+ return tir_mod
diff --git a/src/relax/backend/vm/exec_builder.cc b/src/relax/backend/vm/exec_builder.cc
new file mode 100644
index 0000000000..b5d932137b
--- /dev/null
+++ b/src/relax/backend/vm/exec_builder.cc
@@ -0,0 +1,399 @@
+/*
+ * 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 src/relax/backend/vm/exec_builder.cc
+ */
+#include <tvm/relax/exec_builder.h>
+
+#include <sstream>
+
+namespace tvm {
+namespace relax {
+
+using namespace vm;
+
+TVM_REGISTER_NODE_TYPE(ExecBuilderNode);
+
+ExecBuilder ExecBuilderNode::Create() {
+ ExecBuilder ret(make_object<ExecBuilderNode>());
+ ret->exec_ = make_object<Executable>();
+ return ret;
+}
+
+Executable* ExecBuilderNode::exec() const { return exec_.get(); }
+
+ObjectPtr<Executable> ExecBuilderNode::Get() {
+ this->Formalize();
+ this->CheckExecutable();
+ return exec_;
+}
+
+vm::Instruction::Arg ExecBuilderNode::ConvertConstant_(TVMRetValue cvalue) {
+ // emit constant immediate as immediate.
+ if (cvalue.type_code() == kDLInt) {
+ int64_t val = cvalue.operator int64_t();
+ if (val <= vm::Instruction::kValueMaxLimit && val >= vm::Instruction::kValueMinLimit) {
+ return vm::Instruction::Arg::Immediate(val);
+ }
+ }
+ // convert string to object string
+ if (cvalue.type_code() == kTVMStr) {
+ cvalue = cvalue.operator String();
+ }
+
+ // run dedup for object with structural equality
+ if (cvalue.IsObjectRef<ObjectRef>()) {
+ ObjectRef obj = cvalue.operator ObjectRef();
+ auto it = const_dedup_map_.find(obj);
+ if (it != const_dedup_map_.end()) {
+ return vm::Instruction::Arg::ConstIdx(it->second);
+ }
+ vm::Index idx = exec_->constants.size();
+ exec_->constants.push_back(cvalue);
+ const_dedup_map_[obj] = idx;
+ return vm::Instruction::Arg::ConstIdx(idx);
+ } else {
+ // emit normal constant
+ vm::Index idx = exec_->constants.size();
+ exec_->constants.push_back(cvalue);
+ return vm::Instruction::Arg::ConstIdx(idx);
+ }
+}
+
+void ExecBuilderNode::DeclareFunction(const std::string& func_name, VMFuncInfo::FuncKind kind) {
+ auto it = exec_->func_map.find(func_name);
+ if (it != exec_->func_map.end()) {
+ ICHECK(kind == exec_->func_table[it->second].kind)
+ << "Function " << func_name << "already declared in a different kind";
+ return;
+ }
+ VMFuncInfo vmfunc;
+ vmfunc.kind = kind;
+ vmfunc.name = func_name;
+ // use num args to mark undefined.
+ vmfunc.start_instr = 0;
+ vmfunc.num_args = -2;
+ vmfunc.register_file_size = 0;
+ exec_->func_map[func_name] = exec_->func_table.size();
+ exec_->func_table.push_back(vmfunc);
+}
+
+vm::Instruction::Arg ExecBuilderNode::GetFunction(const std::string& func_name) {
+ auto it = exec_->func_map.find(func_name);
+ ICHECK(it != exec_->func_map.end()) << "Cannot find function " << func_name;
+ return vm::Instruction::Arg::FuncIdx(it->second);
+}
+
+void ExecBuilderNode::EmitFunction(const std::string& func_name, int64_t num_inputs,
+ Optional<Array<String>> param_names,
+ vm::VMFuncInfo::FuncKind kind, int64_t init_register_size) {
+ auto it = exec_->func_map.find(func_name);
+ if (it == exec_->func_map.end()) {
+ this->DeclareFunction(func_name, kind);
+ }
+ auto& vmfunc = exec_->func_table.at(exec_->func_map.at(func_name));
+ ICHECK_EQ(vmfunc.name, func_name);
+ ICHECK_EQ(vmfunc.num_args, -2) << "Function " << func_name << " already defined";
+ vmfunc.num_args = num_inputs;
+ if (param_names.defined()) {
+ std::vector<std::string> names;
+ for (auto name : param_names.value()) {
+ names.push_back(name);
+ }
+ vmfunc.param_names = names;
+ }
+ vmfunc.register_file_size = init_register_size;
+ if (kind == vm::VMFuncInfo::FuncKind::kVMFunc) {
+ vmfunc.start_instr = exec_->instr_offset.size();
+ }
+}
+
+void ExecBuilderNode::EndFunction(const std::string& func_name) {
+ auto it = exec_->func_map.find(func_name);
+ ICHECK(it != exec_->func_map.end());
+ VMFuncInfo& vmfunc = exec_->func_table.at(it->second);
+ ICHECK_EQ(vmfunc.end_instr, 0) << "EndFuncton can only be called once";
+
+ if (vmfunc.kind == vm::VMFuncInfo::FuncKind::kVMFunc) {
+ vmfunc.end_instr = exec_->instr_offset.size();
+ }
+}
+
+void ExecBuilderNode::EmitCall(vm::Instruction::Arg func, std::vector<vm::Instruction::Arg> args,
+ vm::RegName dst) {
+ ICHECK(func.kind() == vm::Instruction::ArgKind::kFuncIdx);
+ // store instruction
+ exec_->instr_offset.push_back(exec_->instr_data.size());
+ exec_->instr_data.push_back(static_cast<ExecWord>(Opcode::Call));
+ exec_->instr_data.push_back(dst);
+ exec_->instr_data.push_back(func.value());
+ exec_->instr_data.push_back(args.size());
+ for (Instruction::Arg arg : args) {
+ exec_->instr_data.push_back(arg.data());
+ }
+}
+
+void ExecBuilderNode::EmitCall(const std::string& func, std::vector<Instruction::Arg> args,
+ RegName dst) {
+ auto it = exec_->func_map.find(func);
+ if (it == exec_->func_map.end()) {
+ this->DeclareFunction(func, VMFuncInfo::FuncKind::kPackedFunc);
+ }
+ Index func_idx = exec_->func_map.at(func);
+ EmitCall(vm::Instruction::Arg::FuncIdx(func_idx), args, dst);
+}
+
+void ExecBuilderNode::EmitRet(vm::Instruction::Arg result) {
+ ICHECK(result.kind() == vm::Instruction::ArgKind::kRegister);
+ exec_->instr_offset.push_back(exec_->instr_data.size());
+ exec_->instr_data.push_back(static_cast<ExecWord>(Opcode::Ret));
+ exec_->instr_data.push_back(result.value());
+}
+
+void ExecBuilderNode::EmitGoto(Index pc_offset) {
+ exec_->instr_offset.push_back(exec_->instr_data.size());
+ exec_->instr_data.push_back(static_cast<ExecWord>(Opcode::Goto));
+ exec_->instr_data.push_back(pc_offset);
+}
+
+void ExecBuilderNode::EmitIf(vm::Instruction::Arg cond, vm::Index false_offset) {
+ ICHECK(cond.kind() == vm::Instruction::ArgKind::kRegister);
+ exec_->instr_offset.push_back(exec_->instr_data.size());
+ exec_->instr_data.push_back(static_cast<ExecWord>(Opcode::If));
+ exec_->instr_data.push_back(cond.value());
+ exec_->instr_data.push_back(false_offset);
+}
+
+void ExecBuilderNode::CheckExecutable() {
+ for (auto it = exec_->func_table.cbegin(); it != exec_->func_table.cend(); ++it) {
+ if (it->kind == VMFuncInfo::FuncKind::kPackedFunc) continue;
+ if (it->kind == VMFuncInfo::FuncKind::kVMTIRFunc) {
+ ICHECK_GE(it->register_file_size, it->num_args + 1)
+ << "Function " << it->name << " do not meet register file constraint.";
+ continue;
+ }
+ Index num_inputs = it->num_args;
+ std::unordered_set<RegName> dst_registers;
+ std::unordered_set<RegName> arg_registers;
+ size_t start_instr = it->start_instr;
+ size_t end_instr = it->end_instr;
+
+ CHECK_LT(start_instr, end_instr)
+ << "Function " << it->name << " EndFunction has not be been called";
+
+ auto check_reg_defined = [&](Instruction::Arg arg) {
+ if (arg.kind() != Instruction::ArgKind::kRegister) return;
+ if (arg.value() >= Instruction::kBeginSpecialReg) return;
+ if (arg.value() < num_inputs) return;
+
+ if (dst_registers.find(arg.value()) == dst_registers.end()) {
+ LOG(FATAL) << "register r(" << arg.value() << ") in VM function \"" << it->name
+ << "\" is used as input while it is never defined"
+ << " as a destination. Dump:\n"
+ << exec_->AsText();
+ }
+ };
+
+ auto check_const_defined = [&](Instruction::Arg arg) {
+ if (arg.kind() != Instruction::ArgKind::kConstIdx) return;
+ CHECK_LT(arg.value(), exec_->constants.size())
+ << "Constant index " << arg.value() << " exceed size of constant pool. Dump:\n"
+ << exec_->AsText();
+ };
+
+ auto check_func_defined = [&](Instruction::Arg arg) {
+ if (arg.kind() != Instruction::ArgKind::kFuncIdx) return;
+ CHECK_LT(arg.value(), exec_->func_table.size())
+ << "Func index " << arg.value() << " exceed size of fun_table. Dump:\n"
+ << exec_->AsText();
+ };
+
+ for (size_t idx = start_instr; idx < end_instr; ++idx) {
+ Instruction instr = exec_->GetInstruction(idx);
+ switch (instr.op) {
+ case Opcode::Call: {
+ check_func_defined(Instruction::Arg::FuncIdx(instr.func_idx));
+ for (int i = 0; i < instr.num_args; ++i) {
+ check_reg_defined(instr.args[i]);
+ check_const_defined(instr.args[i]);
+ check_func_defined(instr.args[i]);
+ arg_registers.emplace(instr.args[i].value());
+ }
+ if (instr.dst != Instruction::kVoidRegister) {
+ dst_registers.emplace(instr.dst);
+ }
+ break;
+ }
+ case Opcode::Ret: {
+ arg_registers.emplace(instr.result);
+ check_reg_defined(Instruction::Arg::Register(instr.result));
+ break;
+ }
+ case Opcode::Goto: {
+ ICHECK_NE(instr.pc_offset, 0);
+ break;
+ }
+ case Opcode::If: {
+ ICHECK_GT(instr.false_offset, 1);
+ check_reg_defined(Instruction::Arg::Register(instr.cond));
+ arg_registers.emplace(instr.cond);
+ break;
+ }
+ default:
+ LOG(FATAL) << "should never hit this case: " << static_cast<int>(instr.op);
+ break;
+ }
+ }
+ }
+}
+
+void ExecBuilderNode::Formalize() {
+ // a pass to formalize user-specified register indexes in the order of use
+ // and decide the number of registers to allocate for each VMFunction in the Executable
+ for (auto it = this->exec_->func_table.begin(); it != this->exec_->func_table.end(); ++it) {
+ if (it->kind == VMFuncInfo::FuncKind::kPackedFunc) continue;
+ if (it->kind == VMFuncInfo::FuncKind::kVMTIRFunc) continue;
+
+ Index num_inputs = it->num_args;
+ RegName register_idx = num_inputs;
+ std::unordered_map<RegName, RegName> register_map;
+ size_t start_instr = it->start_instr;
+ size_t end_instr = it->end_instr;
+
+ for (size_t idx = start_instr; idx < end_instr; ++idx) {
+ Instruction instr = this->exec_->GetInstruction(idx);
+ switch (instr.op) {
+ case Opcode::Call: {
+ // rewrite args
+ for (int i = 0; i < instr.num_args; ++i) {
+ if (instr.args[i].kind() == Instruction::ArgKind::kRegister &&
+ instr.args[i].value() >= num_inputs &&
+ instr.args[i].value() < Instruction::kBeginSpecialReg &&
+ register_map.find(instr.args[i].value()) != register_map.end()) {
+ this->exec_->instr_data[this->exec_->instr_offset[idx] + 4 + i] =
+ register_map[instr.args[i].value()];
+ }
+ }
+ if (instr.dst >= num_inputs && instr.dst < Instruction::kBeginSpecialReg) {
+ auto it = register_map.find(instr.dst);
+ if (it != register_map.end()) {
+ this->exec_->instr_data[this->exec_->instr_offset[idx] + 1] = it->second;
+ } else {
+ this->exec_->instr_data[this->exec_->instr_offset[idx] + 1] = register_idx;
+ register_map[instr.dst] = register_idx++;
+ }
+ }
+ break;
+ }
+ case Opcode::Ret: {
+ if (register_map.find(instr.result) != register_map.end()) {
+ this->exec_->instr_data[this->exec_->instr_offset[idx] + 1] =
+ register_map[instr.result];
+ }
+ break;
+ }
+ case Opcode::Goto: {
+ break;
+ }
+ case Opcode::If: {
+ if (register_map.find(instr.cond) != register_map.end()) {
+ this->exec_->instr_data[this->exec_->instr_offset[idx] + 1] = register_map[instr.cond];
+ }
+ break;
+ }
+ default:
+ LOG(FATAL) << "should never hit this case: " << static_cast<int>(instr.op);
+ break;
+ }
+ }
+ it->register_file_size = register_idx;
+ }
+}
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderCreate").set_body_typed(ExecBuilderNode::Create);
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderConvertConstant")
+ .set_body([](TVMArgs args, TVMRetValue* ret) {
+ ExecBuilder builder = args[0];
+ TVMRetValue rt;
+ rt = args[1];
+ *ret = builder->ConvertConstant(rt).data();
+ });
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitFunction")
+ .set_body_typed([](ExecBuilder builder, String func, int64_t num_inputs,
+ Optional<Array<String>> param_names) {
+ builder->EmitFunction(func, num_inputs, param_names);
+ });
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderEndFunction")
+ .set_body_method<ExecBuilder>(&ExecBuilderNode::EndFunction);
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderDeclareFunction")
+ .set_body_typed([](ExecBuilder builder, String name, int32_t kind) {
+ builder->DeclareFunction(name, static_cast<VMFuncInfo::FuncKind>(kind));
+ });
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitCall")
+ .set_body_typed([](ExecBuilder builder, String name, Array<IntImm> args, int64_t dst) {
+ std::vector<Instruction::Arg> args_;
+ for (size_t i = 0; i < args.size(); ++i) {
+ args_.push_back(Instruction::Arg::FromData(args[i]->value));
+ }
+ auto dst_ = Instruction::Arg::Register(dst);
+ builder->EmitCall(name, args_, dst_.value());
+ });
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitRet")
+ .set_body_typed([](ExecBuilder builder, int64_t data) {
+ builder->EmitRet(Instruction::Arg::FromData(data));
+ });
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitGoto")
+ .set_body_method<ExecBuilder>(&ExecBuilderNode::EmitGoto);
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitIf")
+ .set_body_typed([](ExecBuilder builder, int64_t data, vm::Index false_offset) {
+ builder->EmitIf(Instruction::Arg::FromData(data), false_offset);
+ });
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderR").set_body_typed([](ExecBuilder builder, int64_t value) {
+ return Instruction::Arg::Register(value).data();
+});
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderImm").set_body_typed([](ExecBuilder builder, int64_t value) {
+ return Instruction::Arg::Immediate(value).data();
+});
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderC").set_body_typed([](ExecBuilder builder, int64_t value) {
+ return Instruction::Arg::ConstIdx(value).data();
+});
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderF").set_body_typed([](ExecBuilder builder, String value) {
+ return builder->GetFunction(value).data();
+});
+
+TVM_REGISTER_GLOBAL("relax.ExecBuilderGet").set_body_typed([](ExecBuilder builder) {
+ ObjectPtr<Executable> p_exec = builder->Get();
+ return runtime::Module(p_exec);
+});
+
+} // namespace relax
+} // namespace tvm
diff --git a/src/runtime/relax_vm/builtin.cc b/src/runtime/relax_vm/builtin.cc
new file mode 100644
index 0000000000..0ef63c8a41
--- /dev/null
+++ b/src/runtime/relax_vm/builtin.cc
@@ -0,0 +1,445 @@
+/*
+ * 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 src/runtime/relax_vm/builtin.cc
+ */
+#include <tvm/runtime/container/adt.h>
+#include <tvm/runtime/container/shape_tuple.h>
+#include <tvm/runtime/data_type.h>
+#include <tvm/runtime/device_api.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/runtime/memory.h>
+#include <tvm/runtime/ndarray.h>
+#include <tvm/runtime/packed_func.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/runtime/relax_vm/builtin.h>
+#include <tvm/runtime/relax_vm/bytecode.h>
+#include <tvm/runtime/relax_vm/memory_manager.h>
+#include <tvm/runtime/relax_vm/vm.h>
+
+#include "../runtime_base.h"
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+using tvm::runtime::NDArray;
+
+//-------------------------------------------------
+// Shape/StructInfo handling.
+//-------------------------------------------------
+/*!
+ * \brief Builtin function to allocate shape heap.
+ * \param ctx_ptr The context module pointer.
+ * \param size the size of the heap.
+ * \return An allocate NDArray as shape heap.
+ */
+NDArray AllocShapeHeap(void* ctx_ptr, int64_t size) {
+ VirtualMachine* vm = static_cast<VirtualMachine*>(ctx_ptr);
+ // use host allocator, which is always last element.
+ size_t host_device_index = vm->devices.size() - 1;
+ // specialy handle hexagon on-device RT.
+ // TODO(relax-team): visit and consider other possible choices.
+ if (vm->devices[0].device_type == kDLHexagon) {
+ host_device_index = 0;
+ }
+ auto* alloc = vm->allocators[host_device_index];
+ return alloc->Empty({size}, DLDataType{kDLInt, 64, 1}, vm->devices[host_device_index]);
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.alloc_shape_heap").set_body_typed(AllocShapeHeap);
+
+/*!
+ * \brief Builtin match shape function.
+ * \param args The packed function arguments.
+ * \param rv The return value.
+ *
+ * \sa MatchShapeCode
+ */
+void MatchShape(TVMArgs args, TVMRetValue* rv) {
+ // input shape the first argument can take in tensor or shape.
+ ShapeTuple input_shape;
+ if (args[0].IsObjectRef<NDArray>()) {
+ input_shape = args[0].operator NDArray().Shape();
+ } else {
+ input_shape = args[0];
+ }
+ DLTensor* heap = args[1];
+ int64_t* heap_data = heap == nullptr ? nullptr : static_cast<int64_t*>(heap->data);
+ int64_t size = args[2];
+ const int64_t kBeginCode = 3;
+ ICHECK_LE(kBeginCode + size * 2, args.size());
+ // a function that lazily get context for error reporting
+ const int64_t kErrorContextOffset = kBeginCode + size * 2;
+ Optional<String> err_ctx = args[kErrorContextOffset];
+
+ CHECK_EQ(input_shape.size(), size)
+ << "RuntimeError: " << err_ctx.value_or("") << " match_cast shape size mismatch.";
+
+ for (int64_t i = 0; i < size; ++i) {
+ MatchShapeCode code = static_cast<MatchShapeCode>(args[kBeginCode + i * 2].operator int());
+ int64_t reg = args[kBeginCode + i * 2 + 1];
+
+ if (code == MatchShapeCode::kAssertEqualToImm) {
+ CHECK_EQ(input_shape[i], reg)
+ << "RuntimeError: " << err_ctx.value_or("") << " match_cast error, "
+ << " shape[" << i << "]"
+ << " mismatch to specified constant.";
+ } else if (code == MatchShapeCode::kStoreToHeap) {
+ heap_data[reg] = input_shape[i];
+ } else if (code == MatchShapeCode::kNoOp) {
+ } else {
+ ICHECK(code == MatchShapeCode::kAssertEqualToLoad);
+ CHECK_EQ(input_shape[i], heap_data[reg])
+ << "RuntimeError: " << err_ctx.value_or("") << " match_cast error, "
+ << " shape[" << i << "]"
+ << " mismatch to a previous populated value.";
+ }
+ }
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.match_shape").set_body(MatchShape);
+
+/*!
+ * \brief Builtin make shape function.
+ * \param args The packed function arguments.
+ * \param rv The return value.
+ *
+ * \sa MakeShapeCode
+ */
+void MakeShape(TVMArgs args, TVMRetValue* rv) {
+ // NOTE: heap can be nullptr
+ DLTensor* heap = args[0];
+ int64_t* heap_data = heap == nullptr ? nullptr : static_cast<int64_t*>(heap->data);
+ int64_t size = args[1];
+ const int64_t kBeginCode = 2;
+
+ std::vector<int64_t> shape(size);
+
+ for (int64_t i = 0; i < size; ++i) {
+ MakeShapeCode code = static_cast<MakeShapeCode>(args[kBeginCode + i * 2].operator int());
+ int64_t reg = args[kBeginCode + i * 2 + 1];
+ if (code == MakeShapeCode::kUseImm) {
+ shape[i] = reg;
+ } else {
+ ICHECK(code == MakeShapeCode::kLoadShape);
+ shape[i] = heap_data[reg];
+ }
+ }
+ *rv = ShapeTuple(std::move(shape));
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.make_shape").set_body(MakeShape);
+
+/*!
+ * \brief Builtin function to check if arg is Tensor(dtype, ndim)
+ * \param arg The input argument.
+ * \param ndim Expected ndim of the Tensor, can be -1 (indicate unknown).
+ * \param dtype The expected content data type.
+ * \param err_ctx Additional context if error occurs.
+ */
+void CheckTensorInfo(TVMArgs args, TVMRetValue* rv) {
+ ObjectRef arg = args[0];
+ int ndim = args[1];
+ DataType dtype;
+ Optional<String> err_ctx;
+
+ if (args.size() == 3) {
+ dtype = DataType::Void();
+ err_ctx = args[2].operator Optional<String>();
+ } else {
+ dtype = args[2];
+ err_ctx = args[3].operator Optional<String>();
+ }
+
+ auto* ptr = arg.as<NDArray::ContainerType>();
+ CHECK(ptr != nullptr) << "TypeError: " << err_ctx.value_or("") << " expect a Tensor but get "
+ << arg->GetTypeKey();
+
+ if (ndim != -1) {
+ CHECK(ptr->dl_tensor.ndim == ndim)
+ << "ValueError: " << err_ctx.value_or("") << " expect Tensor with ndim " << ndim
+ << " but get " << ptr->dl_tensor.ndim;
+ }
+
+ if (dtype != DataType::Void()) {
+ CHECK(DataType(ptr->dl_tensor.dtype) == dtype)
+ << "ValueError: " << err_ctx.value_or("") << " expect Tensor with dtype " << dtype
+ << " but get " << ptr->dl_tensor.dtype;
+ }
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.check_tensor_info").set_body(CheckTensorInfo);
+
+/*!
+ * \brief Builtin function to check if arg is Shape(ndim)
+ * \param arg The input argument.
+ * \param ndim Expected size of the shape, can be -1 (indicate unknown).
+ * \param err_ctx Additional context if error occurs.
+ */
+void CheckShapeInfo(ObjectRef arg, int ndim, Optional<String> err_ctx) {
+ // a function that lazily get context for error reporting
+ auto* ptr = arg.as<ShapeTuple::ContainerType>();
+ CHECK(ptr != nullptr) << "TypeError: " << err_ctx.value_or("") << " expect a Shape but get "
+ << arg->GetTypeKey();
+ if (ndim != -1) {
+ CHECK(ptr->size == static_cast<uint64_t>(ndim))
+ << "ValueError: " << err_ctx.value_or("") << " expect Shape with ndim " << ndim
+ << " but get " << ptr->size;
+ }
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.check_shape_info").set_body_typed(CheckShapeInfo);
+
+/*!
+ * \brief Builtin function to check if arg is Tuple with size elements.
+ * \param arg The input argument.
+ * \param size The expected size of the tuple.
+ * \param err_ctx Additional context if error occurs.
+ */
+void CheckTupleInfo(ObjectRef arg, int64_t size, Optional<String> err_ctx) {
+ using Tuple = runtime::ADT;
+ // a function that lazily get context for error reporting
+ auto* ptr = arg.as<Tuple::ContainerType>();
+ CHECK(ptr != nullptr) << "TypeError: " << err_ctx.value_or("") << " expect a Tuple but get "
+ << arg->GetTypeKey();
+ CHECK(static_cast<int64_t>(ptr->size) == size)
+ << "ValueError: " << err_ctx.value_or("") << " expect a Tuple with " << size << " elements, "
+ << " but get a Tuple with " << ptr->size << " elements.";
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.check_tuple_info").set_body_typed(CheckTupleInfo);
+
+/*!
+ * \brief Builtin function to check if arg is a callable function.
+ * \param arg The input argument.
+ * \param err_ctx Additional context if error occurs.
+ */
+void CheckFuncInfo(ObjectRef arg, Optional<String> err_ctx) {
+ // a function that lazily get context for error reporting
+ bool is_func = arg.as<PackedFunc::ContainerType>() || arg.as<VMClosure::ContainerType>();
+ CHECK(is_func) << "TypeError: " << err_ctx.value_or("") << " expect a Function but get "
+ << arg->GetTypeKey();
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.check_func_info").set_body_typed(CheckFuncInfo);
+
+//-------------------------------------------------
+// Storage management.
+//-------------------------------------------------
+Storage VMAllocStorage(void* ctx_ptr, ShapeTuple buffer_size, Index device_index,
+ DLDataType dtype_hint) {
+ VirtualMachine* vm = static_cast<VirtualMachine*>(ctx_ptr);
+
+ ICHECK_EQ(buffer_size.size(), 1);
+ int alignment = runtime::kAllocAlignment;
+ ICHECK_LT(device_index, vm->devices.size())
+ << "The device index is out of VM physical devices list";
+
+ if (device_index == -1) {
+ // Allocate on host. Host is always the last element of vm->devices.
+ device_index = vm->devices.size() - 1;
+ }
+
+ int64_t size_imm = buffer_size[0];
+
+ auto storage_obj = runtime::SimpleObjAllocator().make_object<StorageObj>();
+ auto* alloc = vm->allocators[device_index];
+ ICHECK(alloc) << "Did you forget to init the VirtualMachine with devices?";
+ storage_obj->buffer = alloc->Alloc(size_imm, alignment, dtype_hint);
+ Storage storage(storage_obj);
+ return storage;
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.alloc_storage").set_body_typed(VMAllocStorage);
+
+TVM_REGISTER_GLOBAL("vm.builtin.alloc_tensor").set_body_method<Storage>(&StorageObj::AllocNDArray);
+
+//-------------------------------------------------
+// Closure function handling, calling convention
+//-------------------------------------------------
+TVM_REGISTER_GLOBAL("vm.builtin.make_closure").set_body([](TVMArgs args, TVMRetValue* rv) {
+ VMClosure clo = args[0];
+ std::vector<TVMRetValue> saved_args;
+ saved_args.resize(args.size() - 1);
+ for (size_t i = 0; i < saved_args.size(); ++i) {
+ saved_args[i] = args[i + 1];
+ }
+ auto impl = VMClosure::BindLastArgs(clo->impl, saved_args);
+ *rv = VMClosure(clo->func_name, impl);
+});
+
+TVM_REGISTER_GLOBAL("vm.builtin.invoke_closure").set_body([](TVMArgs args, TVMRetValue* rv) {
+ // args[0]: vm; args[1]: closure; args[2, 3, ...]: function arguments
+ VirtualMachine* vm = VirtualMachine::GetContextPtr(args[0]);
+ ObjectRef vm_closure = args[1];
+ vm->InvokeClosurePacked(vm_closure,
+ TVMArgs(args.values + 2, args.type_codes + 2, args.size() - 2), rv);
+});
+
+TVM_REGISTER_GLOBAL("vm.builtin.call_tir_dyn").set_body([](TVMArgs args, TVMRetValue* rv) {
+ PackedFunc func = args[0];
+ ShapeTuple to_unpack = args[args.size() - 1];
+ size_t num_tensor_args = args.size() - 2;
+
+ std::vector<TVMValue> values(num_tensor_args + to_unpack.size());
+ std::vector<int> tcodes(num_tensor_args + to_unpack.size());
+ runtime::TVMArgsSetter setter(values.data(), tcodes.data());
+
+ std::copy(args.values + 1, args.values + args.size() - 1, values.data());
+ std::copy(args.type_codes + 1, args.type_codes + args.size() - 1, tcodes.data());
+
+ for (size_t i = 0; i < to_unpack.size(); ++i) {
+ setter(i + num_tensor_args, to_unpack[i]);
+ }
+ TVMArgs func_args(values.data(), tcodes.data(), values.size());
+ func.CallPacked(func_args, rv);
+});
+
+//-------------------------------------
+// Builtin runtime operators.
+//-------------------------------------
+TVM_REGISTER_GLOBAL("vm.builtin.shape_of").set_body_method(&NDArray::Shape);
+
+TVM_REGISTER_GLOBAL("vm.builtin.copy").set_body([](TVMArgs args, TVMRetValue* rv) {
+ *rv = args[0];
+});
+
+/*!
+ * \brief Load the scalar value in cond and return the result value.
+ * \param cond The condition
+ * \return Bool
+ */
+bool ReadIfCond(TVMArgValue cond) {
+ if (cond.type_code() == kDLInt) return cond.operator bool();
+ NDArray arr = cond.operator tvm::runtime::NDArray();
+ if (arr->device.device_type != kDLCPU) {
+ arr = arr.CopyTo(DLDevice{kDLCPU, 0});
+ }
+ ICHECK(arr->dtype.code == kDLInt || arr->dtype.code == kDLUInt);
+ int64_t result;
+ switch (arr->dtype.bits) {
+ case 1: {
+ result = reinterpret_cast<int8_t*>(arr->data)[0];
+ break;
+ }
+ case 8: {
+ result = reinterpret_cast<int8_t*>(arr->data)[0];
+ break;
+ }
+ case 16: {
+ result = reinterpret_cast<int16_t*>(arr->data)[0];
+ break;
+ }
+ case 32: {
+ result = reinterpret_cast<int32_t*>(arr->data)[0];
+ break;
+ }
+ case 64: {
+ result = reinterpret_cast<int64_t*>(arr->data)[0];
+ break;
+ }
+ default:
+ LOG(FATAL) << "Unknown scalar int type: " << DLDataType2String(arr->dtype);
+ throw;
+ }
+ return result != 0;
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.read_if_cond").set_body_typed(ReadIfCond);
+
+//-------------------------------------
+// Data structure API
+//-------------------------------------
+TVM_REGISTER_GLOBAL("vm.builtin.tuple_getitem").set_body_typed([](runtime::ADT arr, int64_t index) {
+ return arr[index];
+});
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
+
+//-------------------------------------------------
+// AnyList C runtime API: keep in relax for now.
+//--------------------------------------------------
+extern "C" {
+/*!
+ * \brief Backend function to get anylist item and set into Packed Func call arg stack.
+ *
+ * \param anylist The handle to the anylist, backed by TVMRetValue*
+ * \param int The index.
+ * \param args The args stack.
+ * \param type_codes The type codes stack.
+ * \param arg_offset The offset of argument.
+ * \return 0 when no error is thrown, -1 when failure happens
+ */
+TVM_DLL int TVMBackendAnyListSetPackedArg(void* anylist, int index, TVMValue* args, int* type_codes,
+ int arg_offset);
+/*!
+ * \brief Backend function to get anylist item and set into Packed Func call arg stack.
+ *
+ * \param anylist The handle to the anylist, backed by TVMRetValue*
+ * \param int The index.
+ */
+TVM_DLL int TVMBackendAnyListResetItem(void* anylist, int index);
+
+/*!
+ * \brief Backend function to set anylist item by moving from packed func return.
+ *
+ * \param anylist The handle to the anylist, backed by TVMRetValue*
+ * \param int The index.
+ * \param args The args stack.
+ * \param type_codes The type codes stack.
+ * \param arg_offset The offset of argument.
+ * \return 0 when no error is thrown, -1 when failure happens.
+ */
+TVM_DLL int TVMBackendAnyListMoveFromPackedReturn(void* anylist, int index, TVMValue* args,
+ int* type_codes, int ret_offset);
+
+int TVMBackendAnyListSetPackedArg(void* anylist, int index, TVMValue* args, int* type_codes,
+ int arg_offset) {
+ using namespace tvm::runtime;
+ API_BEGIN();
+ auto* list = static_cast<TVMRetValue*>(anylist);
+ TVMArgsSetter setter(args, type_codes);
+ setter(arg_offset, list[index]);
+ API_END();
+}
+
+int TVMBackendAnyListResetItem(void* anylist, int index) {
+ using namespace tvm::runtime;
+ API_BEGIN();
+ auto* list = static_cast<TVMRetValue*>(anylist);
+ list[index] = nullptr;
+ API_END();
+}
+
+int TVMBackendAnyListMoveFromPackedReturn(void* anylist, int index, TVMValue* args, int* type_codes,
+ int ret_offset) {
+ using namespace tvm::runtime;
+ API_BEGIN();
+ auto* list = static_cast<TVMRetValue*>(anylist);
+ if (type_codes[ret_offset] == kTVMStr || type_codes[ret_offset] == kTVMBytes) {
+ list[index] = TVMArgValue(args[ret_offset], type_codes[ret_offset]);
+ } else {
+ list[index] = TVMRetValue::MoveFromCHost(args[ret_offset], type_codes[ret_offset]);
+ }
+ API_END();
+}
+} // extern "C"
diff --git a/src/runtime/relax_vm/bytecode.cc b/src/runtime/relax_vm/bytecode.cc
new file mode 100644
index 0000000000..9084207848
--- /dev/null
+++ b/src/runtime/relax_vm/bytecode.cc
@@ -0,0 +1,68 @@
+/*
+ * 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 src/runtime/relax_vm/bytecode.cc
+ * \brief The bytecode for Relax virtual machine.
+ */
+
+#include <tvm/runtime/logging.h>
+#include <tvm/runtime/relax_vm/bytecode.h>
+
+#include <functional>
+#include <sstream>
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+Instruction Instruction::Call(Index func_idx, Index num_args, Instruction::Arg* args, RegName dst) {
+ Instruction instr;
+ instr.op = Opcode::Call;
+ instr.dst = dst;
+ instr.func_idx = func_idx;
+ instr.num_args = num_args;
+ instr.args = args;
+ return instr;
+}
+
+Instruction Instruction::Ret(RegName result) {
+ Instruction instr;
+ instr.op = Opcode::Ret;
+ instr.result = result;
+ return instr;
+}
+
+Instruction Instruction::Goto(Index pc_offset) {
+ Instruction instr;
+ instr.op = Opcode::Goto;
+ instr.pc_offset = pc_offset;
+ return instr;
+}
+
+Instruction Instruction::If(RegName cond, Index false_offset) {
+ Instruction instr;
+ instr.op = Opcode::If;
+ instr.cond = cond;
+ instr.false_offset = false_offset;
+ return instr;
+}
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
diff --git a/src/runtime/relax_vm/executable.cc b/src/runtime/relax_vm/executable.cc
new file mode 100644
index 0000000000..b7915d7978
--- /dev/null
+++ b/src/runtime/relax_vm/executable.cc
@@ -0,0 +1,576 @@
+/*
+ * 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 src/runtime/relax_vm/executable.cc
+ */
+
+#include <dmlc/memory_io.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/runtime/relax_vm/executable.h>
+#include <tvm/runtime/relax_vm/vm.h>
+
+#include <functional>
+#include <sstream>
+
+#include "../file_utils.h"
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+/*! \brief The magic number for the serialized VM bytecode file */
+constexpr uint64_t kTVMVMBytecodeMagic = 0xD225DE2F4214151D;
+
+/*! \brief Possible types in the constant pool */
+enum ConstantType : int {
+ kNDArray = 0,
+ kDLDataType = 1,
+ kShapeTuple = 2,
+ kString = 3,
+ kInt = 4,
+};
+
+#define STREAM_CHECK(val, section) \
+ ICHECK(val) << "Invalid VM file format in the " << section << " section." \
+ << "\n";
+
+PackedFunc Executable::GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) {
+ if (name == "stats") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->Stats(); });
+ } else if (name == "as_text") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->AsText(); });
+ } else if (name == "as_python") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->AsPython(); });
+ } else if (name == "vm_load_executable") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ ObjectPtr<VirtualMachine> vm = VirtualMachine::Create();
+ ICHECK(sptr_to_self.get() == this);
+ vm->LoadExecutable(GetObjectPtr<Executable>(this));
+ *rv = Module(vm);
+ });
+ }
+ return nullptr;
+}
+
+std::string Executable::Stats() const {
+ std::ostringstream oss;
+ oss << "Relax VM executable statistics:" << std::endl;
+
+ // Get the number of constants.
+ // If the constant is an NDArray, get the shape of each of them.
+ // If the constant is an DLDataType, get the data type of each of them.
+ oss << " Constant pool (# " << constants.size() << "): [";
+ for (const auto& it : constants) {
+ if (it.IsObjectRef<runtime::NDArray>()) {
+ const auto ndarray = it.operator tvm::runtime::NDArray();
+ const auto& shape = ndarray.Shape();
+ // Scalar
+ if (shape.empty()) {
+ oss << "scalar, ";
+ continue;
+ }
+ oss << "[";
+ for (auto s : shape) {
+ oss << s << ", ";
+ }
+ oss.seekp(-2, oss.cur);
+ oss << "], ";
+ } else if (it.IsObjectRef<ShapeTuple>()) {
+ ShapeTuple shape = it.operator ShapeTuple();
+ oss << "shapetuple[";
+ for (size_t i = 0; i < shape.size(); ++i) {
+ oss << shape.at(i) << ", ";
+ }
+ oss.seekp(-2, oss.cur);
+ oss << "], ";
+ } else if (it.IsObjectRef<String>()) {
+ std::string f = it.AsObjectRef<tvm::runtime::String>().operator std::string();
+ oss << "\"";
+ oss << f;
+ oss << "\", ";
+ } else if (it.type_code() == kDLInt) {
+ oss << static_cast<int64_t>(it);
+ oss << ", ";
+ } else {
+ try {
+ DLDataType dtype = it.operator DLDataType();
+ oss << dtype;
+ oss << ", ";
+ } catch (std::exception& exc) {
+ LOG(FATAL) << "Constant pool can only contain NDArray and DLDataType, but got "
+ << ArgTypeCode2Str(it.type_code());
+ }
+ }
+ }
+ if (!constants.empty()) oss.seekp(-2, oss.cur);
+ oss << "]" << std::endl;
+
+ // Get the number of globals and the name of each of them.
+ oss << " Globals (#" << func_table.size() << "): [";
+ for (const auto& it : func_table) {
+ oss << it.name << ", ";
+ }
+ if (!func_map.empty()) oss.seekp(-2, oss.cur);
+ oss << "]" << std::endl;
+
+ return oss.str();
+}
+
+void Executable::SetInstructionData(Index i, Index j, ExecWord val) {
+ ICHECK_LT(i, instr_offset.size());
+ Index instr_idx = instr_offset[i];
+ ICHECK_LT(instr_idx + j, instr_data.size());
+ instr_data[instr_idx + j] = val;
+}
+
+Instruction Executable::GetInstruction(Index i) const {
+ Index offset = instr_offset[i];
+ Opcode op = static_cast<Opcode>(instr_data[offset]);
+ switch (op) {
+ case Opcode::Call: {
+ RegName dst = instr_data[offset + 1];
+ Index func_idx = instr_data[offset + 2];
+ Index num_args = instr_data[offset + 3];
+ ExecWord* args = const_cast<ExecWord*>(&instr_data[offset + 4]);
+ return Instruction::Call(func_idx, num_args, reinterpret_cast<Instruction::Arg*>(args), dst);
+ }
+ case Opcode::Ret: {
+ RegName result = instr_data[offset + 1];
+ return Instruction::Ret(result);
+ }
+ case Opcode::Goto: {
+ Index pc_offset = instr_data[offset + 1];
+ return Instruction::Goto(pc_offset);
+ }
+ case Opcode::If: {
+ RegName cond = instr_data[offset + 1];
+ Index false_offset = instr_data[offset + 2];
+ return Instruction::If(cond, false_offset);
+ }
+ default:
+ LOG(FATAL) << "should never hit this case: " << static_cast<int>(op);
+ break;
+ }
+ return Instruction();
+}
+
+void SaveHeader(dmlc::Stream* strm) {
+ uint64_t header = kTVMVMBytecodeMagic;
+ strm->Write(header);
+ std::string version = TVM_VERSION;
+ strm->Write(version);
+}
+
+void LoadHeader(dmlc::Stream* strm) {
+ // Check header.
+ uint64_t header;
+ STREAM_CHECK(strm->Read(&header), "header");
+ STREAM_CHECK(header == kTVMVMBytecodeMagic, "header");
+
+ // Check version.
+ std::string version;
+ STREAM_CHECK(strm->Read(&version), "version");
+ STREAM_CHECK(version == TVM_VERSION, "version");
+}
+
+void Executable::SaveToBinary(dmlc::Stream* stream) {
+ std::string code;
+ // Initialize the stream object.
+ dmlc::MemoryStringStream strm(&code);
+
+ // Save header
+ SaveHeader(&strm);
+
+ // Global section.
+ SaveGlobalSection(&strm);
+
+ // Constant section.
+ SaveConstantSection(&strm);
+
+ // Code section.
+ SaveCodeSection(&strm);
+
+ stream->Write(code);
+}
+
+void Executable::SaveToFile(const std::string& file_name, const std::string& format) {
+ std::string data;
+ dmlc::MemoryStringStream writer(&data);
+ dmlc::SeekStream* strm = &writer;
+ Executable::SaveToBinary(strm);
+ runtime::SaveBinaryToFile(file_name, data);
+}
+
+Module Executable::LoadFromBinary(void* stream) {
+ std::string code;
+ static_cast<dmlc::Stream*>(stream)->Read(&code);
+ dmlc::MemoryStringStream strm(&code);
+
+ ObjectPtr<Executable> exec = make_object<Executable>();
+
+ // Load header.
+ LoadHeader(&strm);
+
+ // Global section.
+ exec->LoadGlobalSection(&strm);
+
+ // Constant section.
+ exec->LoadConstantSection(&strm);
+
+ // Code section.
+ exec->LoadCodeSection(&strm);
+
+ return Module(exec);
+}
+
+TVM_REGISTER_GLOBAL("runtime.module.loadbinary_relax.Executable")
+ .set_body_typed(Executable::LoadFromBinary);
+
+Module Executable::LoadFromFile(const std::string& file_name) {
+ std::string data;
+ runtime::LoadBinaryFromFile(file_name, &data);
+ dmlc::MemoryStringStream reader(&data);
+ dmlc::Stream* strm = &reader;
+ return Executable::LoadFromBinary(reinterpret_cast<void*>(strm));
+}
+
+TVM_REGISTER_GLOBAL("runtime.module.loadfile_relax.Executable")
+ .set_body_typed(Executable::LoadFromFile);
+
+void VMFuncInfo::Save(dmlc::Stream* strm) const {
+ int32_t temp_kind = static_cast<int32_t>(kind);
+ strm->Write(temp_kind);
+ strm->Write(name);
+ strm->Write(start_instr);
+ strm->Write(end_instr);
+ strm->Write(num_args);
+ strm->Write(register_file_size);
+ strm->Write(param_names);
+}
+
+bool VMFuncInfo::Load(dmlc::Stream* strm) {
+ int32_t temp_kind;
+ if (!strm->Read(&temp_kind)) return false;
+ this->kind = static_cast<VMFuncInfo::FuncKind>(temp_kind);
+ if (!strm->Read(&name)) return false;
+ if (!strm->Read(&start_instr)) return false;
+ if (!strm->Read(&end_instr)) return false;
+ if (!strm->Read(&num_args)) return false;
+ if (!strm->Read(®ister_file_size)) return false;
+ if (!strm->Read(¶m_names)) return false;
+ return true;
+}
+
+void Executable::SaveGlobalSection(dmlc::Stream* strm) { strm->Write(func_table); }
+
+void Executable::SaveConstantSection(dmlc::Stream* strm) {
+ strm->Write(static_cast<uint64_t>(this->constants.size()));
+ for (const auto& it : this->constants) {
+ if (it.IsObjectRef<runtime::NDArray>()) {
+ strm->Write(ConstantType::kNDArray);
+ runtime::SaveDLTensor(strm, it.operator DLTensor*());
+ } else if (it.IsObjectRef<ShapeTuple>()) {
+ ShapeTuple shape = it.operator ShapeTuple();
+ strm->Write(ConstantType::kShapeTuple);
+ strm->Write(shape.size());
+ for (size_t i = 0; i < shape.size(); ++i) {
+ strm->Write(shape.at(i));
+ }
+ } else if (it.IsObjectRef<String>()) {
+ String str = it.operator String();
+ strm->Write(ConstantType::kString);
+ strm->Write(str.size());
+ for (size_t i = 0; i < str.size(); ++i) {
+ strm->Write(str.at(i));
+ }
+ } else if (it.type_code() == kDLInt) {
+ strm->Write(ConstantType::kInt);
+ strm->Write(it.value());
+ } else {
+ try {
+ strm->Write(ConstantType::kDLDataType);
+ strm->Write(it.operator DLDataType());
+ } catch (std::exception& exc) {
+ LOG(FATAL) << "Constant pool can only contain NDArray, DLDataType, and Integers but got "
+ << ArgTypeCode2Str(it.type_code());
+ }
+ }
+ }
+}
+
+void Executable::SaveCodeSection(dmlc::Stream* strm) {
+ strm->Write(instr_offset);
+ strm->Write(instr_data);
+}
+
+void Executable::LoadGlobalSection(dmlc::Stream* strm) {
+ STREAM_CHECK(strm->Read(&func_table), "Global Section");
+ // setup func map
+ for (size_t i = 0; i < func_table.size(); ++i) {
+ this->func_map[func_table[i].name] = i;
+ }
+}
+
+void Executable::LoadConstantSection(dmlc::Stream* strm) {
+ uint64_t sz;
+ // Load the number of constants.
+ STREAM_CHECK(strm->Read(&sz, sizeof(sz)), "constant");
+
+ size_t size = static_cast<size_t>(sz);
+ runtime::NDArray ndarray;
+ DLDataType dtype;
+ // Load each of the constants.
+ for (size_t i = 0; i < size; i++) {
+ int constant_type;
+ STREAM_CHECK(strm->Read(&constant_type, sizeof(constant_type)), "constant");
+ if (constant_type == ConstantType::kNDArray) {
+ ndarray.Load(strm);
+ TVMRetValue cell;
+ cell = ndarray;
+ this->constants.push_back(cell);
+ } else if (constant_type == ConstantType::kShapeTuple) {
+ uint64_t size;
+ strm->Read(&size);
+ std::vector<ShapeTuple::index_type> data(size);
+ for (size_t i = 0; i < size; ++i) {
+ strm->Read(&(data[i]));
+ }
+ TVMRetValue cell;
+ cell = ShapeTuple(data);
+ this->constants.push_back(cell);
+ } else if (constant_type == ConstantType::kDLDataType) {
+ strm->Read(&dtype);
+ TVMRetValue cell;
+ cell = dtype;
+ this->constants.push_back(cell);
+ } else if (constant_type == ConstantType::kString) {
+ uint64_t size;
+ strm->Read(&size);
+ std::vector<char> data(size);
+ for (size_t i = 0; i < size; ++i) {
+ strm->Read(&(data[i]));
+ }
+ TVMRetValue cell;
+ cell = String(std::string(data.begin(), data.end()));
+ this->constants.push_back(cell);
+ } else if (constant_type == ConstantType::kInt) {
+ int64_t value;
+ strm->Read(&value);
+ TVMRetValue cell;
+ cell = value;
+ this->constants.push_back(cell);
+ } else {
+ LOG(FATAL) << "Constant pool can only contain NDArray and DLDataType, but got "
+ << ArgTypeCode2Str(constant_type) << " when loading the VM constant pool.";
+ }
+ }
+}
+
+void Executable::LoadCodeSection(dmlc::Stream* strm) {
+ STREAM_CHECK(strm->Read(&(this->instr_offset)), "instr offset");
+ STREAM_CHECK(strm->Read(&(this->instr_data)), "instr data");
+}
+
+template <typename T>
+std::string StrJoin(T* items, int offset, int cnt, std::string delim = ", ",
+ std::function<std::string(T)> repr = std::to_string) {
+ if (cnt == 0) {
+ return "";
+ }
+ std::ostringstream oss;
+ oss << repr(items[offset]);
+ for (int i = 1; i < cnt; ++i) {
+ oss << delim << repr(items[offset + i]);
+ }
+ return oss.str();
+}
+
+std::string RegNameToStr(RegName reg) {
+ if (reg == Instruction::kVoidRegister) {
+ return "%void";
+ }
+ if (reg == Instruction::kVMRegister) {
+ return "%vm";
+ }
+ return "%" + std::to_string(reg);
+}
+
+String Executable::AsText() const {
+ auto get_func_name = [&](Index index) -> std::string {
+ if (static_cast<size_t>(index) < func_table.size()) {
+ return func_table[index].name;
+ } else {
+ return "unknown_func_index(" + std::to_string(index) + ")";
+ }
+ };
+
+ auto instr_to_str = [&](Instruction::Arg arg) -> std::string {
+ // only for argument
+ switch (arg.kind()) {
+ case Instruction::ArgKind::kRegister:
+ return RegNameToStr(arg.value());
+ case Instruction::ArgKind::kImmediate:
+ return "i" + std::to_string(arg.value());
+ case Instruction::ArgKind::kConstIdx:
+ return "c[" + std::to_string(arg.value()) + "]";
+ case Instruction::ArgKind::kFuncIdx:
+ return "f[" + get_func_name(arg.value()) + "]";
+ default:
+ LOG(FATAL) << "Wrong instruction kind: " << static_cast<int>(arg.kind());
+ return "";
+ }
+ };
+
+ // print the text format
+ std::ostringstream os;
+ for (size_t fidx = 0; fidx < this->func_table.size(); ++fidx) {
+ const VMFuncInfo& gfunc = this->func_table[fidx];
+ if (gfunc.kind == VMFuncInfo::FuncKind::kPackedFunc) {
+ os << "@" << gfunc.name << " packed_func;\n\n";
+ continue;
+ }
+ if (gfunc.kind == VMFuncInfo::FuncKind::kVMTIRFunc) {
+ os << "@" << gfunc.name << " num_inputs=" << gfunc.num_args << " vm_tir_func;\n\n";
+ continue;
+ }
+ ICHECK(gfunc.kind == VMFuncInfo::FuncKind::kVMFunc);
+ os << "@" << gfunc.name << ":\n";
+ size_t start_instr = gfunc.start_instr;
+ size_t end_instr = gfunc.end_instr;
+
+ for (size_t idx = start_instr; idx < end_instr; ++idx) {
+ os << " ";
+ Instruction instr = this->GetInstruction(idx);
+ switch (instr.op) {
+ case Opcode::Call: {
+ os << std::setw(6) << std::left << "call" << std::setw(16) << std::left
+ << get_func_name(instr.func_idx) << " in: " << std::setw(12) << std::left
+ << StrJoin<Instruction::Arg>(instr.args, 0, instr.num_args, ", ", instr_to_str)
+ << " dst: " << RegNameToStr(instr.dst) << "\n";
+ break;
+ }
+ case Opcode::Ret: {
+ os << std::setw(6) << std::left << "ret " << RegNameToStr(instr.result) << "\n";
+ break;
+ }
+ case Opcode::Goto: {
+ os << std::setw(6) << std::left << "goto" << instr.pc_offset << "\n";
+ break;
+ }
+ case Opcode::If: {
+ os << std::setw(6) << std::left << "If" << RegNameToStr(instr.cond) << ", "
+ << instr.false_offset << "\n";
+ break;
+ }
+ default:
+ LOG(FATAL) << "should never hit this case: " << static_cast<int>(instr.op);
+ break;
+ }
+ }
+ os << "\n";
+ }
+ return String(os.str());
+}
+
+String Executable::AsPython() const {
+ auto get_func_name = [&](Index index) -> std::string {
+ if (static_cast<size_t>(index) < func_table.size()) {
+ return "\"" + func_table[index].name + "\"";
+ } else {
+ return "ib.unknown_func_index(" + std::to_string(index) + ")";
+ }
+ };
+
+ auto arg_to_py_str = [&](Instruction::Arg arg) -> std::string {
+ switch (arg.kind()) {
+ case Instruction::ArgKind::kRegister:
+ if (arg.value() == Instruction::kVMRegister) {
+ return "ib.r(vm)";
+ }
+ return "ib.r(" + std::to_string(arg.value()) + ")";
+ case Instruction::ArgKind::kImmediate:
+ return "ib.imm(" + std::to_string(arg.value()) + ")";
+ case Instruction::ArgKind::kConstIdx:
+ return "ib.c(" + std::to_string(arg.value()) + ")";
+ case Instruction::ArgKind::kFuncIdx: {
+ return "ib.f(" + get_func_name(arg.value()) + ")";
+ }
+ default:
+ LOG(FATAL) << "Wrong instruction kind: " << static_cast<int>(arg.kind());
+ return "";
+ }
+ };
+
+ // print the python format
+ std::ostringstream os;
+ os << "ib = rx.Builder()\n";
+ for (size_t fidx = 0; fidx < this->func_table.size(); ++fidx) {
+ const VMFuncInfo& gfunc = this->func_table[fidx];
+ if (gfunc.kind == VMFuncInfo::FuncKind::kPackedFunc) {
+ continue;
+ }
+ if (gfunc.kind == VMFuncInfo::FuncKind::kVMTIRFunc) {
+ continue;
+ }
+ ICHECK(gfunc.kind == VMFuncInfo::FuncKind::kVMFunc);
+
+ os << "with ib.function(\"" << gfunc.name << "\", num_inputs=" << gfunc.num_args << "):\n";
+ size_t start_instr = gfunc.start_instr;
+ size_t end_instr = gfunc.end_instr;
+
+ for (size_t idx = start_instr; idx < end_instr; ++idx) {
+ Instruction instr = this->GetInstruction(idx);
+ switch (instr.op) {
+ case Opcode::Call: {
+ os << " ib.emit_call(" << get_func_name(instr.func_idx) << ", args=["
+ << StrJoin<Instruction::Arg>(instr.args, 0, instr.num_args, ", ", arg_to_py_str)
+ << "]";
+ if (instr.dst != Instruction::kVoidRegister) os << ", dst=ib.r(" << instr.dst << ")";
+ os << ")\n";
+ break;
+ }
+ case Opcode::Ret: {
+ os << " ib.emit_ret(ib.r(" << instr.result << "))\n";
+ break;
+ }
+ case Opcode::Goto: {
+ os << " ib.emit_goto(" << instr.pc_offset << ")\n";
+ break;
+ }
+ case Opcode::If: {
+ os << " ib.emit_if(ib.r(" << instr.cond << "), " << instr.false_offset << ")\n";
+ break;
+ }
+ default:
+ LOG(FATAL) << "should never hit this case: " << static_cast<int>(instr.op);
+ break;
+ }
+ }
+ }
+ return String(os.str());
+}
+
+TVM_REGISTER_GLOBAL("relax.ExecutableLoadFromFile").set_body_typed(Executable::LoadFromFile);
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
diff --git a/src/runtime/relax_vm/memory_manager.cc b/src/runtime/relax_vm/memory_manager.cc
new file mode 100644
index 0000000000..a017b9c6d9
--- /dev/null
+++ b/src/runtime/relax_vm/memory_manager.cc
@@ -0,0 +1,181 @@
+/*
+ * 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/relax_vm/memory_manager.cc
+ * \brief Allocate and manage memory for the Relay VM.
+ */
+#include <tvm/runtime/relax_vm/memory_manager.h>
+
+#include <memory>
+#include <utility>
+
+#include "naive_allocator.h"
+#include "pooled_allocator.h"
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+static void BufferDeleter(Object* obj) {
+ auto* ptr = static_cast<runtime::NDArray::Container*>(obj);
+ ICHECK(ptr->manager_ctx != nullptr);
+ Buffer* buffer = reinterpret_cast<Buffer*>(ptr->manager_ctx);
+ MemoryManager::GetAllocator(buffer->device)->Free(*(buffer));
+ delete buffer;
+ delete ptr;
+}
+
+void StorageObj::Deleter(Object* obj) {
+ auto* ptr = static_cast<runtime::NDArray::Container*>(obj);
+ // When invoking AllocNDArray we don't own the underlying allocation
+ // and should not delete the buffer, but instead let it be reclaimed
+ // by the storage object's destructor.
+ //
+ // We did bump the reference count by 1 to keep alive the StorageObj
+ // allocation in case this NDArray is the sole owner.
+ //
+ // We decrement the object allowing for the buffer to release our
+ // reference count from allocation.
+ StorageObj* storage = reinterpret_cast<StorageObj*>(ptr->manager_ctx);
+ storage->DecRef();
+ delete ptr;
+}
+
+inline void VerifyDataType(DLDataType dtype) {
+ ICHECK_GE(dtype.lanes, 1);
+ if (dtype.code == kDLFloat) {
+ ICHECK_EQ(dtype.bits % 8, 0);
+ } else {
+ // allow uint1 as a special flag for bool.
+ if (dtype.bits == 1 && dtype.code == kDLUInt) return;
+ ICHECK_EQ(dtype.bits % 8, 0);
+ }
+ ICHECK_EQ(dtype.bits & (dtype.bits - 1), 0);
+}
+
+inline size_t GetDataAlignment(const DLTensor& arr) {
+ size_t align = (arr.dtype.bits / 8) * arr.dtype.lanes;
+ if (align < runtime::kAllocAlignment) return runtime::kAllocAlignment;
+ return align;
+}
+
+runtime::NDArray StorageObj::AllocNDArray(uint64_t offset, ShapeTuple shape, DLDataType dtype) {
+ VerifyDataType(dtype);
+
+ // critical zone: allocate header, cannot throw
+ runtime::NDArray::Container* container =
+ new runtime::NDArray::Container(nullptr, shape, dtype, this->buffer.device);
+
+ container->SetDeleter(StorageObj::Deleter);
+ size_t needed_size = runtime::GetDataSize(container->dl_tensor);
+ this->IncRef();
+ // The manager context pointer must continue to point to the storage object
+ // which owns the backing memory, and keeps track of the reference count.
+ //
+ // When we free a container we extract the storage object, decrement its
+ // reference count, then destroy the container, but leave the underlying
+ // buffer intact.
+ container->manager_ctx = reinterpret_cast<void*>(this);
+
+ // is this UB?
+ // The only change we make w.r.t offset is modifying the data pointer
+ // of the backing tensor to point into the buffer instead of its start.
+ auto offset_ptr = reinterpret_cast<uint8_t*>(this->buffer.data) + offset;
+ container->dl_tensor.data = reinterpret_cast<void*>(offset_ptr);
+
+ runtime::NDArray ret(runtime::GetObjectPtr<Object>(container));
+ // RAII in effect, now run the check.
+
+ ICHECK(offset + needed_size <= this->buffer.size)
+ << "storage allocation failure, attempted to allocate " << needed_size << " at offset "
+ << offset << " in region that is " << this->buffer.size << "bytes";
+
+ return ret;
+}
+
+MemoryManager* MemoryManager::Global() {
+ // NOTE: explicitly use new to avoid exit-time destruction of global state
+ // Global state will be recycled by OS as the process exits.
+ static auto* inst = new MemoryManager();
+ return inst;
+}
+
+Allocator* MemoryManager::GetOrCreateAllocator(Device dev, AllocatorType type) {
+ MemoryManager* m = MemoryManager::Global();
+ std::lock_guard<std::mutex> lock(m->mutex_);
+ if (m->allocators_.find(dev) == m->allocators_.end()) {
+ std::unique_ptr<Allocator> alloc;
+ switch (type) {
+ case kNaive: {
+ DLOG(INFO) << "New naive allocator for " << runtime::DeviceName(dev.device_type) << "("
+ << dev.device_id << ")";
+ alloc.reset(new NaiveAllocator(dev));
+ break;
+ }
+ case kPooled: {
+ DLOG(INFO) << "New pooled allocator for " << runtime::DeviceName(dev.device_type) << "("
+ << dev.device_id << ")";
+ alloc.reset(new PooledAllocator(dev));
+ break;
+ }
+ default:
+ LOG(FATAL) << "Unknown allocator type: " << type;
+ }
+ auto ret = alloc.get();
+ m->allocators_.emplace(dev, std::move(alloc));
+ return ret;
+ }
+ auto alloc = m->allocators_.at(dev).get();
+ if (alloc->type() != type) {
+ LOG(WARNING) << "The type of existing allocator for " << runtime::DeviceName(dev.device_type)
+ << "(" << dev.device_id << ") is different from the request type ("
+ << alloc->type() << " vs " << type << ")";
+ }
+ return alloc;
+}
+
+Allocator* MemoryManager::GetAllocator(Device dev) {
+ MemoryManager* m = MemoryManager::Global();
+ std::lock_guard<std::mutex> lock(m->mutex_);
+ auto it = m->allocators_.find(dev);
+ if (it == m->allocators_.end()) {
+ LOG(FATAL) << "Allocator for " << runtime::DeviceName(dev.device_type) << "(" << dev.device_id
+ << ") has not been created yet.";
+ }
+ return it->second.get();
+}
+
+runtime::NDArray Allocator::Empty(std::vector<int64_t> shape, DLDataType dtype, DLDevice dev) {
+ VerifyDataType(dtype);
+ runtime::NDArray::Container* container =
+ new runtime::NDArray::Container(nullptr, shape, dtype, dev);
+ container->SetDeleter(BufferDeleter);
+ size_t size = runtime::GetDataSize(container->dl_tensor);
+ size_t alignment = GetDataAlignment(container->dl_tensor);
+ Buffer* buffer = new Buffer;
+ *buffer = this->Alloc(size, alignment, dtype);
+ container->manager_ctx = reinterpret_cast<void*>(buffer);
+ container->dl_tensor.data = buffer->data;
+ return runtime::NDArray(runtime::GetObjectPtr<Object>(container));
+}
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
diff --git a/src/runtime/relax_vm/naive_allocator.h b/src/runtime/relax_vm/naive_allocator.h
new file mode 100644
index 0000000000..843a559602
--- /dev/null
+++ b/src/runtime/relax_vm/naive_allocator.h
@@ -0,0 +1,65 @@
+/*
+ * 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/relax_vm/naive_allocator.h
+ */
+#ifndef TVM_RUNTIME_RELAX_VM_NAIVE_ALLOCATOR_H_
+#define TVM_RUNTIME_RELAX_VM_NAIVE_ALLOCATOR_H_
+
+#include <tvm/runtime/device_api.h>
+#include <tvm/runtime/relax_vm/memory_manager.h>
+
+#include <atomic>
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+class NaiveAllocator final : public Allocator {
+ public:
+ explicit NaiveAllocator(Device dev) : Allocator(kNaive), used_memory_(0), device_(dev) {}
+
+ Buffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) override {
+ Buffer buf;
+ buf.device = device_;
+ buf.size = nbytes;
+ buf.data =
+ runtime::DeviceAPI::Get(device_)->AllocDataSpace(device_, nbytes, alignment, type_hint);
+ used_memory_.fetch_add(nbytes, std::memory_order_relaxed);
+ DLOG(INFO) << "allocate " << nbytes << " B, used memory " << used_memory_ << " B";
+ return buf;
+ }
+
+ void Free(const Buffer& buffer) override {
+ runtime::DeviceAPI::Get(device_)->FreeDataSpace(buffer.device, buffer.data);
+ used_memory_.fetch_sub(buffer.size, std::memory_order_relaxed);
+ DLOG(INFO) << "free " << buffer.size << " B, used memory " << used_memory_ << " B";
+ }
+
+ private:
+ std::atomic<size_t> used_memory_;
+ Device device_;
+};
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
+
+#endif // TVM_RUNTIME_RELAX_VM_NAIVE_ALLOCATOR_H_
diff --git a/src/runtime/relax_vm/pooled_allocator.h b/src/runtime/relax_vm/pooled_allocator.h
new file mode 100644
index 0000000000..0dd7d8b027
--- /dev/null
+++ b/src/runtime/relax_vm/pooled_allocator.h
@@ -0,0 +1,111 @@
+/*
+ * 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/relax_vm/pooled_allocator.h
+ */
+#ifndef TVM_RUNTIME_RELAX_VM_POOLED_ALLOCATOR_H_
+#define TVM_RUNTIME_RELAX_VM_POOLED_ALLOCATOR_H_
+
+#include <tvm/runtime/device_api.h>
+#include <tvm/runtime/relax_vm/memory_manager.h>
+
+#include <atomic>
+#include <mutex>
+#include <unordered_map>
+#include <vector>
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+class PooledAllocator final : public Allocator {
+ public:
+ static constexpr size_t kDefaultPageSize = 4096;
+
+ explicit PooledAllocator(Device dev, size_t page_size = kDefaultPageSize)
+ : Allocator(kPooled), page_size_(page_size), used_memory_(0), device_(dev) {}
+
+ ~PooledAllocator() { ReleaseAll(); }
+
+ Buffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) override {
+ std::lock_guard<std::recursive_mutex> lock(mu_);
+ size_t size = ((nbytes + page_size_ - 1) / page_size_) * page_size_;
+ auto&& it = memory_pool_.find(size);
+ if (it != memory_pool_.end() && !it->second.empty()) {
+ auto&& pool = it->second;
+ auto ret = pool.back();
+ pool.pop_back();
+ return ret;
+ }
+ Buffer buf;
+ buf.device = device_;
+ buf.size = size;
+ try {
+ buf.data =
+ runtime::DeviceAPI::Get(device_)->AllocDataSpace(device_, size, alignment, type_hint);
+ } catch (InternalError& err) {
+ LOG(WARNING) << "PooledAllocator got InternalError during allocation: " << err.message();
+ LOG(WARNING) << "Trying to release all unused memory and reallocate...";
+ ReleaseAll();
+ buf.data =
+ runtime::DeviceAPI::Get(device_)->AllocDataSpace(device_, size, alignment, type_hint);
+ }
+
+ used_memory_.fetch_add(size, std::memory_order_relaxed);
+ DLOG(INFO) << "allocate " << size << " B, used memory " << used_memory_ << " B";
+ return buf;
+ }
+
+ void Free(const Buffer& buffer) override {
+ std::lock_guard<std::recursive_mutex> lock(mu_);
+ if (memory_pool_.find(buffer.size) == memory_pool_.end()) {
+ memory_pool_.emplace(buffer.size, std::vector<Buffer>{});
+ }
+ memory_pool_.at(buffer.size).push_back(buffer);
+ DLOG(INFO) << "reclaim buffer " << buffer.size;
+ }
+
+ private:
+ void ReleaseAll() {
+ std::lock_guard<std::recursive_mutex> lock(mu_);
+ for (auto const& it : memory_pool_) {
+ auto const& pool = it.second;
+ for (auto const& buf : pool) {
+ runtime::DeviceAPI::Get(buf.device)->FreeDataSpace(buf.device, buf.data);
+ }
+ }
+ memory_pool_.clear();
+ used_memory_ = 0;
+ DLOG(INFO) << "release all buffers";
+ }
+
+ private:
+ size_t page_size_;
+ std::atomic<size_t> used_memory_;
+ std::unordered_map<size_t, std::vector<Buffer> > memory_pool_;
+ std::recursive_mutex mu_;
+ Device device_;
+};
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
+
+#endif // TVM_RUNTIME_RELAX_VM_POOLED_ALLOCATOR_H_
diff --git a/src/runtime/relax_vm/vm.cc b/src/runtime/relax_vm/vm.cc
new file mode 100644
index 0000000000..3cf65faaa8
--- /dev/null
+++ b/src/runtime/relax_vm/vm.cc
@@ -0,0 +1,811 @@
+/*
+ * 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 src/runtime/relax_vm/vm.cc
+ */
+
+#include <tvm/runtime/container/adt.h>
+#include <tvm/runtime/packed_func.h>
+#include <tvm/runtime/relax_vm/vm.h>
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+//---------------------------------------------
+// VM Closure object
+//---------------------------------------------
+TVM_REGISTER_OBJECT_TYPE(VMClosureObj);
+
+VMClosure::VMClosure(String func_name, PackedFunc impl) {
+ auto ptr = make_object<VMClosureObj>();
+ ptr->func_name = func_name;
+ ptr->impl = std::move(impl);
+ data_ = std::move(ptr);
+}
+
+/*!
+ * \brief Create another PackedFunc with last arguments already bound to last_args.
+ * \param func The input func, can be a VMClosure or PackedFunc.
+ * \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.
+ */
+PackedFunc VMClosure::BindLastArgs(PackedFunc func, std::vector<TVMRetValue> last_args) {
+ return PackedFunc([func, last_args](TVMArgs args, TVMRetValue* rv) {
+ std::vector<TVMValue> values(args.size() + last_args.size());
+ std::vector<int> tcodes(args.size() + last_args.size());
+ runtime::TVMArgsSetter setter(values.data(), tcodes.data());
+ std::copy(args.values, args.values + args.size(), values.data());
+ std::copy(args.type_codes, args.type_codes + args.size(), tcodes.data());
+ for (size_t i = 0; i < last_args.size(); ++i) {
+ setter(i + args.size(), last_args[i]);
+ }
+ func.CallPacked(TVMArgs(values.data(), tcodes.data(), values.size()), rv);
+ });
+}
+
+//-----------------------------------------------------------
+// Utility functions.
+//-----------------------------------------------------------
+// Use the args after `starting_arg_idx` as a series of indices into `obj`,
+// indexing into nested ADTs and returning the final indexed object.
+ObjectRef IndexIntoNestedObject(ObjectRef obj, TVMArgs args, int starting_arg_idx) {
+ for (int i = starting_arg_idx; i < args.size(); i++) {
+ // the object must be an ADT to be able to index into it
+ if (!obj.as<ADTObj>()) {
+ LOG(FATAL) << "ValueError: Attempted to index into an object that is not an ADT.";
+ }
+ int index = args[i];
+ auto adt = Downcast<ADT>(obj);
+ // make sure the index is in bounds
+ if (index >= static_cast<int>(adt.size())) {
+ LOG(FATAL) << "IndexError: Invalid index (" << index << " >= " << adt.size() << ").";
+ }
+ obj = adt[index];
+ }
+ return obj;
+}
+
+NDArray ConvertNDArrayToDevice(NDArray src, const DLDevice& dev) {
+ if (src->device.device_type == dev.device_type && src->device.device_id == dev.device_id) {
+ return src;
+ }
+ return src.CopyTo(dev);
+}
+
+ObjectRef ConvertObjectToDevice(ObjectRef src, const Device& dev) {
+ if (src->IsInstance<NDArray::ContainerType>()) {
+ return ConvertNDArrayToDevice(Downcast<NDArray>(src), dev);
+ } else if (src->IsInstance<ADTObj>()) {
+ std::vector<ObjectRef> ret;
+ ADT adt = Downcast<ADT>(src);
+ for (size_t i = 0; i < adt.size(); i++) {
+ ret.push_back(ConvertObjectToDevice(adt[i], dev));
+ }
+ return ADT(adt->tag, ret.begin(), ret.end());
+ } else {
+ return src;
+ }
+}
+
+TVMRetValue ConvertArgToDevice(TVMArgValue input, Device dev) {
+ // NOTE: NDArray::FromExternalDLTensor is not safe
+ // in terms of memory-behavior.
+ // To be extra careful, we copy DLTensor.
+ // The developer can still explicitly allocate NDArray
+ // in TVM Native API or NDArray::FromDLPack to regain zero copy behavior.
+ TVMRetValue ret;
+
+ if (input.type_code() == kTVMDLTensorHandle) {
+ ret = NDArray::NewFromDLTensor(input, dev);
+ } else if (input.IsObjectRef<ObjectRef>()) {
+ ret = ConvertObjectToDevice(input.operator ObjectRef(), dev);
+ } else {
+ ret = input;
+ }
+ return ret;
+}
+
+TVMRetValue ConvertRegToDevice(TVMRetValue input, Device dev) {
+ TVMRetValue ret;
+ if (input.IsObjectRef<ObjectRef>()) {
+ ret = ConvertObjectToDevice(input.operator ObjectRef(), dev);
+ } else {
+ ret = input;
+ }
+ return ret;
+}
+
+//-----------------------------------------------------------
+// VM implementations.
+//-----------------------------------------------------------
+/*!
+ * \brief The register type.
+ */
+using RegType = TVMRetValue;
+
+/*!
+ * \brief A representation of a stack frame.
+ *
+ * A stack frame is a record containing the information needed
+ * to restore the caller's virtual machine state after returning
+ * from a function call.
+ */
+struct VMFrame {
+ /*! \brief The return program counter. */
+ Index return_pc;
+ /*! \brief Statically allocated space for objects */
+ std::vector<RegType> register_file;
+ /*! \brief Register in caller's frame to put return value */
+ RegName caller_return_register;
+ // The following fields are used for PackedFunc call within
+ // a single function scope. The space is reused across multiple
+ // packed func calls to increase cache locality and avoid re-allocation
+ /*! \brief Temporary argument value stack for packed func call. */
+ std::vector<TVMValue> call_arg_values;
+ /*! \brief Temporary argument tcode stack for packed func call. */
+ std::vector<int> call_arg_tcodes;
+
+ VMFrame(Index pc, Index register_file_size)
+ : return_pc(pc), register_file(register_file_size), caller_return_register(0) {}
+};
+
+class VirtualMachineImpl : public VirtualMachine {
+ public:
+ //---------------------------------------------------
+ // Public facing functions overloading
+ //---------------------------------------------------
+ void LoadExecutable(ObjectPtr<Executable> exec) final;
+
+ void Init(const std::vector<Device>& devices,
+ const std::vector<AllocatorType>& alloc_types) final;
+
+ PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
+
+ VMClosure GetClosure(const String& func_name) final;
+
+ void InvokeClosurePacked(const ObjectRef& closure_or_packedfunc, TVMArgs args,
+ TVMRetValue* rv) final;
+
+ //--------------------------------------------------
+ // Additional support arguments functions for VM
+ //--------------------------------------------------
+ /*!
+ * \brief Set inputs to a function.
+ * \param func_name The function name.
+ * \param args args[offset:] are arguments to the function. If the arguments are not of the
+ * correct device for the function, they will be copied to the device.
+ * \param offset Starting offset of the arguments in \p args.
+ * \note This interface works when using VM over RPC by internally converting NDArray in
+ * the arguments to DLTensor, which is supported in RPC where remote could only have a minimal C
+ * runtime.
+ */
+ void SetInput(std::string func_name, TVMArgs args, int offset);
+
+ /*!
+ * \brief Look up whether the VM has a function by the given name.
+ * \param func_name the function's name
+ * \return The function, if it exists. Logs a fatal error if not.
+ */
+ VMFuncInfo LookupVMFuncInfo(const std::string& func_name);
+
+ /*!
+ * \brief Look up whether the VM has outputs for the given function.
+ * \param func_name the function's name
+ * \return The output, if it exists. Logs a fatal error if not.
+ */
+ RegType LookupVMOutput(const std::string& func_name);
+
+ /*!
+ * \brief Fully bind the argument of a global function and save it in the env.
+ * \param func_name The global function name to be saved.
+ * \param save_name The saved name of the function.
+ * \param include_return Whether forward the return value, set it to false allows
+ * us to ignore forwarding return value, which can be helpful to do benchmarking
+ * in RPC environment when return value is complicated ADT.
+ *
+ * \param args The arguments to bound to the function.
+ * \note This function is used by RPC server to help benchmarking.
+ */
+ void SaveClosure(const String& func_name, const String& save_name, bool include_return,
+ TVMArgs args);
+ /*!
+ * \brief Internal function to invoke a closure.
+ * \param closure_or_packed The closure to be invoked.
+ * \param args The arguments to the function.
+ * \return The result value.
+ */
+ RegType InvokeClosureInternal(const ObjectRef& closure_or_packed,
+ const std::vector<RegType>& args);
+ /*!
+ * \brief Invoke a VM function by interpreting bytecode.
+ * \param fidx The function index.
+ * \param args The arguments to the function.
+ * \return The object representing the result.
+ */
+ RegType InvokeBytecode(Index fidx, const std::vector<RegType>& args);
+
+ protected:
+ /*!
+ * \brief Get function by querying all of the current module's imports.
+ * \param name The name of the function.
+ * \return The result function, can return PackedFunc(nullptr) if nothing is found.
+ */
+ PackedFunc GetFuncFromImports(const String& name) {
+ for (auto& lib : this->imports_) {
+ PackedFunc func = lib->GetFunction(name, true);
+ if (func.defined()) return func;
+ }
+ return PackedFunc(nullptr);
+ }
+ /*!
+ * \brief Initialize function pool.
+ */
+ void InitFuncPool();
+ //-------------------------------------------------
+ // Instruction interpretations.
+ //-------------------------------------------------
+ /*!
+ * \brief Push a call frame onto the call stack.
+ * \param ret_pc The program counter to return to.
+ * \param vm_func The function to be pushed to the call stack.
+ */
+ void PushFrame(Index ret_pc, const VMFuncInfo& vm_func) {
+ frames_.emplace_back(std::make_unique<VMFrame>(ret_pc, vm_func.register_file_size));
+ }
+ /*!
+ * \brief Pop a frame off the call stack.
+ */
+ void PopFrame() {
+ ICHECK_GT(frames_.size(), 0);
+ pc_ = frames_.back()->return_pc;
+ frames_.pop_back();
+ }
+ /*!
+ * \brief Write to a VM register.
+ * \param frame current vm frame.
+ * \param reg The register to write to.
+ * \param obj The object to write to.
+ */
+ void WriteRegister(VMFrame* frame, RegName reg, const RegType& obj) {
+ ICHECK_LT(reg, frame->register_file.size());
+ frame->register_file[reg] = obj;
+ }
+ /*!
+ * \brief Read a VM register.
+ * \param frame current vm frame.
+ * \param reg The register to read from.
+ * \return The value of the register.
+ */
+ RegType ReadRegister(VMFrame* frame, RegName reg) {
+ if (reg < Instruction::kBeginSpecialReg) {
+ return frame->register_file[reg];
+ }
+ RegType ret;
+ if (reg == Instruction::kVoidRegister) {
+ ret = nullptr;
+ } else {
+ ICHECK_EQ(reg, Instruction::kVMRegister);
+ // per convention, ctx ptr must be VirtualMachine* casted to void.
+ // this and VirtualMachine* may or maynot be the same
+ // do first cast to VirtualMachine* then to void*
+ ret = static_cast<void*>(static_cast<VirtualMachine*>(this));
+ }
+ return ret;
+ }
+ /*!
+ * \brief Run call instruction.
+ * \param curr_frame The current frame.
+ * \param inst The call instruction.
+ */
+ inline void RunInstrCall(VMFrame* curr_frame, Instruction inst);
+
+ /*! \brief Run VM dispatch loop. */
+ void RunLoop();
+
+ private:
+ //--------------------------------------------------------
+ // Internal states for execution.
+ //--------------------------------------------------------
+ /*! \brief The loaded executable. */
+ ObjectPtr<Executable> exec_;
+ /*! \brief The global constant pool */
+ std::vector<TVMRetValue> const_pool_;
+ /*!
+ * \brief Function pool to cache functions in func_table
+ */
+ std::vector<TVMRetValue> func_pool_;
+ //--------------------------------------------------------
+ // Executor interface support
+ //--------------------------------------------------------
+ /*! \brief The function name to input register mapping. */
+ std::unordered_map<std::string, std::vector<RegType>> inputs_;
+ /*! \brief The function name to output register. */
+ std::unordered_map<std::string, RegType> outputs_;
+ /*! \brief A store of closures created by `save_function`. */
+ std::unordered_map<std::string, VMClosure> saved_closures_;
+ //------------------------------------------------------------
+ // VM Instruction execution.
+ //------------------------------------------------------------
+ /*!
+ * \brief The current stack of call frames.
+ * \note: Use unique ptr to avoid re-allocation and copy when frames_ get resized.
+ */
+ std::vector<std::unique_ptr<VMFrame>> frames_;
+ /*! \brief The virtual machine PC. */
+ Index pc_{0};
+ /*! \brief The special return register. */
+ RegType return_value_;
+};
+
+void VirtualMachineImpl::LoadExecutable(ObjectPtr<Executable> exec) {
+ this->exec_ = exec;
+ this->imports_ = exec_->imports();
+}
+
+void VirtualMachineImpl::Init(const std::vector<Device>& devices,
+ const std::vector<AllocatorType>& alloc_types) {
+ // TODO(@yuchen): support multi-device heterogeneous execution
+ ICHECK_LT(devices.size(), 3)
+ << "Currently relax vm only supports at most 2 devices (host + device)";
+ ICHECK_EQ(devices.size(), alloc_types.size());
+
+ this->devices.reserve(devices.size());
+ this->allocators.reserve(alloc_types.size());
+ for (size_t i = 0; i < devices.size(); i++) {
+ auto alloc = MemoryManager::GetOrCreateAllocator(devices[i], alloc_types[i]);
+ this->devices.push_back(devices[i]);
+ this->allocators.push_back(alloc);
+ }
+ // Setup constant sections.
+ this->const_pool_.reserve(exec_->constants.size());
+ for (const auto& constant : exec_->constants) {
+ if (constant.type_code() != kTVMNDArrayHandle) {
+ this->const_pool_.push_back(constant);
+ } else {
+ this->const_pool_.push_back(ConvertRegToDevice(constant, devices[0]));
+ }
+ }
+ // Setup function sections.
+ this->InitFuncPool();
+}
+
+VMFuncInfo VirtualMachineImpl::LookupVMFuncInfo(const std::string& func_name) {
+ ICHECK(exec_) << "The executable is not created yet.";
+ auto it = this->exec_->func_map.find(func_name);
+ CHECK(it != this->exec_->func_map.end()) << "ValueError: Unknown function: " << func_name;
+
+ return exec_->func_table[it->second];
+}
+
+RegType VirtualMachineImpl::LookupVMOutput(const std::string& func_name) {
+ if (!outputs_.count(func_name)) {
+ LOG(FATAL) << "ValueError: No output saved for call of \"" << func_name
+ << "\"; use `invoke_stateful` to call it first.";
+ }
+ return outputs_[func_name];
+}
+
+PackedFunc VirtualMachineImpl::GetFunction(const std::string& name,
+ const ObjectPtr<Object>& sptr_to_self) {
+ if (name == "vm_initialization") {
+ // initialize the VirtualMachine, takes variable-length arguments
+ // first argument is a runtime::Module, followed by one or more device_type, device_id,
+ // and the AllocatorType associated with the device.
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ ICHECK_EQ(args.size() % 3, 0);
+ std::vector<Device> devices;
+ std::vector<AllocatorType> alloc_types;
+ for (int i = 0; i < args.size(); i += 3) {
+ Device dev;
+ int device_type = args[i];
+ dev.device_type = DLDeviceType(device_type);
+ dev.device_id = args[i + 1];
+ int type = args[i + 2];
+ devices.push_back(dev);
+ alloc_types.push_back(AllocatorType(type));
+ }
+ this->Init(devices, alloc_types);
+ });
+ } else if (name == "save_function") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ ICHECK_GE(args.size(), 3);
+ this->SaveClosure(args[0], args[1], args[2],
+ TVMArgs(args.values + 3, args.type_codes + 3, args.size() - 3));
+ });
+ } else if (name == "invoke_closure") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ VMClosure clo = args[0];
+ this->InvokeClosurePacked(clo, TVMArgs(args.values + 1, args.type_codes + 1, args.size() - 1),
+ rv);
+ });
+ } else if (name == "invoke_stateful") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string func_name = args[0];
+ const auto& m = this->exec_->func_map;
+ if (m.find(func_name) == m.end()) {
+ LOG(FATAL) << "ValueError: Unknown function: " << func_name;
+ }
+ Index gf_idx = m.at(func_name);
+ if (!inputs_.count(func_name)) {
+ LOG(FATAL) << "ValueError: No inputs set for stateful call of " << func_name
+ << "; use `set_input` first.";
+ return;
+ }
+ outputs_[func_name] = this->InvokeClosureInternal(func_pool_[gf_idx], inputs_[func_name]);
+ });
+ } else if (name == "get_output_arity") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string func_name = args[0];
+ RegType out = LookupVMOutput(func_name);
+ // use remaining args as indices
+ ObjectRef obj = IndexIntoNestedObject(out.AsObjectRef<ObjectRef>(), args, 1);
+ // after chasing through the indices, examine the final object
+ if (const auto* adt = obj.as<ADTObj>()) {
+ *rv = static_cast<int>(adt->size);
+ } else {
+ *rv = -1;
+ }
+ });
+ } else if (name == "get_output") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string func_name = args[0];
+ RegType out = LookupVMOutput(func_name);
+ // use remaining args as indices
+ ObjectRef obj = IndexIntoNestedObject(out.AsObjectRef<ObjectRef>(), args, 1);
+ if (obj.as<ADTObj>()) {
+ LOG(FATAL) << "ValueError: `get_output` cannot return a tuple for RPC compatibility. "
+ "Please specify another index argument.";
+ return;
+ }
+ *rv = obj;
+ });
+ } else if (name == "set_input") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { SetInput(args[0], args, 1); });
+ } else if (name == "get_function_arity") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string func_name = args[0];
+ const VMFuncInfo& vm_func = LookupVMFuncInfo(func_name);
+ *rv = static_cast<int>(vm_func.param_names.size());
+ });
+ } else if (name == "get_function_param_name") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string func_name = args[0];
+ int index = args[1];
+ const VMFuncInfo& vm_func = LookupVMFuncInfo(func_name);
+ if (static_cast<size_t>(index) >= vm_func.param_names.size()) {
+ LOG(FATAL) << "ValueError: Invalid index for " << func_name << " (" << index << " out of "
+ << vm_func.param_names.size() << ")";
+ }
+ *rv = vm_func.param_names[index];
+ });
+ } else {
+ // default case, look up closure in VM.
+ VMClosure clo = this->GetClosure(name);
+ return PackedFunc([sptr_to_self, this, clo](TVMArgs args, TVMRetValue* rv) {
+ this->InvokeClosurePacked(clo, args, rv);
+ });
+ }
+}
+
+void VirtualMachineImpl::SetInput(std::string func_name, TVMArgs args, int offset) {
+ const auto& m = exec_->func_map;
+ if (m.find(func_name) != m.end()) {
+ Index gf_idx = m.at(func_name);
+ const VMFuncInfo& vm_func = exec_->func_table[gf_idx];
+ size_t params_num = vm_func.num_args;
+ ICHECK_EQ(args.size() - offset, params_num)
+ << "The number of provided parameters doesn't match the number of arguments for";
+ std::vector<RegType> func_args(params_num);
+ for (int i = offset; i < args.size(); ++i) {
+ int index = i - offset;
+ func_args[index] = ConvertArgToDevice(args[i], devices[0]);
+ }
+ inputs_.emplace(func_name, func_args);
+ } else {
+ LOG(FATAL) << "ValueError: Unknown function: " << func_name;
+ }
+}
+
+//------------------------------------------
+// Closure handling
+//------------------------------------------
+void VirtualMachineImpl::InvokeClosurePacked(const ObjectRef& closure_or_packedfunc, TVMArgs args,
+ TVMRetValue* rv) {
+ // run packed call if it is a packed func.
+ if (auto* packed = closure_or_packedfunc.as<PackedFunc::ContainerType>()) {
+ packed->CallPacked(args, rv);
+ return;
+ }
+ // run closure call.
+ auto* clo = closure_or_packedfunc.as<VMClosureObj>();
+ ICHECK(clo != nullptr) << "Function expects a closure or PackedFunc ";
+
+ std::vector<TVMValue> values(args.size() + 1);
+ std::vector<int> tcodes(args.size() + 1);
+ runtime::TVMArgsSetter setter(values.data(), tcodes.data());
+ // per convention, ctx ptr must be VirtualMachine* casted to void.
+ // this and VirtualMachine* may or maynot be the same
+ // do first cast to VirtualMachine* then to void*
+ setter(0, static_cast<void*>(static_cast<VirtualMachine*>(this)));
+ std::copy(args.values, args.values + args.size(), values.begin() + 1);
+ std::copy(args.type_codes, args.type_codes + args.size(), tcodes.begin() + 1);
+ clo->impl.CallPacked(TVMArgs(values.data(), tcodes.data(), args.size() + 1), rv);
+}
+
+// internal variant version of invoke closurepacked
+RegType VirtualMachineImpl::InvokeClosureInternal(const ObjectRef& closure_or_packed,
+ const std::vector<RegType>& args) {
+ RegType ret;
+ auto* packed = closure_or_packed.as<PackedFunc::ContainerType>();
+ auto* clo = closure_or_packed.as<VMClosureObj>();
+ int clo_offset = clo != nullptr ? 1 : 0;
+ std::vector<TVMValue> values(args.size() + clo_offset);
+ std::vector<int> tcodes(args.size() + clo_offset);
+ runtime::TVMArgsSetter setter(values.data(), tcodes.data());
+
+ if (clo != nullptr) {
+ setter(0, static_cast<void*>(static_cast<VirtualMachine*>(this)));
+ }
+ for (size_t i = 0; i < args.size(); ++i) {
+ setter(i + clo_offset, args[i]);
+ }
+
+ if (packed != nullptr) {
+ packed->CallPacked(TVMArgs(values.data(), tcodes.data(), values.size()), &ret);
+ } else {
+ ICHECK(clo != nullptr);
+ clo->impl.CallPacked(TVMArgs(values.data(), tcodes.data(), values.size()), &ret);
+ }
+ return ret;
+}
+
+void VirtualMachineImpl::SaveClosure(const String& func_name, const String& save_name,
+ bool include_return, TVMArgs args) {
+ VMClosure clo = this->GetClosure(func_name);
+ std::vector<RegType> inputs(args.size());
+ for (int i = 0; i < args.size(); ++i) {
+ inputs[i] = ConvertArgToDevice(args[i], this->devices[0]);
+ }
+ PackedFunc impl = VMClosure::BindLastArgs(clo->impl, inputs);
+ if (!include_return) {
+ impl = PackedFunc([impl](TVMArgs args, TVMRetValue* rv) {
+ TVMRetValue temp;
+ impl.CallPacked(args, &temp);
+ });
+ }
+ saved_closures_[save_name] = VMClosure(save_name, impl);
+}
+
+VMClosure VirtualMachineImpl::GetClosure(const String& func_name) {
+ // look up saved closures.
+ auto saved_it = saved_closures_.find(func_name);
+ if (saved_it != saved_closures_.end()) {
+ return saved_it->second;
+ }
+ auto it = exec_->func_map.find(func_name);
+ CHECK(it != exec_->func_map.end()) << "ValueError: Unknown function: " << func_name;
+
+ Index gf_idx = it->second;
+ const VMFuncInfo& finfo = exec_->func_table[gf_idx];
+
+ if (finfo.kind == VMFuncInfo::FuncKind::kVMFunc) {
+ // NOTE: should not capture strong ref to self and avoid cyclic ref.
+ auto impl = PackedFunc([gf_idx](TVMArgs args, TVMRetValue* rv) {
+ // Per convention, ctx ptr is a VirtualMachine*
+ VirtualMachine* ctx_ptr = static_cast<VirtualMachine*>(args[0].operator void*());
+
+ std::vector<RegType> inputs(args.size() - 1);
+ for (size_t i = 0; i < inputs.size(); ++i) {
+ inputs[i] = args[i + 1];
+ }
+ *rv = static_cast<VirtualMachineImpl*>(ctx_ptr)->InvokeBytecode(gf_idx, inputs);
+ });
+ return VMClosure(func_name, impl);
+ } else {
+ ICHECK(finfo.kind == VMFuncInfo::FuncKind::kVMTIRFunc)
+ << "Cannot support closure with function kind " << static_cast<int>(finfo.kind);
+ PackedFunc tir_func = GetFuncFromImports("__vmtir__" + finfo.name);
+ ICHECK(tir_func != nullptr) << "Cannot find underlying compiled tir function of VMTIRFunc "
+ << finfo.name;
+ auto impl = PackedFunc([this, finfo, tir_func](TVMArgs args, TVMRetValue* rv) {
+ // Per convention, ctx ptr is a VirtualMachine*
+ VirtualMachine* ctx_ptr = static_cast<VirtualMachine*>(args[0].operator void*());
+ ICHECK(ctx_ptr == this);
+ ICHECK_EQ(args.size() - 1, finfo.num_args)
+ << "Function " << finfo.name << " expects " << finfo.num_args << " arguments";
+ ICHECK_GE(finfo.register_file_size, finfo.num_args + 1);
+ std::vector<TVMRetValue> reg_file(finfo.register_file_size);
+ for (int64_t i = 0; i < finfo.num_args; ++i) {
+ reg_file[i] = args[i + 1];
+ }
+ void* reg_anylist_handle = reg_file.data();
+ void* const_anylist_handle = this->const_pool_.data();
+ void* func_anylist_handle = this->func_pool_.data();
+ tir_func(static_cast<void*>(ctx_ptr), reg_anylist_handle, const_anylist_handle,
+ func_anylist_handle);
+ // Return value always stored after inputs.
+ *rv = reg_file[finfo.num_args];
+ });
+ return VMClosure(func_name, impl);
+ }
+}
+
+//--------------------------------------------------------------------
+// Instruction interpretations.
+//--------------------------------------------------------------------
+RegType VirtualMachineImpl::InvokeBytecode(Index gf_idx, const std::vector<RegType>& args) {
+ const VMFuncInfo& gfunc = exec_->func_table[gf_idx];
+ ICHECK(gfunc.kind == VMFuncInfo::FuncKind::kVMFunc);
+
+ // Get the curr instr which might be a potential caller.
+ Instruction curr_instr = exec_->GetInstruction(pc_);
+ PushFrame(this->pc_, gfunc);
+ // Get new frame and set the caller info.
+ VMFrame* curr_frame = frames_.back().get();
+ if (curr_instr.op == Opcode::Call) {
+ curr_frame->caller_return_register = curr_instr.dst;
+ }
+
+ // load arguments to the register file
+ ICHECK_EQ(static_cast<size_t>(gfunc.num_args), args.size())
+ << "ValueError: Invoking function " << gfunc.name << " requires " << gfunc.num_args
+ << " inputs but only " << args.size() << " inputs are provided.";
+ for (size_t i = 0; i < args.size(); ++i) {
+ WriteRegister(frames_.back().get(), i, args[i]);
+ }
+ // set program counter
+ pc_ = gfunc.start_instr;
+ RunLoop();
+ return return_value_;
+}
+
+void VirtualMachineImpl::InitFuncPool() {
+ func_pool_.resize(exec_->func_table.size());
+
+ for (size_t func_index = 0; func_index < exec_->func_table.size(); ++func_index) {
+ const VMFuncInfo& info = exec_->func_table[func_index];
+ if (info.kind == VMFuncInfo::FuncKind::kPackedFunc) {
+ // only look through imports first
+ PackedFunc func = GetFuncFromImports(info.name);
+ if (!func.defined()) {
+ const PackedFunc* p_func = Registry::Get(info.name);
+ if (p_func != nullptr) func = *(p_func);
+ }
+ ICHECK(func.defined())
+ << "Error: Cannot find PackedFunc " << info.name
+ << " in either Relax VM kernel library, or in TVM runtime PackedFunc registry, or in "
+ "global Relax functions of the VM executable";
+ func_pool_[func_index] = func;
+
+ } else {
+ ICHECK(info.kind == VMFuncInfo::FuncKind::kVMFunc ||
+ info.kind == VMFuncInfo::FuncKind::kVMTIRFunc);
+ auto clo = this->GetClosure(info.name);
+ func_pool_[func_index] = clo;
+ }
+ }
+}
+
+void VirtualMachineImpl::RunInstrCall(VMFrame* curr_frame, Instruction instr) {
+ DLOG(INFO) << "\n pc = " << pc_ << ", execute: " << exec_->func_table[instr.func_idx].name;
+
+ // Use the call arg stack from the current frame to increase reuse
+ // and avoid re-allocation
+ curr_frame->call_arg_values.resize(instr.num_args);
+ curr_frame->call_arg_tcodes.resize(instr.num_args);
+
+ // NOTE: no changes and resize to those vector ref(otherwise can leads to segfault)
+ // in the remainder part of the function.
+ std::vector<TVMValue>& values = curr_frame->call_arg_values;
+ std::vector<int>& tcodes = curr_frame->call_arg_tcodes;
+
+ runtime::TVMArgsSetter setter(values.data(), tcodes.data());
+ for (Index i = 0; i < instr.num_args; ++i) {
+ Instruction::Arg arg = instr.args[i];
+ switch (arg.kind()) {
+ case Instruction::ArgKind::kRegister: {
+ setter(i, ReadRegister(curr_frame, arg.value()));
+ break;
+ }
+ case Instruction::ArgKind::kImmediate: {
+ setter(i, arg.value());
+ break;
+ }
+ case Instruction::ArgKind::kConstIdx: {
+ setter(i, this->const_pool_[arg.value()]);
+ break;
+ }
+ case Instruction::ArgKind::kFuncIdx: {
+ ICHECK_LT(static_cast<size_t>(arg.value()), this->func_pool_.size());
+ setter(i, this->func_pool_[arg.value()]);
+ break;
+ }
+ default: {
+ LOG(FATAL) << "ValueError: Unknown argument kind: " << int(arg.kind());
+ }
+ }
+ }
+ TVMArgs args(values.data(), tcodes.data(), values.size());
+ TVMRetValue ret;
+
+ ICHECK_LT(static_cast<size_t>(instr.func_idx), this->func_pool_.size());
+ this->InvokeClosurePacked(func_pool_[instr.func_idx], args, &ret);
+ // save the return value to the register
+ // saving to special register is a NOP
+ if (instr.dst < Instruction::kBeginSpecialReg) {
+ WriteRegister(curr_frame, instr.dst, ret);
+ }
+ // increment pc
+ pc_++;
+}
+
+void VirtualMachineImpl::RunLoop() {
+ VMFrame* curr_frame = frames_.back().get();
+
+ while (true) {
+ ICHECK_LT(static_cast<size_t>(pc_), exec_->instr_offset.size()) << "run into invalide section";
+ Instruction instr = exec_->GetInstruction(pc_);
+ switch (instr.op) {
+ case Opcode::Call: {
+ this->RunInstrCall(curr_frame, instr);
+ break;
+ }
+ case Opcode::Ret: {
+ // If we have hit the point from which we started
+ // running, we should return to the caller breaking
+ // the dispatch loop.
+ return_value_ = ReadRegister(curr_frame, instr.result);
+ RegName caller_return_register = curr_frame->caller_return_register;
+ PopFrame();
+ if (frames_.size() == 0) {
+ // directly return if no frame in the call stack.
+ } else {
+ // return from a local call.
+ // Update the current frame to be the parent frame.
+ curr_frame = frames_.back().get();
+ WriteRegister(curr_frame, caller_return_register, return_value_);
+ }
+ return;
+ }
+ case Opcode::Goto: {
+ pc_ += instr.pc_offset;
+ break;
+ }
+ case Opcode::If: {
+ int64_t cond_val = ReadRegister(curr_frame, instr.cond);
+ if (cond_val != 0) {
+ pc_++;
+ } else {
+ ICHECK_GT(instr.false_offset, 1);
+ pc_ += instr.false_offset;
+ }
+ break;
+ }
+ }
+ }
+}
+
+ObjectPtr<VirtualMachine> VirtualMachine::Create() { return make_object<VirtualMachineImpl>(); }
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
diff --git a/tests/python/relax/test_vm_execbuilder.py b/tests/python/relax/test_vm_execbuilder.py
new file mode 100644
index 0000000000..9a7cd0c879
--- /dev/null
+++ b/tests/python/relax/test_vm_execbuilder.py
@@ -0,0 +1,262 @@
+# 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.
+"""Lowest level testing VM. Test execbuilder and execution."""
+import tvm
+import pytest
+import numpy as np
+from tvm import relax, TVMError
+from tvm.relax.testing.vm import check_saved_func
+
+
+def test_vm_execute():
+ ib = relax.ExecBuilder()
+ with ib.function("func0", num_inputs=2):
+ ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(1)], dst=ib.r(2))
+ ib.emit_ret(ib.r(2))
+ ex = ib.get()
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ a = tvm.nd.array(
+ np.random.rand(
+ 4,
+ )
+ )
+ b = tvm.nd.array(
+ np.random.rand(
+ 4,
+ )
+ )
+
+ add_res = check_saved_func(vm, "func0", a, b)
+ tvm.testing.assert_allclose(add_res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)
+
+
+def test_vm_multiple_func():
+ ib = relax.ExecBuilder()
+ with ib.function("func0", num_inputs=2):
+ ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(1)], dst=ib.r(2))
+ ib.emit_ret(ib.r(2))
+ with ib.function("func1", num_inputs=2):
+ ib.emit_call("test.vm.mul", args=[ib.r(0), ib.r(1)], dst=ib.r(2))
+ ib.emit_ret(ib.r(2))
+ ex = ib.get()
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ a = tvm.nd.array(
+ np.random.rand(
+ 4,
+ )
+ )
+ b = tvm.nd.array(
+ np.random.rand(
+ 4,
+ )
+ )
+ mul_res = check_saved_func(vm, "func1", a, b)
+ add_res = check_saved_func(vm, "func0", a, b)
+ tvm.testing.assert_allclose(add_res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)
+ tvm.testing.assert_allclose(mul_res.numpy(), a.numpy() * b.numpy(), rtol=1e-7, atol=1e-7)
+
+
+def test_vm_checker():
+ ib = relax.ExecBuilder()
+ with pytest.raises(TVMError):
+ with ib.function("func0", num_inputs=2):
+ ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(2)], dst=ib.r(2))
+ ib.emit_ret(ib.r(2))
+ ib.get()
+
+
+def test_neg_imm():
+ ib = relax.ExecBuilder()
+
+ with ib.function("func0", num_inputs=1):
+ ib.emit_call("test.vm.add_scalar", args=[ib.imm(-3), ib.r(0)], dst=ib.r(1))
+ ib.emit_ret(ib.r(1))
+ ib.get()
+
+ ex = ib.get()
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ assert vm["func0"](1) == -2
+ assert vm["func0"](-3) == -6
+
+
+def test_emit_cache():
+ ib = relax.ExecBuilder()
+
+ with ib.function("func0", num_inputs=1):
+ x0 = ib.convert_constant("str0")
+ x1 = ib.convert_constant("str0")
+ # cache constant str
+ assert x0 == x1
+ s0 = ib.convert_constant(tvm.runtime.container.ShapeTuple([1, 2]))
+ s1 = ib.convert_constant(tvm.runtime.container.ShapeTuple([1, 2]))
+ s2 = ib.convert_constant(tvm.runtime.container.ShapeTuple([1, 3]))
+ assert s0 == s1
+ assert s1 != s2
+ y0 = ib.convert_constant(tvm.nd.array(np.array([1, 2, 3]).astype("int32")))
+ y1 = ib.convert_constant(tvm.nd.array(np.array([1, 2, 3]).astype("int32")))
+ assert y0 == y1
+ ib.emit_ret(ib.r(0))
+
+
+def test_vm_formalize():
+ ib0 = relax.ExecBuilder()
+ ib1 = relax.ExecBuilder()
+ with ib0.function("func0", num_inputs=2):
+ ib0.emit_call("test.vm.add", args=[ib0.r(0), ib0.r(1)], dst=ib0.r(100))
+ ib0.emit_call("test.vm.mul", args=[ib0.r(1), ib0.r(100)], dst=ib0.r(50))
+ ib0.emit_ret(ib0.r(50))
+ with ib1.function("func0", num_inputs=2):
+ ib1.emit_call("test.vm.add", args=[ib1.r(0), ib1.r(1)], dst=ib1.r(2))
+ ib1.emit_call("test.vm.mul", args=[ib1.r(1), ib1.r(2)], dst=ib1.r(3))
+ ib1.emit_ret(ib1.r(3))
+ exec0 = ib0.get()
+ exec1 = ib1.get()
+ assert exec0.as_text() == exec1.as_text()
+
+
+def test_vm_operand():
+ ib0 = relax.ExecBuilder()
+ with ib0.function("func0", num_inputs=2):
+ ib0.emit_call("test.vm.add_scalar", args=[ib0.r(0), ib0.r(1)], dst=ib0.r(2))
+ ib0.emit_ret(ib0.r(2))
+ exec0 = ib0.get()
+ vm = relax.VirtualMachine(exec0, tvm.cpu())
+ res = vm["func0"](2, 3)
+ assert res == 5
+
+ ib1 = relax.ExecBuilder()
+ with ib1.function("func1", num_inputs=1):
+ ib1.emit_call("test.vm.get_device_id", args=[ib1.r(0)], dst=ib1.r(1))
+ ib1.emit_ret(ib1.r(1))
+ exec1 = ib1.get()
+ vm = relax.VirtualMachine(exec1, tvm.cpu())
+ res = vm["func1"](tvm.cpu(3))
+ assert res == 3
+
+
+def test_vm_shapeof():
+ ib = relax.ExecBuilder()
+ shape = (32, 16)
+ arr = tvm.nd.array(np.random.rand(*shape))
+ with ib.function("main", num_inputs=0):
+ ib.emit_call("vm.builtin.shape_of", args=[arr], dst=ib.r(0))
+ ib.emit_ret(ib.r(0))
+ ex = ib.get()
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ res = vm["main"]()
+ for i, s in enumerate(res):
+ assert s == shape[i]
+
+
+def test_vm_storage():
+ dtype = tvm.DataType("float32")
+ shape = (4, 6)
+ ib = relax.ExecBuilder()
+ with ib.function("main", num_inputs=0):
+ ib.emit_call(
+ "vm.builtin.alloc_storage",
+ args=[ib.vm_state(), (24,), ib.convert_constant(0), dtype],
+ dst=ib.r(1),
+ )
+ ib.emit_call(
+ "vm.builtin.alloc_tensor", args=[ib.r(1), ib.imm(0), shape, dtype], dst=ib.r(2)
+ )
+ ib.emit_ret(ib.r(2))
+ ex = ib.get()
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ res = vm["main"]()
+ assert res.device == tvm.cpu()
+ assert res.shape == shape
+
+
+def test_vm_goto():
+ ib = relax.ExecBuilder()
+ with ib.function("main", num_inputs=2):
+ ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(1)], dst=ib.r(2))
+ ib.emit_goto(2)
+ ib.emit_call("test.vm.mul", args=[ib.r(2), ib.r(1)], dst=ib.r(2))
+ ib.emit_ret(ib.r(2))
+ ex = ib.get()
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ a = tvm.nd.array(
+ np.random.rand(
+ 4,
+ )
+ )
+ b = tvm.nd.array(
+ np.random.rand(
+ 4,
+ )
+ )
+ res = check_saved_func(vm, "main", a, b)
+ tvm.testing.assert_allclose(res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)
+
+
+def test_vm_if():
+ ib = relax.ExecBuilder()
+ with ib.function("main", num_inputs=3):
+ ib.emit_if(ib.r(0), 3)
+ ib.emit_call("test.vm.add", args=[ib.r(1), ib.r(2)], dst=ib.r(3))
+ ib.emit_goto(2)
+ ib.emit_call("test.vm.mul", args=[ib.r(1), ib.r(2)], dst=ib.r(3))
+ ib.emit_ret(ib.r(3))
+ ex = ib.get()
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ a = tvm.nd.array(
+ np.random.rand(
+ 4,
+ )
+ )
+ b = tvm.nd.array(
+ np.random.rand(
+ 4,
+ )
+ )
+ res = vm["main"](0, a, b)
+ tvm.testing.assert_allclose(res.numpy(), a.numpy() * b.numpy(), rtol=1e-7, atol=1e-7)
+ res = vm["main"](1, a, b)
+ tvm.testing.assert_allclose(res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)
+
+
+def test_vm_invoke_closure():
+ ib = relax.ExecBuilder()
+ with ib.function("lifted_func_1", num_inputs=4):
+ ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(1)], dst=ib.r(4))
+ ib.emit_call("test.vm.add", args=[ib.r(2), ib.r(4)], dst=ib.r(5))
+ ib.emit_call("test.vm.add", args=[ib.r(3), ib.r(5)], dst=ib.r(6))
+ ib.emit_ret(ib.r(6))
+ with ib.function("main", num_inputs=2):
+ ib.emit_call(
+ "vm.builtin.make_closure", args=[ib.f("lifted_func_1"), ib.r(0), ib.r(1)], dst=ib.r(2)
+ )
+ ib.emit_ret(ib.r(2))
+
+ ex = ib.get()
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ w_inp = tvm.nd.array(np.random.rand(2, 3))
+ x_inp = tvm.nd.array(np.random.rand(2, 3))
+ y_inp = tvm.nd.array([[3.1, 4.0, 5.0], [6.0, 7.1, 9.0]])
+ z_inp = tvm.nd.array(np.random.rand(2, 3))
+ clo = vm["main"](w_inp, x_inp)
+ res = vm.invoke_closure(clo, y_inp, z_inp)
+ tvm.testing.assert_allclose(
+ res.numpy(), w_inp.numpy() + x_inp.numpy() + y_inp.numpy() + z_inp.numpy()
+ )
+
+
+if __name__ == "__main__":
+ tvm.testing.main()