in src/runtime/stackvm/stackvm.cc [241:600]
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";
}
}