bytecode.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#ifndef TVM_RUNTIME_VM_BYTECODE_H_
#define TVM_RUNTIME_VM_BYTECODE_H_

#include <tvm/runtime/data_type.h>
#include <tvm/runtime/logging.h>

#include <iostream>
#include <vector>

namespace tvm {
namespace runtime {
namespace vm {

using RegName = int64_t;

using Index = int64_t;

enum class Opcode {
  Move = 0U,
  Ret = 1U,
  Invoke = 2U,
  InvokeClosure = 3U,
  InvokePacked = 4U,
  AllocTensor = 5U,
  AllocTensorReg = 6U,
  AllocADT = 7U,
  AllocClosure = 8U,
  GetField = 9U,
  If = 10U,
  LoadConst = 11U,
  Goto = 12U,
  GetTag = 13U,
  LoadConsti = 14U,
  Fatal = 15U,
  AllocStorage = 16U,
  ShapeOf = 17U,
  ReshapeTensor = 18U,
  DeviceCopy = 19U,
  KillRegister = 20U,
};

struct Instruction {
  Opcode op;

  RegName dst;

  union {
    struct /* AllocTensor Operands */ {
      RegName storage;
      Index offset;
      uint32_t ndim;
      int64_t* shape;
      DLDataType dtype;
    } alloc_tensor;
    struct /* AllocTensorReg Operands */ {
      RegName storage;
      Index offset;
      RegName shape_register;
      DLDataType dtype;
    } alloc_tensor_reg;
    struct /* InvokeClosure Operands */ {
      RegName closure;
      Index num_closure_args;
      RegName* closure_args;
    };
    struct /* Return Operands */ {
      RegName result;
    };
    struct /* Move Operands */ {
      RegName from;
    };
    struct /* InvokePacked Operands */ {
      Index packed_index;
      Index arity;
      Index output_size;
      RegName* packed_args;
    };
    struct /* If Operands */ {
      RegName test;
      RegName target;
      Index true_offset;
      Index false_offset;
    } if_op;
    struct /* Invoke Operands */ {
      Index func_index;
      Index num_args;
      RegName* invoke_args_registers;
    };
    struct /* LoadConst Operands */ {
      /* \brief The index into the constant pool. */
      Index const_index;
    };
    struct /* LoadConsti Operands */ {
      /* \brief The index into the constant pool. */
      Index val;
    } load_consti;
    struct /* Jump Operands */ {
      Index pc_offset;
    };
    struct /* Proj Operands */ {
      RegName object;
      Index field_index;
    };
    struct /* GetTag Operands */ {
      RegName object;
    } get_tag;
    struct /* AllocADT Operands */ {
      // TODO(mbs): Needs a DeviceAndScope.
      Index constructor_tag;
      Index num_fields;
      RegName* datatype_fields;
    };
    struct /* AllocClosure Operands */ {
      // TODO(mbs): Needs a DeviceAndScope.
      Index clo_index;
      Index num_freevar;
      RegName* free_vars;
    };
    struct /* AllocStorage Operands */ {
      RegName allocation_size;
      Index alignment;
      DLDataType dtype_hint;
      Index device_index;
    } alloc_storage;
    struct /* ShapeOf Operands */ {
      RegName tensor;
    } shape_of;
    struct /* ReshapeTensor Operands */ {
      RegName tensor;
      RegName newshape;
    } reshape_tensor;
    struct /* DeviceCopy Operands */ {
      RegName src;
      Index src_device_index;
      Index dst_device_index;
    } device_copy;
  };

  static Instruction Ret(RegName return_reg);
  static Instruction Fatal();
  static Instruction InvokePacked(Index packed_index, Index arity, Index output_size,
                                  const std::vector<RegName>& args);
  static Instruction AllocTensor(RegName storage, Index offset, const std::vector<int64_t>& shape,
                                 DLDataType dtype, RegName dst);
  static Instruction AllocTensorReg(RegName storage, Index offset, RegName shape_register,
                                    DLDataType dtype, RegName dst);
  static Instruction AllocADT(Index tag, Index num_fields, const std::vector<RegName>& fields,
                              RegName dst);
  static Instruction AllocClosure(Index func_index, Index num_freevar,
                                  const std::vector<RegName>& free_vars, RegName dst);
  static Instruction GetField(RegName object_reg, Index field_index, RegName dst);
  static Instruction GetTag(RegName object_reg, RegName dst);
  static Instruction If(RegName test, RegName target, Index true_branch, Index false_branch);
  static Instruction Goto(Index pc_offset);
  static Instruction Invoke(Index func_index, const std::vector<RegName>& args, RegName dst);
  static Instruction InvokeClosure(RegName closure, const std::vector<RegName>& args, RegName dst);
  static Instruction LoadConst(Index const_index, RegName dst);
  static Instruction LoadConsti(Index val, RegName dst);
  static Instruction Move(RegName src, RegName dst);
  static Instruction AllocStorage(RegName size, Index alignment, DLDataType dtype_hint,
                                  Index device_index, RegName dst);
  static Instruction ShapeOf(RegName tensor, RegName dst);
  static Instruction ReshapeTensor(RegName tensor, RegName newshape, RegName dst);
  static Instruction DeviceCopy(RegName src, Index src_device_index, Index dst_device_index,
                                RegName dst);

  static Instruction KillRegister(RegName dst);

  Instruction();
  Instruction(const Instruction& instr);
  Instruction& operator=(const Instruction& instr);
  ~Instruction();

  friend std::ostream& operator<<(std::ostream& os, const Instruction&);
};

}  // namespace vm
}  // namespace runtime
}  // namespace tvm

#endif  // TVM_RUNTIME_VM_BYTECODE_H_