From 97ea31c8f5d460c5ec401b146cfb16481bef6641 Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Sun, 20 Oct 2019 10:40:10 -0700 Subject: [PATCH] [Runtime] Enable option to use OpenMP thread pool (#4089) --- CMakeLists.txt | 4 ++++ cmake/config.cmake | 4 ++++ cmake/modules/OpenMP.cmake | 48 ++++++++++++++++++++++++++++++++++++++++++++++ src/runtime/thread_pool.cc | 26 +++++++++++++++++++++++++ 4 files changed, 82 insertions(+) create mode 100644 cmake/modules/OpenMP.cmake diff --git a/CMakeLists.txt b/CMakeLists.txt index 10730ac..1b7d5ef 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -33,6 +33,7 @@ tvm_option(USE_LLVM "Build with LLVM, can be set to specific llvm-config path" O tvm_option(USE_STACKVM_RUNTIME "Include stackvm into the runtime" OFF) tvm_option(USE_GRAPH_RUNTIME "Build with tiny graph runtime" ON) tvm_option(USE_GRAPH_RUNTIME_DEBUG "Build with tiny graph runtime debug mode" OFF) +tvm_option(USE_OPENMP "Build with OpenMP thread pool implementation" OFF) tvm_option(USE_RELAY_DEBUG "Building Relay in debug mode..." OFF) tvm_option(USE_SGX "Build with SGX" OFF) tvm_option(USE_RTTI "Build with RTTI" ON) @@ -154,6 +155,7 @@ list(APPEND COMPILER_SRCS ${RELAY_BACKEND_SRCS}) list(APPEND COMPILER_SRCS ${RELAY_IR_SRCS}) list(APPEND COMPILER_SRCS ${RELAY_QNN_SRCS}) + if(USE_VM_PROFILER) message(STATUS "Build compiler with Relay VM profiler support...") file(GLOB BACKEND_VM_PROFILER_SRCS src/relay/backend/vm/profiler/*.cc) @@ -233,6 +235,7 @@ include(cmake/modules/VTA.cmake) include(cmake/modules/CUDA.cmake) include(cmake/modules/OpenCL.cmake) include(cmake/modules/OpenGL.cmake) +include(cmake/modules/OpenMP.cmake) include(cmake/modules/Vulkan.cmake) include(cmake/modules/Metal.cmake) include(cmake/modules/ROCM.cmake) @@ -264,6 +267,7 @@ add_library(tvm SHARED ${COMPILER_SRCS} ${RUNTIME_SRCS}) add_library(tvm_topi SHARED ${TOPI_SRCS}) add_library(tvm_runtime SHARED ${RUNTIME_SRCS}) + if(USE_RELAY_DEBUG) message(STATUS "Building Relay in debug mode...") set_target_properties(tvm PROPERTIES COMPILE_DEFINITIONS "USE_RELAY_DEBUG") diff --git a/cmake/config.cmake b/cmake/config.cmake index d92c215..6a55397 100644 --- a/cmake/config.cmake +++ b/cmake/config.cmake @@ -115,6 +115,10 @@ set(USE_BLAS none) # set(USE_MKL_PATH ) if using `pip install mkl` set(USE_MKL_PATH none) +# Whether use OpenMP thread pool, choices: gnu, intel +# Note: "gnu" uses gomp library, "intel" uses iomp5 library +set(USE_OPENMP none) + # Whether use contrib.random in runtime set(USE_RANDOM OFF) diff --git a/cmake/modules/OpenMP.cmake b/cmake/modules/OpenMP.cmake new file mode 100644 index 0000000..5dd9be5 --- /dev/null +++ b/cmake/modules/OpenMP.cmake @@ -0,0 +1,48 @@ +# 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. + +# OpenMP Module +if(USE_OPENMP STREQUAL "gnu") + find_package(OpenMP) + if(OPENMP_FOUND) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") + list(APPEND TVM_RUNTIME_LINKER_LIBS ${OpenMP_CXX_LIBRARIES}) + add_definitions(-DTVM_THREADPOOL_USE_OPENMP=1) + message(STATUS "Build with OpenMP ${OpenMP_CXX_LIBRARIES}") + else() + add_definitions(-DTVM_THREADPOOL_USE_OPENMP=0) + message(WARNING "OpenMP cannot be found, use TVM threadpool instead.") + endif() +elseif(USE_OPENMP STREQUAL "intel") + find_package(OpenMP) + if(OPENMP_FOUND) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") + if (MSVC) + find_library(OMP_LIBRARY NAMES libiomp5md) + else() + find_library(OMP_LIBRARY NAMES iomp5) + endif() + list(APPEND TVM_RUNTIME_LINKER_LIBS ${OMP_LIBRARY}) + add_definitions(-DTVM_THREADPOOL_USE_OPENMP=1) + message(STATUS "Build with OpenMP " ${OMP_LIBRARY}) + else() + add_definitions(-DTVM_THREADPOOL_USE_OPENMP=0) + message(WARNING "OpenMP cannot be found, use TVM threadpool instead.") + endif() +else() + add_definitions(-DTVM_THREADPOOL_USE_OPENMP=0) +endif() diff --git a/src/runtime/thread_pool.cc b/src/runtime/thread_pool.cc index 2e10136..e9e6d03 100644 --- a/src/runtime/thread_pool.cc +++ b/src/runtime/thread_pool.cc @@ -29,6 +29,9 @@ #include #include #include +#if TVM_THREADPOOL_USE_OPENMP +#include +#endif #include #include #include @@ -394,12 +397,34 @@ int TVMBackendParallelLaunch( FTVMParallelLambda flambda, void* cdata, int num_task) { +#if !TVM_THREADPOOL_USE_OPENMP int res = tvm::runtime::ThreadPool::ThreadLocal()->Launch( flambda, cdata, num_task, 1); return res; +#else + int num_workers = tvm::runtime::threading::MaxConcurrency(); + if (num_task == 0) num_task = num_workers; + omp_set_num_threads(num_workers); + #pragma omp parallel num_threads(num_workers) + { + TVMParallelGroupEnv env; + env.num_task = num_task; + std::atomic* sync_counter = new std::atomic[num_task * tvm::runtime::kSyncStride]; + for (int i = 0; i < num_task; ++i) { + sync_counter[i * tvm::runtime::kSyncStride].store( + 0, std::memory_order_relaxed); + } + env.sync_handle = sync_counter; + (*flambda)(omp_get_thread_num(), &env, cdata); + } + return 0; +#endif } int TVMBackendParallelBarrier(int task_id, TVMParallelGroupEnv* penv) { +#if TVM_THREADPOOL_USE_OPENMP + #pragma omp barrier +#else using tvm::runtime::kSyncStride; int num_task = penv->num_task; std::atomic* sync_counter = @@ -415,5 +440,6 @@ int TVMBackendParallelBarrier(int task_id, TVMParallelGroupEnv* penv) { } } std::atomic_thread_fence(std::memory_order_acquire); +#endif return 0; } -- 2.7.4