saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-24 19:43:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Implement Concat with Strided copy (#8336)

Browse source 
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.
This commit is contained in:
Oliver Rausch 2021-07-27 19:27:56 +03:00 committed by GitHub
parent 4c939e1cb7
commit 1685ab8138
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 806 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
cmake/onnxruntime_unittests.cmake
View file

@ -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

2
onnxruntime/core/providers/cpu/sequence/concat_from_sequence.cc
View file

@ -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

72
onnxruntime/core/providers/cpu/tensor/concat.cc
View file

@ -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<int64_t> StridesForStack(const std::vector<int64_t>& 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<int64_t> 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<int>(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<const uint8_t*>(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<uint8_t*>(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<std::string*>(output)[out + idx_item] =
reinterpret_cast<const std::string*>(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

2
onnxruntime/core/providers/cpu/tensor/concatbase.h
View file

@ -38,7 +38,7 @@ class ConcatBase {
is_stack_ = info.GetAttrOrDefault<int64_t>("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;

347
onnxruntime/core/providers/cpu/tensor/copy.cc Normal file
View file

@ -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<int64_t> StridesForTensor(const Tensor& tensor) {
auto shape = tensor.Shape();
auto strides = std::vector<int64_t>(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<int64_t> dst_strides,
const TensorShape& copy_shape,
const Tensor& src,
const std::vector<int64_t> 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<T>(thread_pool, dst.MutableData<T>() + dst_offset, dst_strides, copy_shape, src.Data<T>(), 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 <typename T>
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 <typename T>
inline void Copy1DContiguous(T* dst, const T* src, std::ptrdiff_t count) {
memcpy(dst, src, count * sizeof(T));
}
template <>
inline void Copy1DContiguous<std::string>(std::string* dst, const std::string* src, std::ptrdiff_t count) {
Copy1DNonContiguous(dst, 1, src, 1, count);
}
template <typename T>
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>(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<int64_t>& 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<std::ptrdiff_t>(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<int64_t> current_index;
const std::vector<int64_t>& 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<std::reference_wrapper<std::vector<int64_t>>>& tensors_strides,
const std::vector<int64_t>& 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<int64_t>& 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<std::reference_wrapper<std::vector<int64_t>>>& tensors_strides,
std::size_t dim, std::size_t ndim) {
for (const auto& strides_ : tensors_strides) {
std::vector<int64_t>& 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<std::reference_wrapper<std::vector<int64_t>>>&& tensors_strides, std::vector<int64_t>& 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<int64_t>& strides = strides_.get();
strides.resize(current_dim + 1);
}
}
template <typename T>
void StridedCopy(concurrency::ThreadPool* thread_pool,
T* dst,
const std::vector<int64_t>& dst_strides_,
const TensorShape& copy_shape_,
const T* src,
const std::vector<int64_t>& src_strides_) {
// Coalesce dimensions
std::vector<int64_t> dst_strides = dst_strides_;
std::vector<int64_t> src_strides = src_strides_;
std::vector<int64_t> 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<float>(sizeof(T)), static_cast<float>(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<std::ptrdiff_t>(elements_to_copy, last - first);
Copy1DContiguous<T>(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<T>(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<T>(dst + dst_idx, src + src_idx, last_span_size);
});
} else {
concurrency::ThreadPool::TryParallelFor(
thread_pool, num_iterations,
{static_cast<float>(sizeof(T)), static_cast<float>(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<T>(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

27
onnxruntime/core/providers/cpu/tensor/copy.h Normal file
View file

@ -0,0 +1,27 @@
#include "core/providers/common.h"
#include "core/platform/threadpool.h"
#include <vector>
namespace onnxruntime {
void CoalesceDimensions(
std::initializer_list<std::reference_wrapper<std::vector<int64_t>>>&& tensors_strides, std::vector<int64_t>& shape);
Status DispatchStridedCopy(concurrency::ThreadPool* thread_pool,
Tensor& dst,
std::ptrdiff_t dst_offset,
const std::vector<int64_t> dst_strides,
const TensorShape& copy_shape,
const Tensor& src,
const std::vector<int64_t> src_strides);
template <typename T>
void StridedCopy(concurrency::ThreadPool* thread_pool,
T* dst,
const std::vector<int64_t>& dst_strides,
const TensorShape& copy_shape,
const T* src,
const std::vector<int64_t>& src_strides);
std::vector<int64_t> StridesForTensor(const Tensor& tensor);
} // namespace onnxruntime

97
onnxruntime/test/onnx/microbenchmark/copy.cc Normal file
View file

@ -0,0 +1,97 @@
#include "common.h"
#include "core/util/thread_utils.h"
#include <core/providers/cpu/tensor/copy.h>
#include <benchmark/benchmark.h>
#include <core/platform/threadpool.h>
using namespace onnxruntime;
using namespace onnxruntime::concurrency;
static void BM_StridedCopy_Memcpy(benchmark::State& state) {
const size_t batch_size = static_cast<size_t>(state.range(0));
const size_t feature_size = static_cast<size_t>(state.range(1));
float* output = (float*)aligned_alloc(sizeof(float) * batch_size * feature_size, 64);
float* data = GenerateArrayWithRandomValue<float>(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<size_t>(state.range(0));
const size_t feature_size = static_cast<size_t>(state.range(1));
float* output = (float*)aligned_alloc(sizeof(float) * batch_size * feature_size, 64);
float* data = GenerateArrayWithRandomValue<float>(batch_size * feature_size, -1, 1);
int64_t ibatch_size = static_cast<int64_t>(batch_size);
int64_t ifeature_size = static_cast<int64_t>(feature_size);
for (auto _ : state) {
// use nullptr threadpool to make it run single threaded
StridedCopy<float>(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<size_t>(state.range(0));
const size_t feature_size = static_cast<size_t>(state.range(1));
float* output = (float*)aligned_alloc(sizeof(float) * batch_size * feature_size, 64);
float* data = GenerateArrayWithRandomValue<float>(batch_size * feature_size, -1, 1);
// setup threadpool
OrtThreadPoolParams tpo;
tpo.auto_set_affinity = true;
std::unique_ptr<concurrency::ThreadPool> tp(
concurrency::CreateThreadPool(&onnxruntime::Env::Default(), tpo, concurrency::ThreadPoolType::INTRA_OP));
int64_t ibatch_size = static_cast<int64_t>(batch_size);
int64_t ifeature_size = static_cast<int64_t>(feature_size);
for (auto _ : state) {
StridedCopy<float>(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<size_t>(state.range(0));
const size_t feature_size = static_cast<size_t>(state.range(1));
float* output = (float*)aligned_alloc(sizeof(float) * batch_size * feature_size * 2, 64);
float* data = GenerateArrayWithRandomValue<float>(batch_size * feature_size, -1, 1);
int64_t ibatch_size = static_cast<int64_t>(batch_size);
int64_t ifeature_size = static_cast<int64_t>(feature_size);
for (auto _ : state) {
StridedCopy<float>(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);

2
onnxruntime/test/providers/cpu/sequence/concat_from_sequence_op_test.cc
View file

@ -133,4 +133,4 @@ TEST(SequenceOpsTest, ConcatFromSequence_Concat_ScalarInputs) {
}
} // namespace test
} // namespace onnxruntime
} // namespace onnxruntime

71
onnxruntime/test/providers/cpu/tensor/concat_op_test.cc
View file

@ -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<int64_t> dims{2, 2};
test.AddInput<double>("input1", dims,
{111.0f, 112.0f,
121.0f, 122.0f});
test.AddInput<double>("input2", dims,
{211.0f, 212.0f,
221.0f, 222.0f});
test.AddOutput<double>("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<float>("input1", {1, 3, 3},
{111.0f, 112.0f, 113.0f,
121.0f, 122.0f, 123.0f,
131.0f, 132.0f, 133.0f});
test.AddInput<float>("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<float>("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<float>("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<float>("input1", {1, 3, 2},
{111.0f, 112.0f,
121.0f, 122.0f,
131.0f, 132.0f});
test.AddInput<float>("input2", {1, 2, 2},
{211.0f, 212.0f,
221.0f, 222.0f});
test.AddInput<float>("input3", {1, 4, 2},
{311.0f, 312.0f,
321.0f, 322.0f,
331.0f, 332.0f,
341.0f, 342.0f});
test.AddOutput<float>("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});

221
onnxruntime/test/providers/cpu/tensor/copy_test.cc Normal file
View file

@ -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<concurrency::ThreadPool> tp;
};
TEST_F(CopyTest, Contiguous1D) {
int src[10];
for (int i = 0; i < 10; i++) {
src[i] = i;
}
int dst[10];
StridedCopy<int>(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<double>(i);
}
double dst[3 * 4 * 5];
StridedCopy<double>(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<double>(i);
}
double* dst = new double[numel];
std::vector<int64_t> dst_strides = {60, 5, 15, 1};
std::vector<int64_t> src_strides = {60, 20, 5, 1};
StridedCopy<double>(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<int64_t> 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<double>(i);
}
double* dst = new double[10 * 5];
for (int i = 0; i < 10 * 5; i++) {
dst[i] = 0;
}
std::vector<int64_t> dst_strides = {5, 1};
std::vector<int64_t> src_strides = {2, 1};
std::ptrdiff_t offset = 3;
StridedCopy<double>(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<int64_t> strides_a{3, 1};
std::vector<int64_t> strides_b{3, 1};
std::vector<int64_t> 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<int64_t> strides_a{3, 3, 1};
std::vector<int64_t> strides_b{3, 3, 1};
std::vector<int64_t> 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<int64_t> strides_a{3, 3, 3, 1};
std::vector<int64_t> strides_b{3, 3, 3, 1};
std::vector<int64_t> 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<int64_t> strides_a{3, 3, 3, 1};
std::vector<int64_t> strides_b{3, 3, 3, 1};
std::vector<int64_t> 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<int64_t> strides_a{320, 1};
std::vector<int64_t> strides_b{320, 1};
std::vector<int64_t> 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<int64_t> strides_a{320, 20, 1};
std::vector<int64_t> strides_b{320, 20, 1};
std::vector<int64_t> 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<int64_t> strides_a{3, 1};
std::vector<int64_t> strides_b{6, 1};
std::vector<int64_t> 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<int64_t> strides_a{3, 1};
std::vector<int64_t> strides_b{6, 1};
std::vector<int64_t> 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<int64_t> strides_a{4, 1};
std::vector<int64_t> strides_b{1, 1};
std::vector<int64_t> 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

2
orttraining/orttraining/training_ops/cpu/tensor/concat.cc
View file

@ -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