| /* |
| * 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. |
| */ |
| |
| /*! |
| * Implementation stack VM. |
| * \file stackvm.cc |
| */ |
| #include "stackvm.h" |
| |
| #include <dmlc/thread_local.h> |
| #include <tvm/runtime/c_backend_api.h> |
| |
| #include <algorithm> |
| |
| namespace tvm { |
| namespace runtime { |
| |
| typedef dmlc::ThreadLocalStore<StackVM::State> StackVMStateStore; |
| |
| StackVM::State* StackVM::ThreadLocalState() { return StackVMStateStore::Get(); } |
| |
| #define STACK_VM_BINOP(OP, FIELD) \ |
| { \ |
| stack[sp - 1].FIELD = stack[sp - 1].FIELD OP stack[sp].FIELD; \ |
| sp -= 1; \ |
| pc += 1; \ |
| } |
| |
| #define STACK_VM_CMPOP(OP, FIELD) \ |
| { \ |
| stack[sp - 1].v_int64 = stack[sp - 1].FIELD OP stack[sp].FIELD; \ |
| sp -= 1; \ |
| pc += 1; \ |
| } |
| |
| #define STACK_VM_LOAD(FIELD, DST_TYPE, SRC_TYPE) \ |
| { \ |
| int index = code[pc + 1].v_int; \ |
| stack[sp] FIELD = static_cast<DST_TYPE>(static_cast<SRC_TYPE*>(stack[sp].v_handle)[index]); \ |
| pc += 2; \ |
| } |
| |
| #define STACK_VM_STORE(FIELD, DST_TYPE) \ |
| { \ |
| int index = code[pc + 1].v_int; \ |
| static_cast<DST_TYPE*>(stack[sp - 1].v_handle)[index] = \ |
| static_cast<DST_TYPE>(stack[sp] FIELD); \ |
| sp -= 2; \ |
| pc += 2; \ |
| } |
| |
| #define STACK_VM_PRINT_CODE0(CODE) \ |
| case CODE: { \ |
| os << "[" << pc << "]\t" << #CODE << std::endl; \ |
| return pc + 1; \ |
| } |
| |
| #define STACK_VM_PRINT_CODE1(CODE) \ |
| case CODE: { \ |
| os << "[" << pc << "]\t" << #CODE << " " << code[pc + 1].v_int << "\n" \ |
| << "[" << pc + 1 << "]" << std::endl; \ |
| return pc + 2; \ |
| } |
| |
| #define STACK_VM_PRINT_CODE2(CODE) \ |
| case CODE: { \ |
| os << "[" << pc << "]\t" << #CODE << " " << code[pc + 1].v_int << " " << code[pc + 2].v_int \ |
| << "\n" \ |
| << "[" << pc + 1 << "]" << std::endl \ |
| << "[" << pc + 2 << "]" << std::endl; \ |
| return pc + 3; \ |
| } |
| |
| #define STACK_VM_PRINT_HEAP_ACCESS(CODE) \ |
| case CODE: { \ |
| os << "[" << pc << "]\t" << #CODE << " " << code[pc + 1].v_int << " " \ |
| << heap_id_name[code[pc + 1].v_int] << "\n" \ |
| << "[" << pc + 1 << "]" << std::endl; \ |
| return pc + 2; \ |
| } |
| |
| #define STACK_VM_PRINT_JUMP(CODE) \ |
| case CODE: { \ |
| os << "[" << pc << "]\t" << #CODE << " rel=" << code[pc + 1].v_int << " to " \ |
| << pc + code[pc + 1].v_int << '\n' \ |
| << "[" << pc + 1 << "]" << std::endl; \ |
| return pc + 2; \ |
| } |
| |
| int64_t StackVM::PrintCode(std::ostream& os, int64_t pc) const { |
| switch (code[pc].op_code) { |
| // int |
| STACK_VM_PRINT_CODE0(ADD_I64); |
| STACK_VM_PRINT_CODE0(SUB_I64); |
| STACK_VM_PRINT_CODE0(MUL_I64); |
| STACK_VM_PRINT_CODE0(MOD_I64); |
| STACK_VM_PRINT_CODE0(DIV_I64); |
| STACK_VM_PRINT_CODE0(EQ_I64); |
| STACK_VM_PRINT_CODE0(LT_I64); |
| STACK_VM_PRINT_CODE0(LE_I64); |
| // floats |
| STACK_VM_PRINT_CODE0(ADD_F64); |
| STACK_VM_PRINT_CODE0(SUB_F64); |
| STACK_VM_PRINT_CODE0(MUL_F64); |
| STACK_VM_PRINT_CODE0(DIV_F64); |
| STACK_VM_PRINT_CODE0(EQ_F64); |
| STACK_VM_PRINT_CODE0(LT_F64); |
| STACK_VM_PRINT_CODE0(LE_F64); |
| // handle. |
| STACK_VM_PRINT_CODE0(EQ_HANDLE); |
| // addressing load |
| STACK_VM_PRINT_CODE1(ARRAY_LOAD_UINT32); |
| STACK_VM_PRINT_CODE1(ARRAY_LOAD_INT32); |
| STACK_VM_PRINT_CODE1(ARRAY_LOAD_INT64); |
| STACK_VM_PRINT_CODE1(ARRAY_LOAD_FP64); |
| STACK_VM_PRINT_CODE1(ARRAY_LOAD_HANDLE); |
| STACK_VM_PRINT_CODE1(ARRAY_LOAD_TVMVALUE); |
| STACK_VM_PRINT_CODE1(ARRAY_STORE_UINT32); |
| STACK_VM_PRINT_CODE1(ARRAY_STORE_INT32); |
| STACK_VM_PRINT_CODE1(ARRAY_STORE_INT64); |
| STACK_VM_PRINT_CODE1(ARRAY_STORE_FP64); |
| STACK_VM_PRINT_CODE1(ARRAY_STORE_HANDLE); |
| STACK_VM_PRINT_CODE1(ARRAY_STORE_TVMVALUE); |
| STACK_VM_PRINT_CODE0(NOT); |
| STACK_VM_PRINT_CODE0(ADDR_ADD); |
| // stack ops |
| STACK_VM_PRINT_CODE1(PUSH_I64); |
| STACK_VM_PRINT_CODE1(PUSH_VALUE); |
| STACK_VM_PRINT_CODE0(POP); |
| STACK_VM_PRINT_CODE0(SELECT); |
| STACK_VM_PRINT_HEAP_ACCESS(STORE_HEAP); |
| STACK_VM_PRINT_HEAP_ACCESS(LOAD_HEAP); |
| STACK_VM_PRINT_CODE1(ASSERT); |
| STACK_VM_PRINT_JUMP(RJUMP_IF_TRUE); |
| STACK_VM_PRINT_JUMP(RJUMP_IF_FALSE); |
| STACK_VM_PRINT_JUMP(RJUMP); |
| STACK_VM_PRINT_CODE1(ASSERT_SP); |
| // Intrinsics |
| STACK_VM_PRINT_CODE2(TVM_STRUCT_GET); |
| STACK_VM_PRINT_CODE2(TVM_STRUCT_SET); |
| // Allocate data by 8 bytes. |
| STACK_VM_PRINT_CODE1(TVM_STACK_ALLOCA_BY_8BYTE); |
| STACK_VM_PRINT_CODE0(TVM_DEVICE_ALLOCA); |
| STACK_VM_PRINT_CODE0(TVM_DEVICE_FREE); |
| STACK_VM_PRINT_CODE0(TVM_THROW_LAST_ERROR); |
| // packed function. |
| case CALL_PACKED_LOWERED: { |
| int call_fid = code[pc + 1].v_int; |
| int begin = code[pc + 2].v_int; |
| int end = code[pc + 3].v_int; |
| os << "[" << pc << "]\tCALL_PACKED_FUNC " |
| << " fid=" << call_fid << " begin=" << begin << " end=" << end; |
| os << '\n'; |
| for (int i = 0; i < 3; ++i) { |
| os << "[" << pc + 1 + i << "]" << std::endl; |
| } |
| return pc + 4; |
| } |
| } |
| LOG(FATAL) << "unknown op code " << code[pc].op_code; |
| } |
| |
| std::ostream& operator<<(std::ostream& os, const StackVM& vm) { // NOLINT(*) |
| int64_t pc = 0; |
| const int64_t code_size = static_cast<int64_t>(vm.code.size()); |
| os << "Program dump: code-size=" << code_size << '\n' << "----------begin-----------------\n"; |
| while (pc < code_size) { |
| pc = vm.PrintCode(os, pc); |
| } |
| os << "----------end--------------------\n"; |
| return os; |
| } |
| |
| void StackVM::Run(const runtime::TVMArgs& args, runtime::ModuleNode* mod_ctx) const { |
| StackVM::State* s = StackVM::ThreadLocalState(); |
| if (s->heap.size() < heap_size) { |
| s->heap.resize(heap_size); |
| } |
| s->sp = 0; |
| s->pc = 0; |
| s->mod_ctx = mod_ctx; |
| s->heap[0].v_handle = (void*)args.values; // NOLINT(*) |
| s->heap[1].v_handle = (void*)args.type_codes; // NOLINT(*) |
| s->heap[2].v_int64 = args.num_args; |
| this->Run(s); |
| } |
| |
| void StackVM::InitCache() { |
| extern_func_cache_.clear(); |
| extern_func_cache_.resize(extern_func_name.size(), PackedFunc(nullptr)); |
| } |
| |
| void StackVM::Save(dmlc::Stream* strm) const { |
| // to be endian invariant. |
| std::vector<int32_t> code_copy(code.size()); |
| std::transform(code.begin(), code.end(), code_copy.begin(), [](Code c) { return c.v_int; }); |
| strm->Write(code_copy); |
| strm->Write(str_data); |
| strm->Write(extern_func_name); |
| strm->Write(heap_id_name); |
| strm->Write(heap_size); |
| strm->Write(stack_size); |
| } |
| |
| bool StackVM::Load(dmlc::Stream* strm) { |
| // to be endian invariant. |
| std::vector<int32_t> code_copy; |
| if (!strm->Read(&code_copy)) return false; |
| code.resize(code_copy.size()); |
| std::transform(code_copy.begin(), code_copy.end(), code.begin(), [](int v) { |
| Code code; |
| code.v_int = v; |
| return code; |
| }); |
| if (!strm->Read(&str_data)) return false; |
| if (!strm->Read(&extern_func_name)) return false; |
| if (!strm->Read(&heap_id_name)) return false; |
| if (!strm->Read(&heap_size)) return false; |
| if (!strm->Read(&stack_size)) return false; |
| this->InitCache(); |
| return true; |
| } |
| |
| void StackVM::Run(State* s) const { |
| int64_t sp = s->sp; |
| int64_t pc = s->pc; |
| int64_t alloca_sp = s->sp; |
| std::vector<TVMValue>& stack = s->stack; |
| std::vector<TVMValue>& heap = s->heap; |
| if (stack.size() < stack_size) { |
| stack.resize(stack_size); |
| } |
| int64_t stack_cap = static_cast<int64_t>(stack_size - 4); |
| if (heap.size() < heap_size) { |
| heap.resize(heap_size); |
| } |
| const int64_t code_size = static_cast<int64_t>(code.size()); |
| while (pc < code_size) { |
| switch (code[pc].op_code) { |
| case ADD_I64: |
| STACK_VM_BINOP(+, v_int64); |
| break; |
| case SUB_I64: |
| STACK_VM_BINOP(-, v_int64); |
| break; |
| case MUL_I64: |
| STACK_VM_BINOP(*, v_int64); |
| break; |
| case DIV_I64: |
| STACK_VM_BINOP(/, v_int64); |
| break; |
| case MOD_I64: |
| STACK_VM_BINOP(%, v_int64); |
| break; |
| case EQ_I64: |
| STACK_VM_CMPOP(==, v_int64); |
| break; |
| case LT_I64: |
| STACK_VM_CMPOP(<, v_int64); |
| break; |
| case LE_I64: |
| STACK_VM_CMPOP(<=, v_int64); |
| break; |
| case ADD_F64: |
| STACK_VM_BINOP(+, v_float64); |
| break; |
| case SUB_F64: |
| STACK_VM_BINOP(-, v_float64); |
| break; |
| case MUL_F64: |
| STACK_VM_BINOP(*, v_float64); |
| break; |
| case DIV_F64: |
| STACK_VM_BINOP(/, v_float64); |
| break; |
| case EQ_F64: |
| STACK_VM_CMPOP(==, v_float64); |
| break; |
| case LT_F64: |
| STACK_VM_CMPOP(<, v_float64); |
| break; |
| case LE_F64: |
| STACK_VM_CMPOP(<=, v_float64); |
| break; |
| case EQ_HANDLE: |
| STACK_VM_CMPOP(==, v_handle); |
| break; |
| // addressing |
| case ARRAY_LOAD_UINT32: |
| STACK_VM_LOAD(.v_int64, int64_t, uint32_t); |
| break; |
| case ARRAY_LOAD_INT32: |
| STACK_VM_LOAD(.v_int64, int64_t, int32_t); |
| break; |
| case ARRAY_LOAD_INT64: |
| STACK_VM_LOAD(.v_int64, int64_t, int64_t); |
| break; |
| case ARRAY_LOAD_FP64: |
| STACK_VM_LOAD(.v_float64, double, double); |
| break; |
| case ARRAY_LOAD_HANDLE: |
| STACK_VM_LOAD(.v_handle, void*, void*); |
| break; |
| case ARRAY_LOAD_TVMVALUE: |
| STACK_VM_LOAD(, TVMValue, TVMValue); |
| break; |
| // store |
| case ARRAY_STORE_UINT32: |
| STACK_VM_STORE(.v_int64, uint32_t); |
| break; |
| case ARRAY_STORE_INT32: |
| STACK_VM_STORE(.v_int64, int32_t); |
| break; |
| case ARRAY_STORE_INT64: |
| STACK_VM_STORE(.v_int64, int64_t); |
| break; |
| case ARRAY_STORE_FP64: |
| STACK_VM_STORE(.v_float64, double); |
| break; |
| case ARRAY_STORE_HANDLE: |
| STACK_VM_STORE(.v_handle, void*); |
| break; |
| case ARRAY_STORE_TVMVALUE: |
| STACK_VM_STORE(, TVMValue); |
| break; |
| // add |
| case ADDR_ADD: { |
| stack[sp - 1].v_handle = (char*)(stack[sp - 1].v_handle) + stack[sp].v_int64; // NOLINT(*) |
| sp = sp - 1; |
| pc = pc + 1; |
| break; |
| } |
| case NOT: { |
| stack[sp].v_int64 = !stack[sp].v_int64; |
| pc += 1; |
| break; |
| } |
| case PUSH_I64: { |
| stack[sp + 1].v_int64 = code[pc + 1].v_int; |
| sp += 1; |
| pc += 2; |
| break; |
| } |
| case PUSH_VALUE: { |
| int relpos = code[pc + 1].v_int; |
| ICHECK_LE(relpos, 0); |
| stack[sp + 1] = stack[sp + relpos]; |
| sp += 1; |
| pc += 2; |
| break; |
| } |
| case POP: { |
| sp -= 1; |
| pc += 1; |
| break; |
| } |
| case SELECT: { |
| stack[sp - 2] = (stack[sp].v_int64 ? stack[sp - 2] : stack[sp - 1]); |
| sp -= 2; |
| pc += 1; |
| break; |
| } |
| case LOAD_HEAP: { |
| stack[sp + 1] = heap[code[pc + 1].v_int]; |
| sp += 1; |
| pc += 2; |
| break; |
| } |
| case STORE_HEAP: { |
| heap[code[pc + 1].v_int] = stack[sp]; |
| sp -= 1; |
| pc += 2; |
| break; |
| } |
| case ASSERT: { |
| ICHECK(stack[sp].v_int64) << str_data[code[pc + 1].v_int]; |
| sp -= 1; |
| pc += 2; |
| break; |
| } |
| case RJUMP_IF_TRUE: { |
| if (stack[sp].v_int64) { |
| pc += code[pc + 1].v_int; |
| } else { |
| pc += 2; |
| } |
| break; |
| } |
| case RJUMP_IF_FALSE: { |
| if (!stack[sp].v_int64) { |
| pc += code[pc + 1].v_int; |
| } else { |
| pc += 2; |
| } |
| break; |
| } |
| case RJUMP: { |
| pc += code[pc + 1].v_int; |
| break; |
| } |
| case ASSERT_SP: { |
| int64_t expected = code[pc + 1].v_int; |
| ICHECK_EQ(sp, expected) << "sp assertion failed, expected=" << expected << " now=" << sp |
| << ", pc=" << pc; |
| pc += 2; |
| break; |
| } |
| case CALL_PACKED_LOWERED: { |
| // call packed function. |
| TVMValue* value_stack = static_cast<TVMValue*>(stack[sp - 1].v_handle); |
| int* type_stack = static_cast<int*>(stack[sp].v_handle); |
| int call_fid = code[pc + 1].v_int; |
| int begin = code[pc + 2].v_int; |
| int end = code[pc + 3].v_int; |
| int num_args = end - begin; |
| static_assert(sizeof(Code) == sizeof(int) && alignof(Code) == alignof(int), "asusmption"); |
| runtime::TVMRetValue rv; |
| GetExtern(s, call_fid) |
| .CallPacked(runtime::TVMArgs(value_stack + begin, type_stack + begin, num_args), &rv); |
| sp = sp - 1; |
| stack[sp] = rv.value(); |
| pc += 4; |
| break; |
| } |
| // intrinsics |
| case TVM_STRUCT_GET: { |
| int index = code[pc + 1].v_int; |
| int kind = code[pc + 2].v_int; |
| DLTensor* arr = static_cast<DLTensor*>(stack[sp].v_handle); |
| switch (kind) { |
| case StackVM::kArrData: { |
| stack[sp].v_handle = arr[index].data; |
| break; |
| } |
| case StackVM::kArrShape: { |
| stack[sp].v_handle = arr[index].shape; |
| break; |
| } |
| case StackVM::kArrStrides: { |
| stack[sp].v_handle = arr[index].strides; |
| break; |
| } |
| case StackVM::kArrNDim: { |
| stack[sp].v_int64 = arr[index].ndim; |
| break; |
| } |
| case StackVM::kArrTypeCode: { |
| stack[sp].v_int64 = static_cast<int64_t>(arr[index].dtype.code); |
| break; |
| } |
| case StackVM::kArrTypeBits: { |
| stack[sp].v_int64 = static_cast<int64_t>(arr[index].dtype.bits); |
| break; |
| } |
| case StackVM::kArrTypeLanes: { |
| stack[sp].v_int64 = static_cast<int64_t>(arr[index].dtype.lanes); |
| break; |
| } |
| case StackVM::kArrByteOffset: { |
| stack[sp].v_int64 = static_cast<int64_t>(arr[index].byte_offset); |
| break; |
| } |
| case StackVM::kArrDeviceId: { |
| stack[sp].v_int64 = arr[index].device.device_id; |
| break; |
| } |
| case StackVM::kArrDeviceType: { |
| stack[sp].v_int64 = static_cast<int64_t>(arr[index].device.device_type); |
| break; |
| } |
| case StackVM::kArrAddr: { |
| stack[sp].v_handle = arr + index; |
| break; |
| } |
| case StackVM::kTVMValueContent: { |
| stack[sp] = static_cast<TVMValue*>(stack[sp].v_handle)[index]; |
| break; |
| } |
| default: |
| LOG(FATAL) << "unhandled get " << kind; |
| } |
| pc = pc + 3; |
| break; |
| } |
| case TVM_STRUCT_SET: { |
| int index = code[pc + 1].v_int; |
| int kind = code[pc + 2].v_int; |
| DLTensor* arr = static_cast<DLTensor*>(stack[sp - 1].v_handle); |
| switch (kind) { |
| case StackVM::kArrData: { |
| arr[index].data = stack[sp].v_handle; |
| break; |
| } |
| case StackVM::kArrShape: { |
| arr[index].shape = static_cast<int64_t*>(stack[sp].v_handle); |
| break; |
| } |
| case StackVM::kArrStrides: { |
| arr[index].strides = static_cast<int64_t*>(stack[sp].v_handle); |
| break; |
| } |
| case StackVM::kArrNDim: { |
| arr[index].ndim = static_cast<int>(stack[sp].v_int64); |
| break; |
| } |
| case StackVM::kArrTypeCode: { |
| arr[index].dtype.code = static_cast<uint8_t>(stack[sp].v_int64); |
| break; |
| } |
| case StackVM::kArrTypeBits: { |
| arr[index].dtype.bits = static_cast<uint8_t>(stack[sp].v_int64); |
| break; |
| } |
| case StackVM::kArrTypeLanes: { |
| arr[index].dtype.lanes = static_cast<uint16_t>(stack[sp].v_int64); |
| break; |
| } |
| case StackVM::kArrByteOffset: { |
| arr[index].byte_offset = static_cast<uint64_t>(stack[sp].v_int64); |
| break; |
| } |
| case StackVM::kArrDeviceId: { |
| arr[index].device.device_id = static_cast<int>(stack[sp].v_int64); |
| break; |
| } |
| case StackVM::kArrDeviceType: { |
| arr[index].device.device_type = static_cast<DLDeviceType>(stack[sp].v_int64); |
| break; |
| } |
| case StackVM::kTVMValueContent: { |
| static_cast<TVMValue*>(stack[sp - 1].v_handle)[index] = stack[sp]; |
| break; |
| } |
| default: |
| LOG(FATAL) << "unhandled tvm_struct_set " << kind; |
| } |
| sp -= 2; |
| pc += 3; |
| break; |
| } |
| // alloca |
| case TVM_STACK_ALLOCA_BY_8BYTE: { |
| static_assert(sizeof(TVMValue) == 8, "invariance"); |
| int num = code[pc + 1].v_int; |
| void* addr = &stack[sp] + 1; |
| sp = sp + num + 1; |
| alloca_sp = sp - 1; |
| stack[sp].v_handle = addr; |
| pc = pc + 2; |
| break; |
| } |
| case TVM_DEVICE_ALLOCA: { |
| int device_type = static_cast<int>(stack[sp - 4].v_int64); |
| int device_id = static_cast<int>(stack[sp - 3].v_int64); |
| size_t nbytes = static_cast<size_t>(stack[sp - 2].v_int64); |
| int dtype_code_hint = static_cast<int>(stack[sp - 1].v_int64); |
| int dtype_bits_hint = static_cast<int>(stack[sp].v_int64); |
| void* ptr = TVMBackendAllocWorkspace(device_type, device_id, nbytes, dtype_code_hint, |
| dtype_bits_hint); |
| stack[sp - 4].v_handle = ptr; |
| sp = sp - 4; |
| pc = pc + 1; |
| break; |
| } |
| case TVM_DEVICE_FREE: { |
| int device_type = static_cast<int>(stack[sp - 2].v_int64); |
| int device_id = static_cast<int>(stack[sp - 1].v_int64); |
| void* ptr = stack[sp].v_handle; |
| int ret = TVMBackendFreeWorkspace(device_type, device_id, ptr); |
| stack[sp - 2].v_int64 = ret; |
| sp = sp - 2; |
| pc = pc + 1; |
| break; |
| } |
| case TVM_THROW_LAST_ERROR: { |
| LOG(FATAL) << TVMGetLastError(); |
| break; |
| } |
| } |
| ICHECK_GE(sp, alloca_sp) << "touch allocated space"; |
| ICHECK_LT(sp, stack_cap) << "Stack overflow"; |
| } |
| } |
| |
| const PackedFunc& StackVM::GetExtern(State* s, int fid) const { |
| ICHECK_LT(static_cast<size_t>(fid), extern_func_cache_.size()); |
| // allow race write in this, since write is idempotent |
| PackedFunc& f = extern_func_cache_[fid]; |
| if (f == nullptr) { |
| ICHECK(s->mod_ctx != nullptr) << "No local context is set in stackvm"; |
| const PackedFunc* pf = s->mod_ctx->GetFuncFromEnv(extern_func_name[fid]); |
| ICHECK(pf != nullptr); |
| f = *pf; |
| } |
| return f; |
| } |
| |
| } // namespace runtime |
| } // namespace tvm |
