mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
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.
This commit is contained in:
Dmitri Smirnov
GitHub
parent
f5e3165cc3
commit
9189ebb415
7 changed files with 155 additions and 61 deletions
|
|
@ -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<const int64_t>& raw_starts,
|
||||
const gsl::span<const int64_t>& raw_ends,
|
||||
const gsl::span<const int64_t>& raw_axes,
|
||||
Status SliceBase__PrepareForCompute(gsl::span<const int64_t> raw_starts,
|
||||
gsl::span<const int64_t> raw_ends,
|
||||
gsl::span<const int64_t> raw_axes,
|
||||
SliceOp__PrepareForComputeMetadata& compute_metadata) override { return SliceBase::PrepareForCompute(raw_starts, raw_ends, raw_axes, reinterpret_cast<SliceOp::PrepareForComputeMetadata&>(compute_metadata)); }
|
||||
|
||||
Status SliceBase__PrepareForCompute(const gsl::span<const int64_t>& raw_starts,
|
||||
const gsl::span<const int64_t>& raw_ends,
|
||||
const gsl::span<const int64_t>& raw_axes,
|
||||
const gsl::span<const int64_t>& raw_steps,
|
||||
Status SliceBase__PrepareForCompute(gsl::span<const int64_t> raw_starts,
|
||||
gsl::span<const int64_t> raw_ends,
|
||||
gsl::span<const int64_t> raw_axes,
|
||||
gsl::span<const int64_t> raw_steps,
|
||||
SliceOp__PrepareForComputeMetadata& compute_metadata) override { return SliceBase::PrepareForCompute(raw_starts, raw_ends, raw_axes, raw_steps, reinterpret_cast<SliceOp::PrepareForComputeMetadata&>(compute_metadata)); }
|
||||
|
||||
Status SliceBase__FillVectorsFromInput(const Tensor& start_tensor,
|
||||
|
|
|
|||
|
|
@ -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<const int64_t>& raw_starts,
|
||||
const gsl::span<const int64_t>& raw_ends,
|
||||
const gsl::span<const int64_t>& raw_axes,
|
||||
virtual Status SliceBase__PrepareForCompute(gsl::span<const int64_t> raw_starts,
|
||||
gsl::span<const int64_t> raw_ends,
|
||||
gsl::span<const int64_t> raw_axes,
|
||||
SliceOp__PrepareForComputeMetadata& compute_metadata) = 0;
|
||||
|
||||
virtual Status SliceBase__PrepareForCompute(const gsl::span<const int64_t>& raw_starts,
|
||||
const gsl::span<const int64_t>& raw_ends,
|
||||
const gsl::span<const int64_t>& raw_axes,
|
||||
const gsl::span<const int64_t>& raw_steps,
|
||||
virtual Status SliceBase__PrepareForCompute(gsl::span<const int64_t> raw_starts,
|
||||
gsl::span<const int64_t> raw_ends,
|
||||
gsl::span<const int64_t> raw_axes,
|
||||
gsl::span<const int64_t> raw_steps,
|
||||
SliceOp__PrepareForComputeMetadata& compute_metadata) = 0;
|
||||
virtual Status SliceBase__FillVectorsFromInput(const Tensor& start_tensor,
|
||||
const Tensor& ends_tensor,
|
||||
|
|
|
|||
|
|
@ -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<const int64_t>& input_dimensions,
|
||||
const gsl::span<const int64_t>& output_dims,
|
||||
static void FlattenOutputDims(gsl::span<const int64_t> input_dimensions,
|
||||
gsl::span<const int64_t> output_dims,
|
||||
TensorShapeVector& starts,
|
||||
TensorShapeVector& ends,
|
||||
TensorShapeVector& steps,
|
||||
|
|
@ -115,9 +115,9 @@ static void FlattenOutputDims(const gsl::span<const int64_t>& input_dimensions,
|
|||
}
|
||||
|
||||
// Slice V1-9 & DynamicSlice
|
||||
Status SliceBase::PrepareForCompute(const gsl::span<const int64_t>& raw_starts,
|
||||
const gsl::span<const int64_t>& raw_ends,
|
||||
const gsl::span<const int64_t>& raw_axes,
|
||||
Status SliceBase::PrepareForCompute(gsl::span<const int64_t> raw_starts,
|
||||
gsl::span<const int64_t> raw_ends,
|
||||
gsl::span<const int64_t> 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<const int64_t>& raw_starts,
|
|||
}
|
||||
|
||||
// DynamicSlice & Slice V10
|
||||
Status SliceBase::PrepareForCompute(const gsl::span<const int64_t>& raw_starts,
|
||||
const gsl::span<const int64_t>& raw_ends,
|
||||
const gsl::span<const int64_t>& raw_axes,
|
||||
const gsl::span<const int64_t>& raw_steps,
|
||||
Status SliceBase::PrepareForCompute(gsl::span<const int64_t> raw_starts,
|
||||
gsl::span<const int64_t> raw_ends,
|
||||
gsl::span<const int64_t> raw_axes,
|
||||
gsl::span<const int64_t> 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<T>& 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_) {
|
||||
|
|
|
|||
|
|
@ -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<const int64_t>& raw_starts,
|
||||
const gsl::span<const int64_t>& raw_ends,
|
||||
const gsl::span<const int64_t>& raw_axes,
|
||||
static Status PrepareForCompute(gsl::span<const int64_t> raw_starts,
|
||||
gsl::span<const int64_t> raw_ends,
|
||||
gsl::span<const int64_t> raw_axes,
|
||||
SliceOp::PrepareForComputeMetadata& compute_metadata);
|
||||
|
||||
// compute output_dims with steps (Slice V10)
|
||||
static Status PrepareForCompute(const gsl::span<const int64_t>& raw_starts,
|
||||
const gsl::span<const int64_t>& raw_ends,
|
||||
const gsl::span<const int64_t>& raw_axes,
|
||||
const gsl::span<const int64_t>& raw_steps,
|
||||
static Status PrepareForCompute(gsl::span<const int64_t> raw_starts,
|
||||
gsl::span<const int64_t> raw_ends,
|
||||
gsl::span<const int64_t> raw_axes,
|
||||
gsl::span<const int64_t> raw_steps,
|
||||
SliceOp::PrepareForComputeMetadata& compute_metadata);
|
||||
|
||||
// Slice V10 & DynamicSlice
|
||||
|
|
|
|||
|
|
@ -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<const int64_t> extents_; // The extents of each axis
|
||||
};
|
||||
|
||||
|
|
@ -177,22 +177,52 @@ struct SliceIteratorBase {
|
|||
|
||||
// Initialize initial skip and inner_extent.
|
||||
void Init(gsl::span<const int64_t> dims, gsl::span<const int64_t> starts, gsl::span<const int64_t> 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<size_t> 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<size_t>(dims[i]);
|
||||
}
|
||||
|
||||
inner_extent_ = static_cast<size_t>(extents_[dims.size() - 1]);
|
||||
//It could be -1
|
||||
inner_step_ = static_cast<ptrdiff_t>(dims.size() == steps.size()
|
||||
? steps[dims.size() - 1]
|
||||
inner_extent_ = static_cast<size_t>(extents_[dims_size - 1]);
|
||||
const auto steps_size = steps.size();
|
||||
// It could be -1
|
||||
inner_step_ = static_cast<ptrdiff_t>(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<int64_t> 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<size_t>(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<byte*>(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<const std::string*>(input_);
|
||||
std::string* out = reinterpret_cast<std::string*>(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<const std::string*>(input_);
|
||||
std::string* out = reinterpret_cast<std::string*>(output);
|
||||
std::copy(input, input + max_copying_elements_block_, out);
|
||||
}
|
||||
input_ += bytes_to_copy;
|
||||
out_bytes = reinterpret_cast<byte*>(output);
|
||||
out_bytes += bytes_to_copy;
|
||||
AdvanceOverExtent(last_batching_axis_);
|
||||
} else {
|
||||
out_bytes = reinterpret_cast<byte*>(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<const int64_t> 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<T*>(new_output);
|
||||
}
|
||||
|
||||
T* CopyContiguousInnermostAxes(void* output) {
|
||||
void* new_output = SliceIteratorBase::CopyContiguousInnermostAxes(output);
|
||||
return static_cast<T*>(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<std::string*>(target)[i] = static_cast<const std::string*>(source)[i];
|
||||
if (src->IsDataTypeString()) {
|
||||
auto src_span = src->DataAsSpan<std::string>();
|
||||
auto* dst_string = tgt->MutableData<std::string>();
|
||||
std::copy(src_span.begin(), src_span.end(), dst_string);
|
||||
} else {
|
||||
memcpy(target, source, static_cast<size_t>(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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<const int64_t>& raw_starts,
|
||||
const gsl::span<const int64_t>& raw_ends,
|
||||
const gsl::span<const int64_t>& raw_axes,
|
||||
Status SliceBase::PrepareForCompute(gsl::span<const int64_t> raw_starts,
|
||||
gsl::span<const int64_t> raw_ends,
|
||||
gsl::span<const int64_t> raw_axes,
|
||||
SliceOp::PrepareForComputeMetadata& compute_metadata) { return g_host_cpu.SliceBase__PrepareForCompute(raw_starts, raw_ends, raw_axes, reinterpret_cast<SliceOp__PrepareForComputeMetadata&>(compute_metadata)); }
|
||||
|
||||
Status SliceBase::PrepareForCompute(const gsl::span<const int64_t>& raw_starts,
|
||||
const gsl::span<const int64_t>& raw_ends,
|
||||
const gsl::span<const int64_t>& raw_axes,
|
||||
const gsl::span<const int64_t>& raw_steps,
|
||||
Status SliceBase::PrepareForCompute(gsl::span<const int64_t> raw_starts,
|
||||
gsl::span<const int64_t> raw_ends,
|
||||
gsl::span<const int64_t> raw_axes,
|
||||
gsl::span<const int64_t> raw_steps,
|
||||
SliceOp::PrepareForComputeMetadata& compute_metadata) { return g_host_cpu.SliceBase__PrepareForCompute(raw_starts, raw_ends, raw_axes, raw_steps, reinterpret_cast<SliceOp__PrepareForComputeMetadata&>(compute_metadata)); }
|
||||
|
||||
Status SliceBase::FillVectorsFromInput(const Tensor& start_tensor,
|
||||
|
|
|
|||
|
|
@ -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<float>({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<float>({0, 6}, // empty dim in shape
|
||||
{},
|
||||
|
|
|
|||
Loading…
Reference in a new issue