saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Move access to intra-op threadpool into OpKernelContext. (#2091)

Browse source
This commit is contained in:
Scott McKay 2019-10-12 03:36:20 +10:00 committed by Changming Sun
parent 368bdfd936
commit bdfff800ea
22 changed files with 159 additions and 166 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
include/onnxruntime/core/framework/op_kernel.h
View file

@ -24,6 +24,9 @@ namespace onnxruntime {
class IExecutionFrame;
class OpKernelContext;
class OpKernelWrapper;
namespace concurrency {
class ThreadPool;
}
class OpKernel {
public:
@ -63,6 +66,7 @@ class OpKernelContext {
explicit OpKernelContext(IExecutionFrame* frame,
const OpKernel* kernel,
concurrency::ThreadPool* threadpool,
const logging::Logger& logger);
virtual ~OpKernelContext() = default;
@ -163,6 +167,11 @@ class OpKernelContext {
**/
const std::string& GetOpDomain() const;
/**
Returns the intra-op threadpool, if available.
*/
_Ret_maybenull_ onnxruntime::concurrency::ThreadPool* GetOperatorThreadPool() const { return threadpool_; }
protected:
onnxruntime::NodeIndex GetNodeIndex() const;
@ -186,6 +195,7 @@ class OpKernelContext {
IExecutionFrame* execution_frame_{nullptr};
const OpKernel* kernel_{nullptr};
concurrency::ThreadPool* threadpool_{nullptr};
const logging::Logger* logger_{nullptr};
// The argument starting index in ExecutionFrame.

5
onnxruntime/contrib_ops/cpu/attnlstm/deep_cpu_attn_lstm.cc
View file

@ -72,9 +72,6 @@ static gsl::span<const T> SecondHalfSpan(const gsl::span<const T>& dspan) {
template <typename T>
Status DeepCpuAttnLstmOp::ComputeImpl(OpKernelContext& context) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(&context);
concurrency::ThreadPool* thread_pool = ctx_internal->GetOperatorThreadPool();
auto& logger = context.Logger();
// original lstm processing
@ -205,6 +202,8 @@ Status DeepCpuAttnLstmOp::ComputeImpl(OpKernelContext& context) const {
}
}
concurrency::ThreadPool* thread_pool = context.GetOperatorThreadPool();
if (direction_ == Direction::kBidirectional) {
// spans for second direction
gsl::span<const T> input_weights_2 = input_weights.subspan(input_weights_size_per_direction,

8
onnxruntime/contrib_ops/cpu/cdist.h
View file

@ -94,7 +94,8 @@ class CDist final : public OpKernel {
private:
typedef void (*DistFunc)(const T* a, const T* b, T* dest, size_t ma, size_t mb, size_t n,
concurrency::ThreadPool* tp);
enum { EUCLIDEAN, SQEUCLIDEAN } mode_;
enum { EUCLIDEAN,
SQEUCLIDEAN } mode_;
public:
CDist(const OpKernelInfo& info) : OpKernel(info) {
@ -109,8 +110,7 @@ class CDist final : public OpKernel {
}
common::Status Compute(OpKernelContext* context) const override {
auto ctx_internal = static_cast<OpKernelContextInternal*>(context);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
assert(context->InputCount() == 2);
const Tensor* A = context->Input<Tensor>(0);
@ -152,4 +152,4 @@ class CDist final : public OpKernel {
}
};
} // namespace contrib
} // namespace onnxruntime
} // namespace onnxruntime

136
onnxruntime/contrib_ops/cpu/crop_and_resize.cc
View file

@ -33,7 +33,7 @@ namespace contrib {
ONNX_OPERATOR_TYPED_KERNEL_EX( \
CropAndResize, \
kMSDomain, \
1, \
1, \
data_type, \
kCpuExecutionProvider, \
KernelDefBuilder() \
@ -46,30 +46,33 @@ ADD_TYPED_CROPANDRESIZE_OP(float);
template <typename T>
static void TryParallelFor(concurrency::ThreadPool* tp, int32_t total, T&& fn) {
if (tp != nullptr)
return tp->ParallelFor(total, fn);
for (int32_t i = 0; i != total; ++i) {
fn(i);
tp->ParallelFor(total, fn);
else {
for (int32_t i = 0; i != total; ++i) {
fn(i);
}
}
}
template <typename T>
void CropAndResizeForward(
int64_t nthreads,
const T* bottom_data,
float extrapolation_value,
int64_t channels,
int64_t height,
int64_t width,
int32_t pooled_height,
int32_t pooled_width,
const T* bottom_rois,
int64_t num_roi_cols,
T* top_data,
const std::string& mode,
const int32_t* batch_indices_ptr,
const ThreadPool* ttp) {
int64_t n_rois = nthreads / channels / pooled_width / pooled_height;
void CropAndResizeForward(const TensorShape& output_shape,
const T* bottom_data,
float extrapolation_value,
int64_t height,
int64_t width,
const T* bottom_rois,
int64_t num_roi_cols,
T* top_data,
const std::string& mode,
const int32_t* batch_indices_ptr,
ThreadPool* ttp) {
int64_t n_rois = output_shape[0];
int64_t channels = output_shape[1];
int64_t pooled_height = output_shape[2];
int64_t pooled_width = output_shape[3];
// TODO: This should do blocks of work based on the number of threads in the threadpool with each block
// being n_rois / num_threads
std::function<void(int32_t)> work_object = [&](int32_t n) {
int64_t index_n = n * channels * pooled_width * pooled_height;
@ -81,36 +84,34 @@ void CropAndResizeForward(
T roi_end_w = offset_bottom_rois[3];
T roi_end_h = offset_bottom_rois[2];
T height_scale =
(pooled_height > 1)
? (roi_end_h - roi_start_h) * (height - 1) / (pooled_height - 1)
: 0;
T width_scale =
(pooled_width > 1) ? (roi_end_w - roi_start_w) * (width - 1) / (pooled_width - 1)
: 0;
T height_scale = (pooled_height > 1)
? (roi_end_h - roi_start_h) * (height - 1) / (pooled_height - 1)
: 0;
T width_scale = (pooled_width > 1)
? (roi_end_w - roi_start_w) * (width - 1) / (pooled_width - 1)
: 0;
for (auto ph = 0; ph < pooled_height; ph++) {
T in_y = static_cast<T>((pooled_height > 1)
? roi_start_h * (height - 1) + ph * height_scale
: 0.5 * (roi_start_h + roi_end_h) * (height - 1));
? roi_start_h * (height - 1) + ph * height_scale
: 0.5 * (roi_start_h + roi_end_h) * (height - 1));
if (ph == pooled_height - 1) {
in_y = static_cast<T>((pooled_height > 1)
? roi_end_h * (height - 1)
: 0.5 * (roi_start_h + roi_end_h) * (height - 1));
? roi_end_h * (height - 1)
: 0.5 * (roi_start_h + roi_end_h) * (height - 1));
}
if (ph == 0) {
in_y = static_cast<T>((pooled_height > 1)
? roi_start_h * (height - 1)
: 0.5 * (roi_start_h + roi_end_h) * (height - 1));
? roi_start_h * (height - 1)
: 0.5 * (roi_start_h + roi_end_h) * (height - 1));
}
if (in_y < 0 || in_y > height - 1) {
for (int64_t pw = 0; pw < pooled_width; pw++) {
for (int64_t c = 0; c < channels; c++)
{
int64_t index_n_c = index_n + c * pooled_width * pooled_height;
int64_t index = index_n_c + ph * pooled_width + pw;
top_data[index] = extrapolation_value;
}
for (int64_t c = 0; c < channels; c++) {
int64_t index_n_c = index_n + c * pooled_width * pooled_height;
int64_t index = index_n_c + ph * pooled_width + pw;
top_data[index] = extrapolation_value;
}
}
continue;
}
@ -121,26 +122,25 @@ void CropAndResizeForward(
for (auto pw = 0; pw < pooled_width; pw++) {
T in_x = static_cast<T>((pooled_width > 1)
? roi_start_w * (width - 1) + pw * width_scale
: 0.5 * (roi_start_w + roi_end_w) * (width - 1));
? roi_start_w * (width - 1) + pw * width_scale
: 0.5 * (roi_start_w + roi_end_w) * (width - 1));
if (pw == pooled_width - 1) {
in_x = static_cast<T>((pooled_width > 1)
? roi_end_w * (width - 1)
: 0.5 * (roi_start_w + roi_end_w) * (width - 1));
in_x = static_cast<T>((pooled_width > 1)
? roi_end_w * (width - 1)
: 0.5 * (roi_start_w + roi_end_w) * (width - 1));
}
if (pw == 0) {
in_x = static_cast<T>((pooled_width > 1)
? roi_start_w * (width - 1)
: 0.5 * (roi_start_w + roi_end_w) * (width - 1));
in_x = static_cast<T>((pooled_width > 1)
? roi_start_w * (width - 1)
: 0.5 * (roi_start_w + roi_end_w) * (width - 1));
}
if (in_x < 0 || in_x > width - 1) {
for (int64_t c = 0; c < channels; c++)
{
int64_t index_n_c = index_n + c * pooled_width * pooled_height;
int64_t index = index_n_c + ph * pooled_width + pw;
top_data[index] = extrapolation_value;
}
continue;
for (int64_t c = 0; c < channels; c++) {
int64_t index_n_c = index_n + c * pooled_width * pooled_height;
int64_t index = index_n_c + ph * pooled_width + pw;
top_data[index] = extrapolation_value;
}
continue;
}
T output_val = extrapolation_value;
@ -153,12 +153,11 @@ void CropAndResizeForward(
auto bottom_left_index = bottom_y_index * width + left_x_index;
auto bottom_right_index = bottom_y_index * width + right_x_index;
for (auto c = 0; c < channels; c++)
{
for (auto c = 0; c < channels; c++) {
int64_t index_n_c = index_n + c * pooled_width * pooled_height;
int64_t index = index_n_c + ph * pooled_width + pw;
const T* offset_bottom_data =
bottom_data + static_cast<int64_t>((roi_batch_ind * channels + c) * height * width);
bottom_data + static_cast<int64_t>((roi_batch_ind * channels + c) * height * width);
const float top_left(static_cast<float>(offset_bottom_data[top_left_index]));
const float top_right(static_cast<float>(offset_bottom_data[top_right_index]));
const float bottom_left(static_cast<float>(offset_bottom_data[bottom_left_index]));
@ -168,25 +167,24 @@ void CropAndResizeForward(
output_val = top + (bottom - top) * y_lerp;
top_data[index] = output_val;
}
}
else { // mode == "nearest"
} else { // mode == "nearest"
const int closest_x_index = static_cast<int>(roundf(static_cast<float>(in_x)));
const int closest_y_index = static_cast<int>(roundf(static_cast<float>(in_y)));
auto closest_index = closest_y_index * width + closest_x_index;
for (auto c = 0; c < channels; c++)
{
for (auto c = 0; c < channels; c++) {
int64_t index_n_c = index_n + c * pooled_width * pooled_height;
int64_t index = index_n_c + ph * pooled_width + pw;
const T* offset_bottom_data =
bottom_data + static_cast<int64_t>((roi_batch_ind * channels + c) * height * width);
bottom_data + static_cast<int64_t>((roi_batch_ind * channels + c) * height * width);
top_data[index] = static_cast<float>(offset_bottom_data[closest_index]);
}
}
} // for pw
} // for ph
}; // for n
TryParallelFor(const_cast<ThreadPool*>(ttp), static_cast<int32_t>(n_rois), work_object);
TryParallelFor(ttp, static_cast<int32_t>(n_rois), work_object);
}
template <typename T>
@ -214,6 +212,7 @@ Status CropAndResize<T>::Compute(OpKernelContext* context) const {
"Number of dimensions for crop size should be exactly 1");
}
auto channels = x_dims[1];
auto num_rois = batch_indices_dims[0];
auto num_roi_cols = rois_dims[1];
auto crop_size_data = crop_size_ptr->Data<int32_t>();
@ -225,23 +224,20 @@ Status CropAndResize<T>::Compute(OpKernelContext* context) const {
return status;
}
auto& Y = *context->Output(0, {num_rois, x_dims[1], crop_height, crop_width});
int64_t output_size = Y.Shape().Size();
TensorShape Y_shape = {num_rois, channels, crop_height, crop_width};
auto& Y = *context->Output(0, Y_shape);
CropAndResizeForward<T>(
output_size, // num threads
Y_shape,
X_ptr->Data<T>(),
extrapolation_value_,
x_dims[1], // num channels
x_dims[2], // height
x_dims[3], // width
crop_height,
crop_width,
rois_ptr->Data<T>(),
num_roi_cols,
Y.template MutableData<T>(),
mode_,
batch_indices_ptr->Data<int32_t>(),
static_cast<OpKernelContextInternal*>(context)->GetOperatorThreadPool());
context->GetOperatorThreadPool());
return Status::OK();
}

9
onnxruntime/contrib_ops/cpu/nchwc_ops.cc
View file

@ -158,7 +158,7 @@ Status NchwcConv::Compute(OpKernelContext* context) const {
y_data,
&activation_,
Sum == nullptr,
const_cast<concurrency::ThreadPool*>(static_cast<OpKernelContextInternal*>(context)->GetOperatorThreadPool()));
context->GetOperatorThreadPool());
return Status::OK();
}
@ -170,7 +170,6 @@ Status NchwcPoolBase::NchwcPool(OpKernelContext* context, MLAS_POOLING_KIND kind
ORT_ENFORCE(X_shape.NumDimensions() == 4);
ORT_ENFORCE((X_shape[1] % MlasNchwcGetBlockSize()) == 0);
std::vector<int64_t> pads = pool_attrs_.pads;
std::vector<int64_t> output_dims = pool_attrs_.SetOutputSize(X_shape, X_shape[1], &pads);
auto* Y = context->Output(0, output_dims);
@ -185,7 +184,7 @@ Status NchwcPoolBase::NchwcPool(OpKernelContext* context, MLAS_POOLING_KIND kind
output_dims.data(),
X->template Data<float>(),
Y->template MutableData<float>(),
const_cast<concurrency::ThreadPool*>(static_cast<OpKernelContextInternal*>(context)->GetOperatorThreadPool()));
context->GetOperatorThreadPool());
return Status::OK();
}
@ -195,8 +194,8 @@ Status NchwcMaxPool::Compute(OpKernelContext* context) const {
}
Status NchwcAveragePool::Compute(OpKernelContext* context) const {
return NchwcPoolBase::NchwcPool(context, pool_attrs_.count_include_pad ? MlasAveragePoolingIncludePad :
MlasAveragePoolingExcludePad);
return NchwcPoolBase::NchwcPool(context, pool_attrs_.count_include_pad ? MlasAveragePoolingIncludePad
: MlasAveragePoolingExcludePad);
}
} // namespace contrib

6
onnxruntime/contrib_ops/cpu/word_conv_embedding.cc
View file

@ -161,9 +161,6 @@ Status WordConvEmbedding::ValidateInputShape(const TensorShape& w_conv_shape, co
}
Status WordConvEmbedding::Compute(OpKernelContext* ctx) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
// original lstm processing
const Tensor& sequence = *(ctx->Input<Tensor>(0)); // sequence: [sequence_length, word_length]
const Tensor& w_conv = *(ctx->Input<Tensor>(1)); // conv weight: [M, C/group, kH, kW]
@ -220,7 +217,8 @@ Status WordConvEmbedding::Compute(OpKernelContext* ctx) const {
char_embedding_size,
filter_width,
filter_size,
Y->MutableData<float>(), tp);
Y->MutableData<float>(),
ctx->GetOperatorThreadPool());
return Status::OK();
}

3
onnxruntime/core/common/threadpool.cc
View file

@ -33,7 +33,8 @@ ThreadPool::ThreadPool(const std::string&, int num_threads) : impl_(num_threads)
void ThreadPool::Schedule(std::function<void()> fn) { impl_.Schedule(fn); }
void ThreadPool::ParallelFor(int32_t total, std::function<void(int32_t)> fn) {
if (total <= 0) return;
if (total <= 0)
return;
if (total == 1) {
fn(0);

2
onnxruntime/core/framework/op_kernel.cc
View file

@ -11,9 +11,11 @@ namespace onnxruntime {
OpKernelContext::OpKernelContext(IExecutionFrame* frame,
const OpKernel* kernel,
concurrency::ThreadPool* threadpool,
const logging::Logger& logger)
: execution_frame_(frame),
kernel_(kernel),
threadpool_(threadpool),
logger_(&logger) {
ORT_ENFORCE(frame != nullptr, "Execution frame was null");
ORT_ENFORCE(kernel != nullptr, "OpKernel was null");

5
onnxruntime/core/framework/op_kernel_context_internal.h
View file

@ -22,7 +22,7 @@ class OpKernelContextInternal : public OpKernelContext {
const OpKernel& kernel,
const logging::Logger& logger,
const bool& terminate_flag)
: OpKernelContext(&frame, &kernel, logger),
: OpKernelContext(&frame, &kernel, session_state.GetThreadPool(), logger),
session_state_(session_state),
terminate_flag_(terminate_flag) {
const auto& implicit_inputs = kernel.Node().ImplicitInputDefs();
@ -60,9 +60,6 @@ class OpKernelContextInternal : public OpKernelContext {
const bool& GetTerminateFlag() const noexcept { return terminate_flag_; }
_Ret_maybenull_ const onnxruntime::concurrency::ThreadPool* GetOperatorThreadPool() const { return session_state_.GetThreadPool(); }
_Ret_maybenull_ onnxruntime::concurrency::ThreadPool* GetOperatorThreadPool() { return session_state_.GetThreadPool(); }
private:
const SessionState& session_state_;
const bool& terminate_flag_;

2
onnxruntime/core/optimizer/constant_folding.cc
View file

@ -50,7 +50,7 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level)
OptimizerExecutionFrame frame(info, fetch_mlvalue_idxs);
auto* kernel = info.GetKernel(node->Index());
OpKernelContext op_kernel_context(&frame, kernel, ::onnxruntime::logging::LoggingManager::DefaultLogger());
OpKernelContext op_kernel_context(&frame, kernel, nullptr, onnxruntime::logging::LoggingManager::DefaultLogger());
ORT_RETURN_IF_ERROR(kernel->Compute(&op_kernel_context));

3
onnxruntime/core/providers/cpu/math/gemm.h
View file

@ -28,8 +28,7 @@ class Gemm : public OpKernel {
}
Status Compute(OpKernelContext* context) const override {
auto ctx_internal = static_cast<OpKernelContextInternal*>(context);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
const auto* X = context->Input<Tensor>(0);
const auto* W = context->Input<Tensor>(1);

3
onnxruntime/core/providers/cpu/math/matmul.cc
View file

@ -68,8 +68,7 @@ ONNX_CPU_OPERATOR_TYPED_KERNEL(
template <typename T>
Status MatMul<T>::Compute(OpKernelContext* ctx) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
concurrency::ThreadPool* thread_pool = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool();
const auto* left_X = ctx->Input<Tensor>(0);
const auto* right_X = ctx->Input<Tensor>(1);

6
onnxruntime/core/providers/cpu/math/matmul_integer.cc
View file

@ -36,8 +36,7 @@ ONNX_OPERATOR_TYPED_KERNEL_EX(
template <>
Status MatMulInteger<uint8_t, uint8_t>::Compute(OpKernelContext* ctx) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
concurrency::ThreadPool* thread_pool = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool();
auto a = ctx->Input<Tensor>(0);
auto b = ctx->Input<Tensor>(1);
@ -82,8 +81,7 @@ Status MatMulInteger<uint8_t, uint8_t>::Compute(OpKernelContext* ctx) const {
template <>
Status MatMulInteger<uint8_t, int8_t>::Compute(OpKernelContext* ctx) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
concurrency::ThreadPool* thread_pool = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool();
auto a = ctx->Input<Tensor>(0);
auto b = ctx->Input<Tensor>(1);

3
onnxruntime/core/providers/cpu/math/softmax.h
View file

@ -86,8 +86,7 @@ class Softmax final : public OpKernel {
Status Compute(OpKernelContext* ctx) const override {
#ifndef USE_OPENMP
auto ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* tp = ctx->GetOperatorThreadPool();
#endif
const auto* tensor_pointer = ctx->Input<Tensor>(0);
if (tensor_pointer == nullptr)

14
onnxruntime/core/providers/cpu/nn/conv.cc
View file

@ -24,13 +24,9 @@ namespace onnxruntime {
template <typename T>
Status Conv<T>::Compute(OpKernelContext* context) const {
size_t num_inputs = OpKernel::Node().InputDefs().size();
auto ctx_internal = static_cast<OpKernelContextInternal*>(context);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
const auto* X = context->Input<Tensor>(0);
const auto* W = context->Input<Tensor>(1);
const Tensor* B = num_inputs == 3 ? context->Input<Tensor>(2) : nullptr;
const Tensor* B = context->Input<Tensor>(2); // optional. nullptr if not provided
const int64_t N = X->Shape()[0];
const int64_t C = X->Shape()[1];
const int64_t M = W->Shape()[0];
@ -84,6 +80,8 @@ Status Conv<T>::Compute(OpKernelContext* context) const {
col_buffer_shape.insert(col_buffer_shape.end(), output_shape.GetDims().begin(),
output_shape.GetDims().end());
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
for (int image_id = 0; image_id < N; ++image_id) {
for (int group_id = 0; group_id < conv_attrs_.group; ++group_id) {
if (Is2DKernel) {
@ -119,6 +117,7 @@ Status Conv<T>::Compute(OpKernelContext* context) const {
col_buffer_data,
&CPUMathUtil::Instance());
}
math::Gemm<T>(
CblasNoTrans,
CblasNoTrans,
@ -147,9 +146,6 @@ Status Conv<T>::Compute(OpKernelContext* context) const {
}
Status Conv<float>::Compute(OpKernelContext* context) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(context);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
size_t num_inputs = OpKernel::Node().InputDefs().size();
const auto* X = context->Input<Tensor>(0);
const auto* W = context->Input<Tensor>(1);
@ -190,6 +186,7 @@ Status Conv<float>::Compute(OpKernelContext* context) const {
auto* Ydata = Y->template MutableData<float>();
const size_t kernel_rank = kernel_shape.size();
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
if (kernel_rank == 2 || kernel_rank == 3) {
MLAS_CONV_PARAMETERS Parameters;
@ -254,6 +251,7 @@ Status Conv<float>::Compute(OpKernelContext* context) const {
static_cast<int>(kernel_shape.size()),
col_buffer_data,
&CPUMathUtil::Instance());
math::Gemm<float>(
CblasNoTrans,
CblasNoTrans,

5
onnxruntime/core/providers/cpu/nn/conv_transpose.cc
View file

@ -42,8 +42,7 @@ Status ConvTranspose<T>::Compute(OpKernelContext* context) const {
template <typename T>
Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(context);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
size_t num_inputs = OpKernel::Node().InputDefs().size();
ConvTransposeAttributes::Prepare p;
@ -55,7 +54,7 @@ Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
const int64_t Y_offset = p.Y->Shape().Size() / p.Y->Shape()[0] / conv_transpose_attrs_.group;
const int64_t W_offset = p.F->Shape().Size() / conv_transpose_attrs_.group;
const int64_t kernel_dim =
p.num_output_channels / conv_transpose_attrs_.group * p.kernel_shape[0] * p.kernel_shape[1];
p.num_output_channels / conv_transpose_attrs_.group * p.kernel_shape[0] * p.kernel_shape[1];
const int64_t output_image_size = p.Y->Shape()[2] * p.Y->Shape()[3];
AllocatorPtr alloc;

5
onnxruntime/core/providers/cpu/nn/pool.cc
View file

@ -190,8 +190,7 @@ Status PoolBase::Compute(OpKernelContext* context, MLAS_POOLING_KIND kind) const
// Get access to the internal threadpool
// Temporarily derive concurrency parameters without access to session state
auto ctx_internal = static_cast<OpKernelContextInternal*>(context);
concurrency::ThreadPool* thread_pool = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool();
MlasPool(kind,
pooling_dims,
@ -202,7 +201,7 @@ Status PoolBase::Compute(OpKernelContext* context, MLAS_POOLING_KIND kind) const
output_dims.data(),
X->template Data<float>(),
Y->template MutableData<float>(),
const_cast<concurrency::ThreadPool*>(thread_pool));
thread_pool);
return Status::OK();
}

87
onnxruntime/core/providers/cpu/object_detection/roialign.cc
View file

@ -45,11 +45,14 @@ namespace {
template <typename T>
void TryParallelFor(concurrency::ThreadPool* tp, int32_t total, T&& fn) {
if (tp != nullptr)
return tp->ParallelFor(total, fn);
for (int32_t i = 0; i != total; ++i) {
fn(i);
tp->ParallelFor(total, fn);
else {
for (int32_t i = 0; i != total; ++i) {
fn(i);
}
}
}
template <typename T>
struct PreCalc {
int64_t pos1;
@ -163,24 +166,25 @@ void pre_calc_for_bilinear_interpolate(
}
template <typename T>
void RoiAlignForward(
int64_t nthreads,
const T* bottom_data,
float spatial_scale,
int64_t channels,
int64_t height,
int64_t width,
int64_t pooled_height,
int64_t pooled_width,
int64_t sampling_ratio,
const T* bottom_rois,
int64_t num_roi_cols,
T* top_data,
RoiAlignMode mode,
const int64_t* batch_indices_ptr,
const ThreadPool* ttp) {
int64_t n_rois = nthreads / channels / pooled_width / pooled_height;
void RoiAlignForward(const TensorShape& output_shape,
const T* bottom_data,
float spatial_scale,
int64_t height,
int64_t width,
int64_t sampling_ratio,
const T* bottom_rois,
int64_t num_roi_cols,
T* top_data,
RoiAlignMode mode,
const int64_t* batch_indices_ptr,
ThreadPool* ttp) {
int64_t n_rois = output_shape[0];
int64_t channels = output_shape[1];
int64_t pooled_height = output_shape[2];
int64_t pooled_width = output_shape[3];
// TODO: This should do blocks of work based on the number of threads in the threadpool with each block
// being n_rois / num_threads
std::function<void(int32_t)> work_object = [&](int32_t n) {
int64_t index_n = n * channels * pooled_width * pooled_height;
@ -258,9 +262,9 @@ void RoiAlignForward(
for (int64_t ix = 0; ix < roi_bin_grid_w; ix++) {
PreCalc<T> pc = pre_calc[pre_calc_index];
T val = std::max(std::max(std::max(pc.w1 * offset_bottom_data[pc.pos1],
pc.w2 * offset_bottom_data[pc.pos2]),
pc.w3 * offset_bottom_data[pc.pos3]),
pc.w4 * offset_bottom_data[pc.pos4]);
pc.w2 * offset_bottom_data[pc.pos2]),
pc.w3 * offset_bottom_data[pc.pos3]),
pc.w4 * offset_bottom_data[pc.pos4]);
if (!max_flag) {
output_val = val;
max_flag = true;
@ -278,7 +282,8 @@ void RoiAlignForward(
} // for ph
} // for c
}; // for n
TryParallelFor(const_cast<ThreadPool*>(ttp), static_cast<int32_t>(n_rois), work_object);
TryParallelFor(ttp, static_cast<int32_t>(n_rois), work_object);
}
} // namespace
@ -334,6 +339,7 @@ Status RoiAlign<T>::Compute(OpKernelContext* context) const {
const auto& rois_dims = rois_ptr->Shape();
const auto& batch_indices_dims = batch_indices_ptr->Shape();
auto num_channels = x_dims[1];
auto num_rois = batch_indices_dims[0];
auto num_roi_cols = rois_dims[1];
@ -342,24 +348,21 @@ Status RoiAlign<T>::Compute(OpKernelContext* context) const {
return status;
}
auto& Y = *context->Output(0, {num_rois, x_dims[1], this->output_height_, this->output_width_});
int64_t output_size = Y.Shape().Size();
RoiAlignForward<T>(
output_size, // num threads
X_ptr->Data<T>(),
this->spatial_scale_,
x_dims[1], // num channels
x_dims[2], // height
x_dims[3], // width
this->output_height_,
this->output_width_,
this->sampling_ratio_,
rois_ptr->Data<T>(),
num_roi_cols,
Y.template MutableData<T>(),
this->mode_,
batch_indices_ptr->Data<int64_t>(),
static_cast<OpKernelContextInternal*>(context)->GetOperatorThreadPool());
TensorShape Y_shape = {num_rois, num_channels, this->output_height_, this->output_width_};
auto& Y = *context->Output(0, Y_shape);
RoiAlignForward<T>(Y_shape,
X_ptr->Data<T>(),
this->spatial_scale_,
x_dims[2], // height
x_dims[3], // width
this->sampling_ratio_,
rois_ptr->Data<T>(),
num_roi_cols,
Y.template MutableData<T>(),
this->mode_,
batch_indices_ptr->Data<int64_t>(),
context->GetOperatorThreadPool());
return Status::OK();
}

3
onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc
View file

@ -268,8 +268,7 @@ Status DeepCpuGruOp::Compute(OpKernelContext* context) const {
template <typename T>
Status DeepCpuGruOp::ComputeImpl(OpKernelContext& context) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(&context);
concurrency::ThreadPool* thread_pool = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* thread_pool = context.GetOperatorThreadPool();
const Tensor& X = *context.Input<Tensor>(0); // inputs. [seq_length, batch_size, input_size]
const Tensor& W = *context.Input<Tensor>(1); // weights. [num_directions, 3*hidden_size, input_size]

3
onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc
View file

@ -314,8 +314,7 @@ DeepCpuLstmOp::Compute(OpKernelContext* context) const {
template <typename T>
Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const {
auto ctx_internal = static_cast<OpKernelContextInternal*>(&context);
concurrency::ThreadPool* mlas_thread_pool = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* mlas_thread_pool = context.GetOperatorThreadPool();
auto& logger = context.Logger();

5
onnxruntime/core/providers/cpu/rnn/rnn.cc
View file

@ -100,8 +100,7 @@ using EigenMatrixMapRowMajor = Eigen::Map<
template <>
Status RNN<float>::Compute(OpKernelContext* ctx) const {
using namespace rnn::detail;
auto ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool();
concurrency::ThreadPool* tp = ctx->GetOperatorThreadPool();
// inputs
const Tensor& X = *ctx->Input<Tensor>(0);
@ -188,7 +187,7 @@ Status RNN<float>::Compute(OpKernelContext* ctx) const {
// the shape of initial_h is [num_directions, batch_size, hidden_size]
// so pick the offset (multiple of Y_frame_size == batch_size * hidden_size_)
// based on the direction
h_prev = initial_h->template Data<float>() + (direction * Y_frame_size);
h_prev = initial_h->template Data<float>() + (direction * Y_frame_size);
}
} else {
if (isReverse)

2
onnxruntime/test/optimizer/optimizer_test.cc
View file

@ -72,7 +72,7 @@ TEST(OptimizerTest, Basic) {
for (auto& node : graph.Nodes()) {
auto* kernel = info.GetKernel(node.Index());
OpKernelContext op_kernel_context(&frame, kernel, logger);
OpKernelContext op_kernel_context(&frame, kernel, nullptr, logger);
auto st = kernel->Compute(&op_kernel_context);
ASSERT_TRUE(st.IsOK()) << st.ErrorMessage();