From 3235d05cbe2a02da77466b9190a8af687d0d6524 Mon Sep 17 00:00:00 2001 From: tqchen Date: Wed, 26 Mar 2025 11:44:49 -0400 Subject: [PATCH] [REFACTOR] Phase out StackVM This PR phases out the stackvm implementation. StackVM historically was used for some host codegen. As we move towards more comprehensive host codegen support, we can leverage LLVM and C for that purpose, so it can be helpful to phase it out to simplify maintainace. --- CMakeLists.txt | 10 - python/tvm/contrib/stackvm.py | 45 -- python/tvm/runtime/module.py | 9 - rust/tvm-sys/src/device.rs | 3 +- rust/tvm-sys/src/value.rs | 2 +- src/runtime/module.cc | 2 - src/runtime/stackvm/stackvm.cc | 615 ------------------ src/runtime/stackvm/stackvm.h | 459 ------------- src/runtime/stackvm/stackvm_module.cc | 149 ----- src/runtime/stackvm/stackvm_module.h | 47 -- src/support/libinfo.cc | 4 - src/target/stackvm/codegen_stackvm.cc | 555 ---------------- src/target/stackvm/codegen_stackvm.h | 165 ----- src/target/target_kind.cc | 3 - .../codegen/test_target_codegen_device.py | 2 +- .../codegen/test_target_codegen_extern.py | 1 - .../codegen/test_target_codegen_vm_basic.py | 143 ---- .../test_tir_transform_lower_tvm_builtin.py | 1 - tests/scripts/task_config_build_gpu.sh | 1 - 19 files changed, 3 insertions(+), 2213 deletions(-) delete mode 100644 python/tvm/contrib/stackvm.py delete mode 100644 src/runtime/stackvm/stackvm.cc delete mode 100644 src/runtime/stackvm/stackvm.h delete mode 100644 src/runtime/stackvm/stackvm_module.cc delete mode 100644 src/runtime/stackvm/stackvm_module.h delete mode 100644 src/target/stackvm/codegen_stackvm.cc delete mode 100644 src/target/stackvm/codegen_stackvm.h delete mode 100644 tests/python/codegen/test_target_codegen_vm_basic.py diff --git a/CMakeLists.txt b/CMakeLists.txt index a86bc4cc332b..caad7fb02b1f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -370,16 +370,6 @@ if(USE_RPC) list(APPEND RUNTIME_SRCS ${RUNTIME_RPC_SRCS}) endif(USE_RPC) -tvm_file_glob(GLOB STACKVM_RUNTIME_SRCS src/runtime/stackvm/*.cc) -tvm_file_glob(GLOB STACKVM_CODEGEN_SRCS src/target/stackvm/*.cc) -list(APPEND COMPILER_SRCS ${STACKVM_CODEGEN_SRCS}) -if(USE_STACKVM_RUNTIME) - message(STATUS "Build with stackvm support in runtime...") - list(APPEND RUNTIME_SRCS ${STACKVM_RUNTIME_SRCS}) -else() - list(APPEND COMPILER_SRCS ${STACKVM_RUNTIME_SRCS}) -endif(USE_STACKVM_RUNTIME) - if(USE_CUDA AND USE_NCCL) message(STATUS "Build with NCCL...") find_nccl(${USE_NCCL}) diff --git a/python/tvm/contrib/stackvm.py b/python/tvm/contrib/stackvm.py deleted file mode 100644 index 458d69235db5..000000000000 --- a/python/tvm/contrib/stackvm.py +++ /dev/null @@ -1,45 +0,0 @@ -# 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. - -"""Dummy StackVM build function.""" -# pylint: disable=invalid-name -from __future__ import absolute_import as _abs -import shutil - - -def build(output, files): - """Simply copy StackVM output to the destination. - - Parameters - ---------- - output : str - The target StackVM file. - - files : list - A single self-contained StackVM module file. - """ - - if len(files) == 0: - raise RuntimeError("StackVM artifact must be provided") - if len(files) > 1: - raise RuntimeError("Unexpected multiple StackVM artifacts") - - shutil.copy(files[0], output) - - -# assign output format -build.output_format = "stackvm" diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py index ca151293bbbd..79e853563037 100644 --- a/python/tvm/runtime/module.py +++ b/python/tvm/runtime/module.py @@ -517,15 +517,6 @@ def export_library( if isinstance(file_name, Path): file_name = str(file_name) - if self.type_key == "stackvm": - if not file_name.endswith(".stackvm"): - raise ValueError( - f"Module[{self.type_key}]: can only be saved as stackvm format." - "did you build with LLVM enabled?" - ) - self.save(file_name) - return - modules = self._collect_dso_modules() if workspace_dir is None: temp = _utils.tempdir() diff --git a/rust/tvm-sys/src/device.rs b/rust/tvm-sys/src/device.rs index 1ebac09bf611..0344983c1622 100644 --- a/rust/tvm-sys/src/device.rs +++ b/rust/tvm-sys/src/device.rs @@ -120,7 +120,6 @@ impl<'a> From<&'a str> for DeviceType { match type_str { "cpu" => DeviceType::CPU, "llvm" => DeviceType::CPU, - "stackvm" => DeviceType::CPU, "cuda" => DeviceType::CUDA, "nvptx" => DeviceType::CUDA, "cl" => DeviceType::OpenCL, @@ -208,7 +207,7 @@ macro_rules! impl_tvm_device { } impl_tvm_device!( - DLDeviceType_kDLCPU: [cpu, llvm, stackvm], + DLDeviceType_kDLCPU: [cpu, llvm], DLDeviceType_kDLCUDA: [cuda, nvptx], DLDeviceType_kDLOpenCL: [cl], DLDeviceType_kDLMetal: [metal], diff --git a/rust/tvm-sys/src/value.rs b/rust/tvm-sys/src/value.rs index f69172f41221..9c987af4cef6 100644 --- a/rust/tvm-sys/src/value.rs +++ b/rust/tvm-sys/src/value.rs @@ -85,7 +85,7 @@ macro_rules! impl_tvm_device { } impl_tvm_device!( - DLDeviceType_kDLCPU: [cpu, llvm, stackvm], + DLDeviceType_kDLCPU: [cpu, llvm], DLDeviceType_kDLCUDA: [cuda, nvptx], DLDeviceType_kDLOpenCL: [cl], DLDeviceType_kDLMetal: [metal], diff --git a/src/runtime/module.cc b/src/runtime/module.cc index 4e60a0d0a285..a21223f142ed 100644 --- a/src/runtime/module.cc +++ b/src/runtime/module.cc @@ -151,8 +151,6 @@ bool RuntimeEnabled(const String& target_str) { f_name = "target.runtime.tflite"; } else if (target == "vulkan") { f_name = "device_api.vulkan"; - } else if (target == "stackvm") { - f_name = "target.build.stackvm"; } else if (target == "rpc") { f_name = "device_api.rpc"; } else if (target == "hexagon") { diff --git a/src/runtime/stackvm/stackvm.cc b/src/runtime/stackvm/stackvm.cc deleted file mode 100644 index 5a4af57b5ec4..000000000000 --- a/src/runtime/stackvm/stackvm.cc +++ /dev/null @@ -1,615 +0,0 @@ -/* - * 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 -#include - -#include - -namespace tvm { -namespace runtime { - -typedef dmlc::ThreadLocalStore 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(static_cast(stack[sp].v_handle)[index]); \ - pc += 2; \ - } - -#define STACK_VM_STORE(FIELD, DST_TYPE) \ - { \ - int index = code[pc + 1].v_int; \ - static_cast(stack[sp - 1].v_handle)[index] = \ - static_cast(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(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 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 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& stack = s->stack; - std::vector& heap = s->heap; - if (stack.size() < stack_size) { - stack.resize(stack_size); - } - int64_t stack_cap = static_cast(stack_size - 4); - if (heap.size() < heap_size) { - heap.resize(heap_size); - } - const int64_t code_size = static_cast(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(stack[sp - 1].v_handle); - int* type_stack = static_cast(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(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(arr[index].dtype.code); - break; - } - case StackVM::kArrTypeBits: { - stack[sp].v_int64 = static_cast(arr[index].dtype.bits); - break; - } - case StackVM::kArrTypeLanes: { - stack[sp].v_int64 = static_cast(arr[index].dtype.lanes); - break; - } - case StackVM::kArrByteOffset: { - stack[sp].v_int64 = static_cast(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(arr[index].device.device_type); - break; - } - case StackVM::kArrAddr: { - stack[sp].v_handle = arr + index; - break; - } - case StackVM::kTVMValueContent: { - stack[sp] = static_cast(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(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(stack[sp].v_handle); - break; - } - case StackVM::kArrStrides: { - arr[index].strides = static_cast(stack[sp].v_handle); - break; - } - case StackVM::kArrNDim: { - arr[index].ndim = static_cast(stack[sp].v_int64); - break; - } - case StackVM::kArrTypeCode: { - arr[index].dtype.code = static_cast(stack[sp].v_int64); - break; - } - case StackVM::kArrTypeBits: { - arr[index].dtype.bits = static_cast(stack[sp].v_int64); - break; - } - case StackVM::kArrTypeLanes: { - arr[index].dtype.lanes = static_cast(stack[sp].v_int64); - break; - } - case StackVM::kArrByteOffset: { - arr[index].byte_offset = static_cast(stack[sp].v_int64); - break; - } - case StackVM::kArrDeviceId: { - arr[index].device.device_id = static_cast(stack[sp].v_int64); - break; - } - case StackVM::kArrDeviceType: { - arr[index].device.device_type = static_cast(stack[sp].v_int64); - break; - } - case StackVM::kTVMValueContent: { - static_cast(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(stack[sp - 4].v_int64); - int device_id = static_cast(stack[sp - 3].v_int64); - size_t nbytes = static_cast(stack[sp - 2].v_int64); - int dtype_code_hint = static_cast(stack[sp - 1].v_int64); - int dtype_bits_hint = static_cast(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(stack[sp - 2].v_int64); - int device_id = static_cast(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(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 diff --git a/src/runtime/stackvm/stackvm.h b/src/runtime/stackvm/stackvm.h deleted file mode 100644 index c967e99dbecb..000000000000 --- a/src/runtime/stackvm/stackvm.h +++ /dev/null @@ -1,459 +0,0 @@ -/* - * 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 stackvm.h - * \brief A simple stack-based virtual machine. - * - * This can be used to interepret host side code - * to setup calls into device functions - * when only Runtime compilation for device is available(via NVRTC or OpenCL). - */ -#ifndef TVM_RUNTIME_STACKVM_STACKVM_H_ -#define TVM_RUNTIME_STACKVM_STACKVM_H_ - -#include -#include -#include - -#include -#include - -namespace tvm { -namespace runtime { - -using runtime::operator<<; - -/*! - * \brief A simple stack-based virtual machine program. - */ -class StackVM { - public: - /*! - * \brief Invoke the StackVM program. - * \param args The arguments to the StackVM. - * \param mod_ctx The module context used in running. - */ - void Run(const TVMArgs& args, runtime::ModuleNode* mod_ctx) const; - /*! - * \brief The opcode of stack vm - * \note Notation - * - sp Stack pointer - * - pc Program pointer - */ - enum OpCode { - // integer ops - ADD_I64, - SUB_I64, - MUL_I64, - DIV_I64, - MOD_I64, - EQ_I64, - LT_I64, - LE_I64, - // floating ops - ADD_F64, - SUB_F64, - MUL_F64, - DIV_F64, - EQ_F64, - LT_F64, - LE_F64, - // Pointer comparison - EQ_HANDLE, - /*! - * \brief Routine to load data from address with const offset. - * \code - * stack[sp].v_int64 = ((DType*)stack[sp].v_handle)[code[pc + 1].v_int]; - * pc = pc + 2; - * \endcode - */ - ARRAY_LOAD_UINT32, - ARRAY_LOAD_INT32, - ARRAY_LOAD_INT64, - ARRAY_LOAD_FP64, - ARRAY_LOAD_HANDLE, - ARRAY_LOAD_TVMVALUE, - /*! - * \brief Routine to store data from constant offset. - * \code - * ((DType*)stack[sp - 1].v_handle)[code[pc + 1].v_int] = stack[sp]; - * pc = pc + 2; - * sp = sp - 2; - * \endcode - */ - ARRAY_STORE_UINT32, - ARRAY_STORE_INT32, - ARRAY_STORE_INT64, - ARRAY_STORE_FP64, - ARRAY_STORE_HANDLE, - ARRAY_STORE_TVMVALUE, - // logical ops - NOT, - /*! - * \brief Add address by an offset. - * \code - * stack[sp - 1].v_handle = ((char*)stack[sp - 1].v_handle + stack[sp].v_int64); - * sp = sp - 1; - * \endcode - */ - ADDR_ADD, - /*! - * \brief push integer fetched from next pc position into stack - * \code - * stack[sp + 1].v_int64 = code[pc + 1].v_int; - * pc = pc + 2; - * sp = sp + 1; - * \endcode - */ - PUSH_I64, - /*! - * \brief push a value given relative index on the stack - * \code - * stack[sp + 1] = stack[sp + code[pc + 1].v_int]; - * pc = pc + 2; - * sp = sp + 1; - * \endcode - */ - PUSH_VALUE, - /*! - * \brief Load data from heap to top of stack - * \code - * stack[sp + 1] = heap[code[pc + 1].v_int]; - * pc = pc + 2; - * sp = sp + 1; - * \endcode - */ - LOAD_HEAP, - /*! - * \brief Store data to heap - * \code - * heap[code[pc + 1].v_int] = stack[sp]; - * sp = sp - 1; - * \endcode - */ - STORE_HEAP, - /*! \brief pop value from top of the stack */ - POP, - /*! - * \brief select based on operands. - * \code - * stack[sp - 2] = stack[sp].v_int64 ? stack[sp - 2] : stack[sp - 1] - * sp = sp - 2; - * \endcode - */ - SELECT, - /*! - * \brief Assert condition is true. - * \code - * ICHECK(stack[sp]) << str_data[code[pc + 1].v_int]; - * sp = sp - 1; - * \endcode - */ - ASSERT, - /*! - * \brief Relative Jump if the condition is true, - * Does not change the stack status. - * \code - * if (stack[sp]) { - * pc += code[pc + 1].v_int - * } else { - * pc = pc + 2; - * } - * \endcode - */ - RJUMP_IF_TRUE, - /*! - * \brief Relative Jump if the condition is true, - * Does not change the stack status. - * \code - * if (stack[sp]) { - * pc += code[pc + 1].v_int - * } else { - * pc = pc + 2; - * } - * \endcode - */ - RJUMP_IF_FALSE, - /*! - * \brief Relative jump to a location. - * \code - * pc += code[pc + 1].v_int; - * \endcode - */ - RJUMP, - /*! - * \brief debug instruction. - * \code - * ICHECK_EQ(sp, code[pc + 1]).v_int; - * pc += 2; - * \code - */ - ASSERT_SP, - /*! - * \brief call an extern packed function - * \code - * value_stack = stack[sp - 1].v_handle; - * type_stack = stack[sp - 0].v_handle; - * call_fid = code[pc + 1].v_int; - * begin = code[pc + 2].v_int; - * end = code[pc + 3].v_int; - * num_args = end - begin - 1; - * f = extern_func[call_fid]; - * stack[sp - 1] = f(&value_stack[begin:end-1], type_stack[begin:end-1], num_args); - * sp = sp - 1; - * // The type codes are hidden in the code space. - * pc = pc + 4 - * \endcode - */ - CALL_PACKED_LOWERED, - // Allocate things on stack - /*! - * \brief allocate data from stack. - * \code - * num = code[pc + 1].v_int; - * void* addr = &stack[sp]; - * sp = sp + num; - * stack[sp].v_handle = addr; - * pc = pc + 1; - * \endcode - */ - TVM_STACK_ALLOCA_BY_8BYTE, - /*! - * \brief allocate data from device. - * \code - * device_type = stack[sp - 2].v_int64; - * device_id = stack[sp - 1].v_int64; - * nbytes = stack[sp].v_int64; - * stack[sp - 2].v_handle = device_alloca(device_type, device_id, nbytes); - * sp = sp - 2; - * pc = pc + 1; - * \endcode - */ - TVM_DEVICE_ALLOCA, - /*! - * \brief free data into device. - * \code - * device_type = stack[sp - 2].v_int64; - * device_id = stack[sp - 1].v_int64; - * ptr = stack[sp].v_handle; - * stack[sp - 2].v_int64 = device_free(device_type, device_id, ptr); - * sp = sp - 2; - * pc = pc + 1; - * \endcode - */ - TVM_DEVICE_FREE, - /*! - * \brief throw last error - */ - TVM_THROW_LAST_ERROR, - /*! - * \brief get data from structure. - * \code - * index = code[pc + 1].v_int; - * field = code[pc + 2].v_int; - * stack[sp] = ((StructType*)stack[sp].v_handle)[index]->field; - * pc = pc + 3 - * \endcode - */ - TVM_STRUCT_GET, - /*! - * \brief set data into structure. - * \code - * index = code[pc + 1].v_int; - * field = code[pc + 2].v_int; - * ((StructType*)stack[sp - 1].v_handle)[index]->field = stack[sp]; - * pc = pc + 3 - * sp = sp - 1 - * \endcode - */ - TVM_STRUCT_SET - }; - /*! \brief The kind of structure field info */ - enum StructFieldKind : int { - // array head address - kArrAddr, - kArrData, - kArrShape, - kArrStrides, - kArrNDim, - kArrTypeCode, - kArrTypeBits, - kArrTypeLanes, - kArrByteOffset, - kArrDeviceId, - kArrDeviceType, - kArrKindBound_, - // TVMValue field - kTVMValueContent, - kTVMValueKindBound_ - }; - /*! \brief The code structure */ - union Code { - OpCode op_code; - int v_int; - }; - /*! \brief The state object of StackVM */ - struct State { - /*! \brief The execution stack */ - std::vector stack; - /*! \brief The global heap space */ - std::vector heap; - /*! \brief stack pointer */ - int64_t sp{0}; - /*! \brief program counter */ - int64_t pc{0}; - /*! \brief The current module context of stackvm */ - runtime::ModuleNode* mod_ctx{nullptr}; - }; - /*! \brief Initialize local cache*/ - void InitCache(); - /*! - * \brief Save stackvm program to an output stream - * \param strm The output stream - */ - void Save(dmlc::Stream* strm) const; - /*! - * \brief Load stackvm program from output stream - * \param strm The output stream - */ - bool Load(dmlc::Stream* strm); - /*! - * \brief Print instruction at location pc - * \param os The ostream - * \param pc The pc - * \return the pc to next instruction. - */ - int64_t PrintCode(std::ostream& os, int64_t pc) const; // NOLINT(*) - /*! \brief Get thread local state of the stack VM */ - static State* ThreadLocalState(); - // The code below are programs - /*! \brief The instructions */ - std::vector code; - /*! \brief constant error messages */ - std::vector str_data; - /*! \brief Extern functions */ - std::vector extern_func_name; - /*! \brief name of each heap id */ - std::vector heap_id_name; - /*! \brief The memory size needed */ - size_t heap_size{0}; - /*! \brief The stack size required */ - size_t stack_size{1024}; - /*! - * \brief Convert I64 opcode to F64 Ones - * \param code The op code. - * \return the F64 op code. - */ - static OpCode CodeI64ToF64(OpCode code) { - switch (code) { - case ADD_I64: - return ADD_F64; - case SUB_I64: - return SUB_F64; - case MUL_I64: - return MUL_F64; - case DIV_I64: - return DIV_F64; - case EQ_I64: - return EQ_F64; - case LT_I64: - return LT_F64; - case LE_I64: - return LE_F64; - case MOD_I64: - LOG(FATAL) << "cannot handle mod for float"; - default: - LOG(FATAL) << "cannot handle op " << code; - } - } - /*! - * \brief Get load opcode for type t - * \param t the type code. - * \return The load opcode - */ - static OpCode GetLoad(DLDataType t) { - ICHECK_EQ(t.lanes, 1U); - if (t.code == kTVMOpaqueHandle) return ARRAY_LOAD_HANDLE; - if (t.code == kDLInt) { - switch (t.bits) { - case 32: - return ARRAY_LOAD_INT32; - case 64: - return ARRAY_LOAD_INT64; - } - } else if (t.code == kDLUInt) { - switch (t.bits) { - case 32: - return ARRAY_LOAD_UINT32; - } - } else if (t.code == kDLFloat) { - switch (t.bits) { - case 64: - return ARRAY_LOAD_FP64; - } - } - LOG(FATAL) << "Cannot load type " << t; - } - /*! - * \brief Get store opcode for type t - * \param t the type code. - * \return The load opcode - */ - static OpCode GetStore(DLDataType t) { - ICHECK_EQ(t.lanes, 1U); - if (t.code == kTVMOpaqueHandle) return ARRAY_STORE_HANDLE; - if (t.code == kDLInt) { - switch (t.bits) { - case 32: - return ARRAY_STORE_INT32; - case 64: - return ARRAY_STORE_INT64; - } - } else if (t.code == kDLUInt) { - switch (t.bits) { - case 32: - return ARRAY_STORE_UINT32; - } - } else if (t.code == kDLFloat) { - switch (t.bits) { - case 64: - return ARRAY_STORE_FP64; - } - } - LOG(FATAL) << "Cannot store type " << t; - } - friend std::ostream& operator<<(std::ostream& os, const StackVM& vm); // NOLINT(*) - - private: - // execute the stack vm with given state - void Run(State* state) const; - // get extern function. - const PackedFunc& GetExtern(State* s, int fid) const; - // cached extern function - mutable std::vector extern_func_cache_; -}; - -} // namespace runtime -} // namespace tvm - -namespace dmlc { -DMLC_DECLARE_TRAITS(has_saveload, ::tvm::runtime::StackVM, true); -} -#endif // TVM_RUNTIME_STACKVM_STACKVM_H_ diff --git a/src/runtime/stackvm/stackvm_module.cc b/src/runtime/stackvm/stackvm_module.cc deleted file mode 100644 index 867ccc8ed082..000000000000 --- a/src/runtime/stackvm/stackvm_module.cc +++ /dev/null @@ -1,149 +0,0 @@ -/* - * 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 stackvm_module.cc - */ -#include "stackvm_module.h" - -#include -#include -#include - -#include -#include -#include - -#include "../file_utils.h" - -namespace tvm { -namespace runtime { - -class StackVMModuleNode : public runtime::ModuleNode { - public: - const char* type_key() const final { return "stackvm"; } - - PackedFunc GetFunction(const String& name, const ObjectPtr& sptr_to_self) final { - if (name == runtime::symbol::tvm_module_main) { - return GetFunction(entry_func_, sptr_to_self); - } - auto it = fmap_.find(name); - if (it == fmap_.end()) return PackedFunc(); - const StackVM& vm = it->second; - // capture sptr_to_self to keep module node alive. - return PackedFunc( - [vm, sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { vm.Run(args, this); }); - } - - String GetSource(const String& format) final { - std::ostringstream os; - for (const auto& kv : fmap_) { - os << "Function: " << kv.first << '\n'; - os << kv.second; - } - return os.str(); - } - - void SaveToFile(const String& file_name, const String& format) final { - std::string data, mblob; - dmlc::MemoryStringStream writer(&data); - dmlc::Stream* strm = &writer; - strm->Write(fmap_); - strm->Write(entry_func_); - // also save imports - uint64_t num_imports = static_cast(imports_.size()); - strm->Write(num_imports); - - for (runtime::Module im : imports_) { - ICHECK_EQ(im->imports().size(), 0U) << "Only support simply one-level hierarchy"; - std::string tkey = im->type_key(); - strm->Write(tkey); - im->SaveToBinary(strm); - } - SaveBinaryToFile(file_name, data); - } - - static Module Create(std::unordered_map fmap, std::string entry_func) { - auto n = make_object(); - n->fmap_ = std::move(fmap); - n->entry_func_ = std::move(entry_func); - return Module(n); - } - - static Module Load(dmlc::Stream* strm) { - std::unordered_map fmap; - std::string entry_func, data; - strm->Read(&fmap); - strm->Read(&entry_func); - auto n = make_object(); - n->fmap_ = std::move(fmap); - n->entry_func_ = std::move(entry_func); - uint64_t num_imports; - strm->Read(&num_imports); - for (uint64_t i = 0; i < num_imports; ++i) { - std::string tkey; - ICHECK(strm->Read(&tkey)); - std::string loadkey = "runtime.module.loadbinary_"; - std::string fkey = loadkey + tkey; - const PackedFunc* f = Registry::Get(fkey); - if (f == nullptr) { - std::string loaders = ""; - for (auto reg_name : Registry::ListNames()) { - std::string name = reg_name; - if (name.rfind(loadkey, 0) == 0) { - if (loaders.size() > 0) { - loaders += ", "; - } - loaders += name.substr(loadkey.size()); - } - } - ICHECK(f != nullptr) - << "Binary was created using " << tkey - << " but a loader of that name is not registered. Available loaders are " << loaders - << ". Perhaps you need to recompile with this runtime enabled."; - } - Module m = (*f)(static_cast(strm)); - n->imports_.emplace_back(std::move(m)); - } - return Module(n); - } - - static Module LoadFromFile(std::string file_name, std::string format) { - std::string data; - LoadBinaryFromFile(file_name, &data); - dmlc::MemoryStringStream reader(&data); - return Load(&reader); - } - - private: - // internal function map - std::unordered_map fmap_; - // entry function. - std::string entry_func_; -}; - -Module StackVMModuleCreate(std::unordered_map fmap, std::string entry_func) { - return StackVMModuleNode::Create(fmap, entry_func); -} - -TVM_REGISTER_GLOBAL("runtime.module.loadfile_stackvm") - .set_body_typed(StackVMModuleNode::LoadFromFile); - -} // namespace runtime -} // namespace tvm diff --git a/src/runtime/stackvm/stackvm_module.h b/src/runtime/stackvm/stackvm_module.h deleted file mode 100644 index 6ae4ae47a92c..000000000000 --- a/src/runtime/stackvm/stackvm_module.h +++ /dev/null @@ -1,47 +0,0 @@ -/* - * 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 stackvm_module.h - * \brief StackVM module - */ -#ifndef TVM_RUNTIME_STACKVM_STACKVM_MODULE_H_ -#define TVM_RUNTIME_STACKVM_STACKVM_MODULE_H_ - -#include - -#include -#include - -#include "stackvm.h" - -namespace tvm { -namespace runtime { -/*! - * \brief create a stackvm module - * - * \param fmap The map from name to function - * \param entry_func The entry function name. - * \return The created module - */ -Module StackVMModuleCreate(std::unordered_map fmap, std::string entry_func); - -} // namespace runtime -} // namespace tvm -#endif // TVM_RUNTIME_STACKVM_STACKVM_MODULE_H_ diff --git a/src/support/libinfo.cc b/src/support/libinfo.cc index b981fcd6d73f..40213e37b650 100644 --- a/src/support/libinfo.cc +++ b/src/support/libinfo.cc @@ -115,10 +115,6 @@ #define TVM_INFO_USE_LLVM "NOT-FOUND" #endif -#ifndef TVM_INFO_USE_STACKVM_RUNTIME -#define TVM_INFO_USE_STACKVM_RUNTIME "NOT-FOUND" -#endif - #ifndef TVM_INFO_USE_OPENMP #define TVM_INFO_USE_OPENMP "NOT-FOUND" #endif diff --git a/src/target/stackvm/codegen_stackvm.cc b/src/target/stackvm/codegen_stackvm.cc deleted file mode 100644 index 36638576d387..000000000000 --- a/src/target/stackvm/codegen_stackvm.cc +++ /dev/null @@ -1,555 +0,0 @@ -/* - * 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 codegen_stackvm.cc - */ -#include "codegen_stackvm.h" - -#include -#include -#include -#include -#include - -#include -#include - -#include "../../runtime/stackvm/stackvm_module.h" - -namespace tvm { -namespace codegen { - -using namespace tir; - -// map struct field kind to runtime variants -// We keep two separate enums to ensure runtime/compiler isolation. -StackVM::StructFieldKind MapFieldKind(int64_t kind) { - auto val = static_cast(kind); - switch (val) { - case builtin::kArrData: - return StackVM::kArrData; - case builtin::kArrShape: - return StackVM::kArrShape; - case builtin::kArrAddr: - return StackVM::kArrAddr; - case builtin::kArrStrides: - return StackVM::kArrStrides; - case builtin::kArrNDim: - return StackVM::kArrNDim; - case builtin::kArrTypeCode: - return StackVM::kArrTypeCode; - case builtin::kArrTypeBits: - return StackVM::kArrTypeBits; - case builtin::kArrTypeLanes: - return StackVM::kArrTypeLanes; - case builtin::kArrByteOffset: - return StackVM::kArrByteOffset; - case builtin::kArrDeviceId: - return StackVM::kArrDeviceId; - case builtin::kArrDeviceType: - return StackVM::kArrDeviceType; - case builtin::kTVMValueContent: - return StackVM::kTVMValueContent; - default: - LOG(FATAL) << "Do not know how to map field " << kind; - } - return StackVM::kArrData; -} - -StackVM CodeGenStackVM::Compile(const PrimFunc& f) { - ICHECK_EQ(f->buffer_map.size(), 0U) - << "Cannot codegen function with buffer_map, please lower them first"; - for (size_t i = 0; i < f->params.size(); ++i) { - Var v = f->params[i]; - int vid = AllocVarID(v.get()); - ICHECK_EQ(static_cast(vid), i); - } - this->Push(f->body); - vm_.InitCache(); - return std::move(vm_); -} - -void CodeGenStackVM::Push(const Stmt& n) { - VisitStmt(n); - if (debug_) { - this->PushOp(StackVM::ASSERT_SP, 0); - } -} - -void CodeGenStackVM::PushOp(StackVM::OpCode opcode) { - StackVM::Code code; - code.op_code = opcode; - vm_.code.push_back(code); -} - -void CodeGenStackVM::SetOperand(int64_t operand_index, int64_t operand) { - ICHECK(operand >= std::numeric_limits::min() && operand <= std::numeric_limits::max()); - vm_.code.at(operand_index).v_int = static_cast(operand); -} - -int64_t CodeGenStackVM::PushOp(StackVM::OpCode opcode, int operand) { - int64_t pc = static_cast(vm_.code.size()); - StackVM::Code code; - code.op_code = opcode; - vm_.code.push_back(code); - code.v_int = operand; - vm_.code.push_back(code); - return pc + 1; -} - -int CodeGenStackVM::GetStrID(const std::string& key) { - auto it = str_idmap_.find(key); - if (it != str_idmap_.end()) return it->second; - int sid = static_cast(vm_.str_data.size()); - vm_.str_data.push_back(key); - str_idmap_[key] = sid; - return sid; -} - -int CodeGenStackVM::AllocVarID(const VarNode* v) { - ICHECK(!var_idmap_.count(v)); - int vid = static_cast(vm_.heap_size); - ICHECK_EQ(vm_.heap_size, var_idmap_.size()); - vm_.heap_id_name.push_back(v->name_hint); - ++vm_.heap_size; - var_idmap_[v] = vid; - return vid; -} - -int CodeGenStackVM::GetVarID(const VarNode* v) const { - auto it = var_idmap_.find(v); - ICHECK(it != var_idmap_.end()) << "Find undefined Variable " << v->name_hint; - return it->second; -} - -void CodeGenStackVM::VisitExpr_(const BufferLoadNode* op) { - ICHECK_EQ(op->indices.size(), 1) << "StackVM expects flat 1-d buffers. " - << "Has FlattenBuffer been run?"; - auto index = op->indices[0]; - - this->Push(op->buffer->data); - StackVM::OpCode code = StackVM::GetLoad(op->dtype); - if (const IntImmNode* int_index = index.as()) { - this->PushOp(code, int_index->value); - } else { - this->Push(index); - this->PushOp(StackVM::PUSH_I64, op->dtype.element_of().bytes()); - this->PushOp(StackVM::MUL_I64); - this->PushOp(StackVM::ADDR_ADD); - this->PushOp(code, 0); - } -} - -void CodeGenStackVM::VisitStmt_(const BufferStoreNode* op) { - ICHECK_EQ(op->indices.size(), 1) << "StackVM expects flat 1-d buffers. " - << "Has FlattenBuffer been run?"; - auto index = op->indices[0]; - - this->Push(op->buffer->data); - StackVM::OpCode code = StackVM::GetStore(op->value.dtype()); - if (const IntImmNode* int_index = index.as()) { - this->Push(op->value); - this->PushOp(code, int_index->value); - } else { - this->Push(index); - this->PushOp(StackVM::PUSH_I64, op->value.dtype().element_of().bytes()); - this->PushOp(StackVM::MUL_I64); - this->PushOp(StackVM::ADDR_ADD); - this->Push(op->value); - this->PushOp(code, 0); - } -} - -void CodeGenStackVM::VisitStmt_(const AllocateNode* op) { - LOG(FATAL) << "Dynamic allocation not supported"; -} - -void CodeGenStackVM::VisitStmt_(const DeclBufferNode* op) { VisitStmt(op->body); } - -void CodeGenStackVM::VisitExpr_(const CallNode* op) { - if (op->op.same_as(builtin::address_of())) { - const BufferLoadNode* load = op->args[0].as(); - ICHECK(op->args.size() == 1 && load); - ICHECK_EQ(load->indices.size(), 1) << "CodeGenStackVM only supports flat memory allocations."; - - this->PushOp(StackVM::LOAD_HEAP, GetVarID(load->buffer->data.get())); - this->Push(load->indices[0]); - this->PushOp(StackVM::PUSH_I64, load->dtype.element_of().bytes()); - this->PushOp(StackVM::MUL_I64); - this->PushOp(StackVM::ADDR_ADD); - } else if (op->op.same_as(builtin::reinterpret())) { - this->Push(op->args[0]); - } else if (op->op.same_as(builtin::tvm_struct_get())) { - ICHECK_EQ(op->args.size(), 3U); - int kind = op->args[2].as()->value; - this->Push(op->args[0]); - const IntImmNode* index = op->args[1].as(); - ICHECK(index != nullptr); - StackVM::Code code; - code.op_code = StackVM::TVM_STRUCT_GET; - vm_.code.push_back(code); - code.v_int = index->value; - vm_.code.push_back(code); - code.v_int = MapFieldKind(kind); - vm_.code.push_back(code); - } else if (op->op.same_as(builtin::tvm_call_packed_lowered())) { - ICHECK_GE(op->args.size(), 5U); - const StringImmNode* s = op->args[0].as(); - ICHECK(s != nullptr) << "tvm_call_global expect first argument as function name"; - this->Push(op->args[1]); - this->Push(op->args[2]); - int begin = op->args[3].as()->value; - int end = op->args[4].as()->value; - // find the fuction id. - const std::string& func_name = s->value; - auto it = extern_fun_idmap_.find(func_name); - int fid; - if (it != extern_fun_idmap_.end()) { - fid = it->second; - } else { - fid = static_cast(vm_.extern_func_name.size()); - vm_.extern_func_name.push_back(func_name); - extern_fun_idmap_[func_name] = fid; - } - // CALL_PACKED_FUNC - StackVM::Code code; - code.op_code = StackVM::CALL_PACKED_LOWERED; - vm_.code.push_back(code); - code.v_int = fid; - vm_.code.push_back(code); - code.v_int = begin; - vm_.code.push_back(code); - code.v_int = end; - vm_.code.push_back(code); - } else if (op->op.same_as(builtin::tvm_stack_alloca())) { - ICHECK_EQ(op->args.size(), 2U); - const std::string& type = op->args[0].as()->value; - const IntImmNode* num = op->args[1].as(); - ICHECK(num != nullptr); - static_assert(alignof(TVMValue) % alignof(DLTensor) == 0, "invariant"); - // static_assert(alignof(TVMValue) % alignof(tvm_index_t) == 0, "invariant"); - size_t unit = sizeof(TVMValue); - size_t size = 0; - if (type == "shape") { - size = (num->value * sizeof(tvm_index_t) + unit - 1) / unit; - } else if (type == "arg_value") { - size = (num->value * sizeof(TVMValue) + unit - 1) / unit; - } else if (type == "arg_tcode") { - size = (num->value * sizeof(int) + unit - 1) / unit; - } else if (type == "array") { - size = (num->value * sizeof(DLTensor) + unit - 1) / unit; - } else { - LOG(FATAL) << "Unknown stack alloca type " << type; - } - // add stack size to be safe. - vm_.stack_size += size; - this->PushOp(StackVM::TVM_STACK_ALLOCA_BY_8BYTE, static_cast(size)); - } else if (op->op.same_as(backend_alloc_workspace_op_)) { - ICHECK_EQ(op->args.size(), 5U); - this->Push(op->args[0]); - this->Push(op->args[1]); - this->Push(op->args[2]); - this->Push(op->args[3]); - this->Push(op->args[4]); - this->PushOp(StackVM::TVM_DEVICE_ALLOCA); - } else if (op->op.same_as(backend_free_workspace_op_)) { - ICHECK_EQ(op->args.size(), 3U); - this->Push(op->args[0]); - this->Push(op->args[1]); - this->Push(op->args[2]); - this->PushOp(StackVM::TVM_DEVICE_FREE); - } else if (op->op.same_as(builtin::tvm_throw_last_error())) { - this->PushOp(StackVM::TVM_THROW_LAST_ERROR); - } else if (op->op.same_as(builtin::isnullptr())) { - ICHECK_EQ(op->args.size(), 1U); - this->Push(op->args[0]); - this->PushOp(StackVM::PUSH_I64, 0); - this->PushOp(StackVM::EQ_HANDLE); - } else if (op->op.same_as(builtin::ret())) { - CHECK(op->args.size() == 1 && op->args[0]->IsInstance() && - op->args[0].as()->value == 0) - << "StackVM does not support return values, " - << "and the return value " << op->args - << " is not special case of returning an error code of zero."; - } else { - LOG(FATAL) << "unknown function call " << op->op; - } -} - -void CodeGenStackVM::PushBinary(StackVM::OpCode op_int64, const PrimExpr& a, const PrimExpr& b) { - this->Push(a); - this->Push(b); - DataType t = a.dtype(); - if (t.is_int()) { - this->PushOp(op_int64); - } else if (t.is_uint()) { - this->PushOp(op_int64); - } else { - this->PushOp(StackVM::CodeI64ToF64(op_int64)); - } -} - -void CodeGenStackVM::PushCast(DataType dst, DataType src) { - if (dst.is_int()) { - if (src.is_int() || src.is_uint()) return; - } else if (dst.is_uint()) { - if (src.is_int() || src.is_uint()) return; - } else if (dst.is_float()) { - if (src.is_float()) return; - } -} - -void CodeGenStackVM::VisitExpr_(const StringImmNode* op) { - int sid = this->GetStrID(op->value); - this->PushOp(StackVM::PUSH_I64, sid); -} - -void CodeGenStackVM::VisitExpr_(const IntImmNode* op) { - ICHECK(op->value >= std::numeric_limits::min() && - op->value <= std::numeric_limits::max()) - << "Int constant exceed bound"; - this->PushOp(StackVM::PUSH_I64, static_cast(op->value)); -} - -void CodeGenStackVM::VisitExpr_(const FloatImmNode* op) { - LOG(FATAL) << "Float Imm is not supported"; -} - -void CodeGenStackVM::VisitExpr_(const VarNode* op) { - int vid = this->GetVarID(op); - this->PushOp(StackVM::LOAD_HEAP, vid); -} - -void CodeGenStackVM::VisitExpr_(const CastNode* op) { - this->Push(op->value); - PushCast(op->dtype, op->value.dtype()); -} - -void CodeGenStackVM::VisitExpr_(const AddNode* op) { PushBinary(StackVM::ADD_I64, op->a, op->b); } - -void CodeGenStackVM::VisitExpr_(const SubNode* op) { PushBinary(StackVM::SUB_I64, op->a, op->b); } - -void CodeGenStackVM::VisitExpr_(const MulNode* op) { PushBinary(StackVM::MUL_I64, op->a, op->b); } - -void CodeGenStackVM::VisitExpr_(const DivNode* op) { PushBinary(StackVM::DIV_I64, op->a, op->b); } - -void CodeGenStackVM::VisitExpr_(const ModNode* op) { PushBinary(StackVM::MOD_I64, op->a, op->b); } - -void CodeGenStackVM::VisitExpr_(const MinNode* op) { - this->Push(op->a); - this->Push(op->b); - this->PushOp(StackVM::PUSH_VALUE, -1); - this->PushOp(StackVM::PUSH_VALUE, -1); - this->PushOp(StackVM::LT_I64); - this->PushOp(StackVM::SELECT); -} - -void CodeGenStackVM::VisitExpr_(const MaxNode* op) { - this->Push(op->a); - this->Push(op->b); - this->PushOp(StackVM::PUSH_VALUE, 0); - this->PushOp(StackVM::PUSH_VALUE, -2); - this->PushOp(StackVM::LT_I64); - this->PushOp(StackVM::SELECT); -} - -void CodeGenStackVM::VisitExpr_(const EQNode* op) { PushBinary(StackVM::EQ_I64, op->a, op->b); } - -void CodeGenStackVM::VisitExpr_(const LENode* op) { PushBinary(StackVM::LE_I64, op->a, op->b); } - -void CodeGenStackVM::VisitExpr_(const NENode* op) { - PushBinary(StackVM::EQ_I64, op->a, op->b); - this->PushOp(StackVM::NOT); -} - -void CodeGenStackVM::VisitExpr_(const LTNode* op) { PushBinary(StackVM::LT_I64, op->a, op->b); } - -void CodeGenStackVM::VisitExpr_(const GENode* op) { - PushBinary(StackVM::LT_I64, op->a, op->b); - this->PushOp(StackVM::NOT); -} - -void CodeGenStackVM::VisitExpr_(const GTNode* op) { - PushBinary(StackVM::LE_I64, op->a, op->b); - this->PushOp(StackVM::NOT); -} - -void CodeGenStackVM::VisitExpr_(const AndNode* op) { - this->Push(op->a); - int64_t pc_jump = this->GetPC(); - int64_t opr_index = this->PushOp(StackVM::RJUMP_IF_FALSE, 0); - this->PushOp(StackVM::POP); - this->Push(op->b); - int64_t diff = this->GetPC() - pc_jump; - this->SetOperand(opr_index, diff); -} - -void CodeGenStackVM::VisitExpr_(const OrNode* op) { - this->Push(op->a); - int64_t pc_jump = this->GetPC(); - int64_t opr_index = this->PushOp(StackVM::RJUMP_IF_TRUE, 0); - this->Push(op->b); - int64_t diff = this->GetPC() - pc_jump; - this->SetOperand(opr_index, diff); -} - -void CodeGenStackVM::VisitExpr_(const NotNode* op) { - this->Push(op->a); - this->PushOp(StackVM::NOT); -} - -void CodeGenStackVM::VisitStmt_(const ForNode* op) { - ICHECK(is_zero(op->min)); - int vid = this->AllocVarID(op->loop_var.get()); - this->PushOp(StackVM::PUSH_I64, 0); - int64_t loop_head = this->GetPC(); - this->PushOp(StackVM::STORE_HEAP, vid); - this->PushOp(StackVM::LOAD_HEAP, vid); - this->Push(op->extent); - this->PushOp(StackVM::LT_I64); - int64_t label_fjump = this->GetPC(); - int64_t foward_jump = this->PushOp(StackVM::RJUMP_IF_FALSE, 0); - this->PushOp(StackVM::POP); - this->Push(op->body); - this->PushOp(StackVM::LOAD_HEAP, vid); - this->PushOp(StackVM::PUSH_I64, 1); - this->PushOp(StackVM::ADD_I64); - int64_t label_bjump = this->GetPC(); - int64_t backward_jump = this->PushOp(StackVM::RJUMP, 0); - int64_t loop_end = this->GetPC(); - this->PushOp(StackVM::POP); - this->SetOperand(foward_jump, loop_end - label_fjump); - this->SetOperand(backward_jump, loop_head - label_bjump); -} - -void CodeGenStackVM::VisitStmt_(const SeqStmtNode* op) { - for (Stmt stmt : op->seq) { - this->Push(stmt); - } -} - -void CodeGenStackVM::VisitStmt_(const EvaluateNode* ev) { - if (is_const_int(ev->value)) return; - const CallNode* op = ev->value.as(); - if (op && op->op.same_as(builtin::tvm_struct_set())) { - ICHECK_EQ(op->args.size(), 4U); - this->Push(op->args[0]); - this->Push(op->args[3]); - const IntImmNode* index = op->args[1].as(); - ICHECK(index != nullptr); - StackVM::Code code; - code.op_code = StackVM::TVM_STRUCT_SET; - vm_.code.push_back(code); - code.v_int = index->value; - vm_.code.push_back(code); - code.v_int = MapFieldKind(op->args[2].as()->value); - vm_.code.push_back(code); - } else { - this->Push(ev->value); - this->PushOp(StackVM::POP); - } -} - -void CodeGenStackVM::VisitStmt_(const IfThenElseNode* op) { - this->Push(op->condition); - int64_t label_ejump = this->GetPC(); - int64_t else_jump = this->PushOp(StackVM::RJUMP_IF_FALSE, 0); - this->PushOp(StackVM::POP); - this->Push(op->then_case); - if (op->else_case) { - int64_t label_then_jump = this->GetPC(); - int64_t then_jump = this->PushOp(StackVM::RJUMP, 0); - int64_t else_begin = this->GetPC(); - this->SetOperand(else_jump, else_begin - label_ejump); - this->PushOp(StackVM::POP); - this->Push(op->else_case.value()); - int64_t if_end = this->GetPC(); - this->SetOperand(then_jump, if_end - label_then_jump); - } else { - int64_t if_end = this->GetPC(); - this->SetOperand(else_jump, if_end - label_ejump); - this->PushOp(StackVM::POP); - } -} - -void CodeGenStackVM::VisitStmt_(const LetStmtNode* op) { - this->Push(op->value); - int64_t vid = this->AllocVarID(op->var.get()); - this->PushOp(StackVM::STORE_HEAP, static_cast(vid)); - this->Push(op->body); -} - -void CodeGenStackVM::VisitExpr_(const RampNode* op) { LOG(FATAL) << "Ramp is not supported"; } - -void CodeGenStackVM::VisitExpr_(const BroadcastNode* op) { - LOG(FATAL) << "Broadcast is not supported"; -} - -void CodeGenStackVM::VisitExpr_(const SelectNode* op) { - this->Push(op->true_value); - this->Push(op->false_value); - this->Push(op->condition); - this->PushOp(StackVM::SELECT); -} - -void CodeGenStackVM::VisitStmt_(const AssertStmtNode* op) { - if (const auto* str = op->message.as()) { - int sid = this->GetStrID(str->value); - this->Push(op->condition); - this->PushOp(StackVM::ASSERT, sid); - } - this->Push(op->body); -} - -void CodeGenStackVM::VisitStmt_(const AttrStmtNode* op) { this->Push(op->body); } - -void CodeGenStackVM::VisitExpr_(const LetNode* op) { - this->Push(op->value); - int64_t vid = this->AllocVarID(op->var.get()); - this->PushOp(StackVM::STORE_HEAP, static_cast(vid)); - this->Push(op->body); -} - -runtime::Module BuildStackVM(IRModule mod, Target target) { - std::unordered_map fmap; - std::string entry_func; - - for (auto kv : mod->functions) { - ICHECK(kv.second->IsInstance()) << "CodeGenStackVM: Can only take PrimFunc"; - auto f = Downcast(kv.second); - auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); - ICHECK(global_symbol.defined()) - << "CodeGenStackVM: Expect PrimFunc to have the global_symbol attribute"; - std::string f_name = global_symbol.value(); - StackVM vm = codegen::CodeGenStackVM().Compile(f); - ICHECK(!fmap.count(f_name)) << "Function name " << f_name << "already exist in list"; - fmap[f_name] = std::move(vm); - - if (f->HasNonzeroAttr(tir::attr::kIsEntryFunc)) { - entry_func = f_name; - } - } - - return runtime::StackVMModuleCreate(fmap, entry_func); -} - -TVM_REGISTER_GLOBAL("target.build.stackvm").set_body_typed(BuildStackVM); -} // namespace codegen -} // namespace tvm diff --git a/src/target/stackvm/codegen_stackvm.h b/src/target/stackvm/codegen_stackvm.h deleted file mode 100644 index 0bac55e3b2af..000000000000 --- a/src/target/stackvm/codegen_stackvm.h +++ /dev/null @@ -1,165 +0,0 @@ -/* - * 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 codegen_stack_vm.h - * \brief Codegen into Simple Stack VM. - */ -#ifndef TVM_TARGET_STACKVM_CODEGEN_STACKVM_H_ -#define TVM_TARGET_STACKVM_CODEGEN_STACKVM_H_ - -#include -#include -#include -#include -#include - -#include -#include -#include - -#include "../../runtime/stackvm/stackvm.h" - -namespace tvm { -namespace codegen { - -using namespace tir; -using runtime::StackVM; - -/*! - * \brief A base class to generate a stack VM. - * This module is used to generate host wrapper - * into device function when only device JIT is available. - */ -class CodeGenStackVM : public ExprFunctor, - public StmtFunctor { - public: - /*! - * \brief Generate a stack VM representing - * \param f The function to be compiled - * \param device_funcs The extern device functions to be linked. - * \note Only call compile once, - * create a new codegen object each time. - */ - StackVM Compile(const PrimFunc& f); - /*! \brief Push stmt to generate new code */ - void Push(const Stmt& n); - /*! \brief Push expr to generate new code */ - void Push(const PrimExpr& n) { VisitExpr(n); } - /*! - * \brief Push the opcode to the code. - * \param opcode The code to be pushed. - */ - void PushOp(StackVM::OpCode opcode); - /*! - * \brief Push the opcode and operand to the code. - * \param opcode The opcode. - * \param operand The operand to be pushed. - * \return operand_index, indicating location of operand - */ - int64_t PushOp(StackVM::OpCode opcode, int operand); - /*! - * \brief Set the relative jump offset to be offset. - * \param operand_index The indexed returned by PushOp. - * \param operand The operand to be set. - */ - void SetOperand(int64_t operand_index, int64_t operand); - /*! \return The current program pointer */ - int64_t GetPC() const { return static_cast(vm_.code.size()); } - /*! - * \brief Get string id in vm - * \param key The string to get id. - * \return the id of the string. - */ - int GetStrID(const std::string& key); - /*! - * \brief Allocate a variable name for a newly defined var. - * \param v The variable. - * \return the heap index of the var. - */ - int AllocVarID(const VarNode* v); - /*! - * \brief Get a variable name. - * \param v The variable. - * \return the heap index of the var. - */ - int GetVarID(const VarNode* v) const; - // Push binary operator - void PushBinary(StackVM::OpCode op_int64, const PrimExpr& a, const PrimExpr& b); - // push cast; - void PushCast(DataType dst, DataType src); - // overloadable functions - // expression - void VisitExpr_(const VarNode* op) final; - void VisitExpr_(const BufferLoadNode* op) final; - void VisitExpr_(const LetNode* op) final; - void VisitExpr_(const CallNode* op) final; - void VisitExpr_(const AddNode* op) final; - void VisitExpr_(const SubNode* op) final; - void VisitExpr_(const MulNode* op) final; - void VisitExpr_(const DivNode* op) final; - void VisitExpr_(const ModNode* op) final; - void VisitExpr_(const MinNode* op) final; - void VisitExpr_(const MaxNode* op) final; - void VisitExpr_(const EQNode* op) final; - void VisitExpr_(const NENode* op) final; - void VisitExpr_(const LTNode* op) final; - void VisitExpr_(const LENode* op) final; - void VisitExpr_(const GTNode* op) final; - void VisitExpr_(const GENode* op) final; - void VisitExpr_(const AndNode* op) final; - void VisitExpr_(const OrNode* op) final; - void VisitExpr_(const CastNode* op) final; - void VisitExpr_(const NotNode* op) final; - void VisitExpr_(const SelectNode* op) final; - void VisitExpr_(const RampNode* op) final; - void VisitExpr_(const BroadcastNode* op) final; - void VisitExpr_(const IntImmNode* op) final; - void VisitExpr_(const FloatImmNode* op) final; - void VisitExpr_(const StringImmNode* op) final; - // statment - void VisitStmt_(const LetStmtNode* op) final; - void VisitStmt_(const BufferStoreNode* op) final; - void VisitStmt_(const ForNode* op) final; - void VisitStmt_(const IfThenElseNode* op) final; - void VisitStmt_(const AllocateNode* op) final; - void VisitStmt_(const DeclBufferNode* op) final; - void VisitStmt_(const AttrStmtNode* op) final; - void VisitStmt_(const AssertStmtNode* op) final; - void VisitStmt_(const EvaluateNode* op) final; - void VisitStmt_(const SeqStmtNode* op) final; - - private: - bool debug_{false}; - /*! \brief The vm to be generated */ - StackVM vm_; - /*! \brief id of each variable */ - std::unordered_map var_idmap_; - /*! \brief id of each string */ - std::unordered_map str_idmap_; - /*! \brief id of each global function */ - std::unordered_map extern_fun_idmap_; - - Op backend_alloc_workspace_op_ = Op::Get("tir.TVMBackendAllocWorkspace"); - Op backend_free_workspace_op_ = Op::Get("tir.TVMBackendFreeWorkspace"); -}; - -} // namespace codegen -} // namespace tvm -#endif // TVM_TARGET_STACKVM_CODEGEN_STACKVM_H_ diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc index 62ba2787a367..a60f63dc2132 100644 --- a/src/target/target_kind.cc +++ b/src/target/target_kind.cc @@ -434,9 +434,6 @@ TVM_REGISTER_TARGET_KIND("hexagon", kDLHexagon) .add_attr_option("vtcm-capacity") .set_default_keys({"hexagon", "cpu"}); -TVM_REGISTER_TARGET_KIND("stackvm", kDLCPU) // line break - .set_default_keys({"cpu"}); - TVM_REGISTER_TARGET_KIND("ext_dev", kDLExtDev); TVM_REGISTER_TARGET_KIND("hybrid", kDLCPU); diff --git a/tests/python/codegen/test_target_codegen_device.py b/tests/python/codegen/test_target_codegen_device.py index 1adb337de0c4..4dad03d7004c 100644 --- a/tests/python/codegen/test_target_codegen_device.py +++ b/tests/python/codegen/test_target_codegen_device.py @@ -89,7 +89,7 @@ def test_add_pipeline(): sch.bind(d_xi, "threadIdx.x") sch.bind(d_xo, "blockIdx.x") - def check_target(device, host="stackvm"): + def check_target(device, host): if not tvm.testing.device_enabled(device) or not tvm.testing.device_enabled(host): return dev = tvm.device(device, 0) diff --git a/tests/python/codegen/test_target_codegen_extern.py b/tests/python/codegen/test_target_codegen_extern.py index 99069b1bd1ed..35227baaff5b 100644 --- a/tests/python/codegen/test_target_codegen_extern.py +++ b/tests/python/codegen/test_target_codegen_extern.py @@ -115,7 +115,6 @@ def check_target(target): f(a, c) tvm.testing.assert_allclose(c.numpy(), a.numpy()) - check_target("stackvm") check_target("llvm") diff --git a/tests/python/codegen/test_target_codegen_vm_basic.py b/tests/python/codegen/test_target_codegen_vm_basic.py deleted file mode 100644 index d1a3c7217aa9..000000000000 --- a/tests/python/codegen/test_target_codegen_vm_basic.py +++ /dev/null @@ -1,143 +0,0 @@ -# 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. -import tvm -import tvm.testing -from tvm import te -from tvm.script import tir as T, ir as I - -import numpy as np - - -def run_jit(fapi, check): - for target in ["llvm", "stackvm"]: - if not tvm.testing.device_enabled(target): - continue - f = tvm.driver.build(fapi, target=target) - s = f.get_source() - check(f) - - -def test_stack_vm_basic(): - a = tvm.nd.array(np.zeros(10, dtype="float32")) - - @tvm.register_func - def tvm_call_back_get_shape(shape0): - print(shape0) - assert shape0 == a.shape[0] - - n = te.size_var("n") - Ab = tvm.tir.decl_buffer((n,), "float32") - stmt = tvm.tir.Evaluate(tvm.tir.call_packed("tvm_call_back_get_shape", Ab.shape[0])) - - mod = tvm.IRModule.from_expr( - tvm.tir.PrimFunc([Ab], stmt).with_attr("global_symbol", "print_shape") - ) - - run_jit(mod, lambda f: f(a)) - - -@tvm.register_func -def tvm_stack_vm_print(*x): - print(x) - - -def test_stack_vm_loop(): - dtype = "int64" - n = te.size_var("n") - Ab = tvm.tir.decl_buffer((n,), dtype) - i = te.size_var("i") - - ib = tvm.tir.ir_builder.create() - A = ib.buffer_ptr(Ab) - with ib.for_range(0, n - 1, "i") as i: - A[i + 1] = A[i] + 1 - ib.emit(tvm.tir.call_packed("tvm_stack_vm_print", i)) - - stmt = ib.get() - mod = tvm.IRModule.from_expr(tvm.tir.PrimFunc([Ab], stmt).with_attr("global_symbol", "ramp")) - a = tvm.nd.array(np.zeros(10, dtype=dtype)) - - def check(f): - f(a) - np.testing.assert_equal(a.numpy(), np.arange(a.shape[0])) - - run_jit(mod, check) - - -def test_stack_vm_cond(): - dtype = "int64" - n = te.size_var("n") - Ab = tvm.tir.decl_buffer((n,), dtype) - - ib = tvm.tir.ir_builder.create() - A = ib.buffer_ptr(Ab) - with ib.for_range(0, n - 1, "i") as i: - with ib.if_scope(tvm.tir.EQ(i, 4)): - A[i + 1] = A[i] + 1 - with ib.else_scope(): - A[i + 1] = A[i] + 2 - - stmt = ib.get() - mod = tvm.IRModule.from_expr(tvm.tir.PrimFunc([Ab], stmt).with_attr("global_symbol", "test")) - - def check(f): - a = tvm.nd.array(np.zeros(10, dtype=dtype)) - f(a) - y = np.arange(a.shape[0]) * 2 - y[5:] -= 1 - np.testing.assert_equal(a.numpy(), y) - - run_jit(mod, check) - - -def test_vm_parallel(): - dtype = "int64" - n = te.size_var("n") - Ab = tvm.tir.decl_buffer((n,), dtype) - i = te.size_var("i") - ib = tvm.tir.ir_builder.create() - A = ib.buffer_ptr(Ab) - with ib.for_range(0, n, "i", kind="parallel") as i: - A[i] = A[i] + 1 - stmt = ib.get() - mod = tvm.IRModule.from_expr(tvm.tir.PrimFunc([Ab], stmt).with_attr("global_symbol", "test")) - - def check(f): - a = tvm.nd.array(np.zeros(10, dtype=dtype)) - f(a) - np.testing.assert_equal(a.numpy(), np.ones(a.shape[0])) - - run_jit(mod, check) - - -def test_codegen_decl_buffer(): - """The codegen should accept DeclBuffer nodes in its input""" - - @I.ir_module - class mod: - @T.prim_func - def kernel(A_data: T.handle("float32")): - T.func_attr({"global_symbol": "kernel"}) - A_buf = T.decl_buffer([256], dtype="float32", scope="global", data=A_data) - - target = tvm.target.Target("stackvm") - stackvm_codegen = tvm.get_global_func("target.build.stackvm") - stackvm_codegen(mod, target) - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py b/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py index 7d7f610123fe..89e8b9e35040 100644 --- a/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py +++ b/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py @@ -148,7 +148,6 @@ def check_packed_func(target="llvm"): def test_lower_packed_func(): check_packed_func("llvm") - check_packed_func("stackvm") @tvm.testing.requires_llvm diff --git a/tests/scripts/task_config_build_gpu.sh b/tests/scripts/task_config_build_gpu.sh index 74bb702a8b08..f306bdf8bf74 100755 --- a/tests/scripts/task_config_build_gpu.sh +++ b/tests/scripts/task_config_build_gpu.sh @@ -32,7 +32,6 @@ echo set\(USE_OPENCL_GTEST \"/googletest\"\) >> config.cmake echo set\(USE_LLVM \"/usr/bin/llvm-config-15 --link-static\"\) >> config.cmake echo set\(USE_RPC ON\) >> config.cmake echo set\(USE_SORT ON\) >> config.cmake -echo set\(USE_STACKVM_RUNTIME ON\) >> config.cmake echo set\(USE_BLAS openblas\) >> config.cmake echo set\(CMAKE_CXX_FLAGS -Werror\) >> config.cmake echo set\(USE_TENSORRT_CODEGEN ON\) >> config.cmake