[Enhancement] Support tvm

.codespell_ignore.txt

@@ -1,2 +1,3 @@
 cann
 CANN
+nd

.github/workflows/backend-tvm.yml

@@ -0,0 +1,41 @@
+name: backend-tvm
+
+on:
+  push:
+    paths-ignore:
+      - "demo/**"
+      - "tools/**"
+
+  pull_request:
+    paths-ignore:
+      - "demo/**"
+      - "tools/**"
+      - "docs/**"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  script_install:
+    runs-on: ubuntu-20.04
+    strategy:
+      matrix:
+        python-version: [3.7]
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+        with:
+          submodules: 'recursive'
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install mmdeploy
+        run: |
+          python3 tools/scripts/build_ubuntu_x64_tvm.py
+          source ~/mmdeploy.env
+          python3 -m pip install torch==1.8.2 torchvision==0.9.2 --extra-index-url https://download.pytorch.org/whl/lts/1.8/cpu
+          python3 -m pip install mmcv-full==1.5.1  -f https://download.openmmlab.com/mmcv/dist/cpu/torch1.8.0/index.html
+          python3 -m pip install decorator psutil scipy attrs tornado pytest
+          python3 -c 'import mmdeploy.apis.tvm as tvm_api; assert tvm_api.is_available()'

.gitmodules

@@ -1,9 +1,9 @@
 [submodule "third_party/cub"]
-	path = third_party/cub
-	url = https://github.com/NVIDIA/cub.git
+path = third_party/cub
+url = https://github.com/NVIDIA/cub.git
 [submodule "third_party/pybind11"]
-	path = third_party/pybind11
-	url = https://github.com/pybind/pybind11.git
+path = third_party/pybind11
+url = https://github.com/pybind/pybind11.git
 [submodule "third_party/spdlog"]
-	path = third_party/spdlog
-	url = https://github.com/gabime/spdlog.git
+path = third_party/spdlog
+url = https://github.com/gabime/spdlog.git

cmake/modules/FindTVM.cmake

@@ -0,0 +1,47 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+if (NOT DEFINED TVM_DIR)
+    set(TVM_DIR $ENV{TVM_DIR})
+endif ()
+if (NOT TVM_DIR)
+    message(FATAL_ERROR "Please set TVM_DIR with cmake -D option.")
+endif()
+
+find_path(
+    TVM_INCLUDE_DIR tvm/runtime/c_runtime_api.h
+    HINTS ${TVM_DIR}
+    PATH_SUFFIXES include)
+
+find_path(
+    DMLC_CORE_INCLUDE_DIR  dmlc/io.h
+    HINTS ${TVM_DIR}/3rdparty/dmlc-core
+    PATH_SUFFIXES include)
+
+find_path(
+    DLPACK_INCLUDE_DIR dlpack/dlpack.h
+    HINTS ${TVM_DIR}/3rdparty/dlpack
+    PATH_SUFFIXES include)
+
+find_library(
+    TVM_LIBRARY_PATH tvm_runtime
+    HINTS ${TVM_DIR}
+    PATH_SUFFIXES build lib build/Release build/Debug)
+if (NOT (TVM_INCLUDE_DIR AND DMLC_CORE_INCLUDE_DIR AND DLPACK_INCLUDE_DIR AND TVM_LIBRARY_PATH))
+    message(FATAL_ERROR "Couldn't find tvm in TVM_DIR: "
+        "${TVM_DIR}, please check if the path is correct.")
+endif()
+
+add_library(tvm_runtime SHARED IMPORTED)
+set_property(TARGET tvm_runtime APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE)
+if (MSVC)
+    set_target_properties(tvm_runtime PROPERTIES
+        IMPORTED_IMPLIB_RELEASE ${TVM_LIBRARY_PATH}
+        INTERFACE_INCLUDE_DIRECTORIES ${TVM_INCLUDE_DIR} ${DMLC_CORE_INCLUDE_DIR} ${DLPACK_INCLUDE_DIR}
+    )
+
+else()
+    set_target_properties(tvm_runtime PROPERTIES
+        IMPORTED_LOCATION_RELEASE ${TVM_LIBRARY_PATH}
+        INTERFACE_INCLUDE_DIRECTORIES ${TVM_INCLUDE_DIR} ${DMLC_CORE_INCLUDE_DIR} ${DLPACK_INCLUDE_DIR}
+    )
+endif()

configs/_base_/backends/tvm.py

@@ -0,0 +1 @@
+backend_config = dict(type='tvm')

configs/mmcls/classification_tvm-ansor_static-224x224.py

@@ -0,0 +1,12 @@
+_base_ = ['./classification_static.py', '../_base_/backends/tvm.py']
+
+onnx_config = dict(input_shape=[224, 224])
+backend_config = dict(model_inputs=[
+    dict(
+        shape=dict(input=[1, 3, 224, 224]),
+        dtype=dict(input='float32'),
+        tuner=dict(
+            type='AutoScheduleTuner',
+            log_file='tvm_tune_log.log',
+            num_measure_trials=2000))
+])

