mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-07 17:15:29 +00:00
Add CPU kernel for Einsum op (#3575)
This commit is contained in:
Hariharan Seshadri
GitHub
parent
c8269e4b89
commit
785b45124d
14 changed files with 2201 additions and 117 deletions
|
|
@ -191,7 +191,11 @@ Scan<9>::Scan(const OpKernelInfo& info) : OpKernel(info) {
|
|||
output_axes_ = std::vector<int64_t>(num_scan_outputs, 0);
|
||||
}
|
||||
|
||||
device_helpers_.transpose_func = TransposeBase::DoTranspose;
|
||||
device_helpers_.transpose_func = [](const std::vector<size_t>& permutations, const Tensor& input,
|
||||
Tensor& output) -> Status {
|
||||
return TransposeBase::DoTranspose(permutations, input, output);
|
||||
};
|
||||
|
||||
device_helpers_.set_data_to_zero_func = [](void* data, size_t size_in_bytes) -> Status {
|
||||
memset(data, 0, size_in_bytes);
|
||||
return Status::OK();
|
||||
|
|
|
|||
|
|
@ -435,8 +435,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain,
|
|||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, int64_t, ReduceMin);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, int8_t, ReduceMin);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, uint8_t, ReduceMin);
|
||||
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, GatherND);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, Einsum);
|
||||
|
||||
Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
|
||||
static const BuildKernelCreateInfoFn function_table[] = {
|
||||
|
|
@ -958,9 +958,9 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
|
|||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint8_t,
|
||||
DynamicQuantizeLinear)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 11, float,
|
||||
ArgMax)>,
|
||||
ArgMax)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 11, int32_t,
|
||||
ArgMax)>,
|
||||
ArgMax)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 11,
|
||||
float, ArgMin)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 11,
|
||||
|
|
@ -1042,21 +1042,20 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
|
|||
int32_t, Resize)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11,
|
||||
uint8_t, Resize)>,
|
||||
|
||||
// OpSet 12
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, float,
|
||||
ArgMax)>,
|
||||
ArgMax)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, int32_t,
|
||||
ArgMax)>,
|
||||
ArgMax)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, float,
|
||||
ArgMin)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, int32_t,
|
||||
ArgMin)>,
|
||||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12,Clip)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, Clip)>,
|
||||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, Min)>,
|
||||
|
||||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, Max)>,
|
||||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, Pow)>,
|
||||
|
|
@ -1084,8 +1083,8 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
|
|||
ReduceMin)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, uint8_t,
|
||||
ReduceMin)>,
|
||||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, GatherND)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, Einsum)>,
|
||||
};
|
||||
|
||||
for (auto& function_table_entry : function_table) {
|
||||
|
|
|
|||
59
onnxruntime/core/providers/cpu/math/einsum.cc
Normal file
59
onnxruntime/core/providers/cpu/math/einsum.cc
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "einsum.h"
|
||||
#include "einsum_utils.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
// Credit: Implementation influenced by Torch's implementation at the time of writing
|
||||
|
||||
ONNX_CPU_OPERATOR_KERNEL(
|
||||
Einsum,
|
||||
12,
|
||||
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllNumericTensorTypes()),
|
||||
Einsum);
|
||||
|
||||
Status Einsum::Compute(OpKernelContext* context) const {
|
||||
int num_inputs = context->InputCount();
|
||||
if (num_inputs == 0) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Einsum op: There must be atleast one input");
|
||||
}
|
||||
|
||||
std::vector<const Tensor*> inputs;
|
||||
inputs.reserve(num_inputs);
|
||||
for (int i = 0; i < num_inputs; ++i) {
|
||||
inputs.push_back(context->Input<Tensor>(i));
|
||||
}
|
||||
|
||||
// Get temp space allocator - we will use this to allocate memory for intermediate tensors
|
||||
AllocatorPtr allocator;
|
||||
auto status = context->GetTempSpaceAllocator(&allocator);
|
||||
if (!status.IsOK()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, RUNTIME_EXCEPTION,
|
||||
"There was a problem acquiring temporary memory allocator in Einsum op");
|
||||
}
|
||||
|
||||
// Instantiate EinsumComputePreprocessor
|
||||
auto einsum_compute_preprocessor = EinsumComputePreprocessor(*einsum_equation_preprocessor_, inputs, allocator);
|
||||
|
||||
// Compute all required metadata to be used at Einsum compute time and return error status code if one was generated
|
||||
ORT_RETURN_IF_ERROR(einsum_compute_preprocessor.Run());
|
||||
|
||||
if (inputs[0]->IsDataType<float>()) {
|
||||
return EinsumTypedComputeProcessor<float>(context, allocator, einsum_compute_preprocessor);
|
||||
} else if (inputs[0]->IsDataType<int32_t>()) {
|
||||
return EinsumTypedComputeProcessor<int32_t>(context, allocator, einsum_compute_preprocessor);
|
||||
} else if (inputs[0]->IsDataType<double>()) {
|
||||
return EinsumTypedComputeProcessor<double>(context, allocator, einsum_compute_preprocessor);
|
||||
} else if (inputs[0]->IsDataType<int64_t>()) {
|
||||
return EinsumTypedComputeProcessor<int64_t>(context, allocator, einsum_compute_preprocessor);
|
||||
}
|
||||
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED,
|
||||
"Einsum op: An implementation for the input type ",
|
||||
inputs[0]->DataType(), " is not supported yet");
|
||||
}
|
||||
|
||||
} // namespace onnxruntime
|
||||
26
onnxruntime/core/providers/cpu/math/einsum.h
Normal file
26
onnxruntime/core/providers/cpu/math/einsum.h
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "core/common/common.h"
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/providers/cpu/math/einsum_utils.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
class Einsum final : public OpKernel {
|
||||
public:
|
||||
Einsum(const OpKernelInfo& info) : OpKernel(info) {
|
||||
ORT_ENFORCE(info.GetAttr<std::string>("equation", &equation_).IsOK(), "Missing 'equation' attribute");
|
||||
einsum_equation_preprocessor_ = onnxruntime::make_unique<EinsumEquationPreprocessor>(equation_);
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* context) const override;
|
||||
|
||||
private:
|
||||
std::string equation_;
|
||||
std::unique_ptr<EinsumEquationPreprocessor> einsum_equation_preprocessor_;
|
||||
};
|
||||
|
||||
} // namespace onnxruntime
|
||||
305
onnxruntime/core/providers/cpu/math/einsum_auxiliary_ops.cc
Normal file
305
onnxruntime/core/providers/cpu/math/einsum_auxiliary_ops.cc
Normal file
|
|
@ -0,0 +1,305 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "einsum_auxiliary_ops.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
namespace EinsumOp {
|
||||
|
||||
std::unique_ptr<Tensor> Transpose(const Tensor& input, const std::vector<int64_t>& input_shape_override,
|
||||
const std::vector<size_t>& permutation, AllocatorPtr allocator) {
|
||||
auto input_rank = input_shape_override.size();
|
||||
ORT_ENFORCE(input_rank == permutation.size(), "Length of permutation must match the rank of the input to be permutated");
|
||||
|
||||
std::vector<int64_t> output_dims;
|
||||
output_dims.reserve(input_rank);
|
||||
|
||||
for (const auto& dim : permutation) {
|
||||
output_dims.push_back(input_shape_override.at(dim));
|
||||
}
|
||||
|
||||
// Pass in allocator as that will be used as an allocator deleter by the framework
|
||||
// and it will de-allocate the memory for this intermediate tensor when it goes out of scope
|
||||
std::unique_ptr<Tensor> output = onnxruntime::make_unique<Tensor>(input.DataType(), output_dims, allocator);
|
||||
|
||||
TensorShape overriden_shape(input_shape_override);
|
||||
TransposeBase::DoTranspose(permutation, input, *output, &overriden_shape);
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::unique_ptr<Tensor> MatMul(const Tensor& input_1, const std::vector<int64_t>& input_shape_1_override,
|
||||
const Tensor& input_2, const std::vector<int64_t>& input_shape_2_override,
|
||||
AllocatorPtr allocator, concurrency::ThreadPool* tp) {
|
||||
// Sanity checks before the actual MatMul
|
||||
ORT_ENFORCE(input_1.DataType() == input_2.DataType(), "Data types of the inputs must match for MatMul");
|
||||
ORT_ENFORCE(input_shape_1_override.size() == 3 && input_shape_2_override.size() == 3, "Only 1 batch dimension is allowed for MatMul");
|
||||
ORT_ENFORCE(input_shape_1_override[0] == input_shape_2_override[0], "Batch dimension should match for MatMul;");
|
||||
ORT_ENFORCE(input_shape_1_override[2] == input_shape_2_override[1], "Incompatible matrix dimensions for matMul");
|
||||
|
||||
size_t batches = static_cast<size_t>(input_shape_1_override[0]);
|
||||
size_t M = static_cast<size_t>(input_shape_1_override[1]);
|
||||
size_t K = static_cast<size_t>(input_shape_1_override[2]);
|
||||
size_t N = static_cast<size_t>(input_shape_2_override[2]);
|
||||
|
||||
size_t left_offset = M * K;
|
||||
size_t right_offset = K * N;
|
||||
size_t output_offset = M * N;
|
||||
|
||||
std::vector<int64_t> output_dims;
|
||||
output_dims.reserve(3);
|
||||
output_dims.push_back(static_cast<int64_t>(batches));
|
||||
output_dims.push_back(static_cast<int64_t>(M));
|
||||
output_dims.push_back(static_cast<int64_t>(N));
|
||||
|
||||
// Pass in allocator as that will be used as an allocator deleter by the framework
|
||||
// and it will de-allocate the memory for this intermediate tensor when it goes out of scope
|
||||
std::unique_ptr<Tensor> output = onnxruntime::make_unique<Tensor>(input_1.DataType(), output_dims, allocator);
|
||||
|
||||
const T* input_1_data = input_1.template Data<T>();
|
||||
const T* input_2_data = input_2.template Data<T>();
|
||||
T* output_data = output->template MutableData<T>();
|
||||
|
||||
// Process each batch
|
||||
// TODO: Currently we parallelize a single MatMul operation, add logic to determine if
|
||||
// we can parallelizing on batches would be more optimal
|
||||
for (size_t i = 0; i < batches; ++i) {
|
||||
math::MatMul<T>(
|
||||
static_cast<int>(M),
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(K),
|
||||
input_1_data + i * left_offset,
|
||||
input_2_data + i * right_offset,
|
||||
output_data + i * output_offset, tp);
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::unique_ptr<Tensor> ReduceSum(const Tensor& input, const std::vector<int64_t>& input_shape_override,
|
||||
const std::vector<int64_t>& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp) {
|
||||
TensorShape overriden_shape(input_shape_override);
|
||||
auto output = onnxruntime::ReduceSum<T>::Impl(input, reduce_axes, allocator, tp, true, &overriden_shape);
|
||||
return onnxruntime::make_unique<Tensor>(std::move(output));
|
||||
}
|
||||
|
||||
// A specific helper just for the Diagonal op
|
||||
static inline bool IsTransposeRequiredForDiagonal(int64_t dim_1, int64_t dim_2, int64_t rank) {
|
||||
// If the input is 2D, we don't need a transpose
|
||||
if (rank == 2)
|
||||
return false;
|
||||
|
||||
// If the two dims are the innermost dims, no transpose is required
|
||||
if ((dim_1 == rank - 1 && dim_2 == rank - 2) ||
|
||||
(dim_1 == rank - 2 && dim_2 == rank - 1))
|
||||
return false;
|
||||
|
||||
// Transpose is required
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void DiagonalDataAssignment(const T* input_data, T* output_data, int64_t batch_size, int64_t base_stride, int64_t inner_stride) {
|
||||
int64_t output_iter = 0;
|
||||
// TODO: Parallelize this operation
|
||||
for (int64_t i = 0; i < batch_size; ++i) {
|
||||
auto base_offset = i * base_stride;
|
||||
for (int64_t j = 0; j < inner_stride; ++j) {
|
||||
output_data[output_iter] = input_data[base_offset + j * inner_stride + j];
|
||||
output_iter++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Parse diagonal elements along the 2 innermost dimensions
|
||||
// E.g.: input_shape = [1, 2, 3, 3]
|
||||
|
||||
// This implementation provides flexibility as to which of the 2 innermost dim values is preserved
|
||||
// via the `preserve_innermost_dim_val` parameter
|
||||
|
||||
// preserve_innermost_dim_val == true,
|
||||
// output_shape = [1, 2, 1, 3] => the diagonal contains 3 elements and the dim value of the innermost dim is preserved
|
||||
|
||||
// preserve_innermost_dim_val == false,
|
||||
// output_shape = [1, 2, 3, 1] => the diagonal contains 3 elements and the dim value of the non-innermost dim is preserved
|
||||
|
||||
static std::unique_ptr<Tensor> DiagonalInnermostDims(const Tensor& input,
|
||||
bool preserve_innermost_dim_val, AllocatorPtr allocator) {
|
||||
const auto& input_dims = input.Shape().GetDims();
|
||||
auto rank = input_dims.size();
|
||||
const size_t element_size_in_bytes = input.DataType()->Size();
|
||||
|
||||
// This is an internal method and we already have finished all validations in the calling method.
|
||||
// We proceed without duplicating all validations again here.
|
||||
|
||||
// We have a minimalistic check here to make sure the innermost dims have the same dim value
|
||||
// as the calling method may have done a transpose before calling this method
|
||||
ORT_ENFORCE(input_dims[rank - 2] == input_dims[rank - 1],
|
||||
"The innermost dims should have the same dim value to parse the diagonal elements");
|
||||
|
||||
std::vector<int64_t> output_dims;
|
||||
output_dims.reserve(rank);
|
||||
|
||||
int64_t batch_size = 1; // Flatten the outermost dims - this will be the number of iterations
|
||||
for (size_t i = 0; i < rank - 2; ++i) {
|
||||
auto input_dim_value = input_dims[i];
|
||||
batch_size *= input_dim_value;
|
||||
output_dims.push_back(input_dim_value);
|
||||
}
|
||||
|
||||
if (preserve_innermost_dim_val) {
|
||||
output_dims.push_back(1);
|
||||
output_dims.push_back(input_dims[rank - 1]);
|
||||
} else {
|
||||
output_dims.push_back(input_dims[rank - 1]);
|
||||
output_dims.push_back(1);
|
||||
}
|
||||
|
||||
int64_t inner_stride = input_dims[rank - 1]; // offset to move over the innermost dim
|
||||
int64_t base_stride = inner_stride * inner_stride; // offset to move over all the axes except the 2 innermost dims
|
||||
|
||||
// Pass in allocator as that will be used as an allocator deleter by the framework
|
||||
// and it will de-allocate the memory for this intermediate tensor when it goes out of scope
|
||||
std::unique_ptr<Tensor> output = onnxruntime::make_unique<Tensor>(input.DataType(), output_dims, allocator);
|
||||
|
||||
switch (element_size_in_bytes) {
|
||||
case 4:
|
||||
DiagonalDataAssignment<float>(reinterpret_cast<const float*>(input.DataRaw()), reinterpret_cast<float*>(output->MutableDataRaw()),
|
||||
batch_size, base_stride, inner_stride);
|
||||
break;
|
||||
case 8:
|
||||
DiagonalDataAssignment<double>(reinterpret_cast<const double*>(input.DataRaw()), reinterpret_cast<double*>(output->MutableDataRaw()),
|
||||
batch_size, base_stride, inner_stride);
|
||||
break;
|
||||
|
||||
default:
|
||||
ORT_THROW("Einsum op: Unsupported data type for Diagonal ", input.DataType());
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
std::unique_ptr<Tensor> Diagonal(const Tensor& input, int64_t dim_1, int64_t dim_2, AllocatorPtr allocator) {
|
||||
const auto& input_shape = input.Shape();
|
||||
const auto& input_dims = input_shape.GetDims();
|
||||
auto rank = static_cast<int64_t>(input_dims.size());
|
||||
|
||||
ORT_ENFORCE(rank >= 2 && dim_1 != dim_2 && input_dims[dim_1] == input_dims[dim_2],
|
||||
"Cannot parse the diagonal elements along dims ", dim_1, " and ", dim_2, " for input shape ", input_shape);
|
||||
|
||||
int64_t first_dim = -1; // first_dim holds the lesser of dim_1 and dim_2
|
||||
int64_t second_dim = -1; // second_dim holds the greater of dim_1 and dim_2
|
||||
if (dim_1 < dim_2) {
|
||||
first_dim = dim_1;
|
||||
second_dim = dim_2;
|
||||
} else {
|
||||
first_dim = dim_2;
|
||||
second_dim = dim_1;
|
||||
}
|
||||
|
||||
std::unique_ptr<Tensor> output;
|
||||
bool preserve_innermost_dim_val = false;
|
||||
|
||||
bool is_transpose_required = IsTransposeRequiredForDiagonal(dim_1, dim_2, rank);
|
||||
if (is_transpose_required) {
|
||||
std::vector<size_t> permutation(rank, 0);
|
||||
int64_t first_dim_axis = -1; // This is the axis eventually occupied by the first_dim
|
||||
|
||||
// If one of the diagonal dimensions is one of the 2 innermost dims, then leave it as such
|
||||
// so as to avoid transpose overhead
|
||||
if (first_dim == rank - 2) { // If rank - 2 is occupied by first_dim, keep it there
|
||||
permutation[rank - 2] = first_dim;
|
||||
first_dim_axis = rank - 2;
|
||||
} else {
|
||||
if (second_dim != rank - 2) { // If rank - 2 is not occupied by second_dim, then put first_dim there
|
||||
permutation[rank - 2] = first_dim;
|
||||
first_dim_axis = rank - 2;
|
||||
} else { // If rank - 2 is occupied by second_dim, then put first_dim in rank - 1
|
||||
permutation[rank - 1] = first_dim;
|
||||
first_dim_axis = rank - 1;
|
||||
preserve_innermost_dim_val = true; // We always want to preserve the dim value of the first_dim
|
||||
}
|
||||
}
|
||||
|
||||
// Put the second_dim in the dim not occupied by the first_dim
|
||||
if (first_dim_axis != rank - 1) {
|
||||
permutation[rank - 1] = second_dim;
|
||||
} else {
|
||||
permutation[rank - 2] = second_dim;
|
||||
}
|
||||
|
||||
int64_t iter = 0;
|
||||
for (int64_t i = 0; i < rank; ++i) {
|
||||
if (i != first_dim && i != second_dim) {
|
||||
permutation[iter++] = i;
|
||||
}
|
||||
}
|
||||
|
||||
// Permutate the input so that the dims from which we need the diagonal forms the innermost dims
|
||||
auto transposed = Transpose(input, input_dims, permutation, allocator);
|
||||
|
||||
// Parse the diagonal from the innermost dims
|
||||
output = DiagonalInnermostDims(*transposed, preserve_innermost_dim_val, allocator);
|
||||
|
||||
// Swap back the dimensions to the original axes ordering using a "reverse permutation"
|
||||
|
||||
// Find the "reverse" permutation
|
||||
iter = 0;
|
||||
std::vector<size_t> reverse_permutation(rank, 0);
|
||||
for (const auto& perm : permutation) {
|
||||
reverse_permutation[perm] = iter++;
|
||||
}
|
||||
|
||||
// Permutate using the reverse permutation to get back the original axes ordering
|
||||
output = Transpose(*output, output->Shape().GetDims(), reverse_permutation, allocator);
|
||||
} else {
|
||||
// No transposing required
|
||||
output = DiagonalInnermostDims(input, preserve_innermost_dim_val, allocator);
|
||||
}
|
||||
|
||||
// Make copy of the output dims
|
||||
auto output_dims = output->Shape().GetDims();
|
||||
|
||||
// Unsqueeze the reduced dim
|
||||
auto iter = output_dims.begin() + second_dim;
|
||||
output_dims.erase(iter);
|
||||
|
||||
output->Reshape(output_dims);
|
||||
return output;
|
||||
}
|
||||
|
||||
// Explicit template instantiation
|
||||
|
||||
// float
|
||||
template std::unique_ptr<Tensor> MatMul<float>(const Tensor& input_1, const std::vector<int64_t>& input_shape_1_override,
|
||||
const Tensor& input_2, const std::vector<int64_t>& input_shape_2_override,
|
||||
AllocatorPtr allocator, concurrency::ThreadPool* tp);
|
||||
template std::unique_ptr<Tensor> ReduceSum<float>(const Tensor& input, const std::vector<int64_t>& input_shape_override,
|
||||
const std::vector<int64_t>& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp);
|
||||
|
||||
// int32_t
|
||||
template std::unique_ptr<Tensor> MatMul<int32_t>(const Tensor& input_1, const std::vector<int64_t>& input_shape_1_override,
|
||||
const Tensor& input_2, const std::vector<int64_t>& input_shape_2_override,
|
||||
AllocatorPtr allocator, concurrency::ThreadPool* tp);
|
||||
template std::unique_ptr<Tensor> ReduceSum<int32_t>(const Tensor& input, const std::vector<int64_t>& input_shape_override,
|
||||
const std::vector<int64_t>& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp);
|
||||
|
||||
// double
|
||||
template std::unique_ptr<Tensor> MatMul<double>(const Tensor& input_1, const std::vector<int64_t>& input_shape_1_override,
|
||||
const Tensor& input_2, const std::vector<int64_t>& input_shape_2_override,
|
||||
AllocatorPtr allocator, concurrency::ThreadPool* tp);
|
||||
template std::unique_ptr<Tensor> ReduceSum<double>(const Tensor& input, const std::vector<int64_t>& input_shape_override,
|
||||
const std::vector<int64_t>& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp);
|
||||
|
||||
// int64_t
|
||||
template std::unique_ptr<Tensor> MatMul<int64_t>(const Tensor& input_1, const std::vector<int64_t>& input_shape_1_override,
|
||||
const Tensor& input_2, const std::vector<int64_t>& input_shape_2_override,
|
||||
AllocatorPtr allocator, concurrency::ThreadPool* tp);
|
||||
template std::unique_ptr<Tensor> ReduceSum<int64_t>(const Tensor& input, const std::vector<int64_t>& input_shape_override,
|
||||
const std::vector<int64_t>& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp);
|
||||
|
||||
} // namespace EinsumOp
|
||||
} // namespace onnxruntime
|
||||
50
onnxruntime/core/providers/cpu/math/einsum_auxiliary_ops.h
Normal file
50
onnxruntime/core/providers/cpu/math/einsum_auxiliary_ops.h
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
// This module hosts implementations and thin wrappers over other onnx operator implementations
|
||||
// that will be called from within the Einsum operator implementation
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "core/common/common.h"
|
||||
#include "core/framework/allocator.h"
|
||||
#include "core/util/math.h"
|
||||
#include "core/providers/cpu/tensor/transpose.h"
|
||||
#include "core/providers/cpu/reduction/reduction_ops.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
namespace EinsumOp {
|
||||
|
||||
// Thin wrapper over the Transpose op
|
||||
std::unique_ptr<Tensor> Transpose(const Tensor& input, const std::vector<int64_t>& input_shape_override,
|
||||
const std::vector<size_t>& permutation, AllocatorPtr allocator);
|
||||
|
||||
// Thin wrapper over the MatMul op
|
||||
// Not using the MatMulHelper to compute output dims as it adds a lot of checking overhead involving transposes of the inputs
|
||||
// In our case, we have a more simplistic version which doesn't need to have those checks
|
||||
template <typename T>
|
||||
std::unique_ptr<Tensor> MatMul(const Tensor& input_1, const std::vector<int64_t>& input_1_shape_override,
|
||||
const Tensor& input_2, const std::vector<int64_t>& input_2_shape_override,
|
||||
AllocatorPtr allocator, concurrency::ThreadPool* tp);
|
||||
|
||||
// Thin wrapper over the ReduceSum op
|
||||
template <typename T>
|
||||
std::unique_ptr<Tensor> ReduceSum(const Tensor& input, const std::vector<int64_t>& input_shape_override,
|
||||
const std::vector<int64_t>& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp);
|
||||
|
||||
// Diagonal - A specialized implementation somewhat similar to Torch's Diagonal op
|
||||
// but is specific enough to what is just required for the Einsum op.
|
||||
// Expects the input to be atleast 2-D and 0 <= dim_1, dim_2 < rank.
|
||||
// input_shape[dim_1] == input_shape[dim_2] and dim_1 cannot be same as dim_2.
|
||||
// The rank of the output is 1 less than the rank of the input and the squeezed dim is the greater of dim_1 and dim_2.
|
||||
|
||||
// Eg. input_shape = [2, 3, 5, 3] and dim_1 = 1 and dim_2 = 3
|
||||
// The output_shape will be [2, 3, 5] and dim_1 will contain the diagonal elements
|
||||
std::unique_ptr<Tensor> Diagonal(const Tensor& input, int64_t dim_1, int64_t dim_2, AllocatorPtr allocator);
|
||||
|
||||
} // namespace EinsumOp
|
||||
|
||||
} // namespace onnxruntime
|
||||
846
onnxruntime/core/providers/cpu/math/einsum_utils.cc
Normal file
846
onnxruntime/core/providers/cpu/math/einsum_utils.cc
Normal file
|
|
@ -0,0 +1,846 @@
|
|||
#include "einsum_utils.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
namespace EinsumOp {
|
||||
|
||||
// This helps decide if we need to apply (and pay the cost) of a Transpose
|
||||
static bool IsTransposeRequired(size_t input_rank, const std::vector<size_t>& permutation) {
|
||||
ORT_ENFORCE(input_rank == permutation.size(), "The rank of the input must match permutation size for Transpose");
|
||||
|
||||
// No transpose required for scalars
|
||||
if (input_rank == 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Weeds out cases where permutation is something like [0, 1, 2] for a 3D input and so on
|
||||
bool transpose_required = false;
|
||||
for (size_t i = 0; i < input_rank; ++i) {
|
||||
if (permutation[i] != i) {
|
||||
transpose_required = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return transpose_required;
|
||||
}
|
||||
// We have an the result in an output "candidate". Now we have to copy the contents in its buffer
|
||||
// into the buffer of the actual output given to us by the execution frame
|
||||
// We need to do this because the buffer owned by the output tensor of the op could be user provided buffer
|
||||
static void CopyOutputCandidateIntoOpOutout(Tensor& output, const Tensor& candidate) {
|
||||
ORT_ENFORCE(output.SizeInBytes() == candidate.SizeInBytes(),
|
||||
"Einsum op: The candidate output does not match the actual output's shape");
|
||||
// There are no string tensors - so safely use memcpy
|
||||
memcpy(output.MutableDataRaw(), candidate.DataRaw(), candidate.SizeInBytes());
|
||||
}
|
||||
// Here we take a "candidate output"(candidate output is a tensor that is a permutation and / or a reshape away from the final output),
|
||||
// and after a few operations to get it to the required output structure, copy it to the op's output
|
||||
// The candidate output might contain dims that may not be part of the op's output (i.e.) the dims will have to be unsqueezed
|
||||
template <typename T>
|
||||
static void FinalizeOutput(const Tensor& candidate_output, const std::vector<int64_t>& ordered_subscript_indices_in_candidate,
|
||||
const std::vector<int64_t>& subscript_indices_to_output_indices,
|
||||
Tensor& output, const TensorShape& output_shape, const AllocatorPtr& allocator) {
|
||||
ORT_ENFORCE(candidate_output.Shape().Size() == output_shape.Size(),
|
||||
"Einsum op: The candidate output cannot be reshaped into the op's output");
|
||||
|
||||
const auto& output_dims = output_shape.GetDims();
|
||||
const auto output_rank = output_dims.size();
|
||||
|
||||
const auto& candidate_output_dims = candidate_output.Shape().GetDims();
|
||||
const auto candidate_output_rank = candidate_output_dims.size();
|
||||
|
||||
// This vector holds the shape of the candidate_output after removing the dims that have
|
||||
// been reduced in the final output
|
||||
std::vector<int64_t> candidate_output_shape_without_reduced_dims;
|
||||
candidate_output_shape_without_reduced_dims.reserve(candidate_output_rank); // reserve upper bound
|
||||
|
||||
// Identify the permutation required by the op's output
|
||||
std::vector<size_t> output_permutation;
|
||||
output_permutation.resize(output_rank, 0);
|
||||
size_t output_iter = 0;
|
||||
|
||||
for (size_t iter = 0; iter < ordered_subscript_indices_in_candidate.size(); ++iter) {
|
||||
auto output_index = subscript_indices_to_output_indices[ordered_subscript_indices_in_candidate[iter]];
|
||||
|
||||
// If output_index is -1, then this dimension does not show up in the op's output and has been reduced along the way
|
||||
if (output_index != -1) {
|
||||
output_permutation[output_index] = output_iter++;
|
||||
candidate_output_shape_without_reduced_dims.push_back(candidate_output_dims[iter]);
|
||||
} else {
|
||||
// This dim doesn't show up in the op's output and hence we check if the dim has been reduced in the candidate output
|
||||
ORT_ENFORCE(candidate_output_dims[iter] == 1, "Not all dimensions to be reduced have been reduced in the candidate output");
|
||||
}
|
||||
}
|
||||
|
||||
// Transpose to the required final output order
|
||||
// (Identify no-op transposes and prevent triggering the transpose)
|
||||
if (IsTransposeRequired(candidate_output_shape_without_reduced_dims.size(), output_permutation)) {
|
||||
auto candidate_output_transposed = Transpose(candidate_output, candidate_output_shape_without_reduced_dims, output_permutation, allocator);
|
||||
CopyOutputCandidateIntoOpOutout(output, *candidate_output_transposed);
|
||||
} else {
|
||||
// Copy the output candidate into the op's output
|
||||
CopyOutputCandidateIntoOpOutout(output, candidate_output);
|
||||
}
|
||||
}
|
||||
|
||||
// Processes Einsum operands in a pair-wise fashion
|
||||
// Employs Transpose, ReduceSum, and MatMul under the hood
|
||||
// to achieve MatMul(a, b) and reduces (by summing) along specified axes
|
||||
template <typename T>
|
||||
static std::unique_ptr<Tensor> PairwiseOperandProcess(const Tensor& left,
|
||||
const TensorShape& left_shape_override,
|
||||
const Tensor& right,
|
||||
const TensorShape& right_shape_override,
|
||||
const std::vector<int64_t>& reduce_dims,
|
||||
concurrency::ThreadPool* tp,
|
||||
const AllocatorPtr& allocator,
|
||||
const EinsumComputePreprocessor& einsum_compute_preprocessor,
|
||||
bool is_final_pair, Tensor& final_output) {
|
||||
// Use the provided dim overrides instead of the actual shapes of the operands
|
||||
ORT_ENFORCE(left.Shape().Size() == left_shape_override.Size(), "The override dims are not compatible with given tensor's shape");
|
||||
ORT_ENFORCE(right.Shape().Size() == right_shape_override.Size(), "The override dims are not compatible with given tensor's shape");
|
||||
|
||||
// Make copy as this may be overridden downstream
|
||||
const auto& left_dims = left_shape_override.GetDims();
|
||||
const auto& right_dims = right_shape_override.GetDims();
|
||||
|
||||
int64_t left_rank = static_cast<int64_t>(left_dims.size());
|
||||
int64_t right_rank = static_cast<int64_t>(right_dims.size());
|
||||
|
||||
std::unique_ptr<Tensor> current_left;
|
||||
std::unique_ptr<Tensor> current_right;
|
||||
|
||||
// If the following error condition is hit, it is most likely a pre-processing bug
|
||||
ORT_ENFORCE(left_rank == right_rank, "Ranks of pair-wise operands must be equal");
|
||||
|
||||
// Following vectors hold:
|
||||
// lro: dim indices that are present in left, right, and reduce_dims
|
||||
// lo: dim indices that are present in left and reduce_dims
|
||||
// ro: dim indices that are present in right and reduce_dims
|
||||
std::vector<size_t> lro;
|
||||
std::vector<size_t> lo;
|
||||
std::vector<size_t> ro;
|
||||
|
||||
// Maintain sizes to create reshaped "views"
|
||||
int64_t lro_size = 1;
|
||||
int64_t lo_size = 1;
|
||||
int64_t ro_size = 1;
|
||||
int64_t reduced_size = 1;
|
||||
|
||||
size_t reduce_dims_iter = 0;
|
||||
size_t reduce_dims_size = reduce_dims.size();
|
||||
|
||||
for (int64_t i = 0; i < left_rank; ++i) {
|
||||
int64_t left_dim = left_dims[i];
|
||||
int64_t right_dim = right_dims[i];
|
||||
|
||||
bool has_left_dim = left_dim > 1; // non-trivial dimension (dim_value != 1)
|
||||
bool has_right_dim = right_dim > 1; // non-trivial dimension (dim_value != 1)
|
||||
|
||||
if (reduce_dims_iter < reduce_dims_size && reduce_dims[reduce_dims_iter] == i) { // This dimension is to be reduced after this pair-wise operation
|
||||
++reduce_dims_iter;
|
||||
if (has_left_dim && has_right_dim) { // Both inputs have non-trivial dim values along this dimension
|
||||
// Both the left and right operands have non-trivial dimension value along this axis
|
||||
// They must be equal
|
||||
ORT_ENFORCE(left_dim == right_dim,
|
||||
"Einsum op: Input dimensions must be equal along an axis to be reduced across all inputs");
|
||||
reduced_size *= left_dim;
|
||||
} else if (has_left_dim) { // if it is only in one of left and right, we can reduce right away
|
||||
current_left = ReduceSum<T>(left, left_dims, {i}, allocator, tp);
|
||||
} else if (has_right_dim) {
|
||||
current_right = ReduceSum<T>(right, right_dims, {i}, allocator, tp);
|
||||
}
|
||||
} else { // This dimension is not reduced (i.e.) it appears in the output after processing these 2 operands
|
||||
// Both the left and right operands have non-trivial dimension value along this axis
|
||||
// They must be equal
|
||||
if (has_left_dim && has_right_dim) {
|
||||
ORT_ENFORCE(left_dim == right_dim, "Einsum op: Input shapes do not align");
|
||||
lro.push_back(i);
|
||||
lro_size *= left_dim;
|
||||
} else if (has_left_dim) {
|
||||
// The left operand has non-trivial dimension value
|
||||
lo.push_back(i);
|
||||
lo_size *= left_dim;
|
||||
} else {
|
||||
// The right operand may or may not have non-trivial dim value
|
||||
// If it has trivial dim value (1),
|
||||
// it will just form a trailing dimension for the right operand
|
||||
ro.push_back(i);
|
||||
ro_size *= right_dim;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Permutate the left operand so that the axes order go like this: [lro, lo, reduce_dims, ro]
|
||||
std::vector<size_t> left_permutation;
|
||||
left_permutation.reserve(lro.size() + lo.size() + reduce_dims.size() + ro.size());
|
||||
left_permutation.insert(left_permutation.end(), lro.begin(), lro.end());
|
||||
left_permutation.insert(left_permutation.end(), lo.begin(), lo.end());
|
||||
left_permutation.insert(left_permutation.end(), reduce_dims.begin(), reduce_dims.end());
|
||||
left_permutation.insert(left_permutation.end(), ro.begin(), ro.end());
|
||||
if (IsTransposeRequired(current_left ? current_left->Shape().GetDims().size() : left_dims.size(),
|
||||
left_permutation)) {
|
||||
current_left = Transpose(current_left ? *current_left : left,
|
||||
current_left ? current_left->Shape().GetDims() : left_dims,
|
||||
left_permutation, allocator);
|
||||
}
|
||||
|
||||
// Permutate the right operand so that the axes order go like this: [lro, reduce_dims, ro, lo]
|
||||
std::vector<size_t> right_permutation;
|
||||
right_permutation.reserve(lro.size() + lo.size() + reduce_dims.size() + ro.size());
|
||||
right_permutation.insert(right_permutation.end(), lro.begin(), lro.end());
|
||||
right_permutation.insert(right_permutation.end(), reduce_dims.begin(), reduce_dims.end());
|
||||
right_permutation.insert(right_permutation.end(), ro.begin(), ro.end());
|
||||
right_permutation.insert(right_permutation.end(), lo.begin(), lo.end());
|
||||
if (IsTransposeRequired(current_right ? current_right->Shape().GetDims().size() : right_dims.size(),
|
||||
right_permutation)) {
|
||||
current_right = Transpose(current_right ? *current_right : right,
|
||||
current_right ? current_right->Shape().GetDims() : right_dims,
|
||||
right_permutation, allocator);
|
||||
}
|
||||
|
||||
// Calculate output size
|
||||
// Output shape will be determined by rules of MatMul:
|
||||
// because we are multiplying two tensors of shapes [lro, lo, reduce_dims] , [lro, reduce_dims, ro]
|
||||
// [dim_value of `lro` dims,
|
||||
// dim_value of `lo` dims,
|
||||
// `1` for each of the `reduce_dims`,
|
||||
// dim_value of `ro` dims]
|
||||
std::vector<int64_t> output_dims;
|
||||
output_dims.reserve(lro.size() + lo.size() + reduce_dims.size() + ro.size());
|
||||
for (size_t i = 0; i < lro.size(); ++i) {
|
||||
output_dims.push_back(left_dims[lro[i]]);
|
||||
}
|
||||
for (size_t i = 0; i < lo.size(); ++i) {
|
||||
output_dims.push_back(left_dims[lo[i]]);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < reduce_dims.size(); ++i) {
|
||||
output_dims.push_back(1); // reduced dimensions will have a value 1 in it
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < ro.size(); ++i) {
|
||||
output_dims.push_back(right_dims[ro[i]]);
|
||||
}
|
||||
|
||||
std::vector<int64_t> current_subscript_order;
|
||||
|
||||
// Calculate output permutation
|
||||
// After the MatMul op, the because the two operands have been permutated,
|
||||
// the output is permutated as well with respect to the original ordering of the axes.
|
||||
// The permutated order will be the dims in: [lro, lo, reduced_dims, ro]
|
||||
// Hence invert the permutation by a permutation that puts the axes in the same ordering
|
||||
std::vector<size_t> output_permutation;
|
||||
if (!is_final_pair) { // If this is not the final pair, we need to permutate the result to match the pre-fixed order for the next iteration
|
||||
output_permutation.resize(lro.size() + lo.size() + reduce_dims.size() + ro.size(), 0);
|
||||
size_t iter = 0;
|
||||
for (size_t i = 0; i < lro.size(); ++i) {
|
||||
output_permutation[lro[i]] = iter++;
|
||||
}
|
||||
for (size_t i = 0; i < lo.size(); ++i) {
|
||||
output_permutation[lo[i]] = iter++;
|
||||
}
|
||||
for (size_t i = 0; i < reduce_dims.size(); ++i) {
|
||||
output_permutation[reduce_dims[i]] = iter++;
|
||||
}
|
||||
for (size_t i = 0; i < ro.size(); ++i) {
|
||||
output_permutation[ro[i]] = iter++;
|
||||
}
|
||||
} else {
|
||||
current_subscript_order.reserve(lro.size() + lo.size() + reduce_dims.size() + ro.size());
|
||||
current_subscript_order.insert(current_subscript_order.end(), lro.begin(), lro.end());
|
||||
current_subscript_order.insert(current_subscript_order.end(), lo.begin(), lo.end());
|
||||
current_subscript_order.insert(current_subscript_order.end(), reduce_dims.begin(), reduce_dims.end());
|
||||
current_subscript_order.insert(current_subscript_order.end(), ro.begin(), ro.end());
|
||||
}
|
||||
|
||||
// Multiply the mutated inputs
|
||||
auto output = MatMul<T>(current_left ? *current_left : left, {lro_size, lo_size, reduced_size},
|
||||
current_right ? *current_right : right, {lro_size, reduced_size, ro_size},
|
||||
allocator, tp);
|
||||
|
||||
output->Reshape(output_dims);
|
||||
|
||||
if (!is_final_pair) { // This is not the final pair - so bring the axes order to what the inputs conformed to
|
||||
if (IsTransposeRequired(output_dims.size(), output_permutation)) {
|
||||
output = Transpose(*output, output_dims, output_permutation, allocator);
|
||||
}
|
||||
} else { // This is the final pair - Transpose directly to the output ordering required and copy the contents to the op's output
|
||||
FinalizeOutput<T>(*output, current_subscript_order,
|
||||
einsum_compute_preprocessor.GetMappedSubscriptIndicesToOutputindices(), final_output,
|
||||
einsum_compute_preprocessor.GetOutputDims(), allocator);
|
||||
}
|
||||
|
||||
return std::move(output);
|
||||
}
|
||||
|
||||
} // namespace EinsumOp
|
||||
|
||||
EinsumComputePreprocessor::EinsumComputePreprocessor(EinsumEquationPreprocessor& einsum_equation_preprocessor,
|
||||
const std::vector<const Tensor*>& inputs,
|
||||
AllocatorPtr allocator)
|
||||
: einsum_equation_preprocessor_(einsum_equation_preprocessor), inputs_(inputs), allocator_(allocator) {
|
||||
letter_to_index_.fill(-1);
|
||||
|
||||
letter_to_count_.fill(0);
|
||||
}
|
||||
|
||||
Status EinsumComputePreprocessor::Run() {
|
||||
ORT_RETURN_IF_ERROR(ProcessSubscripts());
|
||||
|
||||
ORT_RETURN_IF_ERROR(PostProcessBroadcastedDims());
|
||||
|
||||
ORT_RETURN_IF_ERROR(ParseOrCreateOutputSubscript());
|
||||
|
||||
ORT_RETURN_IF_ERROR(CalculateOutputShape());
|
||||
|
||||
ORT_RETURN_IF_ERROR(PreprocessInputs());
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
const std::vector<int64_t>& EinsumComputePreprocessor::GetOutputDims() const {
|
||||
return output_dims_;
|
||||
}
|
||||
|
||||
std::vector<std::unique_ptr<Tensor>>& EinsumComputePreprocessor::GetPreprocessedInputTensors() {
|
||||
return preprocessed_inputs_;
|
||||
}
|
||||
|
||||
const std::vector<const Tensor*>& EinsumComputePreprocessor::GetRawInputTensors() {
|
||||
return inputs_;
|
||||
}
|
||||
|
||||
const std::vector<TensorShape>& EinsumComputePreprocessor::GetHomogenizedInputDims() {
|
||||
return homogenized_input_dims_;
|
||||
}
|
||||
|
||||
const std::vector<int64_t>& EinsumComputePreprocessor::GetMappedSubscriptIndicesToLastInputIndex() const {
|
||||
return subscript_indices_to_last_input_;
|
||||
}
|
||||
|
||||
const std::vector<int64_t>& EinsumComputePreprocessor::GetMappedSubscriptIndicesToOutputindices() const {
|
||||
return subscript_indices_to_output_indices_;
|
||||
}
|
||||
|
||||
int64_t EinsumComputePreprocessor::GetNumSubscriptIndices() const {
|
||||
return num_subscript_indices_;
|
||||
}
|
||||
|
||||
Status EinsumComputePreprocessor::ProcessSubscripts() {
|
||||
const auto& left_equation_split = einsum_equation_preprocessor_.left_equation_split_;
|
||||
if (left_equation_split.size() != inputs_.size()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Number of subscripts in the input equation does not match number of input tensors");
|
||||
}
|
||||
|
||||
int64_t input_index = 0;
|
||||
|
||||
// Holds mapping between input indices to its corresponding subscript labels for each input
|
||||
input_subscript_indices_.reserve(inputs_.size());
|
||||
|
||||
// We arbitrarily reserve space for 10 values as we don't expect to see any input with rank >10
|
||||
// which would make num_subscript_indices_ > 10
|
||||
subscript_indices_to_last_input_.reserve(10);
|
||||
subscript_indices_to_dim_value_.reserve(10);
|
||||
|
||||
for (const auto& subscript : left_equation_split) {
|
||||
const auto& shape = inputs_[input_index]->Shape();
|
||||
const auto& dims = shape.GetDims();
|
||||
size_t rank = dims.size();
|
||||
size_t dim_counter = 0;
|
||||
|
||||
std::vector<int64_t> current_subscript_indices;
|
||||
current_subscript_indices.reserve(rank);
|
||||
|
||||
// Temp variables to deal with "ellipsis" in the input
|
||||
bool is_in_middle_of_ellipsis = false;
|
||||
int64_t ellipsis_char_count = 0;
|
||||
|
||||
// Iterate through all subscript labels in the subscript
|
||||
for (auto subscript_label : subscript) {
|
||||
// Broadcasting based dims
|
||||
if (subscript_label == '.') {
|
||||
is_in_middle_of_ellipsis = true;
|
||||
// Make sure there aren't more than 3 '.'s in the current subscript
|
||||
if (++ellipsis_char_count > 3) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Found a '.' not part of an ellipsis in input: ", input_index);
|
||||
}
|
||||
|
||||
// We have seen all 3 '.'s. We can safely process the ellipsis now.
|
||||
if (ellipsis_char_count == 3) {
|
||||
is_in_middle_of_ellipsis = false;
|
||||
|
||||
// Example for the following line of code
|
||||
// Subscript "...ij" for an input of rank 6
|
||||
// num_of_ellipsis_dims = 6 - 5 + 3 = 4
|
||||
int64_t current_num_of_ellipsis_dims = rank - subscript.length() + 3;
|
||||
if (current_num_of_ellipsis_dims < 0) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Einsum subscripts string contains too many subscript labels when compared to the rank of the input");
|
||||
}
|
||||
|
||||
// Theoretically, current_num_of_ellipsis_dims could be 0
|
||||
// Example: For an input of rank 2 paired with a subscript "...ij"
|
||||
if (current_num_of_ellipsis_dims != 0) {
|
||||
// We have seen a ellipsis before - make sure ranks align as per the ONNX spec -
|
||||
// "Ellipsis must indicate a fixed number of dimensions."
|
||||
if (num_of_ellipsis_dims_ != 0) {
|
||||
if (num_of_ellipsis_dims_ != static_cast<size_t>(current_num_of_ellipsis_dims)) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Ellipsis must indicate a fixed number of dimensions across all inputs");
|
||||
}
|
||||
} else {
|
||||
num_of_ellipsis_dims_ = static_cast<size_t>(current_num_of_ellipsis_dims);
|
||||
}
|
||||
|
||||
// We reserve 'EinsumOp::num_of_letters' for broadcasted dims as we only allow 'a' - 'z' (0 - 25) for non-broadcasted dims
|
||||
// We will assign appropriate indices (based on number of dimensions the ellipsis corresponds to) during broadcasting related post-processing
|
||||
for (size_t i = 0; i < num_of_ellipsis_dims_; ++i) {
|
||||
current_subscript_indices.push_back(EinsumOp::num_of_letters);
|
||||
}
|
||||
|
||||
// Offset 'dim_counter' by number of dimensions the ellipsis corresponds to
|
||||
dim_counter += num_of_ellipsis_dims_;
|
||||
}
|
||||
}
|
||||
} else { // regular letter based dimension -> 'i', 'j', etc.
|
||||
if (is_in_middle_of_ellipsis) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Found '.' not part of an ellipsis in input: ", input_index);
|
||||
}
|
||||
|
||||
if (!(subscript_label >= 'a' && subscript_label <= 'z')) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"The only subscript labels allowed are lowercase letters (a-z)");
|
||||
}
|
||||
|
||||
auto letter_index = static_cast<int64_t>(subscript_label - 'a');
|
||||
auto dim_value = dims[dim_counter];
|
||||
|
||||
// Subscript label not found in global subscript label array
|
||||
// Hence add it to both local and global subscript arrays
|
||||
if (letter_to_count_[letter_index] == 0) {
|
||||
letter_to_index_[letter_index] = num_subscript_indices_++;
|
||||
subscript_indices_to_dim_value_.push_back(dim_value);
|
||||
subscript_indices_to_last_input_.push_back(input_index);
|
||||
} else { // This subscript label has been seen in atleast one other operand's subscript
|
||||
// It must be equal unless one of them is a 1 (Numpy allows this)
|
||||
auto mapped_index = letter_to_index_[letter_index];
|
||||
|
||||
subscript_indices_to_last_input_[mapped_index] = input_index;
|
||||
|
||||
if (subscript_indices_to_dim_value_[mapped_index] != dim_value) {
|
||||
// Set the value to the new dim value if the value is 1 in the map
|
||||
if (subscript_indices_to_dim_value_[mapped_index] == 1) {
|
||||
subscript_indices_to_dim_value_[mapped_index] = dim_value;
|
||||
} else {
|
||||
if (dim_value != 1) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Einsum operands could not be broadcast together. "
|
||||
"Please check input shapes/equation provided."
|
||||
"Input shape of operand ",
|
||||
input_index, " is incompatible in the dimension ", dim_counter,
|
||||
". The shape is: ", shape,
|
||||
"Another operand has a dim value of ", subscript_indices_to_dim_value_[mapped_index],
|
||||
" in the same dimension");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
++letter_to_count_[letter_index];
|
||||
|
||||
current_subscript_indices.push_back(letter_to_index_[letter_index]);
|
||||
if (++dim_counter > rank) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Einsum subscripts string contains too many subscript labels when compared to the rank of the input ",
|
||||
input_index);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If no broadcasting is requested, the number of subscript labels (dim_counter) should match input rank
|
||||
if (num_of_ellipsis_dims_ == 0) {
|
||||
if (dim_counter != rank) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Einsum subscripts does not contain enough subscript labels and there is no ellipsis for input ",
|
||||
input_index);
|
||||
}
|
||||
}
|
||||
|
||||
input_subscript_indices_.push_back(std::move(current_subscript_indices));
|
||||
++input_index;
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status EinsumComputePreprocessor::PostProcessBroadcastedDims() {
|
||||
// Pay the cost of this function only if we saw an ellipsis in any of the inputs
|
||||
if (num_of_ellipsis_dims_ > 0) {
|
||||
// extend the number of subscript labels to include each ellipsis dim as
|
||||
// theoretically each ellipsis dim does correspond to a "virtual" subscript label
|
||||
num_subscript_indices_ += num_of_ellipsis_dims_;
|
||||
|
||||
// We are going to assign the broadcasted dims outermost subscript indices (i.e.) 0 -> num_of_ellipsis_dims_ - 1
|
||||
// as most likely bradcasted dims will be batch dimensions (i.e.) outermost dimensions and hence we don't have to pay
|
||||
// transposing while "homogenizing" the input
|
||||
|
||||
// Hence offset all subscript indices by num_of_ellipsis_dims_
|
||||
for (size_t i = 0; i < EinsumOp::num_of_letters; ++i) {
|
||||
if (letter_to_index_[i] != -1) {
|
||||
letter_to_index_[i] += num_of_ellipsis_dims_;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int64_t> temp_index_to_last_input(num_subscript_indices_, -1);
|
||||
for (size_t i = 0; i < subscript_indices_to_last_input_.size(); ++i) {
|
||||
temp_index_to_last_input[i + num_of_ellipsis_dims_] = subscript_indices_to_last_input_[i];
|
||||
}
|
||||
subscript_indices_to_last_input_ = std::move(temp_index_to_last_input);
|
||||
|
||||
std::vector<int64_t> temp_index_to_dim_value(num_subscript_indices_, -1);
|
||||
for (size_t i = 0; i < subscript_indices_to_dim_value_.size(); ++i) {
|
||||
temp_index_to_dim_value[i + num_of_ellipsis_dims_] = subscript_indices_to_dim_value_[i];
|
||||
}
|
||||
subscript_indices_to_dim_value_ = std::move(temp_index_to_dim_value);
|
||||
|
||||
for (size_t i = 0; i < input_subscript_indices_.size(); ++i) {
|
||||
auto& current_input_dim_indices_to_subscript_indices = input_subscript_indices_[i];
|
||||
std::vector<int64_t> temp_current_input_dim_indices_to_subscript_indices;
|
||||
temp_current_input_dim_indices_to_subscript_indices.reserve(current_input_dim_indices_to_subscript_indices.size());
|
||||
|
||||
const auto& dims = inputs_[i]->Shape().GetDims();
|
||||
auto rank = dims.size();
|
||||
|
||||
size_t dim_iter = 0;
|
||||
size_t num_broadcasted_indices = 0;
|
||||
while (dim_iter < current_input_dim_indices_to_subscript_indices.size()) {
|
||||
auto value = current_input_dim_indices_to_subscript_indices[dim_iter];
|
||||
if (value == EinsumOp::num_of_letters) { //This is a broadcasted dim
|
||||
// Shouldn't hit this error - just a sanity check
|
||||
ORT_ENFORCE(num_broadcasted_indices < num_of_ellipsis_dims_);
|
||||
temp_current_input_dim_indices_to_subscript_indices.push_back(static_cast<int64_t>(num_broadcasted_indices));
|
||||
subscript_indices_to_last_input_[num_broadcasted_indices] = i;
|
||||
|
||||
// This is the first time we are seeing this broadcasted dim
|
||||
if (subscript_indices_to_dim_value_[num_broadcasted_indices] == -1) {
|
||||
subscript_indices_to_dim_value_[num_broadcasted_indices] = dims[dim_iter];
|
||||
} else { // We have seen this broadcasted dim before
|
||||
// Check if the previous value is equal to the current value
|
||||
if (subscript_indices_to_dim_value_[num_broadcasted_indices] != dims[dim_iter]) {
|
||||
// If they are not equal, one of them needs to be 1
|
||||
if (subscript_indices_to_dim_value_[num_broadcasted_indices] == 1) {
|
||||
subscript_indices_to_dim_value_[num_broadcasted_indices] = dims[dim_iter];
|
||||
} else {
|
||||
if (dims[dim_iter] != 1) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"The broadcasted dimensions of the inputs are incompatible");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
++num_broadcasted_indices;
|
||||
} else { // This is a regular dim - offset it by number of broadcasted dims
|
||||
temp_current_input_dim_indices_to_subscript_indices.push_back(value + static_cast<int64_t>(num_of_ellipsis_dims_));
|
||||
}
|
||||
++dim_iter;
|
||||
}
|
||||
// Shouldn't hit this error - just a sanity check
|
||||
ORT_ENFORCE(dim_iter == rank);
|
||||
current_input_dim_indices_to_subscript_indices = std::move(temp_current_input_dim_indices_to_subscript_indices);
|
||||
}
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status EinsumComputePreprocessor::ParseOrCreateOutputSubscript() {
|
||||
// Explicit form - no op as the output would have been parsed while parsing the input
|
||||
if (einsum_equation_preprocessor_.is_explicit_) {
|
||||
// Make sure that the given explicit equation contains an ellipsis if the input contains ellipses in them
|
||||
if (num_of_ellipsis_dims_ > 0) {
|
||||
if (einsum_equation_preprocessor_.right_equation_.find("...") == std::string::npos) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Inputs have ellipses in them but the provided output subscript does not contain an ellipsis");
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
// Implicit form - construct the output subscript
|
||||
std::stringstream output_equation;
|
||||
|
||||
// If the an ellipsis was seen in the input, add it
|
||||
if (num_of_ellipsis_dims_ > 0) {
|
||||
output_equation << "...";
|
||||
}
|
||||
|
||||
// In sorted order of letters, add those letters that were seen only once in the input
|
||||
size_t iter = 0;
|
||||
for (const auto& count : letter_to_count_) {
|
||||
if (count == 1) {
|
||||
output_equation << static_cast<char>('a' + iter);
|
||||
}
|
||||
++iter;
|
||||
}
|
||||
|
||||
einsum_equation_preprocessor_.right_equation_ = output_equation.str();
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status EinsumComputePreprocessor::CalculateOutputShape() {
|
||||
// Iterate through all subscript labels in the output subscript
|
||||
bool is_in_middle_of_ellipsis = false;
|
||||
int64_t ellipsis_char_count = 0;
|
||||
|
||||
subscript_indices_to_output_indices_.resize(num_subscript_indices_, -1);
|
||||
|
||||
std::array<int64_t, EinsumOp::num_of_letters> output_letter_to_count;
|
||||
output_letter_to_count.fill(0);
|
||||
|
||||
// Arbitrarily reserve some space as we don't expect rank of output to be > 10 (pay re-allocation cost if it is)
|
||||
output_dims_.reserve(10);
|
||||
|
||||
int64_t output_dim_counter = 0;
|
||||
for (auto subscript_label : einsum_equation_preprocessor_.right_equation_) {
|
||||
if (subscript_label == '.') {
|
||||
is_in_middle_of_ellipsis = true;
|
||||
// Make sure there aren't more than 3 '.'s in the current subscript
|
||||
if (++ellipsis_char_count > 3) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Found a '.' not part of an ellipsis in the output subscript provided");
|
||||
}
|
||||
|
||||
if (ellipsis_char_count == 3) { // Ellipsis is complete. Process it.
|
||||
is_in_middle_of_ellipsis = false;
|
||||
for (size_t i = 0; i < num_of_ellipsis_dims_; ++i) {
|
||||
output_dims_.push_back(subscript_indices_to_dim_value_[i]);
|
||||
// The ellipsis is seen in the output and hence the corresponding dims are to not be reduced
|
||||
subscript_indices_to_last_input_[i] = -1;
|
||||
subscript_indices_to_output_indices_[i] = output_dim_counter++;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (is_in_middle_of_ellipsis) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Found '.' not part of an ellipsis in the output subscript provided");
|
||||
}
|
||||
|
||||
if (!(subscript_label >= 'a' && subscript_label <= 'z')) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"The only subscript labels allowed in the output subscript "
|
||||
"are lowercase letters (a-z)");
|
||||
}
|
||||
|
||||
auto letter_index = static_cast<int64_t>(subscript_label - 'a');
|
||||
|
||||
if (output_letter_to_count[letter_index] != 0) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Output subscript contains repeated letters");
|
||||
}
|
||||
++output_letter_to_count[letter_index];
|
||||
|
||||
auto mapped_index = letter_to_index_[letter_index];
|
||||
if (mapped_index == -1) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Output subscript contains letters not seen in the inputs");
|
||||
}
|
||||
|
||||
output_dims_.push_back(subscript_indices_to_dim_value_[mapped_index]);
|
||||
|
||||
// Reset the last input index for this subscript label
|
||||
// given that it is seen in the output and hence can't be reduced
|
||||
subscript_indices_to_last_input_[mapped_index] = -1;
|
||||
|
||||
subscript_indices_to_output_indices_[mapped_index] = output_dim_counter++;
|
||||
}
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status EinsumComputePreprocessor::PreprocessInputs() {
|
||||
preprocessed_inputs_.reserve(inputs_.size());
|
||||
homogenized_input_dims_.reserve(inputs_.size());
|
||||
// As part of input preprocessing we "homogenize" them by
|
||||
// 1) Making them all of the same rank
|
||||
// 2) The axes order in all the inputs are to be made the same
|
||||
int64_t input_iter = 0;
|
||||
for (const auto* input : inputs_) {
|
||||
// Eventually will hold the "preprocessed" version of the original input
|
||||
std::unique_ptr<Tensor> preprocessed;
|
||||
|
||||
const auto& input_dims = input->Shape().GetDims();
|
||||
const auto& current_subscript_indices = input_subscript_indices_[input_iter];
|
||||
|
||||
// If all has gone well, we will have a subscript index (subscript label) for each dim of the input
|
||||
if (input_dims.size() != current_subscript_indices.size()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"Rank of the input must match number of subscript labels corresponding to the input");
|
||||
}
|
||||
|
||||
std::vector<int64_t> subscript_indices_to_input_index(num_subscript_indices_, -1);
|
||||
|
||||
// This is the input dims after re-ordering so that all inputs have same axes order
|
||||
std::vector<int64_t> homogenized_input_dims(num_subscript_indices_, 1);
|
||||
|
||||
// Preprocessed dim rank may not be the same as original input rank if we need to parse diagonals along the way
|
||||
// (which reduces rank in the preprocessed input by 1 for each diagonal we parse)
|
||||
int64_t dim_index_in_preprocessed_input = 0;
|
||||
int64_t dim_index_in_original_input = 0;
|
||||
|
||||
// iterate through all subscript indices in this input
|
||||
for (const auto& subscript_index : current_subscript_indices) {
|
||||
if (subscript_indices_to_input_index[subscript_index] == -1) { // This is the first time we are seeing this subscript label in this input
|
||||
subscript_indices_to_input_index[subscript_index] = dim_index_in_preprocessed_input++;
|
||||
homogenized_input_dims[subscript_index] = input_dims[dim_index_in_original_input];
|
||||
} else { // Diagonal needs to be parsed along the repeated axes
|
||||
preprocessed = EinsumOp::Diagonal(preprocessed ? *preprocessed : *inputs_[input_iter],
|
||||
subscript_indices_to_input_index[subscript_index],
|
||||
dim_index_in_preprocessed_input,
|
||||
allocator_);
|
||||
}
|
||||
++dim_index_in_original_input;
|
||||
}
|
||||
|
||||
std::vector<size_t> permutation;
|
||||
permutation.reserve(input_dims.size());
|
||||
for (auto& d : subscript_indices_to_input_index) {
|
||||
if (d != -1) {
|
||||
permutation.push_back(static_cast<size_t>(d));
|
||||
}
|
||||
}
|
||||
|
||||
// (Identify no-op transpose and prevent triggering the transpose)
|
||||
if (EinsumOp::IsTransposeRequired(preprocessed ? preprocessed->Shape().GetDims().size() : inputs_[input_iter]->Shape().GetDims().size(),
|
||||
permutation)) {
|
||||
preprocessed = EinsumOp::Transpose(preprocessed ? *preprocessed : *inputs_[input_iter],
|
||||
preprocessed ? preprocessed->Shape().GetDims() : inputs_[input_iter]->Shape().GetDims(),
|
||||
permutation, allocator_);
|
||||
}
|
||||
|
||||
// pre-processed may be null if the input didn't have need diagonals parsed and didn't need transposing
|
||||
// If the pre-processed inputs are null, we will use raw inputs in conjunction with "homogenized_input_dims" for
|
||||
// downstream compute
|
||||
if (preprocessed) { // If the pre-processed version of the operand exists, reshape it to homogenized_input_dims
|
||||
preprocessed->Reshape(homogenized_input_dims);
|
||||
}
|
||||
preprocessed_inputs_.push_back(std::move(preprocessed));
|
||||
homogenized_input_dims_.emplace_back(homogenized_input_dims);
|
||||
|
||||
++input_iter;
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
// Templated core Einsum logic
|
||||
template <typename T>
|
||||
Status EinsumTypedComputeProcessor(OpKernelContext* context,
|
||||
AllocatorPtr allocator,
|
||||
EinsumComputePreprocessor& einsum_compute_preprocessor) {
|
||||
const auto& mapped_indices_to_last_input_index = einsum_compute_preprocessor.GetMappedSubscriptIndicesToLastInputIndex();
|
||||
|
||||
auto& preprocessed_inputs = einsum_compute_preprocessor.GetPreprocessedInputTensors();
|
||||
|
||||
const auto& raw_inputs = einsum_compute_preprocessor.GetRawInputTensors();
|
||||
|
||||
const auto& homogenized_input_dims = einsum_compute_preprocessor.GetHomogenizedInputDims();
|
||||
|
||||
auto num_subscript_labels = einsum_compute_preprocessor.GetNumSubscriptIndices();
|
||||
|
||||
const auto& output_dims = einsum_compute_preprocessor.GetOutputDims();
|
||||
|
||||
auto* output = context->Output(0, output_dims);
|
||||
|
||||
auto num_inputs = context->InputCount();
|
||||
|
||||
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
|
||||
|
||||
// Pre-process the first input so as to reduce any dims that only it has
|
||||
std::unique_ptr<const Tensor> result;
|
||||
|
||||
{
|
||||
std::vector<int64_t> reduced_dims;
|
||||
std::vector<int64_t> preserved_dims; // dims which were not reduced
|
||||
std::vector<int64_t> preserved_shape; // shape pertaining to only the dims that were preserved (not reduced)
|
||||
reduced_dims.reserve(num_subscript_labels); // num_subscript_labels is the upper bound. No harm in over-reserving.
|
||||
preserved_dims.reserve(num_subscript_labels); // num_subscript_labels is the upper bound. No harm in over-reserving.
|
||||
|
||||
for (int64_t i = 0; i < num_subscript_labels; ++i) {
|
||||
if (mapped_indices_to_last_input_index[i] == 0) {
|
||||
reduced_dims.push_back(i);
|
||||
} else {
|
||||
preserved_dims.push_back(i);
|
||||
}
|
||||
}
|
||||
|
||||
// Reduce the dims that are last seen in the first input alone
|
||||
if (reduced_dims.size() != 0) {
|
||||
result = EinsumOp::ReduceSum<T>(preprocessed_inputs[0] ? *preprocessed_inputs[0] : *raw_inputs[0],
|
||||
homogenized_input_dims[0].GetDims(), reduced_dims, allocator, tp);
|
||||
} else {
|
||||
// Check if there is a pre-processed version of this input
|
||||
// If so assign it to result
|
||||
if (preprocessed_inputs[0]) {
|
||||
result = std::move(preprocessed_inputs[0]);
|
||||
}
|
||||
}
|
||||
|
||||
// Finalize the output at this stage if num_inputs == 1
|
||||
if (num_inputs == 1) {
|
||||
// Finalize the output by applying any transpose required to get it to the required output ordering and move it to the op's output
|
||||
EinsumOp::FinalizeOutput<T>(result ? *result : *raw_inputs[0],
|
||||
preserved_dims,
|
||||
einsum_compute_preprocessor.GetMappedSubscriptIndicesToOutputindices(),
|
||||
*output, output_dims, allocator);
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
}
|
||||
|
||||
// Process the operands in a pair-wise fashion
|
||||
{
|
||||
bool is_final_pair = false;
|
||||
// Keep processing each input pair-wise
|
||||
for (int input = 1; input < num_inputs; ++input) {
|
||||
std::vector<int64_t> reduced_dims;
|
||||
reduced_dims.reserve(num_subscript_labels); // num_subscript_labels is the upper bound. No harm in over-reserving by a small margin.
|
||||
for (int64_t dim = 0; dim < num_subscript_labels; ++dim) {
|
||||
if (mapped_indices_to_last_input_index[dim] == input) {
|
||||
// This is the last input we are seeing this dimension (and it doesn't occur in the output), so reduce along the dimension
|
||||
reduced_dims.push_back(dim);
|
||||
}
|
||||
}
|
||||
if (input == num_inputs - 1) {
|
||||
is_final_pair = true;
|
||||
}
|
||||
// Use either the preprocessed inputs (if it is available) or the corresponding raw inputs
|
||||
result = EinsumOp::PairwiseOperandProcess<T>(result ? *result : *raw_inputs[0],
|
||||
result ? result->Shape() : homogenized_input_dims[0],
|
||||
preprocessed_inputs[input] ? *preprocessed_inputs[input] : *raw_inputs[input],
|
||||
homogenized_input_dims[input],
|
||||
reduced_dims, tp, allocator,
|
||||
einsum_compute_preprocessor, is_final_pair, *output);
|
||||
}
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
// Explicit template instantiation
|
||||
|
||||
// float
|
||||
template Status EinsumTypedComputeProcessor<float>(OpKernelContext* context, AllocatorPtr allocator, EinsumComputePreprocessor& einsum_compute_preprocessor);
|
||||
|
||||
// int32_t
|
||||
template Status EinsumTypedComputeProcessor<int32_t>(OpKernelContext* context, AllocatorPtr allocator, EinsumComputePreprocessor& einsum_compute_preprocessor);
|
||||
|
||||
// double
|
||||
template Status EinsumTypedComputeProcessor<double>(OpKernelContext* context, AllocatorPtr allocator, EinsumComputePreprocessor& einsum_compute_preprocessor);
|
||||
|
||||
// int64_t
|
||||
template Status EinsumTypedComputeProcessor<int64_t>(OpKernelContext* context, AllocatorPtr allocator, EinsumComputePreprocessor& einsum_compute_preprocessor);
|
||||
|
||||
} // namespace onnxruntime
|
||||
200
onnxruntime/core/providers/cpu/math/einsum_utils.h
Normal file
200
onnxruntime/core/providers/cpu/math/einsum_utils.h
Normal file
|
|
@ -0,0 +1,200 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
// This module hosts 3 abstractions -
|
||||
|
||||
// 1) EinsumEquationPreprocessor -
|
||||
// Holds logic to statically pre-process the equation string (i.e.) without input shapes being known
|
||||
// These need not be repeated at Compute() time again
|
||||
|
||||
// 2) EinsumComputePreprocessor -
|
||||
// Holds logic to process the data from EinsumEquationPreprocessor using known input shapes to parse data required
|
||||
// during Einsum Compute(). For example, mapping subscript labels to a dimension value, etc.
|
||||
|
||||
// 3) EinsumTypedComputeProcessor - The core logic of the Einsum operator. Invoked from Einsum Compute().
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "einsum_auxiliary_ops.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
namespace EinsumOp {
|
||||
|
||||
constexpr size_t num_of_letters = 26;
|
||||
|
||||
} // namespace EinsumOp
|
||||
|
||||
struct EinsumEquationPreprocessor {
|
||||
explicit EinsumEquationPreprocessor(const std::string& einsum_equation) {
|
||||
// Make copy of the equation as it will be mutated
|
||||
einsum_preprocessed_equation_ = einsum_equation;
|
||||
|
||||
// Remove space characters in the copy of the Einsum eqution
|
||||
einsum_preprocessed_equation_.erase(std::remove(einsum_preprocessed_equation_.begin(), einsum_preprocessed_equation_.end(), ' '),
|
||||
einsum_preprocessed_equation_.end());
|
||||
|
||||
// Check if the Einsum equation has the output subscript labels
|
||||
auto mid_index = einsum_preprocessed_equation_.find("->");
|
||||
if (mid_index != std::string::npos) {
|
||||
// Separate right and left hand sides of the equation
|
||||
left_equation_ = einsum_preprocessed_equation_.substr(0, mid_index);
|
||||
right_equation_ = einsum_preprocessed_equation_.substr(mid_index + 2);
|
||||
is_explicit_ = true;
|
||||
} else {
|
||||
left_equation_ = einsum_preprocessed_equation_;
|
||||
};
|
||||
|
||||
// Process the left_equation_ by splitting on ','
|
||||
std::string delimiter = ",";
|
||||
size_t pos = 0;
|
||||
std::string token;
|
||||
while ((pos = left_equation_.find(delimiter)) != std::string::npos) {
|
||||
token = left_equation_.substr(0, pos);
|
||||
left_equation_.erase(0, pos + delimiter.length());
|
||||
left_equation_split_.push_back(token); // This copy is done statically at model load, hence should not affect runtime perf
|
||||
}
|
||||
left_equation_split_.push_back(left_equation_); // This holds the portion of the equation after the last ','
|
||||
}
|
||||
|
||||
// Holds the pre-processed equation string
|
||||
// In theory, we could re-write the einsum equation to lower overall cost of intermediate arrays
|
||||
// See numpy.einsum_path for details/examples
|
||||
// These are very advanced optimizations that we don't require for the average use-case
|
||||
std::string einsum_preprocessed_equation_;
|
||||
|
||||
// In explicit form, holds the left side of the einsum equation
|
||||
// (e.g.) Einsum equation = 'i,j->i', then left_equation_ = 'i,j'
|
||||
// In implicit form, holds the entire einsum equation
|
||||
// (e.g.) Einsum equation = 'i,j', then left_equation_ = 'i,j'
|
||||
std::string left_equation_;
|
||||
|
||||
// Holds the strings obtained after splitting left_equation_ on ','
|
||||
std::vector<std::string> left_equation_split_;
|
||||
|
||||
// Holds constructed or parsed output subscript
|
||||
std::string right_equation_;
|
||||
|
||||
// Flag indicating if the Einsum op is being used in explicit form (i.e.) contains '->'
|
||||
bool is_explicit_ = false;
|
||||
};
|
||||
|
||||
// Prologue:
|
||||
// In the sample Einsum string: 'ij, jk'
|
||||
// Subscripts are 'ij' and 'jk'
|
||||
// Subscript labels are 'i', 'j', and 'k'
|
||||
// Subscript labels (letter) and subcript indices (a unique id to the letter) are interchangeable
|
||||
|
||||
// This is a pre-processor class that maps subscript labels to a dimension value, etc.
|
||||
class EinsumComputePreprocessor final {
|
||||
public:
|
||||
explicit EinsumComputePreprocessor(EinsumEquationPreprocessor& equation_preprocessor,
|
||||
const std::vector<const Tensor*>& inputs,
|
||||
AllocatorPtr allocator);
|
||||
|
||||
// The main method that does all the pre-processing - must be invoked before other methods are called
|
||||
// to get relevant metadata
|
||||
Status Run();
|
||||
|
||||
// Get the output dims of the op's output
|
||||
const std::vector<int64_t>& GetOutputDims() const;
|
||||
|
||||
// Pre-process inputs if needed - preprocessing includes -
|
||||
// 1) Parsing diagonals from raw inputs
|
||||
// 2) Transposing some axes to match a chosen fixed ordering
|
||||
// This must be used in conjunction with its corresponding entry in homogenized_input_dims_
|
||||
// (returned by GetHomogenizedInputDims()).
|
||||
// If a particular entry is null, use raw inputs in conjunction with homogenized_input_dims_.
|
||||
std::vector<std::unique_ptr<Tensor>>& GetPreprocessedInputTensors();
|
||||
|
||||
// Get raw inputs to the op
|
||||
const std::vector<const Tensor*>& GetRawInputTensors();
|
||||
|
||||
// Get the "homogenized input dims" for each preprocessed/raw input
|
||||
const std::vector<TensorShape>& GetHomogenizedInputDims();
|
||||
|
||||
// For each subscript index, hold the last input the subscript index was seen in
|
||||
const std::vector<int64_t>& GetMappedSubscriptIndicesToLastInputIndex() const;
|
||||
|
||||
// For each subscript index, hold the index it corresponds to in the output's shape
|
||||
const std::vector<int64_t>& GetMappedSubscriptIndicesToOutputindices() const;
|
||||
|
||||
// Get the number of subscript indices (subscript labels) in the einsum equation
|
||||
int64_t GetNumSubscriptIndices() const;
|
||||
|
||||
private:
|
||||
// Process subscripts of each input and collect metadata along the way
|
||||
Status ProcessSubscripts();
|
||||
|
||||
// A function to process broadcasted dims (ellipsis) of inputs that they occur in
|
||||
Status PostProcessBroadcastedDims();
|
||||
|
||||
// Check if the Einsum equation has an explicit form (equation string contains "->")
|
||||
// If it is of explicit form, parse the output subscript (substring following "->")
|
||||
// If it is of implicit form (equation string does not contain "->"), compose the output subscript
|
||||
// If the output subscript is an empty string, the result is a scalar
|
||||
Status ParseOrCreateOutputSubscript();
|
||||
|
||||
Status CalculateOutputShape();
|
||||
|
||||
Status PreprocessInputs();
|
||||
|
||||
// private members
|
||||
// Instance of EinsumEquationPreprocessor
|
||||
EinsumEquationPreprocessor einsum_equation_preprocessor_;
|
||||
|
||||
// The number of dims that encompasses an "ellipsis"
|
||||
size_t num_of_ellipsis_dims_ = 0;
|
||||
|
||||
// All original inputs to the op
|
||||
const std::vector<const Tensor*>& inputs_;
|
||||
|
||||
// All preprocessed inputs
|
||||
std::vector<std::unique_ptr<Tensor>> preprocessed_inputs_;
|
||||
|
||||
// Holds the preprocessed inputs' homogenized dims
|
||||
std::vector<TensorShape> homogenized_input_dims_;
|
||||
|
||||
// Count of unique subscript labels (subscript indices)
|
||||
// E.g. 1 : With equation -> 'ij, jk -> ik'
|
||||
// num_subscript_indices_ = 3 (i, j, k)
|
||||
// E.g. 2 : With equation -> '...ij', 'jk' -> '...ik'
|
||||
// num_subscript_indices_ = 3 (i, j, k) + number of dims specified by an ellipsis (across all inputs)
|
||||
int64_t num_subscript_indices_ = 0;
|
||||
|
||||
// Hold the count corresponding to the letter seen
|
||||
// `0` means the corresponding letter wasn't seen at all
|
||||
std::array<int64_t, EinsumOp::num_of_letters> letter_to_count_;
|
||||
|
||||
// Hold the assigned index corresponding to the letter seen
|
||||
// `-1` means the corresponding letter wasn't seen at all
|
||||
std::array<int64_t, EinsumOp::num_of_letters> letter_to_index_;
|
||||
|
||||
// Holds the input index of the last input to have the index corresponding to the subscript label
|
||||
// If the value is `-1`, then the subscript label is never seen (or) it appears in the output
|
||||
std::vector<int64_t> subscript_indices_to_last_input_;
|
||||
|
||||
// Hold the dim value of the index corresponding to the subscript label
|
||||
// `-1` means the corresponding label wasn't seen at all
|
||||
std::vector<int64_t> subscript_indices_to_dim_value_;
|
||||
|
||||
// Holds the final calculated output dimensions
|
||||
std::vector<int64_t> output_dims_;
|
||||
|
||||
// All subscript indices in the equation for each input
|
||||
std::vector<std::vector<int64_t>> input_subscript_indices_;
|
||||
|
||||
// Index corresponding to each output dim corresponding to each subscript index
|
||||
// A value of -1 means the corresponding subscript index is not found in the output
|
||||
std::vector<int64_t> subscript_indices_to_output_indices_;
|
||||
|
||||
// Allocator to use for ad-hoc tensor buffer allocation
|
||||
AllocatorPtr allocator_;
|
||||
};
|
||||
|
||||
// This method does the heavy-lifting compute portion of Einsum Compute()
|
||||
template <typename T>
|
||||
Status EinsumTypedComputeProcessor(OpKernelContext* context, AllocatorPtr allocator,
|
||||
EinsumComputePreprocessor& einsum_compute_preprocessor);
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -151,38 +151,45 @@ REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ArgMin, 1, 10);
|
|||
REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ArgMin, 11, 11);
|
||||
REGISTER_UNARY_ELEMENTWISE_KERNEL(ArgMin, 12);
|
||||
|
||||
// When all reduce axises located at the tail of the dims, quite general cases, transpose and extra
|
||||
// When all reduce axes are located at the tail of the dims, quite general cases, transpose and extra
|
||||
// copy could be skipped to improve performance. If required by check_no_transpose = true, then
|
||||
// the calling code will check if the data was transposed and act accordingly.
|
||||
// return value: true means transposedInputData is not created/copied, input tensor data could
|
||||
// be direct use as row major matrix [block_size, blocks], where blocks is the
|
||||
// size of each reduce.
|
||||
// be directly used as row major matrix [block_size, blocks], where blocks is the
|
||||
// size of each reduce.
|
||||
// `input_shape_override` overrides the shape of `input` for compute purposes.
|
||||
template <typename T>
|
||||
bool PrepareForReduce(OpKernelContext* ctx,
|
||||
FastAllocVector<T>& transposedInputData,
|
||||
Tensor** reducedTensor,
|
||||
bool PrepareForReduce(const Tensor* input_tensor_ptr,
|
||||
FastAllocVector<T>& transposed_input_data,
|
||||
int64_t& block_size,
|
||||
int64_t& blocks,
|
||||
const std::vector<int64_t>& axes_,
|
||||
bool keepdims_,
|
||||
bool check_no_transpose = false) {
|
||||
const auto* input_tensor_ptr = ctx->Input<Tensor>(0);
|
||||
ORT_ENFORCE(input_tensor_ptr != nullptr);
|
||||
const Tensor& input = *input_tensor_ptr;
|
||||
/*out*/ std::vector<int64_t>& reduced_dims,
|
||||
bool check_no_transpose = false,
|
||||
const TensorShape* input_shape_override = nullptr) {
|
||||
ORT_ENFORCE(input_tensor_ptr != nullptr, "Input to be reduced is null");
|
||||
|
||||
size_t ndim = input.Shape().NumDimensions();
|
||||
if (input_shape_override) {
|
||||
ORT_ENFORCE(input_tensor_ptr->Shape().Size() == input_shape_override->Size(),
|
||||
"The input shape override's size does not match the input tensor's shape size");
|
||||
}
|
||||
|
||||
const Tensor& input = *input_tensor_ptr;
|
||||
const auto& input_shape = input_shape_override ? *input_shape_override : input.Shape();
|
||||
|
||||
size_t ndim = input_shape.NumDimensions();
|
||||
|
||||
// Scalar tensor
|
||||
if (ndim == 0) {
|
||||
if (!check_no_transpose) {
|
||||
auto size = input.Shape().Size();
|
||||
auto size = input_shape.Size();
|
||||
assert(size == 1);
|
||||
transposedInputData.resize(size, 0);
|
||||
T* to_data = &transposedInputData[0];
|
||||
transposed_input_data.resize(size, 0);
|
||||
T* to_data = &transposed_input_data[0];
|
||||
*to_data = *input.Data<T>();
|
||||
}
|
||||
block_size = blocks = 1;
|
||||
*reducedTensor = ctx->Output(0, input.Shape());
|
||||
return true;
|
||||
}
|
||||
|
||||
|
|
@ -224,11 +231,11 @@ bool PrepareForReduce(OpKernelContext* ctx,
|
|||
|
||||
std::vector<int64_t> new_dims(transposed_axes.size());
|
||||
for (size_t i = 0; i < transposed_axes.size(); ++i) {
|
||||
new_dims[i] = input.Shape().GetDims().at(transposed_axes[i]);
|
||||
new_dims[i] = input_shape.GetDims().at(transposed_axes[i]);
|
||||
}
|
||||
|
||||
int num_axes = static_cast<int>(transposed_axes.size());
|
||||
auto in_dims = input.Shape().GetDims();
|
||||
auto in_dims = input_shape.GetDims();
|
||||
|
||||
// Measure amount of contiguous data we can copy at once
|
||||
int64_t blocksize = 1;
|
||||
|
|
@ -243,11 +250,10 @@ bool PrepareForReduce(OpKernelContext* ctx,
|
|||
}
|
||||
|
||||
const T* from_data = input.template Data<T>();
|
||||
size_t count = input.Shape().Size();
|
||||
size_t count = input_shape.Size();
|
||||
|
||||
//set to-be-reduced axes to one. squeeze is keepdims_ is false
|
||||
int64_t first_dim = 1;
|
||||
std::vector<int64_t> reduced_dims;
|
||||
reduced_dims.reserve(in_dims.size());
|
||||
|
||||
for (size_t i = 0; i < in_dims.size(); i++) {
|
||||
|
|
@ -267,13 +273,12 @@ bool PrepareForReduce(OpKernelContext* ctx,
|
|||
ORT_ENFORCE(in_dim != 0,
|
||||
"Can't reduce on dim with value of 0 if 'keepdims' is false. "
|
||||
"Invalid output shape would be produced. input_shape:",
|
||||
input.Shape());
|
||||
input_shape);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
*reducedTensor = ctx->Output(0, std::move(reduced_dims));
|
||||
auto num_elements = input.Shape().Size();
|
||||
auto num_elements = input_shape.Size();
|
||||
|
||||
// edge case. one or more input dims with value of 0.
|
||||
if (num_elements == 0) {
|
||||
|
|
@ -288,8 +293,8 @@ bool PrepareForReduce(OpKernelContext* ctx,
|
|||
return true;
|
||||
}
|
||||
|
||||
transposedInputData.resize(input.Shape().Size(), 0);
|
||||
T* to_data = &transposedInputData[0];
|
||||
transposed_input_data.resize(input_shape.Size(), 0);
|
||||
T* to_data = &transposed_input_data[0];
|
||||
if (num_axes < 2 || n_shared_idxs == num_axes) {
|
||||
memcpy(to_data, from_data, count * sizeof(T));
|
||||
return false;
|
||||
|
|
@ -356,24 +361,28 @@ bool PrepareForReduce(OpKernelContext* ctx,
|
|||
|
||||
template <typename T>
|
||||
Status ReduceL1<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
std::vector<int64_t> reduced_dims;
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
|
||||
T* output_data = reduced->template MutableData<T>();
|
||||
|
||||
if (no_transpose) {
|
||||
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
|
||||
const T* input_data = input->template Data<T>();
|
||||
|
||||
for (int64_t i = 0; i < block_size; ++i) {
|
||||
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).cwiseAbs().sum();
|
||||
}
|
||||
} else {
|
||||
EigenVectorMap<T> out_vec(output_data, block_size);
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).cwiseAbs().rowwise().sum();
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks).cwiseAbs().rowwise().sum();
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
@ -381,24 +390,28 @@ Status ReduceL1<T>::Compute(OpKernelContext* ctx) const {
|
|||
|
||||
template <typename T>
|
||||
Status ReduceL2<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
std::vector<int64_t> reduced_dims;
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
|
||||
T* output_data = reduced->template MutableData<T>();
|
||||
|
||||
if (no_transpose) {
|
||||
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
|
||||
const T* input_data = input->template Data<T>();
|
||||
|
||||
for (int64_t i = 0; i < block_size; ++i) {
|
||||
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).norm();
|
||||
}
|
||||
} else {
|
||||
EigenVectorMap<T> out_vec(output_data, block_size);
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().norm();
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks).rowwise().norm();
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
@ -406,24 +419,28 @@ Status ReduceL2<T>::Compute(OpKernelContext* ctx) const {
|
|||
|
||||
template <typename T>
|
||||
Status ReduceLogSum<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
std::vector<int64_t> reduced_dims;
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
|
||||
T* output_data = reduced->template MutableData<T>();
|
||||
|
||||
if (no_transpose) {
|
||||
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
|
||||
const T* input_data = input->template Data<T>();
|
||||
|
||||
for (int64_t i = 0; i < block_size; ++i) {
|
||||
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).sum();
|
||||
}
|
||||
} else {
|
||||
EigenVectorMap<T> out_vec(output_data, block_size);
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().sum();
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks).rowwise().sum();
|
||||
}
|
||||
|
||||
for (int j = 0; j < block_size; ++j) {
|
||||
|
|
@ -436,22 +453,26 @@ Status ReduceLogSum<T>::Compute(OpKernelContext* ctx) const {
|
|||
|
||||
template <typename T>
|
||||
Status ReduceLogSumExp<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_);
|
||||
std::vector<int64_t> reduced_dims;
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
|
||||
T* output_data = reduced->template MutableData<T>();
|
||||
|
||||
for (int j = 0; j < block_size; ++j) {
|
||||
T max_value = std::numeric_limits<T>::lowest();
|
||||
for (int i = 0; i < blocks; ++i) {
|
||||
max_value = std::max(max_value, transposedInputData[i * block_size + j]);
|
||||
max_value = std::max(max_value, transposed_input_data[i * block_size + j]);
|
||||
}
|
||||
T scaled_exp_sum = 0;
|
||||
for (int i = 0; i < blocks; ++i) {
|
||||
scaled_exp_sum += static_cast<T>(std::exp(transposedInputData[i * block_size + j] - max_value));
|
||||
scaled_exp_sum += static_cast<T>(std::exp(transposed_input_data[i * block_size + j] - max_value));
|
||||
}
|
||||
*(output_data++) = static_cast<T>(std::log(scaled_exp_sum) + max_value);
|
||||
}
|
||||
|
|
@ -460,23 +481,27 @@ Status ReduceLogSumExp<T>::Compute(OpKernelContext* ctx) const {
|
|||
|
||||
template <typename T>
|
||||
Status ReduceMax<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
std::vector<int64_t> reduced_dims;
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
|
||||
T* output_data = reduced->template MutableData<T>();
|
||||
|
||||
if (no_transpose) {
|
||||
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
|
||||
const T* input_data = input->template Data<T>();
|
||||
|
||||
for (int64_t i = 0; i < block_size; ++i) {
|
||||
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).maxCoeff();
|
||||
}
|
||||
} else {
|
||||
EigenVectorMap<T> out_vec(output_data, block_size);
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().maxCoeff();
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks).rowwise().maxCoeff();
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
@ -484,11 +509,15 @@ Status ReduceMax<T>::Compute(OpKernelContext* ctx) const {
|
|||
|
||||
template <typename T>
|
||||
Status ReduceMean<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
std::vector<int64_t> reduced_dims;
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
|
||||
T* output_data = reduced->template MutableData<T>();
|
||||
|
||||
|
|
@ -500,7 +529,7 @@ Status ReduceMean<T>::Compute(OpKernelContext* ctx) const {
|
|||
concurrency::ThreadPool::TryBatchParallelFor(ctx->GetOperatorThreadPool(), block_size, lambda, 0);
|
||||
} else {
|
||||
EigenVectorMap<T> out_vec(output_data, block_size);
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().mean();
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks).rowwise().mean();
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
@ -508,23 +537,27 @@ Status ReduceMean<T>::Compute(OpKernelContext* ctx) const {
|
|||
|
||||
template <typename T>
|
||||
Status ReduceMin<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
std::vector<int64_t> reduced_dims;
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
|
||||
T* output_data = reduced->template MutableData<T>();
|
||||
|
||||
if (no_transpose) {
|
||||
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
|
||||
const T* input_data = input->template Data<T>();
|
||||
|
||||
for (int64_t i = 0; i < block_size; ++i) {
|
||||
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).minCoeff();
|
||||
}
|
||||
} else {
|
||||
EigenVectorMap<T> out_vec(output_data, block_size);
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().minCoeff();
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks).rowwise().minCoeff();
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
@ -532,75 +565,110 @@ Status ReduceMin<T>::Compute(OpKernelContext* ctx) const {
|
|||
|
||||
template <typename T>
|
||||
Status ReduceProd<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
std::vector<int64_t> reduced_dims;
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
|
||||
T* output_data = reduced->template MutableData<T>();
|
||||
|
||||
if (no_transpose) {
|
||||
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
|
||||
const T* input_data = input->template Data<T>();
|
||||
|
||||
for (int64_t i = 0; i < block_size; ++i) {
|
||||
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).prod();
|
||||
}
|
||||
} else {
|
||||
EigenVectorMap<T> out_vec(output_data, block_size);
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().prod();
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks).rowwise().prod();
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status ReduceSum<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
|
||||
T* output_data = reduced->template MutableData<T>();
|
||||
|
||||
static void ReduceSumCore(const T* input_data, T* output_data, bool no_transpose,
|
||||
int64_t blocks, int64_t block_size, FastAllocVector<T>& transposed_input_data,
|
||||
concurrency::ThreadPool* tp) {
|
||||
if (no_transpose) {
|
||||
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
|
||||
auto lambda = [input_data, blocks, output_data](ptrdiff_t i) {
|
||||
|
||||
// The ConstEigenMatrixMap type is expanded to work around a MS compiler issue
|
||||
output_data[i] = Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1>>(input_data + (i * blocks), blocks).sum();
|
||||
};
|
||||
concurrency::ThreadPool::TryBatchParallelFor(ctx->GetOperatorThreadPool(), block_size, lambda, 0);
|
||||
concurrency::ThreadPool::TryBatchParallelFor(tp, block_size, lambda, 0);
|
||||
} else {
|
||||
EigenVectorMap<T> out_vec(output_data, block_size);
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().sum();
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks).rowwise().sum();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Tensor ReduceSum<T>::Impl(const Tensor& input, const std::vector<int64_t>& reduce_axes,
|
||||
AllocatorPtr allocator, concurrency::ThreadPool* tp, bool keep_dims,
|
||||
const TensorShape* input_shape_override) {
|
||||
FastAllocVector<T> transposed_input_data(allocator);
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
std::vector<int64_t> reduced_dims;
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(&input, transposed_input_data, block_size, blocks,
|
||||
reduce_axes, keep_dims, reduced_dims, true, input_shape_override);
|
||||
|
||||
Tensor output(input.DataType(), reduced_dims, allocator);
|
||||
|
||||
ReduceSumCore(input.template Data<T>(), output.template MutableData<T>(),
|
||||
no_transpose, blocks, block_size, transposed_input_data, tp);
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status ReduceSum<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
std::vector<int64_t> reduced_dims;
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
auto* output = ctx->Output(0, reduced_dims);
|
||||
|
||||
ReduceSumCore(input->template Data<T>(), output->template MutableData<T>(),
|
||||
no_transpose, blocks, block_size, transposed_input_data, ctx->GetOperatorThreadPool());
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status ReduceSumSquare<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
std::vector<int64_t> reduced_dims;
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
|
||||
T* output_data = reduced->template MutableData<T>();
|
||||
|
||||
if (no_transpose) {
|
||||
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
|
||||
const T* input_data = input->template Data<T>();
|
||||
|
||||
for (int64_t i = 0; i < block_size; ++i) {
|
||||
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).squaredNorm();
|
||||
}
|
||||
} else {
|
||||
EigenVectorMap<T> out_vec(output_data, block_size);
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().squaredNorm();
|
||||
out_vec = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks).rowwise().squaredNorm();
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
@ -608,13 +676,16 @@ Status ReduceSumSquare<T>::Compute(OpKernelContext* ctx) const {
|
|||
|
||||
template <typename T>
|
||||
Status ArgMax<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
std::vector<int64_t> reduced_dims;
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
int64_t* output_data = reduced->template MutableData<int64_t>();
|
||||
Eigen::MatrixXf::Index maxIndex;
|
||||
|
||||
|
|
@ -642,7 +713,7 @@ Status ArgMax<T>::Compute(OpKernelContext* ctx) const {
|
|||
}
|
||||
}
|
||||
} else {
|
||||
auto matrixData = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks);
|
||||
auto matrixData = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks);
|
||||
if (select_last_index_) {
|
||||
for (int i = 0; i < block_size; ++i) {
|
||||
int idx = 0;
|
||||
|
|
@ -669,13 +740,16 @@ Status ArgMax<T>::Compute(OpKernelContext* ctx) const {
|
|||
|
||||
template <typename T>
|
||||
Status ArgMin<T>::Compute(OpKernelContext* ctx) const {
|
||||
FastAllocVector<T> transposedInputData(GetAllocator<T>(*ctx));
|
||||
FastAllocVector<T> transposed_input_data(GetAllocator<T>(*ctx));
|
||||
int64_t block_size;
|
||||
int64_t blocks;
|
||||
Tensor* reduced;
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
|
||||
std::vector<int64_t> reduced_dims;
|
||||
const Tensor* input = ctx->Input<Tensor>(0);
|
||||
|
||||
bool no_transpose = PrepareForReduce<T>(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true);
|
||||
|
||||
Tensor* reduced = ctx->Output(0, reduced_dims);
|
||||
int64_t* output_data = reduced->template MutableData<int64_t>();
|
||||
Eigen::MatrixXf::Index minIndex;
|
||||
|
||||
|
|
@ -703,7 +777,7 @@ Status ArgMin<T>::Compute(OpKernelContext* ctx) const {
|
|||
}
|
||||
}
|
||||
} else {
|
||||
auto matrixData = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks);
|
||||
auto matrixData = ConstEigenMatrixMap<T>(&transposed_input_data[0], block_size, blocks);
|
||||
if (select_last_index_) {
|
||||
for (int i = 0; i < block_size; ++i) {
|
||||
int idx = 0;
|
||||
|
|
@ -728,4 +802,13 @@ Status ArgMin<T>::Compute(OpKernelContext* ctx) const {
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
// Explicit template instantiation -
|
||||
// Even though there are kernels registered for ReduceSum op for these types,
|
||||
// these are needed because we seem to get linker errors without these when the linker
|
||||
// tries to resolve symbols in the einsum_auxiliary_ops obj file
|
||||
template class ReduceSum<float>;
|
||||
template class ReduceSum<int32_t>;
|
||||
template class ReduceSum<double>;
|
||||
template class ReduceSum<int64_t>;
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -121,6 +121,12 @@ class ReduceSum final : public ReduceKernel<true> {
|
|||
}
|
||||
|
||||
Status Compute(OpKernelContext* context) const override;
|
||||
|
||||
// For external calls requiring ReduceSum implementation - will return the reduced output.
|
||||
//`input_shape_override` overrides the shape of `input` for compute purposes.
|
||||
static Tensor Impl(const Tensor& input, const std::vector<int64_t>& reduce_axes,
|
||||
AllocatorPtr allocator, concurrency::ThreadPool* tp, bool keep_dims,
|
||||
const TensorShape* input_shape_override = nullptr);
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
|
|
|
|||
|
|
@ -174,8 +174,10 @@ static void DoTransposeEltWise(int64_t num_axes, const std::vector<int64_t>& tar
|
|||
}
|
||||
}
|
||||
|
||||
static Status DoUntypedTranspose(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output) {
|
||||
const auto& input_shape = input.Shape();
|
||||
// `input_shape_override` overrides the shape of `input` for compute purposes.
|
||||
static Status DoUntypedTranspose(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
|
||||
const TensorShape* input_shape_override = nullptr) {
|
||||
const auto& input_shape = input_shape_override ? *input_shape_override : input.Shape();
|
||||
const auto& input_dims = input_shape.GetDims();
|
||||
auto rank = input_shape.NumDimensions();
|
||||
|
||||
|
|
@ -289,11 +291,12 @@ static void SimpleTransposeSingleAxisOutwards(const T* input_data, T* output_dat
|
|||
}
|
||||
}
|
||||
|
||||
// `input_shape_override` overrides the shape of `input` for compute purposes.
|
||||
static void TransposeSingleAxisOutwards(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
|
||||
int64_t from, int64_t to) {
|
||||
int64_t from, int64_t to, const TensorShape* input_shape_override = nullptr) {
|
||||
ORT_UNUSED_PARAMETER(permutations);
|
||||
|
||||
const auto& input_shape = input.Shape();
|
||||
const auto& input_shape = input_shape_override ? *input_shape_override : input.Shape();
|
||||
const auto& input_dims = input_shape.GetDims();
|
||||
|
||||
const auto element_size = input.DataType()->Size();
|
||||
|
|
@ -380,11 +383,12 @@ static void SimpleTransposeSingleAxisInwards(const T* input_data, T* output_data
|
|||
}
|
||||
|
||||
// moving a single axis inwards where the read/write size is a power of 2 and between 8 and 64 bits.
|
||||
// `input_shape_override` overrides the shape of `input` for compute purposes.
|
||||
static void TransposeSingleAxisInwards(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
|
||||
int64_t from, int64_t to) {
|
||||
int64_t from, int64_t to, const TensorShape* input_shape_override = nullptr) {
|
||||
ORT_UNUSED_PARAMETER(permutations);
|
||||
|
||||
const auto& input_shape = input.Shape();
|
||||
const auto& input_shape = input_shape_override ? *input_shape_override : input.Shape();
|
||||
const auto& input_dims = input_shape.GetDims();
|
||||
|
||||
const auto element_size = input.DataType()->Size();
|
||||
|
|
@ -448,12 +452,13 @@ static void TransposeSingleAxisInwards(const std::vector<size_t>& permutations,
|
|||
}
|
||||
}
|
||||
|
||||
// `input_shape_override` overrides the shape of `input` for compute purposes.
|
||||
static void SingleAxisTranspose(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
|
||||
size_t from, size_t to) {
|
||||
size_t from, size_t to, const TensorShape* input_shape_override = nullptr) {
|
||||
if (from > to) {
|
||||
TransposeSingleAxisOutwards(permutations, input, output, from, to);
|
||||
TransposeSingleAxisOutwards(permutations, input, output, from, to, input_shape_override);
|
||||
} else {
|
||||
TransposeSingleAxisInwards(permutations, input, output, from, to);
|
||||
TransposeSingleAxisInwards(permutations, input, output, from, to, input_shape_override);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -526,7 +531,9 @@ static bool IsMovingSingleAxis(const std::vector<size_t>& permutations, size_t&
|
|||
return single_axis_moved;
|
||||
}
|
||||
|
||||
Status TransposeBase::DoTranspose(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output) {
|
||||
//`input_shape_override` overrides the shape of `input` for compute purposes.
|
||||
Status TransposeBase::DoTranspose(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
|
||||
const TensorShape* input_shape_override) {
|
||||
Status status = Status::OK();
|
||||
|
||||
auto input_type = input.DataType();
|
||||
|
|
@ -540,10 +547,10 @@ Status TransposeBase::DoTranspose(const std::vector<size_t>& permutations, const
|
|||
bool moving_single_axis = IsMovingSingleAxis(permutations, from, to);
|
||||
|
||||
if (moving_single_axis && !input.IsDataTypeString()) {
|
||||
SingleAxisTranspose(permutations, input, output, from, to);
|
||||
SingleAxisTranspose(permutations, input, output, from, to, input_shape_override);
|
||||
} else {
|
||||
// fall back to default implementation
|
||||
status = DoUntypedTranspose(permutations, input, output);
|
||||
status = DoUntypedTranspose(permutations, input, output, input_shape_override);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -14,9 +14,10 @@ class TransposeBase {
|
|||
public:
|
||||
/**
|
||||
Transpose the input Tensor into the output Tensor using the provided permutations.
|
||||
Both Tensors must have the same data type.
|
||||
Both Tensors must have the same data type. `input_shape_override` overrides the shape of `input` for compute purposes.
|
||||
*/
|
||||
static Status DoTranspose(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output);
|
||||
static Status DoTranspose(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
|
||||
const TensorShape* input_shape_override = nullptr);
|
||||
|
||||
protected:
|
||||
TransposeBase(const OpKernelInfo& info) {
|
||||
|
|
|
|||
503
onnxruntime/test/providers/cpu/math/einsum_test.cc
Normal file
503
onnxruntime/test/providers/cpu/math/einsum_test.cc
Normal file
|
|
@ -0,0 +1,503 @@
|
|||
// 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/framework/data_types.h"
|
||||
#include "core/util/math.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace test {
|
||||
|
||||
// Tests are aplit up "theme-wise" (i.e.) each kind of operation Einsum can be used for
|
||||
// Within each theme we test "explicit" and "implicit" versions of the Einsum equation (wherever possible)
|
||||
// Some operations are not possible with implicit notation (reordering, reduction, etc.)
|
||||
|
||||
// Theme: Deep copy / No-op
|
||||
|
||||
// Explicit
|
||||
TEST(Einsum, ExplicitEinsumAsIdentity_1D_input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "i->i");
|
||||
test.AddInput<float>("x", {5}, {0.9f, 2.5f, 2.3f, 1.5f, -4.5f});
|
||||
test.AddOutput<float>("y", {5}, {0.9f, 2.5f, 2.3f, 1.5f, -4.5f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Implicit
|
||||
TEST(Einsum, ImplicitEinsumAsIdentity_1D_input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "i");
|
||||
test.AddInput<float>("x", {5}, {0.9f, 2.5f, 2.3f, 1.5f, -4.5f});
|
||||
test.AddOutput<float>("y", {5}, {0.9f, 2.5f, 2.3f, 1.5f, -4.5f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Theme: Transpose/Permutation
|
||||
|
||||
// Explicit
|
||||
TEST(Einsum, ExplicitEinsumAsTransposeOp_2D_input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ji->ij");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("y", {2, 2}, {1.f, 3.f, 2.f, 4.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsTransposeOp_2D_input_With_Broadcasting) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...i->i...");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("y", {2, 2}, {1.f, 3.f, 2.f, 4.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsBatchedTransposeOp_3D_input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...ji->...ij");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("y", {2, 2, 2}, {1.f, 3.f, 2.f, 4.f, 1.f, 3.f, 2.f, 4.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Implicit
|
||||
TEST(Einsum, ImplicitEinsumAsTransposeOp_2D_input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ji");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("y", {2, 2}, {1.f, 3.f, 2.f, 4.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ImplicitEinsumAsBatchedTransposeOp_3D_input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...ji");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("y", {2, 2, 2}, {1.f, 3.f, 2.f, 4.f, 1.f, 3.f, 2.f, 4.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Theme: Axis/Axes reduction
|
||||
|
||||
// Explicit
|
||||
TEST(Einsum, ExplicitEinsumAsReduceOp_2D_input_0) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ij->i");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("y", {2}, {3.f, 7.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsReduceOp_2D_input_1) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ij->j");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("y", {2}, {4.f, 6.f});
|
||||
test.Run();
|
||||
}
|
||||
TEST(Einsum, ExplicitEinsumAsBatchedReduceOp_3D_input_0) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...ji->...j");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("y", {2, 2}, {3.f, 7.f, 3.f, 7.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsBatchedReduceOp_3D_input_1) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...ji->...");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("y", {2}, {10.f, 10.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Implicit
|
||||
// Cannot do implicit reduction
|
||||
|
||||
// Theme: Outer Product
|
||||
|
||||
// Explicit
|
||||
TEST(Einsum, ExplicitEinsumAsOuterProductOp_2D_input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "i,j->ij");
|
||||
test.AddInput<float>("x", {2}, {1.f, 2.f});
|
||||
test.AddInput<float>("y", {2}, {3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {3.f, 4.f, 6.f, 8.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsOuterProductWithTransposeOp_2D_input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "i,j->ji");
|
||||
test.AddInput<float>("x", {2}, {1.f, 2.f});
|
||||
test.AddInput<float>("y", {2}, {3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {3.f, 6.f, 4.f, 8.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsOuterProductWithTransposeOp_Multi_Input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "i,j,k->jik");
|
||||
test.AddInput<float>("x", {2}, {1.f, 2.f});
|
||||
test.AddInput<float>("y", {2}, {3.f, 4.f});
|
||||
test.AddInput<float>("z", {2}, {5.f, 6.f});
|
||||
test.AddOutput<float>("o", {2, 2, 2}, {15.f, 18.f, 30.f, 36.f, 20.f, 24.f, 40.f, 48.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Implicit
|
||||
TEST(Einsum, ImplicitEinsumAsOuterProductOp_2D_input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "i,j,k");
|
||||
test.AddInput<float>("x", {2}, {1.f, 2.f});
|
||||
test.AddInput<float>("y", {2}, {3.f, 4.f});
|
||||
test.AddInput<float>("z", {2}, {5.f, 6.f});
|
||||
test.AddOutput<float>("o", {2, 2, 2}, {15.f, 18.f, 20.f, 24.f, 30.f, 36.f, 40.f, 48.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ImplicitEinsumAsOuterProductOp_Multi_Input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "i,j,k");
|
||||
test.AddInput<float>("x", {2}, {1.f, 2.f});
|
||||
test.AddInput<float>("y", {2}, {3.f, 4.f});
|
||||
test.AddInput<float>("z", {2}, {5.f, 6.f});
|
||||
test.AddOutput<float>("o", {2, 2, 2}, {15.f, 18.f, 20.f, 24.f, 30.f, 36.f, 40.f, 48.f});
|
||||
test.Run();
|
||||
}
|
||||
// Theme: MatMul
|
||||
|
||||
// Explicit
|
||||
TEST(Einsum, ExplicitEinsumAsMatmul) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ij,jk->ik");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {7.f, 10.f, 15.f, 22.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsMatmul_Multi_Input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ij,jk,kl->li");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("z", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {37.f, 81.f, 54.f, 118.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsBatchedMatmul) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "bij,bjk->bik");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("y", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2, 2}, {7.f, 10.f, 15.f, 22.f, 7.f, 10.f, 15.f, 22.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsBatchedMatmulWithBroadcasting_0) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...ij,...jk->...ik");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("y", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2, 2}, {7.f, 10.f, 15.f, 22.f, 7.f, 10.f, 15.f, 22.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsBatchedMatmulWithBroadcasting_1) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...ij,bjk->...ik");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("y", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2, 2}, {14.f, 20.f, 30.f, 44.f, 14.f, 20.f, 30.f, 44.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsMatmul_OutputTransposed) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ij,jk->ki");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {7.f, 15.f, 10.f, 22.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsMatmul_2) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ij,jk->ik");
|
||||
test.AddInput<float>("x", {2, 1}, {2.f, 3.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {8.f, 12.f, 12.f, 18.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Implicit
|
||||
TEST(Einsum, ImplicitEinsumAsMatmul) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ij,jk");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {7.f, 10.f, 15.f, 22.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ImplicitEinsumAsMatmul_Multi_Input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ij,jk,kl");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("z", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {37.f, 54.f, 81.f, 118.f});
|
||||
test.Run();
|
||||
}
|
||||
TEST(Einsum, ImplicitEinsumAsBatchedMatmul) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "bij,bjk");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("y", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {14.f, 20.f, 30.f, 44.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ImplicitEinsumAsBatchedMatmulWithBroadcasting_0) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...ij,...jk");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("y", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2, 2}, {7.f, 10.f, 15.f, 22.f, 7.f, 10.f, 15.f, 22.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ImplicitEinsumAsMatmul_2) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ij,jk");
|
||||
test.AddInput<float>("x", {2, 1}, {2.f, 3.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {8.f, 12.f, 12.f, 18.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, DiagonalWithMatmul) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "iij, jk");
|
||||
test.AddInput<float>("x", {2, 2, 3}, {1.f, 2.f, 3.f, 1.f, 2.f, 3.f, 1.f, 2.f, 3.f, 1.f, 2.f, 3.f});
|
||||
test.AddInput<float>("y", {3, 3}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f});
|
||||
test.AddOutput<float>("o", {3}, {60.f, 72.f, 84.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Theme: Diagonal parsing
|
||||
|
||||
// Explicit
|
||||
TEST(Einsum, ExplicitEinsumAsDiagonalOp) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ii->i");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2}, {1.f, 4.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsDiagonalOp_1) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "iii->i");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2}, {1.f, 4.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsDiagonalOpWithAxisReduced) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "iji->j");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2}, {3.f, 7.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsDiagonalOpWithAxisPreserved) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "iji->ij");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {1.f, 3.f, 2.f, 4.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsDiagonalOpWithTranspose) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "iji->ji");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsDiagonalOpWithTranspose_double) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "iji->ji");
|
||||
test.AddInput<double>("x", {2, 2, 2}, {1., 2., 3., 4., 1., 2., 3., 4.});
|
||||
test.AddOutput<double>("o", {2, 2}, {1., 2., 3., 4.});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsDiagonalOpWithTranspose_int32) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "iji->ji");
|
||||
test.AddInput<int32_t>("x", {2, 2, 2}, {1, 2, 3, 4, 1, 2, 3, 4});
|
||||
test.AddOutput<int32_t>("o", {2, 2}, {1, 2, 3, 4});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsDiagonalOpWithTranspose_int64) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "iji->ji");
|
||||
test.AddInput<int64_t>("x", {2, 2, 2}, {1, 2, 3, 4, 1, 2, 3, 4});
|
||||
test.AddOutput<int64_t>("o", {2, 2}, {1, 2, 3, 4});
|
||||
test.Run();
|
||||
}
|
||||
TEST(Einsum, ExplicitEinsumAsBatchedDiagonalOp) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...ii->...i");
|
||||
test.AddInput<float>("x", {3, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {3, 2}, {1.f, 4.f, 1.f, 4.f, 1.f, 4.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsBatchedDiagonalOp_1) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...iij->...j");
|
||||
test.AddInput<float>("x", {2, 2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {4.f, 6.f, 4.f, 6.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Implicit (Implicit diagonal ops will sum up diagonal values)
|
||||
TEST(Einsum, ImplicitEinsumAsDiagonalOp) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ii");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {}, {5.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ImplicitEinsumAsDiagonalOp_1) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "iii");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {}, {5.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ImplicitEinsumAsDiagonalOpWithAxisReduced) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "iji");
|
||||
test.AddInput<float>("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2}, {3.f, 7.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ImplicitEinsumAsBatchedDiagonalOp) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...ii");
|
||||
test.AddInput<float>("x", {2, 1, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 1}, {5.f, 5.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ImplicitEinsumAsBatchedDiagonalOp_1) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "...iij");
|
||||
test.AddInput<float>("x", {2, 2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {4.f, 6.f, 4.f, 6.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Theme: Scalar inputs and outputs
|
||||
|
||||
// Explicit
|
||||
TEST(Einsum, ExplicitEinsumAsElementwiseMulOpWithOneScalar) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", ",...i->...i");
|
||||
test.AddInput<float>("x", {}, {10.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {10.f, 20.f, 30.f, 40.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumAsElementwiseMulOpWithTwoScalars_Multi_Input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", ",...i,->...i");
|
||||
test.AddInput<float>("x", {}, {10.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("z", {}, {10.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {100.f, 200.f, 300.f, 400.f});
|
||||
test.Run();
|
||||
}
|
||||
TEST(Einsum, ExplicitEinsumAsElementwiseMulOpWithAllScalars) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", ",->");
|
||||
test.AddInput<float>("x", {}, {10.f});
|
||||
test.AddInput<float>("y", {}, {2.f});
|
||||
test.AddOutput<float>("o", {}, {20.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ExplicitEinsumReduceAxesInInputToScalarOutput) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "ij->");
|
||||
test.AddInput<float>("x", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {}, {10.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Implicit
|
||||
TEST(Einsum, ImplicitEinsumAsElementwiseMulOpWithOneScalar) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", ",...i");
|
||||
test.AddInput<float>("x", {}, {10.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {10.f, 20.f, 30.f, 40.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(Einsum, ImplicitEinsumAsElementwiseMulOpWithThreeScalars_Multi_Input) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", ",...i,,");
|
||||
test.AddInput<float>("a", {}, {10.f});
|
||||
test.AddInput<float>("b", {2, 2}, {1.f, 2.f, 3.f, 4.f});
|
||||
test.AddInput<float>("c", {}, {10.f});
|
||||
test.AddInput<float>("d", {}, {10.f});
|
||||
test.AddOutput<float>("o", {2, 2}, {1000.f, 2000.f, 3000.f, 4000.f});
|
||||
test.Run();
|
||||
}
|
||||
TEST(Einsum, ImplicitEinsumAsElementwiseMulOpWithAllScalars) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", ",");
|
||||
test.AddInput<float>("x", {}, {10.f});
|
||||
test.AddInput<float>("y", {}, {2.f});
|
||||
test.AddOutput<float>("o", {}, {20.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Tensor Contraction
|
||||
|
||||
// Explicit
|
||||
TEST(Einsum, ExplicitEinsumAsTensorContraction) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "abcd,ea->bcde");
|
||||
test.AddInput<float>("x", {2, 2, 2, 2}, {1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 1.f, 2.f});
|
||||
test.AddOutput<float>("o", {2, 2, 2, 2}, {3.f, 3.f, 6.f, 6.f, 3.f, 3.f, 6.f, 6.f, 3.f, 3.f, 6.f, 6.f, 3.f, 3.f, 6.f, 6.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
// Implicit
|
||||
TEST(Einsum, ImplicitEinsumAsTensorContraction) {
|
||||
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
|
||||
test.AddAttribute<std::string>("equation", "abcd,ea");
|
||||
test.AddInput<float>("x", {2, 2, 2, 2}, {1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f});
|
||||
test.AddInput<float>("y", {2, 2}, {1.f, 2.f, 1.f, 2.f});
|
||||
test.AddOutput<float>("o", {2, 2, 2, 2}, {3.f, 3.f, 6.f, 6.f, 3.f, 3.f, 6.f, 6.f, 3.f, 3.f, 6.f, 6.f, 3.f, 3.f, 6.f, 6.f});
|
||||
test.Run();
|
||||
}
|
||||
|
||||
} // namespace test
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -7,11 +7,6 @@
|
|||
"^test_batchnorm_epsilon_training_mode",
|
||||
"^test_batchnorm_example_old",
|
||||
"^test_batchnorm_example_training_mode",
|
||||
"^test_einsum_batch_diagonal",
|
||||
"^test_einsum_batch_matmul",
|
||||
"^test_einsum_inner_prod",
|
||||
"^test_einsum_sum",
|
||||
"^test_einsum_transpose",
|
||||
"^test_gathernd_example_int32_batch_dim1",
|
||||
"^test_max_int16",
|
||||
"^test_max_int8",
|
||||
|
|
|
|||
Loading…
Reference in a new issue