Skip to content

Commit ffc11b7

Browse files
icemelonkevinthesun
authored andcommitted
[Runtime] Enable option to use OpenMP thread pool (apache#4089)
1 parent 6f5d9f2 commit ffc11b7

File tree

4 files changed

+82
-0
lines changed

4 files changed

+82
-0
lines changed

CMakeLists.txt

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ tvm_option(USE_LLVM "Build with LLVM, can be set to specific llvm-config path" O
3333
tvm_option(USE_STACKVM_RUNTIME "Include stackvm into the runtime" OFF)
3434
tvm_option(USE_GRAPH_RUNTIME "Build with tiny graph runtime" ON)
3535
tvm_option(USE_GRAPH_RUNTIME_DEBUG "Build with tiny graph runtime debug mode" OFF)
36+
tvm_option(USE_OPENMP "Build with OpenMP thread pool implementation" OFF)
3637
tvm_option(USE_RELAY_DEBUG "Building Relay in debug mode..." OFF)
3738
tvm_option(USE_SGX "Build with SGX" OFF)
3839
tvm_option(USE_RTTI "Build with RTTI" ON)
@@ -155,6 +156,7 @@ list(APPEND COMPILER_SRCS ${RELAY_BACKEND_SRCS})
155156
list(APPEND COMPILER_SRCS ${RELAY_IR_SRCS})
156157
list(APPEND COMPILER_SRCS ${RELAY_QNN_SRCS})
157158

159+
158160
if(USE_VM_PROFILER)
159161
message(STATUS "Build compiler with Relay VM profiler support...")
160162
file(GLOB BACKEND_VM_PROFILER_SRCS src/relay/backend/vm/profiler/*.cc)
@@ -234,6 +236,7 @@ include(cmake/modules/VTA.cmake)
234236
include(cmake/modules/CUDA.cmake)
235237
include(cmake/modules/OpenCL.cmake)
236238
include(cmake/modules/OpenGL.cmake)
239+
include(cmake/modules/OpenMP.cmake)
237240
include(cmake/modules/Vulkan.cmake)
238241
include(cmake/modules/Metal.cmake)
239242
include(cmake/modules/ROCM.cmake)
@@ -267,6 +270,7 @@ add_library(tvm_topi SHARED ${TOPI_SRCS})
267270
add_library(tvm_runtime SHARED ${RUNTIME_SRCS})
268271
add_library(tvm_runtime_static STATIC ${RUNTIME_SRCS})
269272

273+
270274
if(USE_RELAY_DEBUG)
271275
message(STATUS "Building Relay in debug mode...")
272276
set_target_properties(tvm PROPERTIES COMPILE_DEFINITIONS "USE_RELAY_DEBUG")

cmake/config.cmake

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,10 @@ set(USE_BLAS none)
115115
# set(USE_MKL_PATH <path to venv or site-packages directory>) if using `pip install mkl`
116116
set(USE_MKL_PATH none)
117117

118+
# Whether use OpenMP thread pool, choices: gnu, intel
119+
# Note: "gnu" uses gomp library, "intel" uses iomp5 library
120+
set(USE_OPENMP none)
121+
118122
# Whether use contrib.random in runtime
119123
set(USE_RANDOM OFF)
120124

cmake/modules/OpenMP.cmake

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
# OpenMP Module
19+
if(USE_OPENMP STREQUAL "gnu")
20+
find_package(OpenMP)
21+
if(OPENMP_FOUND)
22+
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
23+
list(APPEND TVM_RUNTIME_LINKER_LIBS ${OpenMP_CXX_LIBRARIES})
24+
add_definitions(-DTVM_THREADPOOL_USE_OPENMP=1)
25+
message(STATUS "Build with OpenMP ${OpenMP_CXX_LIBRARIES}")
26+
else()
27+
add_definitions(-DTVM_THREADPOOL_USE_OPENMP=0)
28+
message(WARNING "OpenMP cannot be found, use TVM threadpool instead.")
29+
endif()
30+
elseif(USE_OPENMP STREQUAL "intel")
31+
find_package(OpenMP)
32+
if(OPENMP_FOUND)
33+
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
34+
if (MSVC)
35+
find_library(OMP_LIBRARY NAMES libiomp5md)
36+
else()
37+
find_library(OMP_LIBRARY NAMES iomp5)
38+
endif()
39+
list(APPEND TVM_RUNTIME_LINKER_LIBS ${OMP_LIBRARY})
40+
add_definitions(-DTVM_THREADPOOL_USE_OPENMP=1)
41+
message(STATUS "Build with OpenMP " ${OMP_LIBRARY})
42+
else()
43+
add_definitions(-DTVM_THREADPOOL_USE_OPENMP=0)
44+
message(WARNING "OpenMP cannot be found, use TVM threadpool instead.")
45+
endif()
46+
else()
47+
add_definitions(-DTVM_THREADPOOL_USE_OPENMP=0)
48+
endif()

src/runtime/thread_pool.cc

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@
2929
#include <tvm/runtime/threading_backend.h>
3030
#include <dmlc/thread_local.h>
3131
#include <dmlc/logging.h>
32+
#if TVM_THREADPOOL_USE_OPENMP
33+
#include <omp.h>
34+
#endif
3235
#include <thread>
3336
#include <condition_variable>
3437
#include <mutex>
@@ -394,12 +397,34 @@ int TVMBackendParallelLaunch(
394397
FTVMParallelLambda flambda,
395398
void* cdata,
396399
int num_task) {
400+
#if !TVM_THREADPOOL_USE_OPENMP
397401
int res = tvm::runtime::ThreadPool::ThreadLocal()->Launch(
398402
flambda, cdata, num_task, 1);
399403
return res;
404+
#else
405+
int num_workers = tvm::runtime::threading::MaxConcurrency();
406+
if (num_task == 0) num_task = num_workers;
407+
omp_set_num_threads(num_workers);
408+
#pragma omp parallel num_threads(num_workers)
409+
{
410+
TVMParallelGroupEnv env;
411+
env.num_task = num_task;
412+
std::atomic<int32_t>* sync_counter = new std::atomic<int>[num_task * tvm::runtime::kSyncStride];
413+
for (int i = 0; i < num_task; ++i) {
414+
sync_counter[i * tvm::runtime::kSyncStride].store(
415+
0, std::memory_order_relaxed);
416+
}
417+
env.sync_handle = sync_counter;
418+
(*flambda)(omp_get_thread_num(), &env, cdata);
419+
}
420+
return 0;
421+
#endif
400422
}
401423

402424
int TVMBackendParallelBarrier(int task_id, TVMParallelGroupEnv* penv) {
425+
#if TVM_THREADPOOL_USE_OPENMP
426+
#pragma omp barrier
427+
#else
403428
using tvm::runtime::kSyncStride;
404429
int num_task = penv->num_task;
405430
std::atomic<int>* sync_counter =
@@ -415,5 +440,6 @@ int TVMBackendParallelBarrier(int task_id, TVMParallelGroupEnv* penv) {
415440
}
416441
}
417442
std::atomic_thread_fence(std::memory_order_acquire);
443+
#endif
418444
return 0;
419445
}

0 commit comments

Comments
 (0)