configs/mmcls/classification_tvm-autotvm-int8_static-224x224.py

@@ -0,0 +1,16 @@
+_base_ = ['./classification_tvm-autotvm_static-224x224.py']
+
+calib_config = dict(create_calib=True, calib_file='calib_data.h5')
+backend_config = dict(model_inputs=[
+    dict(
+        shape=dict(input=[1, 3, 224, 224]),
+        dtype=dict(input='float32'),
+        tuner=dict(
+            type='AutoTVMTuner',
+            log_file='tvm_tune_log.log',
+            n_trial=1000,
+            tuner=dict(type='XGBTuner'),
+        ),
+        qconfig=dict(calibrate_mode='kl_divergence', weight_scale='max'),
+    )
+])

configs/mmcls/classification_tvm-autotvm_static-224x224.py

@@ -0,0 +1,13 @@
+_base_ = ['./classification_static.py', '../_base_/backends/tvm.py']
+
+onnx_config = dict(input_shape=[224, 224])
+backend_config = dict(model_inputs=[
+    dict(
+        shape=dict(input=[1, 3, 224, 224]),
+        dtype=dict(input='float32'),
+        tuner=dict(
+            type='AutoTVMTuner',
+            log_file='tvm_tune_log.log',
+            n_trial=1000,
+            tuner=dict(type='XGBTuner')))
+])

configs/mmdet/detection/detection_tvm-ansor_static-800x1344.py

@@ -0,0 +1,13 @@
+_base_ = ['../_base_/base_static.py', '../../_base_/backends/tvm.py']
+
+onnx_config = dict(input_shape=[1344, 800])
+backend_config = dict(model_inputs=[
+    dict(
+        use_vm=True,
+        shape=dict(input=[1, 3, 800, 1344]),
+        dtype=dict(input='float32'),
+        tuner=dict(
+            type='AutoScheduleTuner',
+            log_file='tvm_tune_log.log',
+            num_measure_trials=2000))
+])

configs/mmdet/detection/detection_tvm-autotvm_static-300x300.py

@@ -0,0 +1,15 @@
+_base_ = ['../_base_/base_static.py', '../../_base_/backends/tvm.py']
+
+onnx_config = dict(input_shape=[300, 300])
+backend_config = dict(model_inputs=[
+    dict(
+        use_vm=True,
+        shape=dict(input=[1, 3, 300, 300]),
+        dtype=dict(input='float32'),
+        tuner=dict(
+            type='AutoTVMTuner',
+            log_file='tvm_tune_log.log',
+            n_trial=1000,
+            tuner=dict(type='XGBTuner'),
+        ))
+])

configs/mmdet/detection/detection_tvm-autotvm_static-800x1344.py

@@ -0,0 +1,15 @@
+_base_ = ['../_base_/base_static.py', '../../_base_/backends/tvm.py']
+
+onnx_config = dict(input_shape=[1344, 800])
+backend_config = dict(model_inputs=[
+    dict(
+        use_vm=True,
+        shape=dict(input=[1, 3, 800, 1344]),
+        dtype=dict(input='float32'),
+        tuner=dict(
+            type='AutoTVMTuner',
+            log_file='tvm_tune_log.log',
+            n_trial=1000,
+            tuner=dict(type='XGBTuner'),
+        ))
+])

configs/mmdet/instance-seg/instance-seg_tvm-ansor_static-800x1344.py

@@ -0,0 +1,15 @@
+_base_ = [
+    '../_base_/base_instance-seg_static.py', '../../_base_/backends/tvm.py'
+]
+
+onnx_config = dict(input_shape=[1344, 800])
+backend_config = dict(model_inputs=[
+    dict(
+        use_vm=True,
+        shape=dict(input=[1, 3, 800, 1344]),
+        dtype=dict(input='float32'),
+        tuner=dict(
+            type='AutoScheduleTuner',
+            log_file='tvm_tune_log.log',
+            num_measure_trials=20000))
+])

configs/mmdet/instance-seg/instance-seg_tvm-autotvm_static-800x1344.py

@@ -0,0 +1,17 @@
+_base_ = [
+    '../_base_/base_instance-seg_static.py', '../../_base_/backends/tvm.py'
+]
+
+onnx_config = dict(input_shape=[1344, 800])
+backend_config = dict(model_inputs=[
+    dict(
+        use_vm=True,
+        shape=dict(input=[1, 3, 800, 1344]),
+        dtype=dict(input='float32'),
+        tuner=dict(
+            type='AutoTVMTuner',
+            log_file='tvm_tune_log.log',
+            n_trial=10000,
+            tuner=dict(type='XGBTuner'),
+        ))
+])

configs/mmseg/segmentation_tvm-ansor_static-512x1024.py

