[TRAINING/VM] Instruction, Builder and Executable

CMakeLists.txt

@@ -233,6 +233,7 @@ file(GLOB_RECURSE COMPILER_SRCS
     src/parser/*.cc
     src/printer/*.cc
     src/support/*.cc
+    src/relax/*.cc
     )
 
 file(GLOB CODEGEN_SRCS

include/tvm/relax/builder.h

@@ -0,0 +1,69 @@
+/*
+ * 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/runtime/object.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/node/reflection.h>
+#include <tvm/node/repr_printer.h>
+#include <tvm/ir/expr.h>
+#include "./vm/bytecode.h"
+#include "./vm/executable.h"
+
+namespace tvm {
+namespace relax {
+
+class Builder;
+
+class BuilderNode : public Object {
+ public:
+  ObjectPtr<vm::ExecutableNode> exec; // mutable
+
+  void EmitCall(std::string func, std::vector<vm::InstrArg> args, vm::RegName ret);
+
+  vm::Index EmitConstant(ObjectRef obj);
+
+  vm::Executable Get();
+
+  TVM_DLL static Builder Create();
+
+  void VisitAttrs(AttrVisitor* v) {
+  }
+
+  static constexpr const uint32_t _type_index = TypeIndex::kDynamic;
+  static constexpr const char* _type_key = "relax.Builder"; 
+  TVM_DECLARE_FINAL_OBJECT_INFO(BuilderNode, Object);
+};
+
+class Builder : public ObjectRef {
+ public:
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Builder, ObjectRef, BuilderNode);
+};
+
+
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_BUILDER_H_

include/tvm/relax/vm/bytecode.h

@@ -0,0 +1,125 @@
+/*
+ * 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 relax {
+namespace vm {
+
+
+/*! \brief A register name. */
+using RegName = int64_t;
+
+/*! \brief An alias for the integer type used ubiquitously
+ * in the VM.
+ */
+using Index = int64_t;
+
+using ExecWord = int64_t; 
+
+
+enum ArgKind {
+  kRegister = 0,
+  kImmediate = 1,
+  kConstIdx = 2,
+};
+
+constexpr int64_t kVoidArg = 0xFE0321975A;
+
+struct InstrArg {
+  explicit InstrArg() : data(kVoidArg) {}
+  explicit InstrArg(int64_t data) : data(data) {}
+  InstrArg(ArgKind kind, Index value) {
+    // TODO(ziheng): check value
+    this->data = (uint64_t(kind) << 56) | (value & ((uint64_t(1) << 56)  - 1));
+  }
+  ArgKind kind() {
+    uint8_t kind = (data >> 56) & 0xFF;
+    return ArgKind(kind);
+  }
+  int64_t value() {
+    return data & ((int64_t(1) << 56) - 1);
+  }
+  int64_t data;
+};
+
+/*! \brief An enumeration of Relay's opcodes.
+ *
+ * The opcode is used to implement instruction
+ * as a tagged union.
+ */
+enum class Opcode {
+  Call = 1U,
+};
+
+/*! \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 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;
+
+      ExecWord* args;
+    };
+  };
+
+  static Instruction Call(Index func_idx, Index num_args,
+                          ExecWord* args,
+                          RegName dst);
+
+  Instruction();
+  Instruction(const Instruction& instr);
+  Instruction& operator=(const Instruction& instr);
+  ~Instruction();
+};
+
+}  // namespace vm
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_VM_BYTECODE_H_

include/tvm/relax/vm/executable.h

@@ -0,0 +1,93 @@
+/*
+ * 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/vm/executable.h
+ * \brief 
+ */
+#ifndef TVM_RELAX_VM_EXECUTABLE_H_
+#define TVM_RELAX_VM_EXECUTABLE_H_
+
+#include <tvm/runtime/object.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/node/reflection.h>
+#include <tvm/node/repr_printer.h>
+#include <tvm/ir/expr.h>
+#include "./bytecode.h"
+
+#include <sstream>
+
+namespace tvm {
+namespace relax {
+namespace vm {
+
+class Executable;
+
+class ExecutableNode : public Object {
+ public:
+  std::vector<ExecWord> instr_data;
+  // std::vector<ExecWord> arg_data;
+  std::vector<size_t> instr_offset;
+  std::vector<ObjectRef> constants;
+  std::unordered_map<std::string, Index> func2idx;
+  std::vector<std::string> func_names;
+
+  // magic number, version,  
+  // SaveToBinary(dmlc::Stream* stream);
+  // SaveToFile(const std::string& path, const std::string& format);
+
+  Instruction GetInstruction(size_t i) const {
+    size_t 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];
+        const ExecWord* args = &instr_data[offset + 4];
+        return Instruction::Call(func_idx, num_args, const_cast<ExecWord*>(args), dst);
+      }
+      default:
+        LOG(FATAL) << "should never hit this case: " << static_cast<int>(op);
+        break;
+    }
+    return Instruction();
+  }
+
+  String AsText() const;
+
+  void VisitAttrs(AttrVisitor* v) {
+  }
+
+  static constexpr const uint32_t _type_index = TypeIndex::kDynamic;
+  static constexpr const char* _type_key = "relax.Executable"; 
+  TVM_DECLARE_FINAL_OBJECT_INFO(ExecutableNode, Object);
+};
+
+class Executable : public ObjectRef {
+ public:
+  TVM_DEFINE_OBJECT_REF_METHODS(Executable, ObjectRef, ExecutableNode);
+};
+
+
+}  // namespace vm
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_VM_EXECUTABLE_H_

python/tvm/relax/__init__.py

@@ -0,0 +1 @@
+from .builder import Builder

python/tvm/relax/_ffi_api.py

@@ -0,0 +1,3 @@
+import tvm._ffi
+
+tvm._ffi._init_api("relax", __name__)

python/tvm/relax/builder.py

@@ -0,0 +1,71 @@
+# 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.
+
+from enum import IntEnum
+import tvm
+from tvm.runtime import Object
+from tvm._ffi.base import _LIB, check_call
+from . import _ffi_api
+
+class ArgKind(IntEnum):
+    REGISTER = 0
+    IMMEDIATE = 1
+    CONSTIDX = 2
+
+VOID_ARG_ = 0xFE0321975A
+
+def create_arg(kind, value):
+    return (int(kind) << 56) | (value & ((1 << 56) - 1))
+
+@tvm._ffi.register_object("relax.Executable")
+class Executable(Object):
+    def __init__(self):
+        self.__init_handle_by_constructor__(_ffi_api.Executable)
+
+    def astext(self):
+        return _ffi_api.ExecutableAsText(self) 
+
+
+@tvm._ffi.register_object("relax.Builder")
+class Builder(Object):
+    def __init__(self):
+        self.__init_handle_by_constructor__(_ffi_api.BuilderCreate)
+
+    def r(self, idx):
+        return create_arg(ArgKind.REGISTER, idx)
+
+    def imm(self, value):
+        return create_arg(ArgKind.IMMEDIATE, value)
+
+    def const(self, idx):
+        return create_arg(ArgKind.CONSTIDX, idx)
+
+    def emit_call(self, name, args=[], ret=None):
+        if ret is None:
+            ret = VOID_ARG_
+        args_ = []
+        for arg in args:
+            if isinstance(arg, tvm.nd.NDArray):
+                new_arg = _ffi_api.BuilderEmitConstant(self, arg)
+                args_.append(new_arg)
+            else:
+                args_.append(arg)
+        _ffi_api.BuilderEmitCall(self, name, args_, ret)
+
+    def get(self):
+        return _ffi_api.BuilderGet(self)
+