From bdfff800ea47259d6d72ad12b128a9af5447d49a Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Sat, 12 Oct 2019 03:36:20 +1000 Subject: [PATCH] Move access to intra-op threadpool into OpKernelContext. (#2091) --- .../onnxruntime/core/framework/op_kernel.h | 10 ++ .../cpu/attnlstm/deep_cpu_attn_lstm.cc | 5 +- onnxruntime/contrib_ops/cpu/cdist.h | 8 +- .../contrib_ops/cpu/crop_and_resize.cc | 136 +++++++++--------- onnxruntime/contrib_ops/cpu/nchwc_ops.cc | 9 +- .../contrib_ops/cpu/word_conv_embedding.cc | 6 +- onnxruntime/core/common/threadpool.cc | 3 +- onnxruntime/core/framework/op_kernel.cc | 2 + .../framework/op_kernel_context_internal.h | 5 +- .../core/optimizer/constant_folding.cc | 2 +- onnxruntime/core/providers/cpu/math/gemm.h | 3 +- onnxruntime/core/providers/cpu/math/matmul.cc | 3 +- .../core/providers/cpu/math/matmul_integer.cc | 6 +- onnxruntime/core/providers/cpu/math/softmax.h | 3 +- onnxruntime/core/providers/cpu/nn/conv.cc | 14 +- .../core/providers/cpu/nn/conv_transpose.cc | 5 +- onnxruntime/core/providers/cpu/nn/pool.cc | 5 +- .../cpu/object_detection/roialign.cc | 87 +++++------ .../core/providers/cpu/rnn/deep_cpu_gru.cc | 3 +- .../core/providers/cpu/rnn/deep_cpu_lstm.cc | 3 +- onnxruntime/core/providers/cpu/rnn/rnn.cc | 5 +- onnxruntime/test/optimizer/optimizer_test.cc | 2 +- 22 files changed, 159 insertions(+), 166 deletions(-) diff --git a/include/onnxruntime/core/framework/op_kernel.h b/include/onnxruntime/core/framework/op_kernel.h index f7a1902d3d..323fd19943 100644 --- a/include/onnxruntime/core/framework/op_kernel.h +++ b/include/onnxruntime/core/framework/op_kernel.h @@ -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. diff --git a/onnxruntime/contrib_ops/cpu/attnlstm/deep_cpu_attn_lstm.cc b/onnxruntime/contrib_ops/cpu/attnlstm/deep_cpu_attn_lstm.cc index 50e98f8342..38b33c90be 100644 --- a/onnxruntime/contrib_ops/cpu/attnlstm/deep_cpu_attn_lstm.cc +++ b/onnxruntime/contrib_ops/cpu/attnlstm/deep_cpu_attn_lstm.cc @@ -72,9 +72,6 @@ static gsl::span SecondHalfSpan(const gsl::span& dspan) { template Status DeepCpuAttnLstmOp::ComputeImpl(OpKernelContext& context) const { - auto ctx_internal = static_cast(&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 input_weights_2 = input_weights.subspan(input_weights_size_per_direction, diff --git a/onnxruntime/contrib_ops/cpu/cdist.h b/onnxruntime/contrib_ops/cpu/cdist.h index 9ffaf2fce6..494eb13ee6 100644 --- a/onnxruntime/contrib_ops/cpu/cdist.h +++ b/onnxruntime/contrib_ops/cpu/cdist.h @@ -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(context); - concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool(); + concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); assert(context->InputCount() == 2); const Tensor* A = context->Input(0); @@ -152,4 +152,4 @@ class CDist final : public OpKernel { } }; } // namespace contrib -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/crop_and_resize.cc b/onnxruntime/contrib_ops/cpu/crop_and_resize.cc index c7e8232477..f966c9c821 100644 --- a/onnxruntime/contrib_ops/cpu/crop_and_resize.cc +++ b/onnxruntime/contrib_ops/cpu/crop_and_resize.cc @@ -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 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 -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 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((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((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((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((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((pooled_width > 1) - ? roi_end_w * (width - 1) - : 0.5 * (roi_start_w + roi_end_w) * (width - 1)); + in_x = static_cast((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((pooled_width > 1) - ? roi_start_w * (width - 1) - : 0.5 * (roi_start_w + roi_end_w) * (width - 1)); + in_x = static_cast((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((roi_batch_ind * channels + c) * height * width); + bottom_data + static_cast((roi_batch_ind * channels + c) * height * width); const float top_left(static_cast(offset_bottom_data[top_left_index])); const float top_right(static_cast(offset_bottom_data[top_right_index])); const float bottom_left(static_cast(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(roundf(static_cast(in_x))); const int closest_y_index = static_cast(roundf(static_cast(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((roi_batch_ind * channels + c) * height * width); + bottom_data + static_cast((roi_batch_ind * channels + c) * height * width); top_data[index] = static_cast(offset_bottom_data[closest_index]); } } } // for pw } // for ph }; // for n - TryParallelFor(const_cast(ttp), static_cast(n_rois), work_object); + + TryParallelFor(ttp, static_cast(n_rois), work_object); } template @@ -214,6 +212,7 @@ Status CropAndResize::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(); @@ -225,23 +224,20 @@ Status CropAndResize::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( - output_size, // num threads + Y_shape, X_ptr->Data(), extrapolation_value_, - x_dims[1], // num channels x_dims[2], // height x_dims[3], // width - crop_height, - crop_width, rois_ptr->Data(), num_roi_cols, Y.template MutableData(), mode_, batch_indices_ptr->Data(), - static_cast(context)->GetOperatorThreadPool()); + context->GetOperatorThreadPool()); return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cpu/nchwc_ops.cc b/onnxruntime/contrib_ops/cpu/nchwc_ops.cc index 5e16eda148..325217a1f8 100644 --- a/onnxruntime/contrib_ops/cpu/nchwc_ops.cc +++ b/onnxruntime/contrib_ops/cpu/nchwc_ops.cc @@ -158,7 +158,7 @@ Status NchwcConv::Compute(OpKernelContext* context) const { y_data, &activation_, Sum == nullptr, - const_cast(static_cast(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 pads = pool_attrs_.pads; std::vector 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(), Y->template MutableData(), - const_cast(static_cast(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 diff --git a/onnxruntime/contrib_ops/cpu/word_conv_embedding.cc b/onnxruntime/contrib_ops/cpu/word_conv_embedding.cc index 3213ff4fc1..a0e6d7c33d 100644 --- a/onnxruntime/contrib_ops/cpu/word_conv_embedding.cc +++ b/onnxruntime/contrib_ops/cpu/word_conv_embedding.cc @@ -161,9 +161,6 @@ Status WordConvEmbedding::ValidateInputShape(const TensorShape& w_conv_shape, co } Status WordConvEmbedding::Compute(OpKernelContext* ctx) const { - auto ctx_internal = static_cast(ctx); - concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool(); - // original lstm processing const Tensor& sequence = *(ctx->Input(0)); // sequence: [sequence_length, word_length] const Tensor& w_conv = *(ctx->Input(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(), tp); + Y->MutableData(), + ctx->GetOperatorThreadPool()); return Status::OK(); } diff --git a/onnxruntime/core/common/threadpool.cc b/onnxruntime/core/common/threadpool.cc index 6cdcb3add7..0595f8c56e 100644 --- a/onnxruntime/core/common/threadpool.cc +++ b/onnxruntime/core/common/threadpool.cc @@ -33,7 +33,8 @@ ThreadPool::ThreadPool(const std::string&, int num_threads) : impl_(num_threads) void ThreadPool::Schedule(std::function fn) { impl_.Schedule(fn); } void ThreadPool::ParallelFor(int32_t total, std::function fn) { - if (total <= 0) return; + if (total <= 0) + return; if (total == 1) { fn(0); diff --git a/onnxruntime/core/framework/op_kernel.cc b/onnxruntime/core/framework/op_kernel.cc index 10127b4fb9..62d4e8cd25 100644 --- a/onnxruntime/core/framework/op_kernel.cc +++ b/onnxruntime/core/framework/op_kernel.cc @@ -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"); diff --git a/onnxruntime/core/framework/op_kernel_context_internal.h b/onnxruntime/core/framework/op_kernel_context_internal.h index 44cbb2b532..2c7b91a64b 100644 --- a/onnxruntime/core/framework/op_kernel_context_internal.h +++ b/onnxruntime/core/framework/op_kernel_context_internal.h @@ -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_; diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index ff7882e974..be15e4b2d6 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -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)); diff --git a/onnxruntime/core/providers/cpu/math/gemm.h b/onnxruntime/core/providers/cpu/math/gemm.h index fed4c987e1..a309ad144a 100644 --- a/onnxruntime/core/providers/cpu/math/gemm.h +++ b/onnxruntime/core/providers/cpu/math/gemm.h @@ -28,8 +28,7 @@ class Gemm : public OpKernel { } Status Compute(OpKernelContext* context) const override { - auto ctx_internal = static_cast(context); - concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool(); + concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); const auto* X = context->Input(0); const auto* W = context->Input(1); diff --git a/onnxruntime/core/providers/cpu/math/matmul.cc b/onnxruntime/core/providers/cpu/math/matmul.cc index 3eda03219a..cff1b1066a 100644 --- a/onnxruntime/core/providers/cpu/math/matmul.cc +++ b/onnxruntime/core/providers/cpu/math/matmul.cc @@ -68,8 +68,7 @@ ONNX_CPU_OPERATOR_TYPED_KERNEL( template Status MatMul::Compute(OpKernelContext* ctx) const { - auto ctx_internal = static_cast(ctx); - concurrency::ThreadPool* thread_pool = ctx_internal->GetOperatorThreadPool(); + concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool(); const auto* left_X = ctx->Input(0); const auto* right_X = ctx->Input(1); diff --git a/onnxruntime/core/providers/cpu/math/matmul_integer.cc b/onnxruntime/core/providers/cpu/math/matmul_integer.cc index 68ddf891d3..374c49fc6f 100644 --- a/onnxruntime/core/providers/cpu/math/matmul_integer.cc +++ b/onnxruntime/core/providers/cpu/math/matmul_integer.cc @@ -36,8 +36,7 @@ ONNX_OPERATOR_TYPED_KERNEL_EX( template <> Status MatMulInteger::Compute(OpKernelContext* ctx) const { - auto ctx_internal = static_cast(ctx); - concurrency::ThreadPool* thread_pool = ctx_internal->GetOperatorThreadPool(); + concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool(); auto a = ctx->Input(0); auto b = ctx->Input(1); @@ -82,8 +81,7 @@ Status MatMulInteger::Compute(OpKernelContext* ctx) const { template <> Status MatMulInteger::Compute(OpKernelContext* ctx) const { - auto ctx_internal = static_cast(ctx); - concurrency::ThreadPool* thread_pool = ctx_internal->GetOperatorThreadPool(); + concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool(); auto a = ctx->Input(0); auto b = ctx->Input(1); diff --git a/onnxruntime/core/providers/cpu/math/softmax.h b/onnxruntime/core/providers/cpu/math/softmax.h index 484c6057d5..bf4e7ca433 100644 --- a/onnxruntime/core/providers/cpu/math/softmax.h +++ b/onnxruntime/core/providers/cpu/math/softmax.h @@ -86,8 +86,7 @@ class Softmax final : public OpKernel { Status Compute(OpKernelContext* ctx) const override { #ifndef USE_OPENMP - auto ctx_internal = static_cast(ctx); - concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool(); + concurrency::ThreadPool* tp = ctx->GetOperatorThreadPool(); #endif const auto* tensor_pointer = ctx->Input(0); if (tensor_pointer == nullptr) diff --git a/onnxruntime/core/providers/cpu/nn/conv.cc b/onnxruntime/core/providers/cpu/nn/conv.cc index 56430eb44b..4ea54ae89e 100644 --- a/onnxruntime/core/providers/cpu/nn/conv.cc +++ b/onnxruntime/core/providers/cpu/nn/conv.cc @@ -24,13 +24,9 @@ namespace onnxruntime { template Status Conv::Compute(OpKernelContext* context) const { - size_t num_inputs = OpKernel::Node().InputDefs().size(); - auto ctx_internal = static_cast(context); - concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool(); - const auto* X = context->Input(0); const auto* W = context->Input(1); - const Tensor* B = num_inputs == 3 ? context->Input(2) : nullptr; + const Tensor* B = context->Input(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::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::Compute(OpKernelContext* context) const { col_buffer_data, &CPUMathUtil::Instance()); } + math::Gemm( CblasNoTrans, CblasNoTrans, @@ -147,9 +146,6 @@ Status Conv::Compute(OpKernelContext* context) const { } Status Conv::Compute(OpKernelContext* context) const { - auto ctx_internal = static_cast(context); - concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool(); - size_t num_inputs = OpKernel::Node().InputDefs().size(); const auto* X = context->Input(0); const auto* W = context->Input(1); @@ -190,6 +186,7 @@ Status Conv::Compute(OpKernelContext* context) const { auto* Ydata = Y->template MutableData(); 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::Compute(OpKernelContext* context) const { static_cast(kernel_shape.size()), col_buffer_data, &CPUMathUtil::Instance()); + math::Gemm( CblasNoTrans, CblasNoTrans, diff --git a/onnxruntime/core/providers/cpu/nn/conv_transpose.cc b/onnxruntime/core/providers/cpu/nn/conv_transpose.cc index 08b99c0d77..91739f6e7d 100644 --- a/onnxruntime/core/providers/cpu/nn/conv_transpose.cc +++ b/onnxruntime/core/providers/cpu/nn/conv_transpose.cc @@ -42,8 +42,7 @@ Status ConvTranspose::Compute(OpKernelContext* context) const { template Status ConvTranspose::DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const { - auto ctx_internal = static_cast(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::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; diff --git a/onnxruntime/core/providers/cpu/nn/pool.cc b/onnxruntime/core/providers/cpu/nn/pool.cc index 6582bd965e..a049f2c044 100644 --- a/onnxruntime/core/providers/cpu/nn/pool.cc +++ b/onnxruntime/core/providers/cpu/nn/pool.cc @@ -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(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(), Y->template MutableData(), - const_cast(thread_pool)); + thread_pool); return Status::OK(); } diff --git a/onnxruntime/core/providers/cpu/object_detection/roialign.cc b/onnxruntime/core/providers/cpu/object_detection/roialign.cc index 64389d08d2..e68258cb3b 100644 --- a/onnxruntime/core/providers/cpu/object_detection/roialign.cc +++ b/onnxruntime/core/providers/cpu/object_detection/roialign.cc @@ -45,11 +45,14 @@ namespace { template 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 struct PreCalc { int64_t pos1; @@ -163,24 +166,25 @@ void pre_calc_for_bilinear_interpolate( } template -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 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 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(ttp), static_cast(n_rois), work_object); + + TryParallelFor(ttp, static_cast(n_rois), work_object); } } // namespace @@ -334,6 +339,7 @@ Status RoiAlign::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::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( - output_size, // num threads - X_ptr->Data(), - 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(), - num_roi_cols, - Y.template MutableData(), - this->mode_, - batch_indices_ptr->Data(), - static_cast(context)->GetOperatorThreadPool()); + TensorShape Y_shape = {num_rois, num_channels, this->output_height_, this->output_width_}; + auto& Y = *context->Output(0, Y_shape); + + RoiAlignForward(Y_shape, + X_ptr->Data(), + this->spatial_scale_, + x_dims[2], // height + x_dims[3], // width + this->sampling_ratio_, + rois_ptr->Data(), + num_roi_cols, + Y.template MutableData(), + this->mode_, + batch_indices_ptr->Data(), + context->GetOperatorThreadPool()); return Status::OK(); } diff --git a/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc b/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc index c2fda16f25..bb26bddc76 100644 --- a/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc +++ b/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc @@ -268,8 +268,7 @@ Status DeepCpuGruOp::Compute(OpKernelContext* context) const { template Status DeepCpuGruOp::ComputeImpl(OpKernelContext& context) const { - auto ctx_internal = static_cast(&context); - concurrency::ThreadPool* thread_pool = ctx_internal->GetOperatorThreadPool(); + concurrency::ThreadPool* thread_pool = context.GetOperatorThreadPool(); const Tensor& X = *context.Input(0); // inputs. [seq_length, batch_size, input_size] const Tensor& W = *context.Input(1); // weights. [num_directions, 3*hidden_size, input_size] diff --git a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc index 05d6753cd7..af0bdbce19 100644 --- a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc +++ b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc @@ -314,8 +314,7 @@ DeepCpuLstmOp::Compute(OpKernelContext* context) const { template Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const { - auto ctx_internal = static_cast(&context); - concurrency::ThreadPool* mlas_thread_pool = ctx_internal->GetOperatorThreadPool(); + concurrency::ThreadPool* mlas_thread_pool = context.GetOperatorThreadPool(); auto& logger = context.Logger(); diff --git a/onnxruntime/core/providers/cpu/rnn/rnn.cc b/onnxruntime/core/providers/cpu/rnn/rnn.cc index d26b02f81a..82675b4d14 100644 --- a/onnxruntime/core/providers/cpu/rnn/rnn.cc +++ b/onnxruntime/core/providers/cpu/rnn/rnn.cc @@ -100,8 +100,7 @@ using EigenMatrixMapRowMajor = Eigen::Map< template <> Status RNN::Compute(OpKernelContext* ctx) const { using namespace rnn::detail; - auto ctx_internal = static_cast(ctx); - concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool(); + concurrency::ThreadPool* tp = ctx->GetOperatorThreadPool(); // inputs const Tensor& X = *ctx->Input(0); @@ -188,7 +187,7 @@ Status RNN::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() + (direction * Y_frame_size); + h_prev = initial_h->template Data() + (direction * Y_frame_size); } } else { if (isReverse) diff --git a/onnxruntime/test/optimizer/optimizer_test.cc b/onnxruntime/test/optimizer/optimizer_test.cc index 48c84f47bc..097ba35821 100644 --- a/onnxruntime/test/optimizer/optimizer_test.cc +++ b/onnxruntime/test/optimizer/optimizer_test.cc @@ -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();