Relax Virtual Machine

Co-Authored-By: Yuchen Jin <[email protected]>

.github/workflows/main.yml

@@ -1,3 +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.
+
 # GH actions.
 # We use it to cover windows and mac builds
 # Jenkins is still the primary CI
@@ -62,8 +79,7 @@ jobs:
           python -m pytest -v tests/python/contrib/test_rpc_server_device.py
 
   Windows:
-    if: ${{ github.repository == 'apache/tvm' }}
-    runs-on: windows-2019
+    runs-on: windows-2016
     steps:
       - uses: actions/checkout@v2
         with:

CMakeLists.txt

@@ -291,6 +291,7 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
     src/printer/*.cc
     src/support/*.cc
     src/script/*.cc
+    src/relax/*.cc
     )
 
 tvm_file_glob(GLOB CODEGEN_SRCS

include/tvm/relax/builder.h

@@ -0,0 +1,105 @@
+/*
+ * 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/builder.h
+ * \brief
+ */
+#ifndef TVM_RELAX_BUILDER_H_
+#define TVM_RELAX_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 "./vm/bytecode.h"
+#include "./vm/executable.h"
+
+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 The mutable internal executable node. */
+  ObjectPtr<vm::ExecutableNode> exec;  // mutable
+  /*!
+   * \brief To annotate the start of a vm function.
+   * \param func The function name.
+   * \param num_inputs The number of inputs.
+   */
+  void Function(std::string func, int64_t num_inputs);
+  /*!
+   * \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(std::string func, std::vector<vm::Instruction::Arg> args, vm::RegName ret);
+  /*!
+   * \brief Emit a ret instruction.
+   * \param result The return result.
+   */
+  void EmitRet(vm::RegName result);
+  /*!
+   * \brief Emit a constant value to the constant pool.
+   * \return The index that represents the constant.
+   */
+  vm::Index EmitConstant(ObjectRef obj);
+  /*!
+   * \brief Get the built executable.
+   * \return The built executable.
+   */
+  vm::Executable Get();
+  /*!
+   * \brief Create a 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 Formalize the executable.
+   */
+  void Formalize();
+};
+
+class ExecBuilder : public ObjectRef {
+ public:
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ExecBuilder, ObjectRef, ExecBuilderNode);
+};
+
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_BUILDER_H_

include/tvm/relax/vm/bytecode.h

@@ -0,0 +1,162 @@
+/*
+ * 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/vm/bytecode.h
+ * \brief The bytecode for the virtual machine.
+ */
+#ifndef TVM_RELAX_VM_BYTECODE_H_
+#define TVM_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,
+};
+
+
+/*! \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 Random magic number that represents void argument. */
+  static constexpr RegName kVoidArg = 0x00EC66FE0321975A;
+  /*! \brief Random magic number that represents the VM state. */
+  static constexpr RegName kVMStateRegister = 0x008D14FA4379015C;
+  /*!
+   * \brief The kind of instruction's argument.
+   */
+  enum ArgKind {
+    kRegister = 0,
+    kImmediate = 1,
+    kConstIdx = 2,
+  };
+  /*!
+   * \brief The auxiliary data structure for instruction argument.
+   */
+  struct Arg {
+    /*! \brief The number of bit for storing value. */
+    static constexpr ExecWord kValueBit = sizeof(ExecWord) * 8 - 8;
+    /*! \brief The bit mask of the value part. */
+    static constexpr ExecWord kValueMask = (static_cast<ExecWord>(1) << kValueBit) - 1;
+    /*! \brief Construct a void argument. */
+    explicit Arg() : data(Instruction::kVoidArg) {}
+    /*! \brief Construct from the data. */
+    explicit Arg(ExecWord data) : data(data) {}
+    /*! \brief Construct from the kind and value. */
+    Arg(ArgKind kind, Index value) {
+      // TODO(ziheng): check value?
+      this->data = (static_cast<ExecWord>(kind) << kValueBit) |
+                   (value & kValueMask);
+    }
+    /*!
+     * \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.
+     */
+    ExecWord value() const {
+      return data & ((static_cast<ExecWord>(1) << kValueBit) - 1);
+    }
+    /*! \brief The underlying stored data. */
+    ExecWord data;
+  };
+  /*! \brief The instruction opcode. */
+  Opcode op;
+  /*! \brief The destination register. */
+  RegName dst; 
+  union {
+    struct /* Call */ {
+      /*! \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;
+    };
+  };
+  /*!
+   * \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);
+};
+
+}  // namespace relax_vm
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RELAX_VM_BYTECODE_H_