odla torchscript

ODLA/platforms/CMakeLists.txt

@@ -86,3 +86,8 @@ option(ODLA_BUILD_TF_Wrapper "Build ODLA Tensorflow Wrapper" OFF)
 if (ODLA_BUILD_TF_Wrapper)
  add_subdirectory(tensorflow)
 endif()
+
+option(ODLA_BUILD_TORCH_Wrapper "Build ODLA TORCHSCRIPT Wrapper" OFF)
+if (ODLA_BUILD_TORCH_Wrapper)
+ add_subdirectory(torch)
+endif()

ODLA/platforms/torch/CMakeLists.txt

@@ -0,0 +1,65 @@
+# ==============================================================================
+# Copyright (C) 2022 Alibaba Group Holding Limited.
+#
+# Licensed 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
+# ==============================================================================
+add_odla_library(odla_torch SHARED odla_torchscript.cc)
+
+if (CMAKE_PREFIX_PATH)
+ find_package(Torch REQUIRED)
+endif()
+if (TORCH_INSTALL_PREFIX)
+ set(TORCH_PATH ${TORCH_INSTALL_PREFIX})
+else()
+ execute_process(
+ COMMAND python3 -c "import torch; print(torch.__path__[0])"
+ OUTPUT_VARIABLE TORCH_PATH
+ OUTPUT_STRIP_TRAILING_WHITESPACE
+ RESULT_VARIABLE retcode)
+endif()
+if (NOT TORCH_PATH)
+ message(FATAL_ERROR "torch install is not found.")
+else()
+ message(STATUS "torch install path: ${TORCH_PATH}")
+endif()
+
+set(PYTORCH_INC_DIR
+ ${TORCH_PATH}/include
+ ${TORCH_PATH}/torch/csrc/api/include
+ ${TORCH_PATH}/include/TH
+ ${TORCH_PATH}/include/THC
+)
+message(STATUS "torch include dirs: ${PYTORCH_INC_DIR}")
+target_include_directories(odla_torch PRIVATE ${PYTORCH_INC_DIR})
+
+set(PYTORCH_LIBS
+ -lc10
+ -ltorch
+ -ltorch_cpu
+ -ltorch_python)
+
+if (NOT TORCH_CXX_FLAGS)
+ execute_process(
+ COMMAND python3 -c "import torch; print(torch._C._GLIBCXX_USE_CXX11_ABI)"
+ OUTPUT_VARIABLE TORCH_CXX11_ABI
+ OUTPUT_STRIP_TRAILING_WHITESPACE
+ RESULT_VARIABLE retcode)
+ message(STATUS "Torch CXX flags:-D_GLIBCXX_USE_CXX11_ABI=${TORCH_CXX11_ABI}")
+ target_compile_definitions(odla_torch PRIVATE _GLIBCXX_USE_CXX11_ABI=${TORCH_CXX11_ABI})
+else()
+ message(STATUS "Torch CXX flags:${TORCH_CXX_FLAGS}")
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
+endif()
+target_link_options(odla_torch PUBLIC -Wl,-rpath,${TORCH_PATH}/lib -L${TORCH_PATH}/lib)
+target_link_libraries(odla_torch PUBLIC ${PYTORCH_LIBS})
+target_link_libraries(odla_torch PUBLIC ODLA)

ODLA/platforms/torch/odla_torchscript.cc

