diff --git a/onnxruntime/core/providers/cpu/ml/scaler.cc b/onnxruntime/core/providers/cpu/ml/scaler.cc index 360f76ca07..2c62bd6e9b 100644 --- a/onnxruntime/core/providers/cpu/ml/scaler.cc +++ b/onnxruntime/core/providers/cpu/ml/scaler.cc @@ -83,15 +83,30 @@ common::Status ScalerOp::Compute(OpKernelContext* context) const { size_t x_size = x_shape.Size(); int64_t stride = x_dims.size() == 1 ? x_dims[0] : x_dims[1]; + auto* ttp = context->GetOperatorThreadPool(); + auto num_threads = std::min(concurrency::ThreadPool::DegreeOfParallelism(ttp), static_cast(x_size)); + if (static_cast(offset_.size()) == stride && static_cast(scale_.size()) == stride) { - for (size_t i = 0; i < x_size; i++) { - y_data[i] = static_cast((x_data[i] - offset_[i % stride]) * scale_[i % stride]); - } + concurrency::ThreadPool::TrySimpleParallelFor( + ttp, + num_threads, + [this, num_threads, y_data, x_data, stride, x_size](ptrdiff_t batch_num) { + auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, x_size); + for (auto i = work.start; i < work.end; ++i) { + y_data[i] = static_cast((x_data[i] - offset_[i % stride]) * scale_[i % stride]); + } + }); } else if (offset_.size() == 1 && scale_.size() == 1) { - for (size_t i = 0; i < x_size; i++) { - y_data[i] = static_cast((x_data[i] - offset_[0]) * scale_[0]); - } + concurrency::ThreadPool::TrySimpleParallelFor( + ttp, + num_threads, + [this, num_threads, y_data, x_data, x_size](ptrdiff_t batch_num) { + auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, x_size); + for (auto i = work.start; i < work.end; ++i) { + y_data[i] = static_cast((x_data[i] - offset_[0]) * scale_[0]); + } + }); } else { std::ostringstream err_msg; err_msg << "Either both scale and offset can be of feature size (" << stride << ") or 1"; diff --git a/onnxruntime/core/providers/cpu/tensor/gather_elements.cc b/onnxruntime/core/providers/cpu/tensor/gather_elements.cc index b8238d9a56..b5fd0bfcf8 100644 --- a/onnxruntime/core/providers/cpu/tensor/gather_elements.cc +++ b/onnxruntime/core/providers/cpu/tensor/gather_elements.cc @@ -86,42 +86,15 @@ static inline void increment_over_inner_dim(std::vector& current_dims, } } -// parse indices_tensor and along the way validate its shape and contents -static std::vector parse_and_validate_indices_tensor(const Tensor* indices_tensor, - int64_t axis, const TensorShape& input_shape) { - // first parse 'indices' data - auto num_elements = indices_tensor->Shape().Size(); - std::vector indices_data; - // reserving memory ahead as we know the size of the container - indices_data.reserve(num_elements); - if (indices_tensor->IsDataType()) { - const auto* data = indices_tensor->Data(); - for (int64_t i = 0; i < num_elements; ++i) - indices_data.push_back(data[i]); - } else if (indices_tensor->IsDataType()) { - const auto* data = indices_tensor->Data(); - for (int64_t i = 0; i < num_elements; ++i) - indices_data.push_back(data[i]); +template +static inline int64_t GetNegativeIndexAdjustedValue(const Tin* indices_data, Tin index, int64_t axis, const TensorShape& input_shape) { + int64_t retval = -1; + if (indices_data[index] < 0) { + retval = static_cast(indices_data[index] + input_shape[axis]); } else { - ORT_THROW("GatherElements op: Data type for 'indices' tensor must be 'int32_t' and 'int64_t'"); + retval = static_cast(indices_data[index]); } - - // validate 'indices' data - // along the way 'fix' negative index values if within bounds - int64_t lower_index_limit = -input_shape[axis]; - int64_t upper_index_limit = input_shape[axis] - 1; - - for (int64_t i = 0; i < num_elements; ++i) { - auto indices_val = indices_data[i]; - if (indices_val < lower_index_limit || indices_val > upper_index_limit) - ORT_THROW("GatherElements op: Value in indices must be within bounds [", - lower_index_limit, " , ", upper_index_limit, "]. Actual value is ", indices_val); - - if (indices_val < 0) - indices_data[i] += input_shape[axis]; - } - - return indices_data; + return retval; } #ifdef __GNUC__ @@ -130,7 +103,7 @@ static std::vector parse_and_validate_indices_tensor(const Tensor* indi #pragma GCC diagnostic ignored "-Wclass-memaccess" #endif #endif -template +template static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor, Tensor* output_tensor, int64_t axis) { // get pointer to input data @@ -154,15 +127,27 @@ static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor, const int64_t input_rank = static_cast(input_tensor->Shape().NumDimensions()); const TensorPitches input_shape_pitches(*input_tensor); - const std::vector& indices_data = parse_and_validate_indices_tensor(indices_tensor, axis, input_tensor->Shape()); + const auto& input_shape = input_tensor->Shape(); const TensorShape& indices_shape = indices_tensor->Shape(); + const Tin* indices_data = indices_tensor->Data(); + + // validate indices + auto num_elements = indices_tensor->Shape().Size(); + int64_t lower_index_limit = -input_shape[axis]; + int64_t upper_index_limit = input_shape[axis] - 1; + for (int64_t i = 0; i < num_elements; ++i) { + auto indices_val = indices_data[i]; + if (indices_val < lower_index_limit || indices_val > upper_index_limit) + ORT_THROW("GatherElements op: Value in indices must be within bounds [", + lower_index_limit, " , ", upper_index_limit, "]. Actual value is ", indices_val); + } int64_t num_inner_dim = calculate_num_inner_dim(indices_shape); int64_t inner_dim_size = indices_shape[input_rank - 1]; bool processing_inner_dim = (axis == input_rank - 1) ? true : false; int64_t base_offset = 0; - int64_t indices_counter = -1; + Tin indices_counter = -1; int64_t output_counter = -1; size_t element_size = input_tensor->DataType()->Size(); @@ -173,17 +158,24 @@ static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor, base_offset = compute_base_offset(process_dims, input_shape_pitches, axis); // process 1 chunk of 'inner dimension' length - for (int64_t i = 0; i < inner_dim_size; ++i) { - // optimizer will remove the redundant if/else block based on 'is_string' template parameter - if (is_string) { - output_data[++output_counter] = input_data[base_offset + (indices_data[++indices_counter] * input_shape_pitches[axis]) + i]; - } else { + // optimizer will remove the redundant if/else block based on 'is_string' template parameter + if (is_string) { + for (int64_t i = 0; i < inner_dim_size; ++i) { + output_data[++output_counter] = + input_data[base_offset + + (GetNegativeIndexAdjustedValue(indices_data, ++indices_counter, axis, input_shape) * + input_shape_pitches[axis]) + + i]; + } + } else { + for (int64_t i = 0; i < inner_dim_size; ++i) { + // optimizer will remove the redundant if/else block based on 'is_string' template parameter memcpy(output_data, - input_data + (base_offset + (indices_data[++indices_counter] * input_shape_pitches[axis]) + i) * element_size, element_size); + input_data + (base_offset + (GetNegativeIndexAdjustedValue(indices_data, ++indices_counter, axis, input_shape) * input_shape_pitches[axis]) + i) * element_size, + element_size); output_data += element_size; } } - increment_over_inner_dim(process_dims, indices_shape); } } @@ -193,13 +185,23 @@ static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor, base_offset = compute_base_offset(process_dims, input_shape_pitches, axis); // process 1 chunk of 'inner dimension' length - for (int64_t i = 0; i < inner_dim_size; ++i) { - // for innermost axis, input_shape_pitches[axis] = 1 (so no need to multiply) - // optimizer will remove the redundant if/else block based on 'is_string' template parameter - if (is_string) { - output_data[++output_counter] = input_data[base_offset + indices_data[++indices_counter]]; - } else { - memcpy(output_data, input_data + (base_offset + indices_data[++indices_counter]) * element_size, element_size); + // optimizer will remove the redundant if/else block based on 'is_string' template parameter + if (is_string) { + for (int64_t i = 0; i < inner_dim_size; ++i) { + // for innermost axis, input_shape_pitches[axis] = 1 (so no need to multiply) + output_data[++output_counter] = + input_data[base_offset + + GetNegativeIndexAdjustedValue(indices_data, ++indices_counter, axis, input_shape)]; + } + } else { + for (int64_t i = 0; i < inner_dim_size; ++i) { + // for innermost axis, input_shape_pitches[axis] = 1 (so no need to multiply) + // optimizer will remove the redundant if/else block based on 'is_string' template parameter + memcpy(output_data, + input_data + (base_offset + + GetNegativeIndexAdjustedValue(indices_data, ++indices_counter, axis, input_shape)) * + element_size, + element_size); output_data += element_size; } } @@ -270,10 +272,17 @@ Status GatherElements::Compute(OpKernelContext* context) const { if (indices_shape.Size() == 0) return Status::OK(); - if (input_tensor->IsDataTypeString()) - core_impl(input_tensor, indices_tensor, output_tensor, axis); - else - core_impl(input_tensor, indices_tensor, output_tensor, axis); + if (input_tensor->IsDataTypeString()) { + if (indices_tensor->IsDataType()) + core_impl(input_tensor, indices_tensor, output_tensor, axis); + else + core_impl(input_tensor, indices_tensor, output_tensor, axis); + } else { + if (indices_tensor->IsDataType()) + core_impl(input_tensor, indices_tensor, output_tensor, axis); + else + core_impl(input_tensor, indices_tensor, output_tensor, axis); + } return Status::OK(); }