From b7128e9c012f098487060b5938e8f99f1ba6acf4 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 29 Jan 2025 11:32:56 -0800 Subject: [PATCH] Use Eigen for Round Retrofit NearestSample Fix up casting Use InlinedVector --- onnxruntime/core/providers/cpu/math/round.cc | 21 +- .../core/providers/cpu/tensor/upsample.cc | 206 ++++++++++-------- 2 files changed, 127 insertions(+), 100 deletions(-) diff --git a/onnxruntime/core/providers/cpu/math/round.cc b/onnxruntime/core/providers/cpu/math/round.cc index 86be7cce43..c437d298ee 100644 --- a/onnxruntime/core/providers/cpu/math/round.cc +++ b/onnxruntime/core/providers/cpu/math/round.cc @@ -9,6 +9,7 @@ #include #include "core/providers/cpu/math/element_wise_ops.h" #include "core/util/math.h" +#include "core/util/math_cpuonly.h" namespace onnxruntime { @@ -24,24 +25,28 @@ template Status Round::Compute(OpKernelContext* ctx) const { const auto& X = *ctx->Input(0); auto& Y = *ctx->Output(0, X.Shape()); - auto* input = X.Data(); + const auto* input = X.Data(); auto* output = Y.MutableData(); const auto size = X.Shape().Size(); - for (int64_t i = 0; i < size; ++i, ++output, ++input) { - *output = ::rint(*input); - } + + EigenArrayMap Y_arr(output, 1, size); + ConstEigenArrayMap X_arr(input, 1, size); + Y_arr = X_arr.rint(); + return Status::OK(); } template <> Status Round::Compute(OpKernelContext* ctx) const { const auto& X = *ctx->Input(0); auto& Y = *ctx->Output(0, X.Shape()); - auto* input = X.Data(); + const auto* input = X.Data(); auto* output = Y.MutableData(); const auto size = X.Shape().Size(); - for (int64_t i = 0; i < size; ++i, ++output, ++input) { - *output = MLFloat16(static_cast(::rint(input->ToFloat()))); - } + + ConstEigenArrayMap X_arr(reinterpret_cast(input), 1, size); + EigenArrayMap Y_arr(reinterpret_cast(output), 1, size); + Y_arr = X_arr.rint(); + return Status::OK(); } diff --git a/onnxruntime/core/providers/cpu/tensor/upsample.cc b/onnxruntime/core/providers/cpu/tensor/upsample.cc index babbac0b7b..2ee859c900 100644 --- a/onnxruntime/core/providers/cpu/tensor/upsample.cc +++ b/onnxruntime/core/providers/cpu/tensor/upsample.cc @@ -10,11 +10,10 @@ #include "core/platform/threadpool.h" #include "core/providers/cpu/tensor/upsample_antialias.h" -using namespace onnxruntime::common; -using namespace std; -using onnxruntime::narrow; namespace onnxruntime { +using namespace common; + #define REGISTER_VERSIONED_TYPED_KERNEL(T, start, end) \ ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL( \ Upsample, \ @@ -102,18 +101,19 @@ void UpsampleNearest2x(int64_t batch_size, } } -static std::vector UpsampleNearestSetupRank1InputMapping( +static InlinedVector UpsampleNearestSetupRank1InputMapping( int64_t length_original, - int64_t length_resized, + int64_t length_res, float x_scale, float roi_start, float roi_end, bool extrapolation_enabled, const GetOriginalCoordinateFunc& get_original_coordinate, const GetNearestPixelFunc& get_nearest_pixel) { - std::vector input_mapping(onnxruntime::narrow(length_resized)); + const auto length_resized = narrow(length_res); + InlinedVector input_mapping(length_resized); - for (int64_t output_dim0_idx = 0; output_dim0_idx < length_resized; ++output_dim0_idx) { + for (size_t output_dim0_idx = 0; output_dim0_idx < length_resized; ++output_dim0_idx) { float original_0_idx = get_original_coordinate(static_cast(output_dim0_idx), x_scale, static_cast(length_resized), @@ -124,57 +124,67 @@ static std::vector UpsampleNearestSetupRank1InputMapping( // leave as -1 to indicate the extrapolation value should be used } else { input_dim0_idx = get_nearest_pixel(original_0_idx, x_scale < 1); - if (input_dim0_idx > length_original - 1) input_dim0_idx = length_original - 1; - if (input_dim0_idx < 0) input_dim0_idx = 0; + if (input_dim0_idx < 0) { + input_dim0_idx = 0; + } else if (input_dim0_idx > length_original - 1) { + input_dim0_idx = length_original - 1; + } } - input_mapping[narrow(output_dim0_idx)] = input_dim0_idx; + input_mapping[output_dim0_idx] = input_dim0_idx; } return input_mapping; } -static std::vector> -UpsampleNearestSetupInputMappings(int64_t n_dim, +static std::vector> +UpsampleNearestSetupInputMappings(size_t n_dims, const TensorShape& input_shape, const TensorShape& output_shape, - const std::vector& input_dim_factor, + gsl::span input_dim_factor, gsl::span scales, gsl::span roi, bool extrapolation_enabled, const GetOriginalCoordinateFunc& get_original_coordinate, const GetNearestPixelFunc& get_nearest_pixel) { - std::vector> input_mappings(narrow(n_dim)); + std::vector> input_mappings; + input_mappings.reserve(n_dims); - for (int64_t axis = 0; axis < n_dim; ++axis) { - std::vector& input_mapping = input_mappings[narrow(axis)]; - input_mapping.resize(narrow(output_shape[narrow(axis)])); + for (size_t axis = 0; axis < n_dims; ++axis) { + const auto output_axis_val = narrow(output_shape[axis]); + auto& input_mapping = input_mappings.emplace_back(); + input_mapping.reserve(output_axis_val); // When scale is 1.0, there is a one-to-one mapping between the dimension // in the input and the output and there is no need to apply the co-ordinate // transformation which should only be done when there is "resizing" required - if (scales[narrow(axis)] == 1.0f) { - for (int64_t dim = 0; dim < output_shape[narrow(axis)]; dim++) { - input_mapping[narrow(dim)] = dim * input_dim_factor[narrow(axis)]; + if (scales[axis] == 1.0f) { + for (int64_t dim = 0; dim < output_shape[axis]; dim++) { + input_mapping.push_back(SafeInt(dim) * input_dim_factor[axis]); } continue; } // scale != 1.0 - const int64_t input_size = input_dim_factor[0] * input_shape[0]; - for (int64_t dim = 0; dim < output_shape[narrow(axis)]; dim++) { + const int64_t input_size = SafeInt(input_dim_factor[0]) * input_shape[0]; + for (size_t dim = 0; dim < output_axis_val; dim++) { float original_dim = get_original_coordinate(static_cast(dim), - scales[narrow(axis)], - static_cast(output_shape[narrow(axis)]), - static_cast(input_shape[narrow(axis)]), - roi[narrow(axis)], roi[SafeInt(n_dim) + axis]); + scales[axis], + static_cast(output_shape[axis]), + static_cast(input_shape[axis]), + roi[axis], roi[SafeInt(n_dims) + axis]); - bool need_extrapolation = (extrapolation_enabled && (original_dim < 0 || original_dim > input_shape[narrow(axis)] - 1)); - int64_t input_dim = get_nearest_pixel(original_dim, scales[narrow(axis)] < 1); - if (input_dim >= input_shape[narrow(axis)]) input_dim = input_shape[narrow(axis)] - 1; - if (input_dim < 0) input_dim = 0; + bool need_extrapolation = (extrapolation_enabled && (original_dim < 0 || original_dim > input_shape[axis] - 1)); + int64_t input_dim = get_nearest_pixel(original_dim, scales[axis] < 1); + if (input_dim < 0) { + input_dim = 0; + } else if (input_dim >= input_shape[axis]) { + input_dim = input_shape[axis] - 1; + } - input_mapping[narrow(dim)] = need_extrapolation ? (-input_size) : (input_dim * input_dim_factor[narrow(axis)]); + input_mapping.push_back(need_extrapolation + ? (-input_size) + : (SafeInt(input_dim) * input_dim_factor[axis])); } } @@ -192,65 +202,68 @@ static Status UpsampleNearestImpl(const T* input, const T extrapolation_value, const GetOriginalCoordinateFunc& get_original_coordinate, const GetNearestPixelFunc& get_nearest_pixel) { - int64_t n_dim = static_cast(input_shape.NumDimensions()); + const size_t n_dim = input_shape.NumDimensions(); - std::vector input_dim_counters(narrow(n_dim)); - std::vector input_dim_factor(narrow(n_dim)); - input_dim_factor[SafeInt(n_dim) - 1] = 1; // initialize dimension factor - for (int64_t dim_idx = n_dim - 2; dim_idx >= 0; dim_idx--) { - input_dim_factor[narrow(dim_idx)] = input_dim_factor[SafeInt(dim_idx) + 1] * input_shape[SafeInt(dim_idx) + 1]; + TensorShapeVector input_dim_factor(n_dim); + input_dim_factor[n_dim - 1] = 1; // initialize dimension factor + for (size_t dim_idx = n_dim - 1; dim_idx > 0; dim_idx--) { + input_dim_factor[dim_idx - 1] = input_dim_factor[dim_idx] * + input_shape[dim_idx]; } - int64_t output_idx = 0; - int64_t input_idx = 0; - if (n_dim == 1) { - std::vector input_mapping = UpsampleNearestSetupRank1InputMapping(input_shape[0], - output_shape[0], - scales[0], - roi[0], roi[narrow(n_dim + 0)], - extrapolation_enabled, - get_original_coordinate, - get_nearest_pixel); + auto input_mapping = UpsampleNearestSetupRank1InputMapping(input_shape[0], + output_shape[0], + scales[0], + roi[0], roi[n_dim], + extrapolation_enabled, + get_original_coordinate, + get_nearest_pixel); - for (int64_t output_dim0_idx = 0; output_dim0_idx < output_shape[0]; output_dim0_idx++) { - int64_t input_dim0_idx = input_mapping[narrow(output_dim0_idx)]; - output[narrow(output_dim0_idx)] = input_dim0_idx < 0 ? extrapolation_value : input[input_dim0_idx]; + for (size_t output_dim0_idx = 0, lim = narrow(output_shape[0]); output_dim0_idx < lim; output_dim0_idx++) { + SafeInt input_dim0_idx = input_mapping[output_dim0_idx]; + output[output_dim0_idx] = input_dim0_idx < 0 ? extrapolation_value : input[static_cast(input_dim0_idx)]; } return Status::OK(); } - std::vector> input_mappings = + auto input_mappings = UpsampleNearestSetupInputMappings(n_dim, input_shape, output_shape, input_dim_factor, scales, roi, extrapolation_enabled, get_original_coordinate, get_nearest_pixel); + size_t output_idx = 0; if (n_dim == 2) { - const std::vector& input_mapping_0 = input_mappings[0]; - const std::vector& input_mapping_1 = input_mappings[1]; + const auto& input_mapping_0 = input_mappings[0]; + const auto& input_mapping_1 = input_mappings[1]; - for (int64_t output_dim0_inx = 0; output_dim0_inx < output_shape[0]; output_dim0_inx++) { - int64_t input_idx_0 = input_mapping_0[narrow(output_dim0_inx)]; - for (int64_t output_dim1_inx = 0; output_dim1_inx < output_shape[1]; output_dim1_inx++) { - int64_t input_idx_1 = input_idx_0 + input_mapping_1[narrow(output_dim1_inx)]; - output[output_idx++] = (input_idx_1 < 0) ? extrapolation_value : input[input_idx_1]; + const size_t lim = narrow(output_shape[0]); + const size_t limin = narrow(output_shape[1]); + for (size_t output_dim0_inx = 0; output_dim0_inx < lim; output_dim0_inx++) { + SafeInt input_idx_0 = input_mapping_0[output_dim0_inx]; + for (size_t output_dim1_inx = 0; output_dim1_inx < limin; output_dim1_inx++) { + auto input_idx_1 = input_idx_0 + input_mapping_1[output_dim1_inx]; + output[output_idx++] = (input_idx_1 < 0) ? extrapolation_value : input[static_cast(input_idx_1)]; } } return Status::OK(); } if (n_dim == 3) { - const std::vector& input_mapping_0 = input_mappings[0]; - const std::vector& input_mapping_1 = input_mappings[1]; - const std::vector& input_mapping_2 = input_mappings[2]; + const auto& input_mapping_0 = input_mappings[0]; + const auto& input_mapping_1 = input_mappings[1]; + const auto& input_mapping_2 = input_mappings[2]; - for (int64_t output_dim0_inx = 0; output_dim0_inx < output_shape[0]; output_dim0_inx++) { - int64_t input_idx_0 = input_mapping_0[narrow(output_dim0_inx)]; - for (int64_t output_dim1_inx = 0; output_dim1_inx < output_shape[1]; output_dim1_inx++) { - int64_t input_idx_1 = input_idx_0 + input_mapping_1[narrow(output_dim1_inx)]; - for (int64_t output_dim2_inx = 0; output_dim2_inx < output_shape[2]; output_dim2_inx++) { - int64_t input_idx_2 = input_idx_1 + input_mapping_2[narrow(output_dim2_inx)]; - output[output_idx++] = (input_idx_2 < 0) ? extrapolation_value : input[input_idx_2]; + const size_t output0_lim = narrow(output_shape[0]); + const size_t output1_lim = narrow(output_shape[1]); + const size_t output2_lim = narrow(output_shape[2]); + for (size_t output_dim0_inx = 0; output_dim0_inx < output0_lim; output_dim0_inx++) { + SafeInt input_idx_0 = input_mapping_0[output_dim0_inx]; + for (size_t output_dim1_inx = 0; output_dim1_inx < output1_lim; output_dim1_inx++) { + auto input_idx_1 = input_idx_0 + input_mapping_1[output_dim1_inx]; + for (size_t output_dim2_inx = 0; output_dim2_inx < output2_lim; output_dim2_inx++) { + auto input_idx_2 = input_idx_1 + input_mapping_2[output_dim2_inx]; + output[output_idx++] = (input_idx_2 < 0) ? extrapolation_value : input[static_cast(input_idx_2)]; } } } @@ -258,20 +271,27 @@ static Status UpsampleNearestImpl(const T* input, } if (n_dim == 4) { - const std::vector& input_mapping_0 = input_mappings[0]; - const std::vector& input_mapping_1 = input_mappings[1]; - const std::vector& input_mapping_2 = input_mappings[2]; - const std::vector& input_mapping_3 = input_mappings[3]; + const auto& input_mapping_0 = input_mappings[0]; + const auto& input_mapping_1 = input_mappings[1]; + const auto& input_mapping_2 = input_mappings[2]; + const auto& input_mapping_3 = input_mappings[3]; - for (int64_t output_dim0_inx = 0; output_dim0_inx < output_shape[0]; output_dim0_inx++) { - int64_t input_idx_0 = input_mapping_0[narrow(output_dim0_inx)]; - for (int64_t output_dim1_inx = 0; output_dim1_inx < output_shape[1]; output_dim1_inx++) { - int64_t input_idx_1 = input_idx_0 + input_mapping_1[narrow(output_dim1_inx)]; - for (int64_t output_dim2_inx = 0; output_dim2_inx < output_shape[2]; output_dim2_inx++) { - int64_t input_idx_2 = input_idx_1 + input_mapping_2[narrow(output_dim2_inx)]; - for (int64_t output_dim3_inx = 0; output_dim3_inx < output_shape[3]; output_dim3_inx++) { - int64_t input_idx_3 = input_idx_2 + input_mapping_3[narrow(output_dim3_inx)]; - output[output_idx++] = (input_idx_3 < 0) ? static_cast(extrapolation_value) : input[narrow(input_idx_3)]; + const size_t output0_lim = narrow(output_shape[0]); + const size_t output1_lim = narrow(output_shape[1]); + const size_t output2_lim = narrow(output_shape[2]); + const size_t output3_lim = narrow(output_shape[3]); + + for (size_t output_dim0_inx = 0; output_dim0_inx < output0_lim; output_dim0_inx++) { + SafeInt input_idx_0 = input_mapping_0[output_dim0_inx]; + for (size_t output_dim1_inx = 0; output_dim1_inx < output1_lim; output_dim1_inx++) { + auto input_idx_1 = input_idx_0 + input_mapping_1[output_dim1_inx]; + for (size_t output_dim2_inx = 0; output_dim2_inx < output2_lim; output_dim2_inx++) { + auto input_idx_2 = input_idx_1 + input_mapping_2[output_dim2_inx]; + for (size_t output_dim3_inx = 0; output_dim3_inx < output3_lim; output_dim3_inx++) { + auto input_idx_3 = input_idx_2 + input_mapping_3[output_dim3_inx]; + output[output_idx++] = (input_idx_3 < 0) + ? static_cast(extrapolation_value) + : input[static_cast(input_idx_3)]; } } } @@ -279,22 +299,24 @@ static Status UpsampleNearestImpl(const T* input, return Status::OK(); } - std::vector output_dim_counter(onnxruntime::narrow(n_dim)); - for (int64_t dim_idx = 0; dim_idx < n_dim; dim_idx++) { - input_idx += input_mappings[narrow(dim_idx)][0 /* output_dim_counter[narrow(dim_idx)] */]; + SafeInt input_idx = 0; + TensorShapeVector output_dim_counter(n_dim); + for (size_t dim_idx = 0; dim_idx < n_dim; dim_idx++) { + input_idx += input_mappings[dim_idx][0 /* output_dim_counter[narrow(dim_idx)] */]; } - for (int64_t output_size = output_shape.Size(); output_idx < output_size; output_idx++) { - output[narrow(output_idx)] = (input_idx < 0) ? extrapolation_value : input[narrow(input_idx)]; + for (size_t output_size = narrow(output_shape.Size()); output_idx < output_size; output_idx++) { + output[output_idx] = (input_idx < 0) ? extrapolation_value : input[static_cast(input_idx)]; - for (int64_t dim_idx = n_dim - 1; dim_idx >= 0; dim_idx--) { - input_idx -= input_mappings[narrow(dim_idx)][narrow(output_dim_counter[narrow(dim_idx)])]; - if (++output_dim_counter[narrow(dim_idx)] < output_shape[narrow(dim_idx)]) { - input_idx += input_mappings[narrow(dim_idx)][narrow(output_dim_counter[narrow(dim_idx)])]; + for (size_t dim_idx = n_dim; dim_idx > 0; dim_idx--) { + const auto idx = dim_idx - 1; + input_idx -= input_mappings[idx][narrow(output_dim_counter[idx])]; + if (++output_dim_counter[idx] < output_shape[idx]) { + input_idx += input_mappings[idx][narrow(output_dim_counter[idx])]; break; } - output_dim_counter[narrow(dim_idx)] = 0; - input_idx += input_mappings[narrow(dim_idx)][0 /* output_dim_counter[narrow(dim_idx)] */]; + output_dim_counter[idx] = 0; + input_idx += input_mappings[idx][0 /* output_dim_counter[dim_idx] */]; } }