@@ -0,0 +1,264 @@
+//===- odla_torchscript.cc ------------------------------------------------===//
+//
+// Copyright (C) 2022 Alibaba Group Holding Limited.
+//
+// Licensed 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.
+// =============================================================================
+#include <c10/util/ArrayRef.h>
+#include <torch/csrc/jit/serialization/import.h>
+
+#include <ODLA/odla.h>
+
+#include <sstream>
+
+const uint32_t MAX_OUTPUT_TENSORS = 10;
+const uint32_t MAX_INPUT_TENSORS = 20;
+
+struct _odla_device {
+ c10::DeviceType device_t_;
+};
+
+struct _odla_value {
+ _odla_value(uint32_t v):id_(v) {}
+ uint32_t id_;
+};
+
+struct _odla_executable {
+ torch::jit::Module module_;
+ std::vector<odla_value> odla_inputs_outputs_;
+ odla_uint32 num_inputs_;
+};
+
+struct _odla_context {
+ _odla_context();
+ std::vector<torch::jit::IValue> inputs_;
+ std::vector<odla_value_type> input_types_;
+ torch::jit::IValue output_;
+ std::vector<at::Tensor> output_tensors_;
+ odla_uint32 num_output_tensors_;
+};
+
+_odla_context::_odla_context() {
+ inputs_.resize(MAX_INPUT_TENSORS);
+ input_types_.resize(MAX_INPUT_TENSORS);
+}
+
+size_t static getElementCount(const odla_value_shape& dims) {
+ return dims.size == 0 ? 1
+ : std::accumulate(dims.dims, dims.dims + dims.size, 1,
+ std::multiplies<size_t>());
+}
+
+c10::IntArrayRef static toTensorDim(odla_value_shape& dims) {
+ return dims.size == 0 ? c10::IntArrayRef(1) : 
+ c10::IntArrayRef(dims.dims, dims.size);
+}
+
+c10::ScalarType static toTensorDataType(odla_element_type dt) {
+ static const std::unordered_map<odla_element_type, c10::ScalarType> dt_map = {
+ {ODLA_FLOAT32, c10::ScalarType::Float},
+ {ODLA_INT32, c10::ScalarType::Int},
+ {ODLA_BOOL, c10::ScalarType::Bool}
+ };
+ auto it = dt_map.find(dt);
+ return it == dt_map.end() ? c10::ScalarType::Float : it->second;
+}
+
+odla_element_type static toODLADataType(const c10::ScalarType& st) {
+ static const std::unordered_map<c10::ScalarType, odla_element_type> dt_map = {
+ {c10::ScalarType::Float, ODLA_FLOAT32},
+ {c10::ScalarType::Int, ODLA_INT32},
+ {c10::ScalarType::Bool, ODLA_BOOL}
+ };
+ auto it = dt_map.find(st);
+ return it == dt_map.end() ? ODLA_FLOAT32 : it->second;
+}
+
+odla_value_type static toODLAValueType(const c10::ScalarType& dt, at::IntArrayRef dims) {
+ odla_value_type ty;
+ ty.element_type = toODLADataType(dt);
+ ty.shape.size = dims.size();
+ int i = 0;
+ for (auto d : dims) {
+ ty.shape.dims[i++] = d;
+ }
+ return ty;
+}
+
+static std::unordered_map<odla_context, std::unique_ptr<_odla_context>> g_ctxs;
+static std::unordered_map<odla_executable, std::unique_ptr<_odla_executable>>
+ g_executables;
+
+static _odla_device g_device{c10::kCUDA};
+
+odla_status odla_AllocateDevice(const odla_vendor vendor,
+ const odla_device_name device_name,
+ odla_device* device,
+ const char* config) {
+ *device = &g_device;
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_LoadExecutable(odla_resource_location location, 
+ odla_device device,
+ odla_executable* computation) {
+ *computation = nullptr;
+ if (location.location_type != ODLA_LOCATION_MEMORY &&
+ location.location_type != ODLA_LOCATION_PATH) {
+ return ODLA_FAILURE;
+ }
+ auto comp = std::make_unique<_odla_executable>();
+ if (location.location_type == ODLA_LOCATION_MEMORY) {
+ std::istringstream s;
+ s.rdbuf()->pubsetbuf(const_cast<char*>(
+ reinterpret_cast<const char*>(location.location)),
+ location.size);
+ comp->module_ = torch::jit::load(s, c10::Device(g_device.device_t_));
+ } else {
+ comp->module_ = torch::jit::load(reinterpret_cast<const char*>(
+ location.location),
+ c10::Device(g_device.device_t_));
+ }
+ auto schema = comp->module_.get_method("forward").function().getSchema();
+ assert(!schema.is_vararg());
+ assert(!schema.is_varret());
+ auto num_inputs = comp->module_.get_method("forward").function().num_inputs();
+ comp->num_inputs_ = num_inputs - 1;
+ for (uint32_t idx = 0; idx < std::max(comp->num_inputs_, MAX_OUTPUT_TENSORS); ++idx) {
+ auto v = std::make_unique<_odla_value>(idx);
+ comp->odla_inputs_outputs_.push_back(v.get());
+ }
+ *computation = comp.get();
+ g_executables[*computation] = std::move(comp);
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_GetArgFromExecutableByIdx(odla_executable comp, 
+ odla_uint32 idx, 
+ odla_value* value) {
+ if (idx > comp->num_inputs_) {
+ *value = nullptr;
+ return ODLA_FAILURE;
+ }
+ *value = comp->odla_inputs_outputs_[idx];
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_GetOutputFromExecutableByIdx(const odla_executable comp,
+ const odla_uint32 output_idx,
+ odla_value* output_value) {
+ if (output_idx > comp->odla_inputs_outputs_.size()) {
+ *output_value = nullptr;
+ return ODLA_FAILURE;
+ }
+ *output_value = comp->odla_inputs_outputs_[output_idx];
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_CreateContext(odla_context* context) {
+ *context = nullptr;
+ auto ctx = std::make_unique<_odla_context>();
+ *context = ctx.get();
+ g_ctxs[*context] = std::move(ctx);
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_SetRuntimeValueType(odla_context context, odla_value v, odla_value_type ty) {
+ assert(v->id_ < MAX_INPUT_TENSORS);
+ context->input_types_[v->id_] = std::move(ty);
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_GetRuntimeValueType(odla_context context, odla_value value, odla_value_type* ty) {
+ assert(value->id_ <= context->num_output_tensors_);
+ auto t = context->output_tensors_[value->id_];
+ *ty = toODLAValueType(t.scalar_type(), t.sizes());
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_BindToArgument(odla_value value, const odla_void* data_ptr,
+ odla_context context) {
+ assert(value->id_ < MAX_INPUT_TENSORS);
+ auto ty = context->input_types_[value->id_];
+ auto options = c10::TensorOptions()
+ .dtype(toTensorDataType(ty.element_type))
+ .device(c10::kCPU);
+ auto t = at::from_blob(const_cast<void*>(data_ptr), toTensorDim(ty.shape), options);
+ if (g_device.device_t_ == c10::kCUDA) {
+ t = t.to(c10::device(c10::kCUDA));
+ }
+ context->inputs_[value->id_] = c10::IValue(t);
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_BindToOutput(odla_value value, odla_void* data_ptr,
+ odla_context context) {
+ assert(value->id_ < context->num_output_tensors_);
+ auto t = context->output_tensors_[value->id_];
+ auto ty = toODLAValueType(t.scalar_type(), t.sizes());
+ void* raw_data = t.storage().data();
+ int len = at::elementSize(t.scalar_type()) * getElementCount(ty.shape);
+ if (g_device.device_t_ == c10::kCPU) {
+ memcpy(data_ptr, raw_data, len);
+ } else {
+ // cudaMemcpy(data_ptr, raw_data, len, cudaMemcpyDeviceToHost);
+ t = t.to(c10::Device(c10::kCPU));
+ memcpy(data_ptr, t.storage().data(), len);
+ }
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_GetRuntimeNumOfOutputs(odla_context context, 
+ odla_uint32 *num_output_ptr) {
+ *num_output_ptr = (odla_uint32)context->num_output_tensors_;
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_LaunchExecutable(const odla_executable computation,
+ const odla_context context) {
+ context->inputs_.resize(computation->num_inputs_);
+ context->input_types_.resize(computation->num_inputs_);
+ context->output_ = computation->module_.forward(context->inputs_);
+
+ if (context->output_.isTensor()) {
+ context->output_tensors_.push_back(context->output_.toTensor());
+ } else {
+ assert(context->output_.isTuple());
+ for (const auto& item : context->output_.toTuple()->elements()) {
+ assert (item.isTensor());
+ context->output_tensors_.push_back(item.toTensor());
+ }
+ }
+ context->num_output_tensors_ = context->output_tensors_.size();
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_DestroyContext(odla_context context) {
+ auto it = g_ctxs.find(context);
+ if (it == g_ctxs.end()) {
+ return ODLA_FAILURE;
+ }
+ g_ctxs.erase(it);
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_DestroyExecutable(odla_executable computation) {
+ auto it = g_executables.find(computation);
+ if (it == g_executables.end()) {
+ return ODLA_FAILURE;
+ }
+ g_executables.erase(it);
+ return ODLA_SUCCESS;
+}
+
+odla_status odla_DestroyDevice(odla_device device) { return ODLA_SUCCESS; }