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