From 1685ab81381fdc4d6948bcd19c5729c06a0cd49c Mon Sep 17 00:00:00 2001 From: Oliver Rausch Date: Tue, 27 Jul 2021 19:27:56 +0300 Subject: [PATCH] Implement Concat with Strided copy (#8336) Adds a StridedCopy function that implements a copy from strided tensor to another. This parallelizes the Concat operator, and can also be used in the future to parallelize many other data movement operators (e.g. Transpose, Split, etc.). This operation is also required for the proposed data layout extensions to ORT. --- cmake/onnxruntime_unittests.cmake | 1 + .../cpu/sequence/concat_from_sequence.cc | 2 +- .../core/providers/cpu/tensor/concat.cc | 72 ++-- .../core/providers/cpu/tensor/concatbase.h | 2 +- onnxruntime/core/providers/cpu/tensor/copy.cc | 347 ++++++++++++++++++ onnxruntime/core/providers/cpu/tensor/copy.h | 27 ++ onnxruntime/test/onnx/microbenchmark/copy.cc | 97 +++++ .../sequence/concat_from_sequence_op_test.cc | 2 +- .../providers/cpu/tensor/concat_op_test.cc | 71 ++++ .../test/providers/cpu/tensor/copy_test.cc | 221 +++++++++++ .../training_ops/cpu/tensor/concat.cc | 2 +- 11 files changed, 806 insertions(+), 38 deletions(-) create mode 100644 onnxruntime/core/providers/cpu/tensor/copy.cc create mode 100644 onnxruntime/core/providers/cpu/tensor/copy.h create mode 100644 onnxruntime/test/onnx/microbenchmark/copy.cc create mode 100644 onnxruntime/test/providers/cpu/tensor/copy_test.cc diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index d75b1ee4af..b3c2628279 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -826,6 +826,7 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) ${BENCHMARK_DIR}/batchnorm2.cc ${BENCHMARK_DIR}/tptest.cc ${BENCHMARK_DIR}/eigen.cc + ${BENCHMARK_DIR}/copy.cc ${BENCHMARK_DIR}/gelu.cc ${BENCHMARK_DIR}/activation.cc ${BENCHMARK_DIR}/quantize.cc diff --git a/onnxruntime/core/providers/cpu/sequence/concat_from_sequence.cc b/onnxruntime/core/providers/cpu/sequence/concat_from_sequence.cc index 3b8b3e52e4..352cdb83e0 100644 --- a/onnxruntime/core/providers/cpu/sequence/concat_from_sequence.cc +++ b/onnxruntime/core/providers/cpu/sequence/concat_from_sequence.cc @@ -40,7 +40,7 @@ Status ConcatFromSequence::Compute(OpKernelContext* ctx) const { return Status::OK(); // Compute values to be placed in the output tensor - return ComputeImpl(p); + return ComputeImpl(p, ctx); } } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/tensor/concat.cc b/onnxruntime/core/providers/cpu/tensor/concat.cc index 7e346c20e2..102ee1acc8 100644 --- a/onnxruntime/core/providers/cpu/tensor/concat.cc +++ b/onnxruntime/core/providers/cpu/tensor/concat.cc @@ -4,6 +4,7 @@ #include "core/providers/cpu/tensor/concat.h" #include "core/providers/common.h" #include "core/framework/TensorSeq.h" +#include "core/providers/cpu/tensor/copy.h" namespace onnxruntime { @@ -210,11 +211,31 @@ Status ConcatBase::PrepareForCompute(OpKernelContext* ctx, return Status::OK(); } +namespace { +std::vector StridesForStack(const std::vector& full_strides, uint64_t axis) { + // if we are stacking, skip the dimension that will be stacked along in the output strides + // (the striding for that dimension is handled by the initial_output_offset) + auto num_dims = full_strides.size(); + + std::vector strides(num_dims - 1); + + for (size_t i = 0; i < num_dims - 1; i++) { + auto read_i = (i >= axis) ? i + 1 : i; + strides[i] = full_strides[read_i]; + } + return strides; +} +} // namespace + // This method computes the output tensor for Concat/ConcatFromSequence ops -Status ConcatBase::ComputeImpl(Prepare& p) const { +Status ConcatBase::ComputeImpl(Prepare& p, OpKernelContext* ctx) const { int input_count = static_cast(p.inputs.size()); int64_t initial_output_offset = 0; // initial offset for each input - auto element_bytes = p.output_tensor->DataType()->Size(); + + auto output_strides_full = StridesForTensor(*p.output_tensor); + // Note that output_strides_full is only used later when is_stack_ is true, so it's safe to move + auto output_strides_for_copy = is_stack_ ? StridesForStack(output_strides_full, p.axis) : std::move(output_strides_full); + for (int input_index = 0; input_index < input_count; input_index++) { const auto& prep = p.inputs[input_index]; @@ -222,39 +243,22 @@ Status ConcatBase::ComputeImpl(Prepare& p) const { if (prep.num_elements == 0) continue; - auto input_axis_pitch = prep.axis_pitch; - const uint8_t* input = static_cast(prep.tensor->DataRaw()); + // parallel copy the data across + auto status = DispatchStridedCopy(ctx->GetOperatorThreadPool(), + *p.output_tensor, + initial_output_offset, + output_strides_for_copy, + prep.tensor->Shape(), + *prep.tensor, + StridesForTensor(*prep.tensor)); + ORT_RETURN_IF_ERROR(status); - auto input_size = prep.num_elements; - - // Copy the data across. For every 'input_axis_pitch' values copied, we move over by the 'output_axis_pitch' - // TODO: Optimization possibility: There are cases where we simply need to "merge" raw buffers and this - // could be done without the pointer house-keeping as below. Some scenarios whether this is possible are: - // 1) Concatenating on input axis = 0 - // 2) Stacking on output axis = 0 - // 3) Stacking scalars - uint8_t* output = static_cast(p.output_tensor->MutableDataRaw()); - int64_t cur_out_offset = 0; - int64_t cur_in_offset = 0; - for (size_t idx_copy = 0, end = input_size / input_axis_pitch; idx_copy < end; ++idx_copy) { - if (p.is_string_type) { - size_t out = initial_output_offset + cur_out_offset; - for (int idx_item = 0; idx_item < input_axis_pitch; ++idx_item) { - reinterpret_cast(output)[out + idx_item] = - reinterpret_cast(input)[cur_in_offset + idx_item]; - } - } else { - memcpy( - output + (initial_output_offset + cur_out_offset) * element_bytes, - input + cur_in_offset * element_bytes, - input_axis_pitch * element_bytes); - } - - cur_out_offset += p.output_axis_pitch; - cur_in_offset += input_axis_pitch; + // advance along the axis that we are concatenating on (by the size of the axis of the tensor that we just copied) + if (is_stack_) { + initial_output_offset += output_strides_full[p.axis]; + } else { + initial_output_offset += prep.tensor->Shape()[p.axis] * output_strides_for_copy[p.axis]; } - - initial_output_offset += input_axis_pitch; } return Status::OK(); @@ -283,7 +287,7 @@ Status Concat::Compute(OpKernelContext* ctx) const { return Status::OK(); // Compute values to be placed in the output tensor - return ComputeImpl(p); + return ComputeImpl(p, ctx); } } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/tensor/concatbase.h b/onnxruntime/core/providers/cpu/tensor/concatbase.h index fb8ebd8fdf..2f62ec9f5a 100644 --- a/onnxruntime/core/providers/cpu/tensor/concatbase.h +++ b/onnxruntime/core/providers/cpu/tensor/concatbase.h @@ -38,7 +38,7 @@ class ConcatBase { is_stack_ = info.GetAttrOrDefault("new_axis", 0) == 0 ? false : true; } } - Status ComputeImpl(Prepare& p) const; + Status ComputeImpl(Prepare& p, OpKernelContext* ctx) const; int64_t axis_; bool is_stack_ = false; diff --git a/onnxruntime/core/providers/cpu/tensor/copy.cc b/onnxruntime/core/providers/cpu/tensor/copy.cc new file mode 100644 index 0000000000..2eb1629554 --- /dev/null +++ b/onnxruntime/core/providers/cpu/tensor/copy.cc @@ -0,0 +1,347 @@ +// +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include "core/common/common.h" +#include "core/providers/common.h" +#include "core/providers/cpu/tensor/copy.h" +#include "core/platform/threadpool.h" + +namespace onnxruntime { + +std::vector StridesForTensor(const Tensor& tensor) { + auto shape = tensor.Shape(); + auto strides = std::vector(shape.NumDimensions()); + int64_t running_size = 1; + for (auto i = shape.NumDimensions(); i > 0; i--) { + strides[i - 1] = running_size; + running_size *= shape[i - 1]; + } + + return strides; +} + +Status DispatchStridedCopy(concurrency::ThreadPool* thread_pool, + Tensor& dst, + std::ptrdiff_t dst_offset, + const std::vector dst_strides, + const TensorShape& copy_shape, + const Tensor& src, + const std::vector src_strides) { + ORT_ENFORCE(dst.DataType() == src.DataType(), "src and dst types must match"); + + // Manual dispatching: DispatchOnTensorType doesn't work here because we need to pass the type to the MutableData call +#define CALL_FOR_TYPE(T) \ + StridedCopy(thread_pool, dst.MutableData() + dst_offset, dst_strides, copy_shape, src.Data(), src_strides); \ + break + + auto tensor_type = dst.DataType()->AsPrimitiveDataType()->GetDataType(); + switch (tensor_type) { + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + CALL_FOR_TYPE(float); + case ONNX_NAMESPACE::TensorProto_DataType_BOOL: + CALL_FOR_TYPE(bool); + case ONNX_NAMESPACE::TensorProto_DataType_DOUBLE: + CALL_FOR_TYPE(double); + case ONNX_NAMESPACE::TensorProto_DataType_STRING: + CALL_FOR_TYPE(std::string); + case ONNX_NAMESPACE::TensorProto_DataType_INT8: + CALL_FOR_TYPE(int8_t); + case ONNX_NAMESPACE::TensorProto_DataType_UINT8: + CALL_FOR_TYPE(uint8_t); + case ONNX_NAMESPACE::TensorProto_DataType_INT16: + CALL_FOR_TYPE(int16_t); + case ONNX_NAMESPACE::TensorProto_DataType_UINT16: + CALL_FOR_TYPE(uint16_t); + case ONNX_NAMESPACE::TensorProto_DataType_INT32: + CALL_FOR_TYPE(int32_t); + case ONNX_NAMESPACE::TensorProto_DataType_UINT32: + CALL_FOR_TYPE(uint32_t); + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + CALL_FOR_TYPE(int64_t); + case ONNX_NAMESPACE::TensorProto_DataType_UINT64: + CALL_FOR_TYPE(uint64_t); + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: + CALL_FOR_TYPE(MLFloat16); + case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: + CALL_FOR_TYPE(BFloat16); + default: + ORT_ENFORCE(false, "Unknown tensor type of ", tensor_type); + } + return Status::OK(); +} + +namespace { + +template +inline void Copy1DNonContiguous(T* dst, int64_t dst_stride, const T* src, int64_t src_stride, std::ptrdiff_t count) { + for (std::ptrdiff_t i = 0; i < count; i++) { + dst[0] = src[0]; + dst += dst_stride; + src += src_stride; + } +} + +template +inline void Copy1DContiguous(T* dst, const T* src, std::ptrdiff_t count) { + memcpy(dst, src, count * sizeof(T)); +} +template <> +inline void Copy1DContiguous(std::string* dst, const std::string* src, std::ptrdiff_t count) { + Copy1DNonContiguous(dst, 1, src, 1, count); +} + +template +inline void Copy1D(T* dst, int64_t dst_stride, const T* src, int64_t src_stride, std::ptrdiff_t count) { + if (dst_stride == 1 && src_stride == 1) { + Copy1DContiguous(dst, src, count); + } else { + Copy1DNonContiguous(dst, dst_stride, src, src_stride, count); + } +} + +template <> +inline void Copy1D(std::string* dst, int64_t dst_stride, const std::string* src, int64_t src_stride, std::ptrdiff_t count) { + // strings should always be copied using the for loop + Copy1DNonContiguous(dst, dst_stride, src, src_stride, count); +} + +struct NdCounter { + NdCounter(const std::vector& shape, std::ptrdiff_t first, std::ptrdiff_t last) + : dims(shape.size()), + last_dim_size(shape[dims - 1]), + current_offset(first), + last(last), + current_index(dims), + shape(shape) { + // compute the initial n-dimensional index + int64_t remaining_index = first; + // Iterate from dims to 1 so we don't roll over to positive on the bounds check + for (std::size_t dim = dims; dim > 0; dim--) { + auto shape_val = shape[dim - 1]; + current_index[dim - 1] = remaining_index % shape_val; + remaining_index /= shape_val; + } + } + + /* + Return the size of the largest step in the last dimension. + */ + std::ptrdiff_t NextStepSize() const { + auto elements_in_dimension = last_dim_size - current_index[dims - 1]; + std::ptrdiff_t span_end = std::min(last, current_offset + elements_in_dimension); + return span_end - current_offset; + } + + /* + Advance the counter by step_size elements. + */ + void Step(std::ptrdiff_t step_size) { + current_offset += step_size; + current_index[dims - 1] += step_size; + + // update the current_nd_idx if needed + std::size_t dim = dims - 1; + while (dim > 0 && current_index[dim] >= shape[dim]) { + current_index[dim] = 0; + dim--; + current_index[dim]++; + } + } + + const std::size_t dims; + const int64_t last_dim_size; + ptrdiff_t current_offset; + const ptrdiff_t last; + std::vector current_index; + const std::vector& shape; +}; + +/* + Check if we can coalesce dim with dim + 1. + + We can do this if: + * either of the dims have shape 1 + * strides[dim + 1] * shape[dim + 1] = strides[dim] (for all tensors) +*/ +inline bool CanCoalesce( + std::initializer_list>>& tensors_strides, + const std::vector& shape, + std::size_t dim, + std::size_t ndim) { + auto shape_dim = shape[dim]; + auto shape_ndim = shape[ndim]; + if (shape_dim == 1 || shape_ndim == 1) { + return true; + } + + for (const auto& strides_ : tensors_strides) { + std::vector& strides = strides_.get(); + if (shape_ndim * strides[ndim] != strides[dim]) { + return false; + } + } + return true; +} + +/* + Copy the stride from ndim to dim in all tensors. +*/ +inline void CopyStride( + std::initializer_list>>& tensors_strides, + std::size_t dim, std::size_t ndim) { + for (const auto& strides_ : tensors_strides) { + std::vector& strides = strides_.get(); + strides[dim] = strides[ndim]; + } +} + +} // namespace + +/* + Coalesce contiguous dimensions in the tensors. Operates inplace on the function arguments. +*/ +void CoalesceDimensions(std::initializer_list>>&& tensors_strides, std::vector& shape) { + const std::size_t dims = shape.size(); + + // the current dimension is the one we are attempting to "coalesce onto" + std::size_t current_dim = 0; + + for (std::size_t dim = 1; dim < dims; dim++) { + // check if dim can be coalesced with current_dim + if (CanCoalesce(tensors_strides, shape, current_dim, dim)) { + if (shape[dim] != 1) { + CopyStride(tensors_strides, current_dim, dim); + } + shape[current_dim] *= shape[dim]; + } else { + current_dim++; + + if (current_dim != dim) { + // we have coaleseced at least one value before this: bump forward the values into the correct place + CopyStride(tensors_strides, current_dim, dim); + shape[current_dim] = shape[dim]; + } + } + } + + shape.resize(current_dim + 1); + for (const auto& strides_ : tensors_strides) { + std::vector& strides = strides_.get(); + strides.resize(current_dim + 1); + } +} + +template +void StridedCopy(concurrency::ThreadPool* thread_pool, + T* dst, + const std::vector& dst_strides_, + const TensorShape& copy_shape_, + const T* src, + const std::vector& src_strides_) { + // Coalesce dimensions + std::vector dst_strides = dst_strides_; + std::vector src_strides = src_strides_; + std::vector copy_shape(copy_shape_.GetDims()); + + CoalesceDimensions({dst_strides, src_strides}, copy_shape); + ORT_ENFORCE(dst_strides.size() == src_strides.size() && src_strides.size() == copy_shape.size(), "src and dst must have same shape"); + + const std::size_t dims = copy_shape.size(); + // We will iterate over the output dimensions + int64_t num_iterations = 1; + for (std::size_t dim = 0; dim < dims; dim++) { + num_iterations *= copy_shape[dim]; + } + + if (num_iterations <= 1) { + // scalar edge case + dst[0] = src[0]; + return; + } + + // TODOs for when we have strided tensors: + // - Reorder dimensions so that we iterate along the smallest strides first + + ORT_ENFORCE(dims > 0); + + if (dims <= 2 && src_strides[dims - 1] == 1 && dst_strides[dims - 1] == 1) { + // Fast path for 2D copies that skips the NdCounter required in the general case. + // This avoids overhead which becomes noticable at smaller iteration sizes. + // + // After coalescing, the case is actually quite common since all tensors in ORT are contiguous + + int64_t dst_stride = dims == 2 ? dst_strides[0] : 0; + int64_t src_stride = dims == 2 ? src_strides[0] : 0; + + // the size of contiguous spans that we can copy before having to advance the non-contiguous stride + int64_t contiguous_span_size = dims == 2 ? copy_shape[1] : copy_shape[0]; + + concurrency::ThreadPool::TryParallelFor( + thread_pool, num_iterations, + {static_cast(sizeof(T)), static_cast(sizeof(T)), 1.0F}, + [src_stride, dst_stride, dst, src, contiguous_span_size](std::ptrdiff_t first, std::ptrdiff_t last) { + // get the current inner and outer index + int64_t inner = first % contiguous_span_size; + int64_t outer = first / contiguous_span_size; + + std::ptrdiff_t dst_idx = outer * dst_stride + inner; + std::ptrdiff_t src_idx = outer * src_stride + inner; + + // Step 1: if there is anything left in the contiguous span that we are starting in, finish copying it + if (inner != 0) { + auto elements_to_copy = contiguous_span_size - inner; + // never copy more than what is in our partition + elements_to_copy = std::min(elements_to_copy, last - first); + Copy1DContiguous(dst + dst_idx, src + src_idx, elements_to_copy); + inner = 0; + outer++; + first += elements_to_copy; + + // reset the dst and src idx now that we are aligned to the start of a contiguous span + dst_idx = outer * dst_stride; + src_idx = outer * src_stride; + } + + // Step 2: copy contiguous span by contiguous span until we reach the penultimate span + while (first < last - contiguous_span_size) { + Copy1DContiguous(dst + dst_idx, src + src_idx, contiguous_span_size); + dst_idx += dst_stride; + src_idx += src_stride; + first += contiguous_span_size; + } + // Step 3: finish off the last (possibly partial) span manually, making sure that we don't go past the last + // element in our partition + ORT_ENFORCE(last >= first); + auto last_span_size = last - first; + Copy1DContiguous(dst + dst_idx, src + src_idx, last_span_size); + }); + } else { + concurrency::ThreadPool::TryParallelFor( + thread_pool, num_iterations, + {static_cast(sizeof(T)), static_cast(sizeof(T)), 1.0F}, + [copy_shape, dst_strides, dst, src, src_strides, dims](std::ptrdiff_t first, std::ptrdiff_t last) { + NdCounter counter(copy_shape, first, last); + + auto last_dst_stride = dst_strides[dims - 1]; + auto last_src_stride = src_strides[dims - 1]; + + auto iter_size = counter.NextStepSize(); + while (iter_size > 0) { + // Compute the src and dst addresses + std::ptrdiff_t dst_idx = 0; + std::ptrdiff_t src_idx = 0; + for (std::size_t dim = 0; dim < dims; dim++) { + dst_idx += counter.current_index[dim] * dst_strides[dim]; + src_idx += counter.current_index[dim] * src_strides[dim]; + } + // we can copy until the current dimension is done (or until we hit the last element we are trying to copy) + Copy1D(dst + dst_idx, last_dst_stride, src + src_idx, last_src_stride, iter_size); + + counter.Step(iter_size); + iter_size = counter.NextStepSize(); + } + ORT_ENFORCE(counter.current_offset == last); + }); + } +} +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/tensor/copy.h b/onnxruntime/core/providers/cpu/tensor/copy.h new file mode 100644 index 0000000000..fabbd36cb6 --- /dev/null +++ b/onnxruntime/core/providers/cpu/tensor/copy.h @@ -0,0 +1,27 @@ +#include "core/providers/common.h" +#include "core/platform/threadpool.h" +#include + +namespace onnxruntime { + +void CoalesceDimensions( + std::initializer_list>>&& tensors_strides, std::vector& shape); + +Status DispatchStridedCopy(concurrency::ThreadPool* thread_pool, + Tensor& dst, + std::ptrdiff_t dst_offset, + const std::vector dst_strides, + const TensorShape& copy_shape, + const Tensor& src, + const std::vector src_strides); + +template +void StridedCopy(concurrency::ThreadPool* thread_pool, + T* dst, + const std::vector& dst_strides, + const TensorShape& copy_shape, + const T* src, + const std::vector& src_strides); + +std::vector StridesForTensor(const Tensor& tensor); +} // namespace onnxruntime diff --git a/onnxruntime/test/onnx/microbenchmark/copy.cc b/onnxruntime/test/onnx/microbenchmark/copy.cc new file mode 100644 index 0000000000..a56669fe65 --- /dev/null +++ b/onnxruntime/test/onnx/microbenchmark/copy.cc @@ -0,0 +1,97 @@ +#include "common.h" +#include "core/util/thread_utils.h" + +#include +#include +#include + +using namespace onnxruntime; +using namespace onnxruntime::concurrency; + +static void BM_StridedCopy_Memcpy(benchmark::State& state) { + const size_t batch_size = static_cast(state.range(0)); + const size_t feature_size = static_cast(state.range(1)); + + float* output = (float*)aligned_alloc(sizeof(float) * batch_size * feature_size, 64); + float* data = GenerateArrayWithRandomValue(batch_size * feature_size, -1, 1); + + for (auto _ : state) { + memcpy(output, data, batch_size * feature_size * sizeof(float)); + } + aligned_free(data); + aligned_free(output); +} + +static void BM_StridedCopy_SingleThread(benchmark::State& state) { + const size_t batch_size = static_cast(state.range(0)); + const size_t feature_size = static_cast(state.range(1)); + + float* output = (float*)aligned_alloc(sizeof(float) * batch_size * feature_size, 64); + float* data = GenerateArrayWithRandomValue(batch_size * feature_size, -1, 1); + + int64_t ibatch_size = static_cast(batch_size); + int64_t ifeature_size = static_cast(feature_size); + for (auto _ : state) { + // use nullptr threadpool to make it run single threaded + StridedCopy(nullptr, output, {ifeature_size, 1}, {ibatch_size, ifeature_size}, data, {ifeature_size, 1}); + } + aligned_free(data); + aligned_free(output); +} + +static void BM_StridedCopy_Parallel(benchmark::State& state) { + const size_t batch_size = static_cast(state.range(0)); + const size_t feature_size = static_cast(state.range(1)); + + float* output = (float*)aligned_alloc(sizeof(float) * batch_size * feature_size, 64); + float* data = GenerateArrayWithRandomValue(batch_size * feature_size, -1, 1); + + // setup threadpool + OrtThreadPoolParams tpo; + tpo.auto_set_affinity = true; + std::unique_ptr tp( + concurrency::CreateThreadPool(&onnxruntime::Env::Default(), tpo, concurrency::ThreadPoolType::INTRA_OP)); + + int64_t ibatch_size = static_cast(batch_size); + int64_t ifeature_size = static_cast(feature_size); + for (auto _ : state) { + StridedCopy(tp.get(), output, {ifeature_size, 1}, {ibatch_size, ifeature_size}, data, {ifeature_size, 1}); + } + aligned_free(data); + aligned_free(output); +} + +static void BM_StridedCopy_SingleThread_Axis_1(benchmark::State& state) { + const size_t batch_size = static_cast(state.range(0)); + const size_t feature_size = static_cast(state.range(1)); + + float* output = (float*)aligned_alloc(sizeof(float) * batch_size * feature_size * 2, 64); + float* data = GenerateArrayWithRandomValue(batch_size * feature_size, -1, 1); + + int64_t ibatch_size = static_cast(batch_size); + int64_t ifeature_size = static_cast(feature_size); + for (auto _ : state) { + StridedCopy(nullptr, output, {ifeature_size * 2, 1}, {ibatch_size, ifeature_size}, data, {ifeature_size, 1}); + } + aligned_free(data); + aligned_free(output); +} + +#define SC_BENCHMARK(name) \ + BENCHMARK(name) \ + ->UseRealTime() \ + ->UseRealTime() \ + ->Unit(benchmark::TimeUnit::kNanosecond) \ + ->Args({32, 64}) \ + ->Args({64, 64}) \ + ->Args({128, 64}) \ + ->Args({256, 64}) \ + ->Args({10000, 64}) \ + ->Args({20000, 64}) \ + ->Args({40000, 64}) \ + ->Args({400000, 64}); + +SC_BENCHMARK(BM_StridedCopy_Memcpy); +SC_BENCHMARK(BM_StridedCopy_SingleThread); +SC_BENCHMARK(BM_StridedCopy_Parallel); +SC_BENCHMARK(BM_StridedCopy_SingleThread_Axis_1); diff --git a/onnxruntime/test/providers/cpu/sequence/concat_from_sequence_op_test.cc b/onnxruntime/test/providers/cpu/sequence/concat_from_sequence_op_test.cc index a0e796c8ce..cca69effd5 100644 --- a/onnxruntime/test/providers/cpu/sequence/concat_from_sequence_op_test.cc +++ b/onnxruntime/test/providers/cpu/sequence/concat_from_sequence_op_test.cc @@ -133,4 +133,4 @@ TEST(SequenceOpsTest, ConcatFromSequence_Concat_ScalarInputs) { } } // namespace test -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/tensor/concat_op_test.cc b/onnxruntime/test/providers/cpu/tensor/concat_op_test.cc index 435049d94c..c58475ee2c 100644 --- a/onnxruntime/test/providers/cpu/tensor/concat_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/concat_op_test.cc @@ -131,6 +131,25 @@ TEST(ConcatOpTest, Concat2D_4) { {kTensorrtExecutionProvider}); //TensorRT: no support for dynamic shape tensor } +TEST(ConcatOpTest, Concat2D_5) { + OpTester test("Concat"); + test.AddAttribute("axis", int64_t{0}); + + std::vector dims{2, 2}; + test.AddInput("input1", dims, + {111.0f, 112.0f, + 121.0f, 122.0f}); + test.AddInput("input2", dims, + {211.0f, 212.0f, + 221.0f, 222.0f}); + test.AddOutput("concat_result", {4, 2}, + {111.0f, 112.0f, + 121.0f, 122.0f, + 211.0f, 212.0f, + 221.0f, 222.0f}); + test.Run(); +} + TEST(ConcatOpTest, Concat3D_1) { OpTester test("Concat"); test.AddAttribute("axis", int64_t{0}); @@ -257,6 +276,58 @@ TEST(ConcatOpTest, Concat3D_3) { test.Run(); } +TEST(ConcatOpTest, Concat3D_4) { + OpTester test("Concat"); + test.AddAttribute("axis", int64_t{2}); + + test.AddInput("input1", {1, 3, 3}, + {111.0f, 112.0f, 113.0f, + 121.0f, 122.0f, 123.0f, + 131.0f, 132.0f, 133.0f}); + test.AddInput("input2", {1, 3, 4}, + {211.0f, 212.0f, 213.0f, 214.0f, + 221.0f, 222.0f, 223.0f, 224.0f, + 231.0f, 232.0f, 233.0f, 234.0f}); + test.AddInput("input3", {1, 3, 5}, + {311.0f, 312.0f, 313.0f, 314.0f, 315.0f, + 321.0f, 322.0f, 323.0f, 324.0f, 325.0f, + 331.0f, 332.0f, 333.0f, 334.0f, 335.0f}); + test.AddOutput("concat_result", {1, 3, 12}, + {111.0f, 112.0f, 113.0f, 211.0f, 212.0f, 213.0f, 214.0f, 311.0f, 312.0f, 313.0f, 314.0f, 315.0f, + 121.0f, 122.0f, 123.0f, 221.0f, 222.0f, 223.0f, 224.0f, 321.0f, 322.0f, 323.0f, 324.0f, 325.0f, + 131.0f, 132.0f, 133.0f, 231.0f, 232.0f, 233.0f, 234.0f, 331.0f, 332.0f, 333.0f, 334.0f, 335.0f}); + test.Run(); +} + +TEST(ConcatOpTest, Concat3D_5) { + OpTester test("Concat"); + test.AddAttribute("axis", int64_t{1}); + + test.AddInput("input1", {1, 3, 2}, + {111.0f, 112.0f, + 121.0f, 122.0f, + 131.0f, 132.0f}); + test.AddInput("input2", {1, 2, 2}, + {211.0f, 212.0f, + 221.0f, 222.0f}); + test.AddInput("input3", {1, 4, 2}, + {311.0f, 312.0f, + 321.0f, 322.0f, + 331.0f, 332.0f, + 341.0f, 342.0f}); + test.AddOutput("concat_result", {1, 9, 2}, + {111.0f, 112.0f, + 121.0f, 122.0f, + 131.0f, 132.0f, + 211.0f, 212.0f, + 221.0f, 222.0f, + 311.0f, 312.0f, + 321.0f, 322.0f, + 331.0f, 332.0f, + 341.0f, 342.0f}); + test.Run(); +} + TEST(ConcatOpTest, Concat4D_1) { OpTester test("Concat"); test.AddAttribute("axis", int64_t{1}); diff --git a/onnxruntime/test/providers/cpu/tensor/copy_test.cc b/onnxruntime/test/providers/cpu/tensor/copy_test.cc new file mode 100644 index 0000000000..91c496a8ed --- /dev/null +++ b/onnxruntime/test/providers/cpu/tensor/copy_test.cc @@ -0,0 +1,221 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" +#include "test/providers/provider_test_utils.h" +#include "core/providers/cpu/tensor/copy.h" +#include "core/platform/threadpool.h" + +namespace onnxruntime { +namespace test { + +class CopyTest : public ::testing::Test { + protected: + void SetUp() override { + OrtThreadPoolParams tpo; + tpo.auto_set_affinity = true; + tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), tpo, concurrency::ThreadPoolType::INTRA_OP); + } + std::unique_ptr tp; +}; + +TEST_F(CopyTest, Contiguous1D) { + int src[10]; + for (int i = 0; i < 10; i++) { + src[i] = i; + } + + int dst[10]; + + StridedCopy(tp.get(), dst, {1}, {10}, src, {1}); + + for (int i = 0; i < 10; i++) { + EXPECT_EQ(src[i], dst[i]); + } +} + +TEST_F(CopyTest, Contiguous3D) { + double src[3 * 4 * 5]; + for (int i = 0; i < 3 * 4 * 5; i++) { + src[i] = static_cast(i); + } + + double dst[3 * 4 * 5]; + + StridedCopy(tp.get(), dst, {20, 5, 1}, {3, 4, 5}, src, {20, 5, 1}); + + for (int i = 0; i < 3 * 4 * 5; i++) { + EXPECT_EQ(src[i], dst[i]); + } +} + +TEST_F(CopyTest, Transpose4D) { + // Test performing a transpose using a strided copy + int64_t numel = 2 * 3 * 4 * 5; + double* src = new double[numel]; + for (int i = 0; i < numel; i++) { + src[i] = static_cast(i); + } + + double* dst = new double[numel]; + + std::vector dst_strides = {60, 5, 15, 1}; + std::vector src_strides = {60, 20, 5, 1}; + StridedCopy(tp.get(), dst, dst_strides, {2, 3, 4, 5}, src, src_strides); + + // stride to access the dst tensor as if it were contiguous + std::vector contig_dst_strides = {60, 15, 5, 1}; + + for (int i0 = 0; i0 < 2; i0++) { + for (int i1 = 0; i1 < 3; i1++) { + for (int i2 = 0; i2 < 4; i2++) { + for (int i3 = 0; i3 < 5; i3++) { + size_t src_access = src_strides[0] * i0 + src_strides[1] * i1 + src_strides[2] * i2 + src_strides[3] * i3; + size_t dst_access = contig_dst_strides[0] * i0 + contig_dst_strides[1] * i2 + contig_dst_strides[2] * i1 + contig_dst_strides[3] * i3; + + EXPECT_EQ(src[src_access], dst[dst_access]); + } + } + } + } + delete[] src; + delete[] dst; +} + +TEST_F(CopyTest, Concat2D) { + // test performing a concat using a strided copy + double* src = new double[6 * 2]; + for (int i = 0; i < 6 * 2; i++) { + src[i] = static_cast(i); + } + + double* dst = new double[10 * 5]; + for (int i = 0; i < 10 * 5; i++) { + dst[i] = 0; + } + + std::vector dst_strides = {5, 1}; + std::vector src_strides = {2, 1}; + std::ptrdiff_t offset = 3; + StridedCopy(tp.get(), dst + offset, dst_strides, {6, 2}, src, src_strides); + + for (int i0 = 0; i0 < 10; i0++) { + for (int i1 = 0; i1 < 5; i1++) { + size_t dst_access = dst_strides[0] * i0 + dst_strides[1] * i1; + if (3 <= i1 && 0 <= i0 && i0 < 6) { + size_t src_access = src_strides[0] * i0 + src_strides[1] * (i1 - 3); + EXPECT_EQ(src[src_access], dst[dst_access]); + } else { + EXPECT_EQ(0, dst[dst_access]); + } + } + } + delete[] src; + delete[] dst; +} + +TEST_F(CopyTest, CoalesceTensorsTest) { + { + std::vector strides_a{3, 1}; + std::vector strides_b{3, 1}; + std::vector shape{5, 3}; + + CoalesceDimensions({strides_a, strides_b}, shape); + + ASSERT_THAT(strides_a, testing::ElementsAre(1)); + ASSERT_THAT(strides_b, testing::ElementsAre(1)); + ASSERT_THAT(shape, testing::ElementsAre(15)); + } + + { + std::vector strides_a{3, 3, 1}; + std::vector strides_b{3, 3, 1}; + std::vector shape{5, 1, 3}; + + CoalesceDimensions({strides_a, strides_b}, shape); + + ASSERT_THAT(strides_a, testing::ElementsAre(1)); + ASSERT_THAT(strides_b, testing::ElementsAre(1)); + ASSERT_THAT(shape, testing::ElementsAre(15)); + } + { + std::vector strides_a{3, 3, 3, 1}; + std::vector strides_b{3, 3, 3, 1}; + std::vector shape{1, 5, 1, 3}; + + CoalesceDimensions({strides_a, strides_b}, shape); + + ASSERT_THAT(strides_a, testing::ElementsAre(1)); + ASSERT_THAT(strides_b, testing::ElementsAre(1)); + ASSERT_THAT(shape, testing::ElementsAre(15)); + } + { + std::vector strides_a{3, 3, 3, 1}; + std::vector strides_b{3, 3, 3, 1}; + std::vector shape{1, 5, 1, 3}; + + CoalesceDimensions({strides_a, strides_b}, shape); + + ASSERT_THAT(strides_a, testing::ElementsAre(1)); + ASSERT_THAT(strides_b, testing::ElementsAre(1)); + ASSERT_THAT(shape, testing::ElementsAre(15)); + } + { + std::vector strides_a{320, 1}; + std::vector strides_b{320, 1}; + std::vector shape{20, 10}; + + CoalesceDimensions({strides_a, strides_b}, shape); + + ASSERT_THAT(strides_a, testing::ElementsAre(320, 1)); + ASSERT_THAT(strides_b, testing::ElementsAre(320, 1)); + ASSERT_THAT(shape, testing::ElementsAre(20, 10)); + } + { + std::vector strides_a{320, 20, 1}; + std::vector strides_b{320, 20, 1}; + std::vector shape{10, 2, 20}; + + CoalesceDimensions({strides_a, strides_b}, shape); + + ASSERT_THAT(strides_a, testing::ElementsAre(320, 1)); + ASSERT_THAT(strides_b, testing::ElementsAre(320, 1)); + ASSERT_THAT(shape, testing::ElementsAre(10, 40)); + } + { + std::vector strides_a{3, 1}; + std::vector strides_b{6, 1}; + std::vector shape{5, 3}; + + CoalesceDimensions({strides_a, strides_b}, shape); + + ASSERT_THAT(strides_a, testing::ElementsAre(3, 1)); + ASSERT_THAT(strides_b, testing::ElementsAre(6, 1)); + ASSERT_THAT(shape, testing::ElementsAre(5, 3)); + } + { + std::vector strides_a{3, 1}; + std::vector strides_b{6, 1}; + std::vector shape{5, 3}; + + CoalesceDimensions({strides_a, strides_b}, shape); + + ASSERT_THAT(strides_a, testing::ElementsAre(3, 1)); + ASSERT_THAT(strides_b, testing::ElementsAre(6, 1)); + ASSERT_THAT(shape, testing::ElementsAre(5, 3)); + } + { + std::vector strides_a{4, 1}; + std::vector strides_b{1, 1}; + std::vector shape{4, 1}; + + CoalesceDimensions({strides_a, strides_b}, shape); + + ASSERT_THAT(strides_a, testing::ElementsAre(4)); + ASSERT_THAT(strides_b, testing::ElementsAre(1)); + ASSERT_THAT(shape, testing::ElementsAre(4)); + } +} + +} // namespace test +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/tensor/concat.cc b/orttraining/orttraining/training_ops/cpu/tensor/concat.cc index 80b296ae85..54ed679657 100644 --- a/orttraining/orttraining/training_ops/cpu/tensor/concat.cc +++ b/orttraining/orttraining/training_ops/cpu/tensor/concat.cc @@ -48,7 +48,7 @@ Status ConcatTraining::Compute(OpKernelContext* ctx) const { } } // Compute values to be placed in the output tensor - return ComputeImpl(p); + return ComputeImpl(p, ctx); } } // namespace contrib