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_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 {
 
using ExecWord = int64_t;
 
using RegName = ExecWord;
 
using Index = ExecWord;
 
enum class Opcode {
  Call = 1U,
  Ret = 2U,
  Goto = 3U,
  If = 4U,
};
 
struct Instruction {
  static constexpr ExecWord kKindBit = 8;
  static constexpr ExecWord kValueBit = sizeof(ExecWord) * 8 - kKindBit;
  static constexpr ExecWord kValueMask = (static_cast<ExecWord>(1) << kValueBit) - 1;
  static constexpr ExecWord kValueMaxLimit = (static_cast<ExecWord>(1) << (kValueBit - 1)) - 1;
  static constexpr ExecWord kValueMinLimit = -kValueMaxLimit;
  static constexpr RegName kBeginSpecialReg = static_cast<ExecWord>(1) << 54;
  static constexpr RegName kVoidRegister = kBeginSpecialReg + 0;
  static constexpr RegName kVMRegister = kBeginSpecialReg + 1;
  enum class ArgKind : int { kRegister = 0, kImmediate = 1, kConstIdx = 2, kFuncIdx = 3 };
 
  friend std::ostream& operator<<(std::ostream& os, const ArgKind& kind) {
    switch (kind) {
      case ArgKind::kRegister:
        os << "kRegister";
        break;
      case ArgKind::kImmediate:
        os << "kImmediate";
        break;
      case ArgKind::kConstIdx:
        os << "kConstIdx";
        break;
      case ArgKind::kFuncIdx:
        os << "kFuncIdx";
        break;
      default:
        LOG(FATAL) << "Internal error: "
                   << "Invalid ArgKind with integer value " << static_cast<int>(kind);
    }
    return os;
  }
 
  class Arg {
   public:
    Arg() : data_(Instruction::kVoidRegister) {}
    static Arg FromData(ExecWord data) { return Arg(data); }
    static Arg Register(RegName reg) { return Arg(ArgKind::kRegister, reg); }
    static Arg ConstIdx(Index index) { return Arg(ArgKind::kConstIdx, index); }
    static Arg Immediate(int64_t imm_value) { return Arg(ArgKind::kImmediate, imm_value); }
    static Arg FuncIdx(Index index) { return Arg(ArgKind::kFuncIdx, index); }
    ArgKind kind() const {
      uint8_t kind = (data_ >> kValueBit) & 0xFF;
      return Instruction::ArgKind(kind);
    }
    ExecWord value() const { return ((data_ & kValueMask) << kKindBit) >> kKindBit; }
    ExecWord data() const { return data_; }
 
   private:
    explicit Arg(ExecWord data) : data_(data) {}
    Arg(ArgKind kind, Index value) {
      ICHECK_LE(value, kValueMaxLimit);
      ICHECK_GE(value, kValueMinLimit);
      data_ = (static_cast<ExecWord>(kind) << kValueBit) | (value & kValueMask);
    }
    ExecWord data_;
  };
  Opcode op;
  union {
    struct /* Call */ {
      RegName dst;
      Index func_idx;
      Index num_args;
      Arg* args;
    };
    struct /* Ret */ {
      RegName result;
    };
    struct /* Goto */ {
      Index pc_offset;
    };
    struct /* If */ {
      RegName cond;
      Index false_offset;
    };
  };
  static Instruction Call(Index func_idx, Index num_args, Arg* args, RegName dst);
  static Instruction Ret(RegName result);
  static Instruction Goto(RegName pc_offset);
  static Instruction If(RegName cond, Index false_offset);
};
 
}  // namespace relax_vm
}  // namespace runtime
}  // namespace tvm
 
#endif  // TVM_RUNTIME_RELAX_VM_BYTECODE_H_