From 9189ebb41591fccfd8c4c90af2db36f9b53b4ffe Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Tue, 18 Oct 2022 14:41:46 -0700 Subject: [PATCH] Optimize slicing when possible by copying bigger blocks at once (#13261) ### Description Currently, SliceIterator copies inner dimension size at once at best. However, there are many slices when several inner dimensions can be copied at once. Furthermore, even if a dimension is sliced, it may employ step 1 and, therefore, has a continuous block of inner dimensions that can be copied at once. ### Motivation and Context For example, `[N, C, H, W]` with slice `[:, :, i:, :]` and `[N, C, H-i, W]`. Meaning, we slice along single axis, with step = 1. Current implementation does `C * (H-i) memcpy` with W elements each. With this change we can do `C memcpy with (H-i)*W` elements each. The optimization produces ~11% savings on certain internal models. --- .../core/providers/cpu/cpu_provider_shared.cc | 14 +-- .../core/providers/cpu/cpu_provider_shared.h | 14 +-- .../core/providers/cpu/tensor/slice.cc | 30 +++--- onnxruntime/core/providers/cpu/tensor/slice.h | 14 +-- onnxruntime/core/providers/cpu/tensor/utils.h | 99 ++++++++++++++++--- .../provider_bridge_provider.cc | 14 +-- .../providers/cpu/tensor/slice_op.test.cc | 31 ++++++ 7 files changed, 155 insertions(+), 61 deletions(-) diff --git a/onnxruntime/core/providers/cpu/cpu_provider_shared.cc b/onnxruntime/core/providers/cpu/cpu_provider_shared.cc index 8a17058d0d..bd9c6b1793 100644 --- a/onnxruntime/core/providers/cpu/cpu_provider_shared.cc +++ b/onnxruntime/core/providers/cpu/cpu_provider_shared.cc @@ -102,15 +102,15 @@ struct ProviderHostCPUImpl : ProviderHostCPU { Status PrepareOutputShape(const Tensor* indices, const int64_t depth_val, const int64_t axis, int64_t& prefix_dim_size, int64_t& suffix_dim_size, TensorShapeVector& output_shape) override { return onnxruntime::PrepareOutputShape(indices, depth_val, axis, prefix_dim_size, suffix_dim_size, output_shape); } // From cpu/tensor/slice.h (direct) - Status SliceBase__PrepareForCompute(const gsl::span& raw_starts, - const gsl::span& raw_ends, - const gsl::span& raw_axes, + Status SliceBase__PrepareForCompute(gsl::span raw_starts, + gsl::span raw_ends, + gsl::span raw_axes, SliceOp__PrepareForComputeMetadata& compute_metadata) override { return SliceBase::PrepareForCompute(raw_starts, raw_ends, raw_axes, reinterpret_cast(compute_metadata)); } - Status SliceBase__PrepareForCompute(const gsl::span& raw_starts, - const gsl::span& raw_ends, - const gsl::span& raw_axes, - const gsl::span& raw_steps, + Status SliceBase__PrepareForCompute(gsl::span raw_starts, + gsl::span raw_ends, + gsl::span raw_axes, + gsl::span raw_steps, SliceOp__PrepareForComputeMetadata& compute_metadata) override { return SliceBase::PrepareForCompute(raw_starts, raw_ends, raw_axes, raw_steps, reinterpret_cast(compute_metadata)); } Status SliceBase__FillVectorsFromInput(const Tensor& start_tensor, diff --git a/onnxruntime/core/providers/cpu/cpu_provider_shared.h b/onnxruntime/core/providers/cpu/cpu_provider_shared.h index 48860490cd..cad55e4829 100644 --- a/onnxruntime/core/providers/cpu/cpu_provider_shared.h +++ b/onnxruntime/core/providers/cpu/cpu_provider_shared.h @@ -61,15 +61,15 @@ struct ProviderHostCPU { virtual Status PrepareOutputShape(const Tensor* indices, const int64_t depth_val, const int64_t axis, int64_t& prefix_dim_size, int64_t& suffix_dim_size, TensorShapeVector& output_shape) = 0; // From cpu/tensor/slice.h - virtual Status SliceBase__PrepareForCompute(const gsl::span& raw_starts, - const gsl::span& raw_ends, - const gsl::span& raw_axes, + virtual Status SliceBase__PrepareForCompute(gsl::span raw_starts, + gsl::span raw_ends, + gsl::span raw_axes, SliceOp__PrepareForComputeMetadata& compute_metadata) = 0; - virtual Status SliceBase__PrepareForCompute(const gsl::span& raw_starts, - const gsl::span& raw_ends, - const gsl::span& raw_axes, - const gsl::span& raw_steps, + virtual Status SliceBase__PrepareForCompute(gsl::span raw_starts, + gsl::span raw_ends, + gsl::span raw_axes, + gsl::span raw_steps, SliceOp__PrepareForComputeMetadata& compute_metadata) = 0; virtual Status SliceBase__FillVectorsFromInput(const Tensor& start_tensor, const Tensor& ends_tensor, diff --git a/onnxruntime/core/providers/cpu/tensor/slice.cc b/onnxruntime/core/providers/cpu/tensor/slice.cc index 7e4bb3c3a3..6b404e1f48 100644 --- a/onnxruntime/core/providers/cpu/tensor/slice.cc +++ b/onnxruntime/core/providers/cpu/tensor/slice.cc @@ -74,8 +74,8 @@ ONNX_CPU_OPERATOR_KERNEL( // Updates starts and steps to match flattened_output_dims if it is. // e.g. if input shape is { 2, 2, 2 }, output shape is { 1, 2, 2 }, and the 'steps' value for the last two dims is 1, // we are keeping all the data of the inner most two dimensions so can combine those into dims of { 1, 4 } -static void FlattenOutputDims(const gsl::span& input_dimensions, - const gsl::span& output_dims, +static void FlattenOutputDims(gsl::span input_dimensions, + gsl::span output_dims, TensorShapeVector& starts, TensorShapeVector& ends, TensorShapeVector& steps, @@ -115,9 +115,9 @@ static void FlattenOutputDims(const gsl::span& input_dimensions, } // Slice V1-9 & DynamicSlice -Status SliceBase::PrepareForCompute(const gsl::span& raw_starts, - const gsl::span& raw_ends, - const gsl::span& raw_axes, +Status SliceBase::PrepareForCompute(gsl::span raw_starts, + gsl::span raw_ends, + gsl::span raw_axes, SliceOp::PrepareForComputeMetadata& compute_metadata) { ORT_RETURN_IF_ERROR(SliceOp::PrepareForComputeHelper(raw_starts, raw_ends, raw_axes, compute_metadata)); FlattenOutputDims(compute_metadata.input_dimensions_, compute_metadata.output_dims_, compute_metadata.starts_, @@ -126,10 +126,10 @@ Status SliceBase::PrepareForCompute(const gsl::span& raw_starts, } // DynamicSlice & Slice V10 -Status SliceBase::PrepareForCompute(const gsl::span& raw_starts, - const gsl::span& raw_ends, - const gsl::span& raw_axes, - const gsl::span& raw_steps, +Status SliceBase::PrepareForCompute(gsl::span raw_starts, + gsl::span raw_ends, + gsl::span raw_axes, + gsl::span raw_steps, SliceOp::PrepareForComputeMetadata& compute_metadata) { ORT_RETURN_IF_ERROR(SliceOp::PrepareForComputeHelper(raw_starts, raw_ends, raw_axes, raw_steps, compute_metadata)); FlattenOutputDims(compute_metadata.input_dimensions_, compute_metadata.output_dims_, compute_metadata.starts_, @@ -223,17 +223,11 @@ static Status SliceImpl(OpKernelContext* ctx, const auto* output_end = output + output_tensor.Shape().Size(); auto create_output = [&output, &output_end](SliceIterator& slice_input_iterator) { - if (slice_input_iterator.SolitaryInnerStep()) { - while (output < output_end) { - output = slice_input_iterator.CopyInnermostAxisSolitaryInnerStep(output); - } - } else { - while (output < output_end) { - output = slice_input_iterator.CopyInnermostAxisNonSolitaryInnerStep(output); - } + while (output < output_end) { + output = slice_input_iterator.CopyContiguousInnermostAxes(output); } - ORT_ENFORCE(output == output_end); + ORT_ENFORCE(output == output_end); }; if (compute_metadata.p_flattened_output_dims_) { diff --git a/onnxruntime/core/providers/cpu/tensor/slice.h b/onnxruntime/core/providers/cpu/tensor/slice.h index 0ff724aae1..28e76aca4e 100644 --- a/onnxruntime/core/providers/cpu/tensor/slice.h +++ b/onnxruntime/core/providers/cpu/tensor/slice.h @@ -16,16 +16,16 @@ class SliceBase { // static methods that can be used from other ops if needed public: // compute output_dims without steps (Slice V1-9 & DynamicSlice) - static Status PrepareForCompute(const gsl::span& raw_starts, - const gsl::span& raw_ends, - const gsl::span& raw_axes, + static Status PrepareForCompute(gsl::span raw_starts, + gsl::span raw_ends, + gsl::span raw_axes, SliceOp::PrepareForComputeMetadata& compute_metadata); // compute output_dims with steps (Slice V10) - static Status PrepareForCompute(const gsl::span& raw_starts, - const gsl::span& raw_ends, - const gsl::span& raw_axes, - const gsl::span& raw_steps, + static Status PrepareForCompute(gsl::span raw_starts, + gsl::span raw_ends, + gsl::span raw_axes, + gsl::span raw_steps, SliceOp::PrepareForComputeMetadata& compute_metadata); // Slice V10 & DynamicSlice diff --git a/onnxruntime/core/providers/cpu/tensor/utils.h b/onnxruntime/core/providers/cpu/tensor/utils.h index c519278a59..fd81045405 100644 --- a/onnxruntime/core/providers/cpu/tensor/utils.h +++ b/onnxruntime/core/providers/cpu/tensor/utils.h @@ -129,7 +129,7 @@ struct ExtentAxisCounters { private: bool running_{true}; size_t axis_; - TensorShapeVector indices_; // There is no index for innermost axis since it's a special case + TensorShapeVector indices_; // There is no index for innermost axis since it's a special case gsl::span extents_; // The extents of each axis }; @@ -177,22 +177,52 @@ struct SliceIteratorBase { // Initialize initial skip and inner_extent. void Init(gsl::span dims, gsl::span starts, gsl::span steps) { - ORT_ENFORCE(dims.size() == starts.size() && - dims.size() == extents_.size() && - dims.size() >= steps.size()); + const auto dims_size = dims.size(); + ORT_ENFORCE(dims_size == starts.size() && + dims_size == extents_.size() && + dims_size >= steps.size()); SafeInt pitch = 1; // Initial skip, so that input_ points to the first element to copy - for (size_t i = dims.size(); i-- > 0;) { + for (size_t i = dims_size; i-- > 0;) { input_ += pitch * starts[i] * element_size_; pitch *= static_cast(dims[i]); } - inner_extent_ = static_cast(extents_[dims.size() - 1]); - //It could be -1 - inner_step_ = static_cast(dims.size() == steps.size() - ? steps[dims.size() - 1] + inner_extent_ = static_cast(extents_[dims_size - 1]); + const auto steps_size = steps.size(); + // It could be -1 + inner_step_ = static_cast(dims_size == steps_size + ? steps[steps_size - 1] : 1); + + // This block serves to maximize the amount of data that is copied + // in one shot when trailing axis are not being sliced and, therefore, + // more data can be copied in one shot. The inner dimension with step = 1 + // can always copy with more than one element unless its extent is 1. + // However, when more than one trailing dimensions have steps one and copied + // in its entirety, we can copy bigger blocks at the same time and start skipping + // them as the skips are expected to be zero then for them. + SafeInt max_copyable_elements = (inner_step_ == 1) ? inner_extent_ : 1; + last_batching_axis_ = dims_size - 1; + + // Check if inner dimension is copied as a block in its entirety + if (dims_size > 1 && inner_step_ == 1 && inner_extent_ == gsl::narrow(dims[dims_size - 1])) { + for (size_t dim = dims_size - 2;; dim--) { + if (dim < steps_size && steps[dim] != 1) { + break; + } + max_copyable_elements *= extents_[dim]; + last_batching_axis_ = dim; + if (extents_[dim] != dims[dim]) { + break; + } + if (dim == 0) { + break; + } + } + } + max_copying_elements_block_ = max_copyable_elements; } protected: @@ -226,6 +256,10 @@ struct SliceIteratorBase { void AdvanceOverInnerExtent() { size_t axis = skips_.size() - 1; + AdvanceOverExtent(axis); + } + + void AdvanceOverExtent(size_t axis) { input_ += skips_[axis] * element_size_; while (axis-- && ++indices_[axis] == extents_[axis]) { indices_[axis] = 0; @@ -247,11 +281,12 @@ struct SliceIteratorBase { // Assumes SolitaryInnerStep() == true void* CopyInnermostAxisSolitaryInnerStep(void* output) { + byte* out_bytes = reinterpret_cast(output); auto bytes_to_copy = inner_extent_ * element_size_; if (!is_string_tensor_) { - std::copy(input_, input_ + bytes_to_copy, out_bytes); + memcpy(output, input_, bytes_to_copy); } else { const std::string* input = reinterpret_cast(input_); std::string* out = reinterpret_cast(output); @@ -265,6 +300,27 @@ struct SliceIteratorBase { return out_bytes; } + void* CopyContiguousInnermostAxes(void* output) { + byte* out_bytes = nullptr; + const auto bytes_to_copy = max_copying_elements_block_ * element_size_; + if (SolitaryInnerStep()) { + if (!is_string_tensor_) { + memcpy(output, input_, bytes_to_copy); + } else { + const std::string* input = reinterpret_cast(input_); + std::string* out = reinterpret_cast(output); + std::copy(input, input + max_copying_elements_block_, out); + } + input_ += bytes_to_copy; + out_bytes = reinterpret_cast(output); + out_bytes += bytes_to_copy; + AdvanceOverExtent(last_batching_axis_); + } else { + out_bytes = reinterpret_cast(CopyInnermostAxisNonSolitaryInnerStep(output)); + } + return out_bytes; + } + // Assumes generic inner_step_ void* CopyInnermostAxisNonSolitaryInnerStep(void* output) { // need to special case std::string so the copy works correctly @@ -322,6 +378,12 @@ struct SliceIteratorBase { gsl::span extents_; size_t inner_counter_{}, inner_extent_; ptrdiff_t inner_step_; + // Max copyable continuous block in elements + int64_t max_copying_elements_block_; + // The last slicing axis with step 1 that can combines continuous copyable blocks. + // Used for skipping the max copyable block. If it is equal to the most inner dimension + // then no extents can be combined for copying + size_t last_batching_axis_; SliceSkips skips_; TensorShapeVector indices_; // There is no index for innermost axis since it's a special case }; @@ -371,6 +433,11 @@ struct SliceIterator : public SliceIteratorBase { void* new_output = SliceIteratorBase::CopyInnermostAxisNonSolitaryInnerStep(output); return static_cast(new_output); } + + T* CopyContiguousInnermostAxes(void* output) { + void* new_output = SliceIteratorBase::CopyContiguousInnermostAxes(output); + return static_cast(new_output); + } }; inline void CopyCpuTensor(const Tensor* src, Tensor* tgt) { @@ -378,12 +445,14 @@ inline void CopyCpuTensor(const Tensor* src, Tensor* tgt) { const void* source = src->DataRaw(); if (target != source) { - auto is_string_type = utils::IsDataTypeString(src->DataType()); - if (is_string_type) { - for (int64_t i = 0; i < src->Shape().Size(); ++i) - static_cast(target)[i] = static_cast(source)[i]; + if (src->IsDataTypeString()) { + auto src_span = src->DataAsSpan(); + auto* dst_string = tgt->MutableData(); + std::copy(src_span.begin(), src_span.end(), dst_string); } else { - memcpy(target, source, static_cast(src->Shape().Size() * src->DataType()->Size())); + const auto element_size = src->DataType()->Size(); + const auto elements = src->Shape().Size(); + memcpy(target, source, elements * element_size); } } } diff --git a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc index 0824a82c2b..f92967d4dc 100644 --- a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc +++ b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc @@ -485,15 +485,15 @@ bool TileOp::IsTileMemcpy(const TensorShape& input_shape, const int64_t* repeats return g_host_cpu.TileOp__IsTileMemcpy(input_shape, repeats, rank, is_batched_memcpy, num_of_elements_per_batch, num_of_copies_per_batch, num_of_batch_copies); } -Status SliceBase::PrepareForCompute(const gsl::span& raw_starts, - const gsl::span& raw_ends, - const gsl::span& raw_axes, +Status SliceBase::PrepareForCompute(gsl::span raw_starts, + gsl::span raw_ends, + gsl::span raw_axes, SliceOp::PrepareForComputeMetadata& compute_metadata) { return g_host_cpu.SliceBase__PrepareForCompute(raw_starts, raw_ends, raw_axes, reinterpret_cast(compute_metadata)); } -Status SliceBase::PrepareForCompute(const gsl::span& raw_starts, - const gsl::span& raw_ends, - const gsl::span& raw_axes, - const gsl::span& raw_steps, +Status SliceBase::PrepareForCompute(gsl::span raw_starts, + gsl::span raw_ends, + gsl::span raw_axes, + gsl::span raw_steps, SliceOp::PrepareForComputeMetadata& compute_metadata) { return g_host_cpu.SliceBase__PrepareForCompute(raw_starts, raw_ends, raw_axes, raw_steps, reinterpret_cast(compute_metadata)); } Status SliceBase::FillVectorsFromInput(const Tensor& start_tensor, diff --git a/onnxruntime/test/providers/cpu/tensor/slice_op.test.cc b/onnxruntime/test/providers/cpu/tensor/slice_op.test.cc index faf8657c82..c70bef43cb 100644 --- a/onnxruntime/test/providers/cpu/tensor/slice_op.test.cc +++ b/onnxruntime/test/providers/cpu/tensor/slice_op.test.cc @@ -650,6 +650,37 @@ TEST(SliceTest, Slice5D_LargeStep) { {}); } +TEST(SliceTest, Slice5D_CopyAxis2LargeBlock) { + // We are trying to accomplish two things + // 1) we still need to copy multiple slices because of the axis 1 (slices 0,1) + // 2) we are combining axis 2 slices(1,2) in one copy because its step is 1 and + // dims below it are copied as a whole + RunSliceTest({1, 3, 4, 2, 2}, + {1.f, 2.f, 3.f, 4.f, + 5.f, 6.f, 7.f, 8.f, + -1.f, -2.f, -3.f, -4.f, + -5.f, -6.f, -7.f, -8.f, + 1.f, 2.f, 3.f, 4.f, + 5.f, 6.f, 7.f, 8.f, + -1.f, -2.f, -3.f, -4.f, + -5.f, -6.f, -7.f, -8.f, + 1.f, 2.f, 3.f, 4.f, + 5.f, 6.f, 7.f, 8.f, + -1.f, -2.f, -3.f, -4.f, + -5.f, -6.f, -7.f, -8.f}, + {0, 1}, // starts + {2, 3}, // ends + {1, 2}, // axis + {}, // steps defaults to 1 + {1, 2, 2, 2, 2}, + {5.f, 6.f, 7.f, 8.f, + -1.f, -2.f, -3.f, -4.f, + 5.f, 6.f, 7.f, 8.f, + -1.f, -2.f, -3.f, -4.f}, + true, + {}); +} + TEST(SliceTest, EmptyDim) { RunSliceTest({0, 6}, // empty dim in shape {},