@@ -0,0 +1,12 @@
+_base_ = ['./segmentation_static.py', '../_base_/backends/tvm.py']
+
+onnx_config = dict(input_shape=[1024, 512])
+backend_config = dict(model_inputs=[
+    dict(
+        shape=dict(input=[1, 3, 512, 1024]),
+        dtype=dict(input='float32'),
+        tuner=dict(
+            type='AutoScheduleTuner',
+            log_file='tvm_tune_log.log',
+            num_measure_trials=2000))
+])

configs/mmseg/segmentation_tvm-autotvm_static-512x1024.py

@@ -0,0 +1,13 @@
+_base_ = ['./segmentation_static.py', '../_base_/backends/tvm.py']
+
+onnx_config = dict(input_shape=[1024, 512])
+backend_config = dict(model_inputs=[
+    dict(
+        shape=dict(input=[1, 3, 512, 1024]),
+        dtype=dict(input='float32'),
+        tuner=dict(
+            type='AutoTVMTuner',
+            log_file='tvm_tune_log.log',
+            n_trial=1000,
+            tuner=dict(type='XGBTuner')))
+])

csrc/mmdeploy/core/tensor.h

@@ -14,6 +14,7 @@ namespace mmdeploy {
 namespace framework {
 
 using TensorShape = std::vector<int64_t>;
+
 struct TensorDesc {
   Device device;
   DataType data_type{DataType::kFLOAT};

csrc/mmdeploy/device/cpu/cpu_device.cpp

@@ -11,13 +11,16 @@ class CpuHostMemory : public NonCopyable {
  public:
   CpuHostMemory() : size_(), data_(), owned_data_{false} {}
   Result<void> Init(size_t size, size_t alignment) {
-    if (alignment != 1) {
-      return Status(eNotSupported);
-    }
-    data_ = std::malloc(size);
+    size_t space = (size + alignment - 1) / alignment * alignment;
+#ifdef _MSC_VER
+    data_ = _aligned_malloc(space, alignment);
+#else
+    data_ = std::aligned_alloc(alignment, space);
+#endif
     if (!data_) {
       return Status(eOutOfMemory);
     }
+    aligned_data_ = data_;
     size_ = size;
     owned_data_ = true;
     return success();
@@ -38,7 +41,11 @@ class CpuHostMemory : public NonCopyable {
   ~CpuHostMemory() {
     if (data_) {
       if (owned_data_) {
+#ifdef _MSC_VER
+        _aligned_free(data_);
+#else
         std::free(data_);
+#endif
         owned_data_ = false;
       }
       data_ = nullptr;
@@ -47,11 +54,12 @@ class CpuHostMemory : public NonCopyable {
     size_ = 0;
   }
   size_t size() const { return size_; }
-  void* data() const { return data_; }
+  void* data() const { return owned_data_ ? aligned_data_ : data_; }
 
  private:
   size_t size_;
   void* data_;
+  void* aligned_data_{nullptr};
   bool owned_data_;
   std::shared_ptr<void> external_;
 };

csrc/mmdeploy/device/cuda/cuda_device.cpp

@@ -69,7 +69,7 @@ class CudaDeviceMemory : public NonCopyable {
  public:
   explicit CudaDeviceMemory(int device_id) : device_id_(device_id), size_(), owned_block_() {}
   Result<void> Init(size_t size, Allocator allocator, size_t alignment, uint64_t flags) {
-    if (alignment != 1) {
+    if (alignment > 256 || 256 % alignment != 0) {
       return Status(eNotSupported);
     }
     allocator_ = std::move(allocator);

csrc/mmdeploy/net/CMakeLists.txt

@@ -38,6 +38,10 @@ if ("coreml" IN_LIST MMDEPLOY_TARGET_BACKENDS)
     add_subdirectory(coreml)
 endif ()
 
+if ("tvm" IN_LIST MMDEPLOY_TARGET_BACKENDS)
+    add_subdirectory(tvm)
+endif ()
+
 if ("rknn" IN_LIST MMDEPLOY_TARGET_BACKENDS)
     add_subdirectory(rknn)
 endif ()

csrc/mmdeploy/net/tvm/CMakeLists.txt

@@ -0,0 +1,11 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+project(mmdeploy_tvm_net)
+
+include(${CMAKE_SOURCE_DIR}/cmake/modules/FindTVM.cmake)
+
+mmdeploy_add_module(${PROJECT_NAME} tvm_net.cpp)
+target_include_directories(${PROJECT_NAME} PRIVATE ${TVM_INCLUDE_DIR} ${DLPACK_INCLUDE_DIR} ${DMLC_CORE_INCLUDE_DIR})
+target_link_libraries(${PROJECT_NAME} PRIVATE tvm_runtime mmdeploy_dlpack_utils)
+
+add_library(mmdeploy::tvm_net ALIAS ${PROJECT_NAME})