mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Use Eigen ThreadPool in OnnxRuntime (#323)
* switch to nonblocking threadpool in inference session and sessions state * switch to eigen threadpool - first draft * refine * refine * add a switch to easily revert back to windows thread pool * switch thread pool in test runner and turn on leak checker * remove unncessary files * fix build error * more build fixes * catch exceptions in parallel executor * fix mac build error * fix mac build error * more build fixes * more mac build fixes * fix cv issue * change macro to include cuda compiler for disabled compiler warning * try switching the macro to win32 only * test #error * move #disable warning to the top * Update onnxruntime_framework.cmake * move eigen include to public scope * turn off eigenthreadpool by default and add todo comment
This commit is contained in:
Ashwin Kumar
GitHub
parent
139abda393
commit
492d9fd6cc
23 changed files with 282 additions and 192 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -33,3 +33,4 @@ onnxruntime_profile*.json
|
|||
/csharp/**/bin/
|
||||
docs/python/*.onnx
|
||||
*.onnx
|
||||
onnxprofile_profile_test_*.json
|
||||
|
|
|
|||
|
|
@ -11,8 +11,7 @@ source_group(TREE ${REPO_ROOT} FILES ${onnxruntime_framework_srcs})
|
|||
|
||||
add_library(onnxruntime_framework ${onnxruntime_framework_srcs})
|
||||
|
||||
#TODO: remove ${eigen_INCLUDE_DIRS} from here
|
||||
target_include_directories(onnxruntime_framework PRIVATE ${ONNXRUNTIME_ROOT} ${eigen_INCLUDE_DIRS} PUBLIC ${CMAKE_CURRENT_BINARY_DIR})
|
||||
target_include_directories(onnxruntime_framework PRIVATE ${ONNXRUNTIME_ROOT} PUBLIC ${CMAKE_CURRENT_BINARY_DIR} ${eigen_INCLUDE_DIRS})
|
||||
onnxruntime_add_include_to_target(onnxruntime_framework onnxruntime_common gsl onnx onnx_proto protobuf::libprotobuf)
|
||||
set_target_properties(onnxruntime_framework PROPERTIES FOLDER "ONNXRuntime")
|
||||
# need onnx to build to create headers that this project includes
|
||||
|
|
@ -23,3 +22,7 @@ if (WIN32)
|
|||
# Add Code Analysis properties to enable C++ Core checks. Have to do it via a props file include.
|
||||
set_target_properties(onnxruntime_framework PROPERTIES VS_USER_PROPS ${PROJECT_SOURCE_DIR}/ConfigureVisualStudioCodeAnalysis.props)
|
||||
endif()
|
||||
|
||||
if(onnxruntime_USE_EIGEN_THREADPOOL)
|
||||
target_compile_definitions(onnxruntime_framework PUBLIC USE_EIGEN_THREADPOOL)
|
||||
endif()
|
||||
|
|
|
|||
|
|
@ -98,3 +98,7 @@ endif()
|
|||
if (onnxruntime_ENABLE_MICROSOFT_INTERNAL)
|
||||
include(onnxruntime_providers_internal.cmake)
|
||||
endif()
|
||||
|
||||
if(onnxruntime_USE_EIGEN_THREADPOOL)
|
||||
target_compile_definitions(onnxruntime_providers PUBLIC USE_EIGEN_THREADPOOL)
|
||||
endif()
|
||||
|
|
|
|||
|
|
@ -12,8 +12,10 @@ source_group(TREE ${REPO_ROOT} FILES ${onnxruntime_session_srcs})
|
|||
add_library(onnxruntime_session ${onnxruntime_session_srcs})
|
||||
install(DIRECTORY ${PROJECT_SOURCE_DIR}/../include/onnxruntime/core/session DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/onnxruntime/core)
|
||||
onnxruntime_add_include_to_target(onnxruntime_session onnxruntime_common onnxruntime_framework gsl onnx onnx_proto protobuf::libprotobuf)
|
||||
target_include_directories(onnxruntime_session PRIVATE ${ONNXRUNTIME_ROOT})
|
||||
target_include_directories(onnxruntime_session PRIVATE ${ONNXRUNTIME_ROOT} ${eigen_INCLUDE_DIRS})
|
||||
add_dependencies(onnxruntime_session ${onnxruntime_EXTERNAL_DEPENDENCIES})
|
||||
set_target_properties(onnxruntime_session PROPERTIES FOLDER "ONNXRuntime")
|
||||
|
||||
|
||||
if(onnxruntime_USE_EIGEN_THREADPOOL)
|
||||
target_compile_definitions(onnxruntime_session PUBLIC USE_EIGEN_THREADPOOL)
|
||||
endif()
|
||||
|
|
|
|||
|
|
@ -351,7 +351,8 @@ set(onnx_test_runner_common_srcs
|
|||
${onnx_test_runner_src_dir}/TestCase.cc
|
||||
${onnx_test_runner_src_dir}/TestCase.h
|
||||
${onnx_test_runner_src_dir}/path_lib.h
|
||||
${onnx_test_runner_src_dir}/sync_api.h)
|
||||
${onnx_test_runner_src_dir}/sync_api.h
|
||||
${onnx_test_runner_src_dir}/sync_api.cc)
|
||||
|
||||
if(WIN32)
|
||||
set(wide_get_opt_src_dir ${TEST_SRC_DIR}/win_getopt/wide)
|
||||
|
|
@ -363,11 +364,12 @@ if(WIN32)
|
|||
target_include_directories(win_getopt_mb INTERFACE ${mb_get_opt_src_dir}/include)
|
||||
set_target_properties(win_getopt_mb PROPERTIES FOLDER "ONNXRuntimeTest")
|
||||
|
||||
set(onnx_test_runner_common_srcs ${onnx_test_runner_common_srcs} ${onnx_test_runner_src_dir}/sync_api_win.cc)
|
||||
set(onnx_test_runner_common_srcs ${onnx_test_runner_common_srcs})
|
||||
|
||||
set(GETOPT_LIB_WIDE win_getopt_wide)
|
||||
set(GETOPT_LIB win_getopt_mb)
|
||||
else()
|
||||
set(onnx_test_runner_common_srcs ${onnx_test_runner_common_srcs} ${onnx_test_runner_src_dir}/onnxruntime_event.h ${onnx_test_runner_src_dir}/simple_thread_pool.h ${onnx_test_runner_src_dir}/sync_api_linux.cc)
|
||||
set(onnx_test_runner_common_srcs ${onnx_test_runner_common_srcs} ${onnx_test_runner_src_dir}/onnxruntime_event.h ${onnx_test_runner_src_dir}/simple_thread_pool.h)
|
||||
if(HAS_FILESYSTEM_H OR HAS_EXPERIMENTAL_FILESYSTEM_H)
|
||||
set(FS_STDLIB stdc++fs)
|
||||
endif()
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@
|
|||
|
||||
#include "core/common/common.h"
|
||||
#include "core/common/logging/logging.h"
|
||||
#include "core/common/task_thread_pool.h"
|
||||
#include "core/framework/allocator.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
|
@ -340,16 +339,16 @@ static Status ValidateRnnInputsWithExtraInputFromState(
|
|||
W_shape[1] != hidden_size * WRB_dim_1_multipler ||
|
||||
W_shape[2] != input_size)
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input W must have shape {",
|
||||
num_directions, ",", WRB_dim_1_multipler, "*", hidden_size, ",",
|
||||
input_size, "}. Actual:", W_shape);
|
||||
num_directions, ",", WRB_dim_1_multipler, "*", hidden_size, ",",
|
||||
input_size, "}. Actual:", W_shape);
|
||||
|
||||
if (R_shape.NumDimensions() != 3 ||
|
||||
R_shape[0] != num_directions ||
|
||||
R_shape[1] != hidden_size * WRB_dim_1_multipler ||
|
||||
R_shape[2] != hidden_size)
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input R must have shape {",
|
||||
num_directions, ",", WRB_dim_1_multipler, "*", hidden_size, ",",
|
||||
hidden_size, "}. Actual:", R_shape);
|
||||
num_directions, ",", WRB_dim_1_multipler, "*", hidden_size, ",",
|
||||
hidden_size, "}. Actual:", R_shape);
|
||||
|
||||
if (B != nullptr) {
|
||||
auto& B_shape = B->Shape();
|
||||
|
|
@ -357,7 +356,7 @@ static Status ValidateRnnInputsWithExtraInputFromState(
|
|||
B_shape[0] != num_directions ||
|
||||
B_shape[1] != 2 * WRB_dim_1_multipler * hidden_size)
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input B must have shape {",
|
||||
num_directions, ",", 2 * WRB_dim_1_multipler, "*", hidden_size, "}. Actual:", B_shape);
|
||||
num_directions, ",", 2 * WRB_dim_1_multipler, "*", hidden_size, "}. Actual:", B_shape);
|
||||
}
|
||||
|
||||
if (sequence_lens != nullptr) {
|
||||
|
|
@ -365,7 +364,7 @@ static Status ValidateRnnInputsWithExtraInputFromState(
|
|||
if (sequence_lens_shape.NumDimensions() != 1 ||
|
||||
sequence_lens_shape[0] != batch_size) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input sequence_lens must have shape {",
|
||||
batch_size, "}. Actual:", sequence_lens_shape);
|
||||
batch_size, "}. Actual:", sequence_lens_shape);
|
||||
}
|
||||
|
||||
auto sequence_len_entries = sequence_lens->DataAsSpan<int>();
|
||||
|
|
@ -387,7 +386,7 @@ static Status ValidateRnnInputsWithExtraInputFromState(
|
|||
initial_h_shape[2] != hidden_size)
|
||||
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Input initial_h must have shape {",
|
||||
num_directions, ",", batch_size, ",", hidden_size, "}. Actual:", initial_h_shape);
|
||||
num_directions, ",", batch_size, ",", hidden_size, "}. Actual:", initial_h_shape);
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
@ -403,8 +402,8 @@ Status DeepCpuAttnLstmOp::ValidateInputs(
|
|||
auto memory_shape = attn_memory.Shape();
|
||||
if (memory_shape.NumDimensions() != 3 || memory_shape[0] != batch_size) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Attention mechanism memory shape error! Expected: {", batch_size,
|
||||
"}, actural: ", memory_shape);
|
||||
"Attention mechanism memory shape error! Expected: {", batch_size,
|
||||
"}, actural: ", memory_shape);
|
||||
}
|
||||
const int max_memory_step = gsl::narrow<int>(memory_shape[1]);
|
||||
const int memory_depth = gsl::narrow<int>(memory_shape[2]);
|
||||
|
|
@ -412,8 +411,8 @@ Status DeepCpuAttnLstmOp::ValidateInputs(
|
|||
auto memory_seq_lens_shape = attn_memory_seq_lens->Shape();
|
||||
if (memory_seq_lens_shape.NumDimensions() != 1 || memory_seq_lens_shape[0] != batch_size) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Attention mechanism memory sequence lengths must have shape {", batch_size,
|
||||
"}, actural: ", memory_seq_lens_shape);
|
||||
"Attention mechanism memory sequence lengths must have shape {", batch_size,
|
||||
"}, actural: ", memory_seq_lens_shape);
|
||||
}
|
||||
const gsl::span<const int> mem_seq_lens_span = attn_memory_seq_lens->DataAsSpan<int>();
|
||||
auto item_not_in_range = std::find_if(
|
||||
|
|
@ -421,8 +420,8 @@ Status DeepCpuAttnLstmOp::ValidateInputs(
|
|||
[max_memory_step](int len) { return len <= 0 || len > max_memory_step; });
|
||||
if (item_not_in_range != mem_seq_lens_span.cend()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Attention mechanism memory sequence lengths value must in (0, ",
|
||||
max_memory_step, "], while ", *item_not_in_range, " found!");
|
||||
"Attention mechanism memory sequence lengths value must in (0, ",
|
||||
max_memory_step, "], while ", *item_not_in_range, " found!");
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -432,8 +431,8 @@ Status DeepCpuAttnLstmOp::ValidateInputs(
|
|||
memory_layer_shape[0] != num_directions_ ||
|
||||
memory_layer_shape[1] != memory_depth) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Attention memory layer weight shape error! Expected:{",
|
||||
num_directions_, ",", memory_depth, ", am_attn_size}, Got:", memory_layer_shape);
|
||||
"Attention memory layer weight shape error! Expected:{",
|
||||
num_directions_, ",", memory_depth, ", am_attn_size}, Got:", memory_layer_shape);
|
||||
}
|
||||
const int am_attn_size = gsl::narrow<int>(memory_layer_shape[2]);
|
||||
|
||||
|
|
@ -444,8 +443,8 @@ Status DeepCpuAttnLstmOp::ValidateInputs(
|
|||
query_layer_shape[1] != hidden_size_ ||
|
||||
query_layer_shape[2] != am_attn_size) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Attention query layer weight shape error! Expected:{",
|
||||
num_directions_, ", ", hidden_size_, ", ", am_attn_size, "}, Got: ", query_layer_shape);
|
||||
"Attention query layer weight shape error! Expected:{",
|
||||
num_directions_, ", ", hidden_size_, ", ", am_attn_size, "}, Got: ", query_layer_shape);
|
||||
}
|
||||
|
||||
// check attention v for [num_directions, am_attn_size]
|
||||
|
|
@ -454,8 +453,8 @@ Status DeepCpuAttnLstmOp::ValidateInputs(
|
|||
v_shape[0] != num_directions_ ||
|
||||
v_shape[1] != am_attn_size) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Attention v weight shape error! Expected:{", num_directions_, ", ", am_attn_size,
|
||||
"}. Got: ", v_shape);
|
||||
"Attention v weight shape error! Expected:{", num_directions_, ", ", am_attn_size,
|
||||
"}. Got: ", v_shape);
|
||||
}
|
||||
|
||||
// Check attention layer weights for [num_directions, memory_depth+cell_hidden_size, aw_attn_size]
|
||||
|
|
@ -467,8 +466,8 @@ Status DeepCpuAttnLstmOp::ValidateInputs(
|
|||
attn_layer_shape[0] != num_directions_ ||
|
||||
attn_layer_shape[1] != memory_depth + hidden_size_) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Attention layer weight shape error! Expected: {", num_directions_, ", ",
|
||||
memory_depth + hidden_size_, ", aw_attn_size}. Got:", attn_layer_shape);
|
||||
"Attention layer weight shape error! Expected: {", num_directions_, ", ",
|
||||
memory_depth + hidden_size_, ", aw_attn_size}. Got:", attn_layer_shape);
|
||||
}
|
||||
aw_attn_size = gsl::narrow<int>(attn_layer_shape[2]);
|
||||
}
|
||||
|
|
@ -486,7 +485,7 @@ Status DeepCpuAttnLstmOp::ValidateInputs(
|
|||
initial_c_shape[2] != hidden_size_)
|
||||
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Input initial_c must have shape {",
|
||||
num_directions_, ",", batch_size, ",", hidden_size_, "}. Actual:", initial_c_shape);
|
||||
num_directions_, ",", batch_size, ",", hidden_size_, "}. Actual:", initial_c_shape);
|
||||
}
|
||||
|
||||
if (P != nullptr) {
|
||||
|
|
@ -497,7 +496,7 @@ Status DeepCpuAttnLstmOp::ValidateInputs(
|
|||
p_shape[1] != 3 * hidden_size_)
|
||||
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Input P must have shape {",
|
||||
num_directions_, ",", 3 * hidden_size_, "}. Actual:", p_shape);
|
||||
num_directions_, ",", 3 * hidden_size_, "}. Actual:", p_shape);
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@
|
|||
#include "attention_wrapper.h"
|
||||
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/common/task_thread_pool.h"
|
||||
#include "core/providers/cpu/rnn/rnn_helpers.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
|
@ -93,11 +92,15 @@ class DeepCpuAttnLstmOp final : public OpKernel {
|
|||
|
||||
ActivationFuncs activation_funcs_;
|
||||
|
||||
// Threadpool for operator. If concurrent Compute calls are possible, it will be shared
|
||||
// across them. mutable due to this.
|
||||
// The alternative would be to create a threadpool in each call to Compute but that would incur thread creation
|
||||
// cost on every call.
|
||||
// Threadpool for operator. If concurrent Compute calls are possible, it will be shared
|
||||
// across them. mutable due to this.
|
||||
// The alternative would be to create a threadpool in each call to Compute but that would incur thread creation
|
||||
// cost on every call.
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
mutable Eigen::NonBlockingThreadPool ttp_{static_cast<int>(std::thread::hardware_concurrency())};
|
||||
#else
|
||||
mutable TaskThreadPool ttp_{std::thread::hardware_concurrency()};
|
||||
#endif
|
||||
};
|
||||
|
||||
} // namespace contrib
|
||||
|
|
|
|||
|
|
@ -45,7 +45,11 @@ UniDirectionalAttnLstm<T>::UniDirectionalAttnLstm(AllocatorPtr allocator,
|
|||
const ActivationFuncs::Entry& activation_func_g,
|
||||
const ActivationFuncs::Entry& activation_func_h,
|
||||
const float clip,
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool& ttp)
|
||||
#else
|
||||
TaskThreadPool& ttp)
|
||||
#endif
|
||||
: allocator_(allocator),
|
||||
logger_(logger),
|
||||
seq_length_(seq_length),
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@
|
|||
|
||||
#include "core/common/common.h"
|
||||
#include "core/common/logging/logging.h"
|
||||
#include "core/common/task_thread_pool.h"
|
||||
#include "core/framework/allocator.h"
|
||||
|
||||
#include <gsl/span>
|
||||
|
|
@ -22,7 +21,6 @@ namespace logging = ::onnxruntime::logging;
|
|||
|
||||
using ::onnxruntime::AllocatorPtr;
|
||||
using ::onnxruntime::IAllocatorUniquePtr;
|
||||
using ::onnxruntime::TaskThreadPool;
|
||||
using ::onnxruntime::contrib::detail::ActivationInfo;
|
||||
using ::onnxruntime::rnn::detail::ActivationFuncs;
|
||||
using ::onnxruntime::rnn::detail::Direction;
|
||||
|
|
@ -53,7 +51,11 @@ class UniDirectionalAttnLstm {
|
|||
const ActivationFuncs::Entry& activation_func_g,
|
||||
const ActivationFuncs::Entry& activation_func_h,
|
||||
const float clip,
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool& ttp);
|
||||
#else
|
||||
TaskThreadPool& ttp);
|
||||
#endif
|
||||
|
||||
void Compute(const gsl::span<const T>& inputs,
|
||||
const gsl::span<const int>& sequence_lengths,
|
||||
|
|
@ -153,7 +155,11 @@ class UniDirectionalAttnLstm {
|
|||
|
||||
AttentionWrapper<T>& attention_wrapper_;
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool& ttp_;
|
||||
#else
|
||||
TaskThreadPool& ttp_;
|
||||
#endif
|
||||
};
|
||||
|
||||
} // namespace detail
|
||||
|
|
|
|||
|
|
@ -9,7 +9,11 @@
|
|||
#include <vector>
|
||||
#include "core/common/common.h"
|
||||
#include "core/common/logging/logging.h"
|
||||
|
||||
#ifndef USE_EIGEN_THREADPOOL
|
||||
#include "core/common/task_thread_pool.h"
|
||||
#endif
|
||||
|
||||
#include "core/framework/allocation_planner.h"
|
||||
#include "core/framework/execution_frame.h"
|
||||
#include "core/framework/session_state.h"
|
||||
|
|
@ -238,9 +242,19 @@ void ParallelExecutor::EnqueueNode(size_t p_node_index, const SessionState& sess
|
|||
std::unique_lock<std::mutex> lock(complete_mutex_);
|
||||
out_standings_++;
|
||||
}
|
||||
//std::cout << "Enqueue async node: " << p_node_index << ", out_standings: " << out_standings_ << std::endl;
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
session_state.GetThreadPool()->Schedule([this, p_node_index, &session_state, &logger]() {
|
||||
try {
|
||||
ParallelExecutor::RunNodeAsync(p_node_index, std::cref(session_state), std::cref(logger));
|
||||
} catch (...) {
|
||||
// catch node processing failure exceptions here to prevent app crash.
|
||||
}
|
||||
});
|
||||
#else
|
||||
std::packaged_task<void()> task{std::bind(&ParallelExecutor::RunNodeAsync, this, p_node_index, std::cref(session_state), std::cref(logger))};
|
||||
session_state.GetThreadPool()->RunTask(std::move(task));
|
||||
#endif
|
||||
}
|
||||
|
||||
Status ParallelExecutor::FetchOutput(const MLValueNameIdxMap& name_idx_map,
|
||||
|
|
|
|||
|
|
@ -21,15 +21,22 @@
|
|||
#include "core/graph/graph_viewer.h"
|
||||
#include "core/framework/fuse_nodes_funcs.h"
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
#include <unsupported/Eigen/CXX11/ThreadPool>
|
||||
#endif
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
class ExecutionProviders;
|
||||
class KernelDef;
|
||||
class OpKernel;
|
||||
class TaskThreadPool;
|
||||
struct SequentialExecutionPlan;
|
||||
struct MemoryPatternGroup;
|
||||
|
||||
#ifndef USE_EIGEN_THREADPOOL
|
||||
class TaskThreadPool;
|
||||
#endif
|
||||
|
||||
// SessionState should be modified by the inference session class only.
|
||||
// It is supposed to be passed by const-ref only to all the executors.
|
||||
class SessionState {
|
||||
|
|
@ -146,8 +153,13 @@ class SessionState {
|
|||
/// Return SessionState for the given Node index and attribute name if found.
|
||||
const SessionState* GetSubgraphSessionState(onnxruntime::NodeIndex index, const std::string& attribute_name) const;
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool* GetThreadPool() const { return thread_pool_; }
|
||||
void SetThreadPool(Eigen::NonBlockingThreadPool* p_pool) { thread_pool_ = p_pool; }
|
||||
#else
|
||||
TaskThreadPool* GetThreadPool() const { return thread_pool_; }
|
||||
void SetThreadPool(TaskThreadPool* p_pool) { thread_pool_ = p_pool; }
|
||||
#endif
|
||||
|
||||
bool ExportDll() const { return export_fused_dll_; }
|
||||
void SetExportDllFlag(bool flag) { export_fused_dll_ = flag; }
|
||||
|
|
@ -187,7 +199,12 @@ class SessionState {
|
|||
std::unordered_map<onnxruntime::NodeIndex,
|
||||
std::unordered_map<std::string, gsl::not_null<const SessionState*>>>;
|
||||
SubgraphSessionStateMap subgraph_session_states_;
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool* thread_pool_ = nullptr;
|
||||
#else
|
||||
TaskThreadPool* thread_pool_ = nullptr;
|
||||
#endif
|
||||
|
||||
bool export_fused_dll_ = false;
|
||||
FuncManager fused_funcs_mgr_;
|
||||
|
|
|
|||
|
|
@ -231,11 +231,7 @@ Memory_LeakCheck::~Memory_LeakCheck() {
|
|||
else {
|
||||
// If we're on the command line (like on a build machine), output to the console and exit(-1)
|
||||
std::cout << "\n----- MEMORY LEAKS: " << string.c_str() << "\n";
|
||||
#if 0
|
||||
// There is currently a memory leak due to a static thread_local variable not being destroyed on exit in mkldnn_common.h
|
||||
// The bug is caused by sync_api.h using the windows thread pool functions instead of C++ std::async libraries.
|
||||
exit(-1);
|
||||
#endif
|
||||
}
|
||||
|
||||
} else {
|
||||
|
|
|
|||
|
|
@ -177,7 +177,11 @@ class UniDirectionalGru {
|
|||
const ActivationFuncs::Entry& activation_func_f,
|
||||
const ActivationFuncs::Entry& activation_func_g,
|
||||
const float clip,
|
||||
TaskThreadPool& ttp);
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool& ttp_);
|
||||
#else
|
||||
TaskThreadPool& ttp_);
|
||||
#endif
|
||||
|
||||
void Compute(const gsl::span<const T>& inputs,
|
||||
const gsl::span<const int>& sequence_lengths,
|
||||
|
|
@ -192,7 +196,12 @@ class UniDirectionalGru {
|
|||
private:
|
||||
AllocatorPtr allocator_;
|
||||
const logging::Logger& logger_;
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool& ttp_;
|
||||
#else
|
||||
TaskThreadPool& ttp_;
|
||||
#endif
|
||||
|
||||
int seq_length_;
|
||||
int batch_size_;
|
||||
|
|
@ -375,9 +384,14 @@ Status DeepCpuGruOp::ComputeImpl(OpKernelContext& context) const {
|
|||
hidden_output_size_per_direction);
|
||||
|
||||
#ifndef USE_MKLDNN
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
auto fn =
|
||||
#else
|
||||
std::packaged_task<void()> task_fw{
|
||||
#endif
|
||||
[&]() {
|
||||
#endif // ! USE_MKLDNN
|
||||
#endif // ! USE_MKLDNN
|
||||
std::unique_ptr<detail::UniDirectionalGru<T>> fw = std::make_unique<detail::UniDirectionalGru<T>>(
|
||||
alloc, logger,
|
||||
seq_length, batch_size, input_size, hidden_size_, linear_before_reset_, Direction::kForward,
|
||||
|
|
@ -387,43 +401,67 @@ Status DeepCpuGruOp::ComputeImpl(OpKernelContext& context) const {
|
|||
clip_, ttp_);
|
||||
fw->Compute(input, sequence_lens_span, num_directions_, input_weights_1, recurrent_weights_1, output_1, hidden_output_1);
|
||||
#ifndef USE_MKLDNN
|
||||
}};
|
||||
auto task_results_fw = task_fw.get_future();
|
||||
ttp_.RunTask(std::move(task_fw));
|
||||
#endif // ! USE_MKLDNN
|
||||
|
||||
std::unique_ptr<detail::UniDirectionalGru<T>> bw = std::make_unique<detail::UniDirectionalGru<T>>(
|
||||
alloc, logger,
|
||||
seq_length, batch_size, input_size, hidden_size_, linear_before_reset_, Direction::kReverse,
|
||||
bias_2, initial_hidden_2,
|
||||
activation_funcs_.Entries()[2],
|
||||
activation_funcs_.Entries()[3],
|
||||
clip_, ttp_);
|
||||
bw->Compute(input, sequence_lens_span, num_directions_, input_weights_2, recurrent_weights_2, output_2, hidden_output_2);
|
||||
#ifndef USE_EIGEN_THREADPOOL
|
||||
}
|
||||
};
|
||||
auto task_results_fw = task_fw.get_future();
|
||||
ttp_.RunTask(std::move(task_fw));
|
||||
#else
|
||||
};
|
||||
std::condition_variable cv;
|
||||
std::mutex lock;
|
||||
bool done = false;
|
||||
|
||||
ttp_.Schedule([&]() {
|
||||
fn();
|
||||
auto ul = std::unique_lock<std::mutex>(lock);
|
||||
done = true;
|
||||
cv.notify_one();
|
||||
});
|
||||
#endif // USE_EIGEN_THREADPOOL
|
||||
|
||||
#endif // ! USE_MKLDNN
|
||||
|
||||
std::unique_ptr<detail::UniDirectionalGru<T>> bw = std::make_unique<detail::UniDirectionalGru<T>>(
|
||||
alloc, logger,
|
||||
seq_length, batch_size, input_size, hidden_size_, linear_before_reset_, Direction::kReverse,
|
||||
bias_2, initial_hidden_2,
|
||||
activation_funcs_.Entries()[2],
|
||||
activation_funcs_.Entries()[3],
|
||||
clip_, ttp_);
|
||||
bw->Compute(input, sequence_lens_span, num_directions_, input_weights_2, recurrent_weights_2, output_2, hidden_output_2);
|
||||
|
||||
#ifndef USE_MKLDNN
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
auto ul = std::unique_lock<std::mutex>(lock);
|
||||
if (!done) cv.wait(ul);
|
||||
#else
|
||||
task_results_fw.get();
|
||||
#endif // ! USE_MKLDNN
|
||||
} else {
|
||||
std::unique_ptr<detail::UniDirectionalGru<T>> gru_p = std::make_unique<detail::UniDirectionalGru<T>>(
|
||||
alloc, logger,
|
||||
seq_length, batch_size, input_size, hidden_size_, linear_before_reset_, direction_,
|
||||
bias_1, initial_hidden_1,
|
||||
activation_funcs_.Entries()[0],
|
||||
activation_funcs_.Entries()[1],
|
||||
clip_, ttp_);
|
||||
#endif
|
||||
|
||||
gru_p->Compute(input, sequence_lens_span, num_directions_, input_weights_1, recurrent_weights_1, output_1, hidden_output_1);
|
||||
}
|
||||
#endif // ! USE_MKLDNN
|
||||
} // namespace onnxruntime
|
||||
else {
|
||||
std::unique_ptr<detail::UniDirectionalGru<T>> gru_p = std::make_unique<detail::UniDirectionalGru<T>>(
|
||||
alloc, logger,
|
||||
seq_length, batch_size, input_size, hidden_size_, linear_before_reset_, direction_,
|
||||
bias_1, initial_hidden_1,
|
||||
activation_funcs_.Entries()[0],
|
||||
activation_funcs_.Entries()[1],
|
||||
clip_, ttp_);
|
||||
|
||||
if (!output.empty())
|
||||
DumpMatrix("Y", output.data(), seq_length * num_directions_ * batch_size, hidden_size_);
|
||||
|
||||
DumpMatrix("Y_h", hidden_output.data(), num_directions_ * batch_size, hidden_size_);
|
||||
|
||||
return Status::OK();
|
||||
gru_p->Compute(input, sequence_lens_span, num_directions_, input_weights_1, recurrent_weights_1, output_1, hidden_output_1);
|
||||
}
|
||||
|
||||
if (!output.empty())
|
||||
DumpMatrix("Y", output.data(), seq_length* num_directions_* batch_size, hidden_size_);
|
||||
|
||||
DumpMatrix("Y_h", hidden_output.data(), num_directions_* batch_size, hidden_size_);
|
||||
|
||||
return Status::OK();
|
||||
} // namespace onnxruntime
|
||||
|
||||
//
|
||||
// Implementation of internal helper code
|
||||
namespace detail {
|
||||
|
|
@ -442,7 +480,11 @@ UniDirectionalGru<T>::UniDirectionalGru(AllocatorPtr allocator,
|
|||
const ActivationFuncs::Entry& activation_func_f,
|
||||
const ActivationFuncs::Entry& activation_func_g,
|
||||
const float clip,
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool& ttp)
|
||||
#else
|
||||
TaskThreadPool& ttp)
|
||||
#endif
|
||||
: allocator_(allocator),
|
||||
logger_(logger),
|
||||
ttp_(ttp),
|
||||
|
|
@ -490,8 +532,8 @@ UniDirectionalGru<T>::UniDirectionalGru(AllocatorPtr allocator,
|
|||
// replicate what we just wrote to the start of the output span so we have batch_size_ copies
|
||||
auto values = output.cbegin();
|
||||
ORT_IGNORE_RETURN_VALUE(RepeatVectorToConstructArray(values, values + hidden_size_,
|
||||
output.begin() + hidden_size_, // skip the first batch
|
||||
batch_size_ - 1)); // and replicate batch size - 1 times
|
||||
output.begin() + hidden_size_, // skip the first batch
|
||||
batch_size_ - 1)); // and replicate batch size - 1 times
|
||||
};
|
||||
|
||||
// we can always combine the z and r weights
|
||||
|
|
@ -603,7 +645,7 @@ void UniDirectionalGru<T>::Compute(const gsl::span<const T>& inputs_arg,
|
|||
if (batch_size_ % hidden_num_threads_ != 0)
|
||||
fused_hidden_rows++;
|
||||
|
||||
// lambda executed by TaskThreadPool
|
||||
// lambda executed by Eigen::NonBlockingThreadPool
|
||||
auto hidden_gemm_and_activations = [&](const int row) {
|
||||
//handling boundaries
|
||||
int local_fused_hidden_rows = fused_hidden_rows;
|
||||
|
|
@ -815,8 +857,7 @@ void UniDirectionalGru<T>::Compute(const gsl::span<const T>& inputs_arg,
|
|||
}
|
||||
};
|
||||
|
||||
ExecuteLambdaInParallel("Processing batch", hidden_gemm_and_activations, batch_size_, fused_hidden_rows,
|
||||
ttp_, logger_);
|
||||
ExecuteLambdaInParallel("Processing batch", hidden_gemm_and_activations, batch_size_, fused_hidden_rows, ttp_, logger_);
|
||||
} else {
|
||||
size_t out_added_offset;
|
||||
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@
|
|||
|
||||
#include <limits>
|
||||
|
||||
#include "core/common/task_thread_pool.h"
|
||||
#include "core/framework/allocator.h"
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/providers/cpu/rnn/rnn_helpers.h"
|
||||
|
|
@ -71,7 +70,11 @@ class DeepCpuGruOp final : public OpKernel {
|
|||
// across them. mutable due to this.
|
||||
// The alternative would be to create a threadpool in each call to Compute but that would incur thread creation
|
||||
// cost on every call.
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
mutable Eigen::NonBlockingThreadPool ttp_{static_cast<int>(std::thread::hardware_concurrency())};
|
||||
#else
|
||||
mutable TaskThreadPool ttp_{std::thread::hardware_concurrency()};
|
||||
#endif
|
||||
|
||||
template <typename T>
|
||||
Status ComputeImpl(OpKernelContext& context) const;
|
||||
|
|
|
|||
|
|
@ -13,7 +13,6 @@
|
|||
|
||||
#include "core/common/common.h"
|
||||
#include "core/common/logging/logging.h"
|
||||
#include "core/common/task_thread_pool.h"
|
||||
#include "core/framework/allocator.h"
|
||||
|
||||
#ifdef _MSC_VER
|
||||
|
|
@ -204,7 +203,11 @@ class UniDirectionalLstm {
|
|||
const ActivationFuncs::Entry& activation_func_g,
|
||||
const ActivationFuncs::Entry& activation_func_h,
|
||||
const float clip,
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool& ttp);
|
||||
#else
|
||||
TaskThreadPool& ttp);
|
||||
#endif
|
||||
|
||||
void Compute(const gsl::span<const T>& inputs,
|
||||
const gsl::span<const int>& sequence_lengths,
|
||||
|
|
@ -296,7 +299,11 @@ class UniDirectionalLstm {
|
|||
ActivationInfo<deepcpu::ActivationFuncPtr> activation_g_;
|
||||
ActivationInfo<deepcpu::LstmMergeGatesFuncPtr> activation_h_;
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool& ttp_;
|
||||
#else
|
||||
TaskThreadPool& ttp_;
|
||||
#endif
|
||||
};
|
||||
|
||||
} // namespace detail
|
||||
|
|
@ -518,7 +525,7 @@ Status DeepCpuLstmOp::ValidateInputs(const Tensor& X, const Tensor& W, const Ten
|
|||
initial_c_shape[2] != hidden_size_)
|
||||
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Input initial_c must have shape {",
|
||||
num_directions_, ",", batch_size, ",", hidden_size_, "}. Actual:", initial_c_shape);
|
||||
num_directions_, ",", batch_size, ",", hidden_size_, "}. Actual:", initial_c_shape);
|
||||
}
|
||||
|
||||
if (P != nullptr) {
|
||||
|
|
@ -529,7 +536,7 @@ Status DeepCpuLstmOp::ValidateInputs(const Tensor& X, const Tensor& W, const Ten
|
|||
p_shape[1] != 3 * hidden_size_)
|
||||
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Input P must have shape {",
|
||||
num_directions_, ",", 3 * hidden_size_, "}. Actual:", p_shape);
|
||||
num_directions_, ",", 3 * hidden_size_, "}. Actual:", p_shape);
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
@ -559,7 +566,11 @@ UniDirectionalLstm<T>::UniDirectionalLstm(AllocatorPtr allocator,
|
|||
const ActivationFuncs::Entry& activation_func_g,
|
||||
const ActivationFuncs::Entry& activation_func_h,
|
||||
const float clip,
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool& ttp)
|
||||
#else
|
||||
TaskThreadPool& ttp)
|
||||
#endif
|
||||
: allocator_(allocator),
|
||||
logger_(logger),
|
||||
seq_length_(seq_length),
|
||||
|
|
@ -882,8 +893,7 @@ void UniDirectionalLstm<T>::Compute(const gsl::span<const T>& inputs_arg,
|
|||
}
|
||||
};
|
||||
|
||||
ExecuteLambdaInParallel("Processing batch", hidden_gemm_and_activations, batch_size_, fused_hidden_rows,
|
||||
ttp_, logger_);
|
||||
ExecuteLambdaInParallel("Processing batch", hidden_gemm_and_activations, batch_size_, fused_hidden_rows, ttp_, logger_);
|
||||
|
||||
} else {
|
||||
span_T_iter c_prev = batched_internal_state_prev_one_step.begin();
|
||||
|
|
|
|||
|
|
@ -6,9 +6,12 @@
|
|||
#include <limits>
|
||||
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/common/task_thread_pool.h"
|
||||
#include "core/providers/cpu/rnn/rnn_helpers.h"
|
||||
|
||||
#ifndef USE_EIGEN_THREADPOOL
|
||||
#include "core/common/task_thread_pool.h"
|
||||
#endif
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
/// The class represents DeepCPU implementation of a long short term memory (LSTM) operator.
|
||||
|
|
@ -82,7 +85,12 @@ class DeepCpuLstmOp final : public OpKernel {
|
|||
// across them. mutable due to this.
|
||||
// The alternative would be to create a threadpool in each call to Compute but that would incur thread creation
|
||||
// cost on every call.
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
mutable Eigen::NonBlockingThreadPool ttp_{static_cast<int>(std::thread::hardware_concurrency())};
|
||||
#else
|
||||
mutable TaskThreadPool ttp_{std::thread::hardware_concurrency()};
|
||||
#endif
|
||||
};
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -3,6 +3,10 @@
|
|||
|
||||
#pragma once
|
||||
|
||||
#ifdef _WIN32
|
||||
#pragma warning(disable : 4267)
|
||||
#endif
|
||||
|
||||
#include <algorithm>
|
||||
#include <functional>
|
||||
#include <future>
|
||||
|
|
@ -13,12 +17,17 @@
|
|||
#include "gsl/gsl_algorithm"
|
||||
|
||||
#include "core/common/common.h"
|
||||
#include "core/common/task_thread_pool.h"
|
||||
#include "core/common/logging/logging.h"
|
||||
#include "core/framework/allocator.h"
|
||||
#include "core/util/math.h"
|
||||
#include "core/util/math_cpuonly.h"
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
#include <unsupported/Eigen/CXX11/ThreadPool>
|
||||
#else
|
||||
#include "core/common/task_thread_pool.h"
|
||||
#endif
|
||||
|
||||
namespace onnxruntime {
|
||||
class Tensor;
|
||||
class OpKernelContext;
|
||||
|
|
@ -209,7 +218,12 @@ T* SafeRawPointer(typename gsl::span<T> span, size_t offset, size_t size) {
|
|||
|
||||
template <typename TLambda>
|
||||
void ExecuteLambdaInParallel(const std::string& name, TLambda lambda, int max, int step,
|
||||
TaskThreadPool& ttp, const ::onnxruntime::logging::Logger& logger) {
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
Eigen::NonBlockingThreadPool& ttp,
|
||||
#else
|
||||
TaskThreadPool& ttp,
|
||||
#endif
|
||||
const ::onnxruntime::logging::Logger& logger) {
|
||||
// #define NOTHREADS to execute the lambdas directly and in order if you need to do that to debug
|
||||
|
||||
#ifdef NOTHREADS
|
||||
|
|
@ -221,6 +235,23 @@ void ExecuteLambdaInParallel(const std::string& name, TLambda lambda, int max, i
|
|||
std::bind(lambda, i)();
|
||||
}
|
||||
#else
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
ORT_UNUSED_PARAMETER(name);
|
||||
ORT_UNUSED_PARAMETER(logger);
|
||||
|
||||
std::atomic<int> done(0);
|
||||
for (int i = 0; i < max; i += step) {
|
||||
ttp.Schedule([lambda, i, &done]() {
|
||||
lambda(i);
|
||||
++done;
|
||||
});
|
||||
}
|
||||
|
||||
int totalTasks = (int)max / (step > 0 ? step : 1) + (max % step > 0 ? 1 : 0);
|
||||
while (done != totalTasks) {
|
||||
}
|
||||
#else
|
||||
std::vector<std::future<void> > task_results{};
|
||||
task_results.reserve(static_cast<size_t>(std::ceil(max / step)));
|
||||
|
||||
|
|
@ -229,7 +260,6 @@ void ExecuteLambdaInParallel(const std::string& name, TLambda lambda, int max, i
|
|||
task_results.emplace_back(task.get_future());
|
||||
ttp.RunTask(std::move(task));
|
||||
}
|
||||
|
||||
try {
|
||||
// wait for all and propagate any exceptions
|
||||
for (auto& future : task_results)
|
||||
|
|
@ -238,7 +268,8 @@ void ExecuteLambdaInParallel(const std::string& name, TLambda lambda, int max, i
|
|||
LOGS(logger, ERROR) << name << " - exception running tasks: " << ex.what();
|
||||
throw;
|
||||
}
|
||||
#endif
|
||||
#endif // else part of #ifdef USE_EIGEN_THREADPOOLs
|
||||
#endif // else part of #ifdef NOTHREADS
|
||||
}
|
||||
|
||||
void DumpMatrixImpl(const std::string& name, const float* src, int row, int col,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,10 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#ifdef _WIN32
|
||||
#pragma warning(disable : 4267)
|
||||
#endif
|
||||
|
||||
#include "core/session/inference_session.h"
|
||||
|
||||
#include <memory>
|
||||
|
|
@ -40,6 +44,10 @@
|
|||
#include "core/session/CustomOpsLoader.h"
|
||||
#include "core/session/IOBinding.h"
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
#include <unsupported/Eigen/CXX11/ThreadPool>
|
||||
#endif
|
||||
|
||||
using namespace ONNX_NAMESPACE;
|
||||
|
||||
namespace onnxruntime {
|
||||
|
|
@ -63,7 +71,12 @@ class InferenceSession::Impl {
|
|||
int pool_size = session_options_.session_thread_pool_size == 0
|
||||
? std::thread::hardware_concurrency() / 2
|
||||
: session_options_.session_thread_pool_size;
|
||||
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
thread_pool_ = std::make_unique<Eigen::NonBlockingThreadPool>(pool_size);
|
||||
#else
|
||||
thread_pool_ = std::make_unique<TaskThreadPool>(pool_size);
|
||||
#endif
|
||||
}
|
||||
|
||||
session_state_.SetThreadPool(thread_pool_.get());
|
||||
|
|
@ -1130,14 +1143,19 @@ class InferenceSession::Impl {
|
|||
|
||||
// Threadpool for this session
|
||||
//thread::ThreadPool thread_pool_; // not used for now; will add it later when implementing RunAsync
|
||||
#ifdef USE_EIGEN_THREADPOOL
|
||||
std::unique_ptr<Eigen::NonBlockingThreadPool> thread_pool_;
|
||||
#else
|
||||
std::unique_ptr<TaskThreadPool> thread_pool_;
|
||||
#endif
|
||||
|
||||
// Number of concurrently running executors
|
||||
std::atomic<int> current_num_runs_;
|
||||
std::atomic<int>
|
||||
current_num_runs_;
|
||||
|
||||
mutable onnxruntime::OrtMutex session_mutex_; // to ensure only one thread can invoke Load/Initialize
|
||||
bool is_model_loaded_ = false; // GUARDED_BY(session_mutex_)
|
||||
bool is_inited_ = false; // GUARDED_BY(session_mutex_)
|
||||
bool is_model_loaded_ = false; // GUARDED_BY(session_mutex_)
|
||||
bool is_inited_ = false; // GUARDED_BY(session_mutex_)
|
||||
|
||||
std::map<OrtAllocatorInfo, BufferUniquePtr> weights_buffers_;
|
||||
InsertCastTransformer insert_cast_transformer_;
|
||||
|
|
|
|||
|
|
@ -1,14 +1,15 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
//TODO: switch to nsync mutex and cv
|
||||
#include <mutex>
|
||||
#include <pthread.h>
|
||||
#include <condition_variable>
|
||||
#include <core/common/common.h>
|
||||
|
||||
struct OnnxRuntimeEvent {
|
||||
public:
|
||||
pthread_mutex_t finish_event_mutex = PTHREAD_MUTEX_INITIALIZER;
|
||||
pthread_cond_t finish_event_data = PTHREAD_COND_INITIALIZER;
|
||||
std::mutex finish_event_mutex;
|
||||
std::condition_variable finish_event_data;
|
||||
bool finished = false;
|
||||
OnnxRuntimeEvent() = default;
|
||||
|
||||
|
|
|
|||
|
|
@ -3,10 +3,13 @@
|
|||
|
||||
#include "sync_api.h"
|
||||
#include <mutex>
|
||||
|
||||
#if defined(_MSC_VER)
|
||||
#pragma warning(disable : 4267)
|
||||
#endif
|
||||
#include <unsupported/Eigen/CXX11/ThreadPool>
|
||||
#include <core/common/common.h>
|
||||
#include <core/common/logging/logging.h>
|
||||
#include "simple_thread_pool.h"
|
||||
#include "onnxruntime_event.h"
|
||||
|
||||
using onnxruntime::common::Status;
|
||||
|
|
@ -25,11 +28,12 @@ class OnnxRuntimeCallbackInstance {
|
|||
Status WaitAndCloseEvent(ORT_EVENT finish_event) {
|
||||
if (finish_event == nullptr)
|
||||
return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::INVALID_ARGUMENT, "");
|
||||
pthread_mutex_lock(&finish_event->finish_event_mutex);
|
||||
while (!finish_event->finished) {
|
||||
pthread_cond_wait(&finish_event->finish_event_data, &finish_event->finish_event_mutex);
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(finish_event->finish_event_mutex);
|
||||
while (!finish_event->finished) {
|
||||
finish_event->finish_event_data.wait(lock);
|
||||
}
|
||||
}
|
||||
pthread_mutex_unlock(&finish_event->finish_event_mutex);
|
||||
delete finish_event;
|
||||
return Status::OK();
|
||||
}
|
||||
|
|
@ -51,14 +55,14 @@ Status CreateAndSubmitThreadpoolWork(ORT_CALLBACK_FUNCTION callback, void* data,
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
using DefaultThreadPoolType = onnxruntime::SimpleThreadPoolTempl<onnxruntime::Env>;
|
||||
using DefaultThreadPoolType = Eigen::NonBlockingThreadPool;
|
||||
static std::unique_ptr<DefaultThreadPoolType> default_pool;
|
||||
static std::once_flag default_pool_init;
|
||||
|
||||
PThreadPool GetDefaultThreadPool(const onnxruntime::Env& env) {
|
||||
std::call_once(default_pool_init, [&env] {
|
||||
int core_num = env.GetNumCpuCores();
|
||||
default_pool.reset(new DefaultThreadPoolType(core_num, env));
|
||||
default_pool.reset(new DefaultThreadPoolType(core_num));
|
||||
});
|
||||
return default_pool.get();
|
||||
}
|
||||
|
|
@ -68,16 +72,12 @@ Status OnnxRuntimeSetEventWhenCallbackReturns(ORT_CALLBACK_INSTANCE pci, ORT_EVE
|
|||
return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::INVALID_ARGUMENT, "");
|
||||
|
||||
if (pci == nullptr) {
|
||||
if (pthread_mutex_lock(&finish_event->finish_event_mutex)) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "lock failed");
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(finish_event->finish_event_mutex);
|
||||
finish_event->finished = true;
|
||||
}
|
||||
finish_event->finished = true;
|
||||
if (pthread_mutex_unlock(&finish_event->finish_event_mutex))
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "unlock failed");
|
||||
if (!pthread_cond_broadcast(&finish_event->finish_event_data))
|
||||
return Status::OK();
|
||||
else
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "pthread_cond_broadcast failed");
|
||||
finish_event->finish_event_data.notify_all();
|
||||
return Status::OK();
|
||||
} else {
|
||||
pci->AddEvent(finish_event);
|
||||
return Status::OK();
|
||||
|
|
@ -90,14 +90,11 @@ void OnnxRuntimeCallbackInstance::AddEvent(ORT_EVENT event) {
|
|||
|
||||
Status OnnxRuntimeCallbackInstance::SignalAllEvents() {
|
||||
for (ORT_EVENT finish_event : events_to_signal_) {
|
||||
if (pthread_mutex_lock(&finish_event->finish_event_mutex)) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "lock failed");
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(finish_event->finish_event_mutex);
|
||||
finish_event->finished = true;
|
||||
}
|
||||
finish_event->finished = true;
|
||||
if (pthread_mutex_unlock(&finish_event->finish_event_mutex))
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "unlock failed");
|
||||
if (pthread_cond_broadcast(&finish_event->finish_event_data))
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "pthread_cond_broadcast failed");
|
||||
finish_event->finish_event_data.notify_all();
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
|
@ -3,27 +3,11 @@
|
|||
|
||||
#pragma once
|
||||
|
||||
#ifdef _WIN32
|
||||
#include <Windows.h>
|
||||
#else
|
||||
#include <vector>
|
||||
#endif
|
||||
#include <core/common/status.h>
|
||||
#include <core/common/common.h>
|
||||
#include <core/platform/env.h>
|
||||
|
||||
#ifdef _WIN32
|
||||
using ORT_CALLBACK_INSTANCE = PTP_CALLBACK_INSTANCE;
|
||||
using ORT_EVENT = HANDLE;
|
||||
#define ORT_CALLBACK __stdcall
|
||||
using ORT_WORK = PTP_WORK;
|
||||
using PThreadPool = PTP_CALLBACK_ENVIRON;
|
||||
using ORT_CALLBACK_FUNCTION = PTP_WORK_CALLBACK;
|
||||
#define OnnxRuntimeCloseThreadpoolWork CloseThreadpoolWork
|
||||
inline PThreadPool GetDefaultThreadPool(const ::onnxruntime::Env&) {
|
||||
return nullptr;
|
||||
}
|
||||
#else
|
||||
#define ORT_CALLBACK
|
||||
namespace Eigen {
|
||||
class ThreadPoolInterface;
|
||||
|
|
@ -38,7 +22,6 @@ using ORT_CALLBACK_INSTANCE = OnnxRuntimeCallbackInstance*;
|
|||
using ORT_CALLBACK_FUNCTION = void ORT_CALLBACK (*)(ORT_CALLBACK_INSTANCE pci, void* context, ORT_WORK work);
|
||||
//Do nothing
|
||||
inline void OnnxRuntimeCloseThreadpoolWork(ORT_WORK) {}
|
||||
#endif
|
||||
|
||||
//The returned value will be used with CreateAndSubmitThreadpoolWork function
|
||||
PThreadPool GetDefaultThreadPool(const ::onnxruntime::Env& env);
|
||||
|
|
|
|||
|
|
@ -1,55 +0,0 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "sync_api.h"
|
||||
#include <core/common/common.h>
|
||||
|
||||
using ::onnxruntime::common::Status;
|
||||
|
||||
Status CreateAndSubmitThreadpoolWork(ORT_CALLBACK_FUNCTION callback, void* data, PThreadPool pool) {
|
||||
PTP_WORK work = CreateThreadpoolWork(callback, data, pool);
|
||||
if (!work) {
|
||||
return Status(::onnxruntime::common::ONNXRUNTIME, ::onnxruntime::common::FAIL, "create thread pool task failed");
|
||||
}
|
||||
SubmitThreadpoolWork(work);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status WaitAndCloseEvent(ORT_EVENT finish_event) {
|
||||
DWORD dwWaitResult = WaitForSingleObject(finish_event, INFINITE);
|
||||
(void)CloseHandle(finish_event);
|
||||
if (dwWaitResult != WAIT_OBJECT_0) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "WaitForSingleObject failed");
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status CreateOnnxRuntimeEvent(ORT_EVENT* out) {
|
||||
if (out == nullptr)
|
||||
return Status(::onnxruntime::common::ONNXRUNTIME, ::onnxruntime::common::INVALID_ARGUMENT, "");
|
||||
HANDLE finish_event = CreateEvent(
|
||||
NULL, // default security attributes
|
||||
TRUE, // manual-reset event
|
||||
FALSE, // initial state is nonsignaled
|
||||
NULL);
|
||||
if (finish_event == NULL) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "unable to create finish event");
|
||||
}
|
||||
*out = finish_event;
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status OnnxRuntimeSetEventWhenCallbackReturns(ORT_CALLBACK_INSTANCE pci, ORT_EVENT finish_event) {
|
||||
if (finish_event == nullptr)
|
||||
return Status(::onnxruntime::common::ONNXRUNTIME, ::onnxruntime::common::INVALID_ARGUMENT, "");
|
||||
if (pci)
|
||||
SetEventWhenCallbackReturns(pci, finish_event);
|
||||
else if (!SetEvent(finish_event)) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "SetEvent failed");
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
void OrtCloseEvent(ORT_EVENT finish_event) {
|
||||
(void)CloseHandle(finish_event);
|
||||
}
|
||||
|
|
@ -108,6 +108,7 @@ Use the individual flags to only run the specified stages.
|
|||
parser.add_argument("--use_tvm", action="store_true", help="Build with tvm")
|
||||
parser.add_argument("--use_openmp", action='store_true', help="Build with OpenMP.")
|
||||
parser.add_argument("--use_llvm", action="store_true", help="Build tvm with llvm")
|
||||
parser.add_argument("--use_eigenthreadpool", action="store_true", help="Build with eigenthreadpool")
|
||||
parser.add_argument("--enable_msinternal", action="store_true", help="Enable for Microsoft internal builds only.")
|
||||
parser.add_argument("--llvm_path", help="Path to llvm dir")
|
||||
parser.add_argument("--azure_sas_key", help="Azure storage sas key, starts with '?'")
|
||||
|
|
@ -293,6 +294,7 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
|
|||
"-Donnxruntime_ENABLE_MICROSOFT_INTERNAL=" + ("ON" if args.enable_msinternal else "OFF"),
|
||||
"-Donnxruntime_USE_BRAINSLICE=" + ("ON" if args.use_brainslice else "OFF"),
|
||||
"-Donnxruntime_USE_NUPHAR=" + ("ON" if args.use_nuphar else "OFF"),
|
||||
"-Donnxruntime_USE_EIGEN_THREADPOOL=" + ("ON" if args.use_eigenthreadpool else "OFF"),
|
||||
]
|
||||
if args.use_brainslice:
|
||||
bs_pkg_name = args.brain_slice_package_name.split('.', 1)
|
||||
|
|
|
|||
Loading…
Reference in a new issue