diff --git a/onnxruntime/core/providers/cpu/controlflow/scan_9.cc b/onnxruntime/core/providers/cpu/controlflow/scan_9.cc index 1564905d59..660a3699c1 100644 --- a/onnxruntime/core/providers/cpu/controlflow/scan_9.cc +++ b/onnxruntime/core/providers/cpu/controlflow/scan_9.cc @@ -191,7 +191,11 @@ Scan<9>::Scan(const OpKernelInfo& info) : OpKernel(info) { output_axes_ = std::vector(num_scan_outputs, 0); } - device_helpers_.transpose_func = TransposeBase::DoTranspose; + device_helpers_.transpose_func = [](const std::vector& 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(); diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index 279f1f225e..f7cb7c048f 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -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, BuildKernelCreateInfo, + ArgMax)>, BuildKernelCreateInfo, + ArgMax)>, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - // OpSet 12 BuildKernelCreateInfo, + ArgMax)>, BuildKernelCreateInfo, + ArgMax)>, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, - + BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -1084,8 +1083,8 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { ReduceMin)>, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/core/providers/cpu/math/einsum.cc b/onnxruntime/core/providers/cpu/math/einsum.cc new file mode 100644 index 0000000000..aec5a6c3de --- /dev/null +++ b/onnxruntime/core/providers/cpu/math/einsum.cc @@ -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 inputs; + inputs.reserve(num_inputs); + for (int i = 0; i < num_inputs; ++i) { + inputs.push_back(context->Input(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()) { + return EinsumTypedComputeProcessor(context, allocator, einsum_compute_preprocessor); + } else if (inputs[0]->IsDataType()) { + return EinsumTypedComputeProcessor(context, allocator, einsum_compute_preprocessor); + } else if (inputs[0]->IsDataType()) { + return EinsumTypedComputeProcessor(context, allocator, einsum_compute_preprocessor); + } else if (inputs[0]->IsDataType()) { + return EinsumTypedComputeProcessor(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 diff --git a/onnxruntime/core/providers/cpu/math/einsum.h b/onnxruntime/core/providers/cpu/math/einsum.h new file mode 100644 index 0000000000..31ad20f4eb --- /dev/null +++ b/onnxruntime/core/providers/cpu/math/einsum.h @@ -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("equation", &equation_).IsOK(), "Missing 'equation' attribute"); + einsum_equation_preprocessor_ = onnxruntime::make_unique(equation_); + } + + Status Compute(OpKernelContext* context) const override; + + private: + std::string equation_; + std::unique_ptr einsum_equation_preprocessor_; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/einsum_auxiliary_ops.cc b/onnxruntime/core/providers/cpu/math/einsum_auxiliary_ops.cc new file mode 100644 index 0000000000..7aca74cfa8 --- /dev/null +++ b/onnxruntime/core/providers/cpu/math/einsum_auxiliary_ops.cc @@ -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 Transpose(const Tensor& input, const std::vector& input_shape_override, + const std::vector& 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 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 output = onnxruntime::make_unique(input.DataType(), output_dims, allocator); + + TensorShape overriden_shape(input_shape_override); + TransposeBase::DoTranspose(permutation, input, *output, &overriden_shape); + + return output; +} + +template +std::unique_ptr MatMul(const Tensor& input_1, const std::vector& input_shape_1_override, + const Tensor& input_2, const std::vector& 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(input_shape_1_override[0]); + size_t M = static_cast(input_shape_1_override[1]); + size_t K = static_cast(input_shape_1_override[2]); + size_t N = static_cast(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 output_dims; + output_dims.reserve(3); + output_dims.push_back(static_cast(batches)); + output_dims.push_back(static_cast(M)); + output_dims.push_back(static_cast(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 output = onnxruntime::make_unique(input_1.DataType(), output_dims, allocator); + + const T* input_1_data = input_1.template Data(); + const T* input_2_data = input_2.template Data(); + T* output_data = output->template MutableData(); + + // 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( + static_cast(M), + static_cast(N), + static_cast(K), + input_1_data + i * left_offset, + input_2_data + i * right_offset, + output_data + i * output_offset, tp); + } + + return output; +} + +template +std::unique_ptr ReduceSum(const Tensor& input, const std::vector& input_shape_override, + const std::vector& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp) { + TensorShape overriden_shape(input_shape_override); + auto output = onnxruntime::ReduceSum::Impl(input, reduce_axes, allocator, tp, true, &overriden_shape); + return onnxruntime::make_unique(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 +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 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 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 output = onnxruntime::make_unique(input.DataType(), output_dims, allocator); + + switch (element_size_in_bytes) { + case 4: + DiagonalDataAssignment(reinterpret_cast(input.DataRaw()), reinterpret_cast(output->MutableDataRaw()), + batch_size, base_stride, inner_stride); + break; + case 8: + DiagonalDataAssignment(reinterpret_cast(input.DataRaw()), reinterpret_cast(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 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(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 output; + bool preserve_innermost_dim_val = false; + + bool is_transpose_required = IsTransposeRequiredForDiagonal(dim_1, dim_2, rank); + if (is_transpose_required) { + std::vector 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 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 MatMul(const Tensor& input_1, const std::vector& input_shape_1_override, + const Tensor& input_2, const std::vector& input_shape_2_override, + AllocatorPtr allocator, concurrency::ThreadPool* tp); +template std::unique_ptr ReduceSum(const Tensor& input, const std::vector& input_shape_override, + const std::vector& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp); + +// int32_t +template std::unique_ptr MatMul(const Tensor& input_1, const std::vector& input_shape_1_override, + const Tensor& input_2, const std::vector& input_shape_2_override, + AllocatorPtr allocator, concurrency::ThreadPool* tp); +template std::unique_ptr ReduceSum(const Tensor& input, const std::vector& input_shape_override, + const std::vector& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp); + +// double +template std::unique_ptr MatMul(const Tensor& input_1, const std::vector& input_shape_1_override, + const Tensor& input_2, const std::vector& input_shape_2_override, + AllocatorPtr allocator, concurrency::ThreadPool* tp); +template std::unique_ptr ReduceSum(const Tensor& input, const std::vector& input_shape_override, + const std::vector& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp); + +// int64_t +template std::unique_ptr MatMul(const Tensor& input_1, const std::vector& input_shape_1_override, + const Tensor& input_2, const std::vector& input_shape_2_override, + AllocatorPtr allocator, concurrency::ThreadPool* tp); +template std::unique_ptr ReduceSum(const Tensor& input, const std::vector& input_shape_override, + const std::vector& reduce_axes, AllocatorPtr allocator, concurrency::ThreadPool* tp); + +} // namespace EinsumOp +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/einsum_auxiliary_ops.h b/onnxruntime/core/providers/cpu/math/einsum_auxiliary_ops.h new file mode 100644 index 0000000000..47b77c377f --- /dev/null +++ b/onnxruntime/core/providers/cpu/math/einsum_auxiliary_ops.h @@ -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 + +namespace onnxruntime { + +namespace EinsumOp { + +// Thin wrapper over the Transpose op +std::unique_ptr Transpose(const Tensor& input, const std::vector& input_shape_override, + const std::vector& 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 +std::unique_ptr MatMul(const Tensor& input_1, const std::vector& input_1_shape_override, + const Tensor& input_2, const std::vector& input_2_shape_override, + AllocatorPtr allocator, concurrency::ThreadPool* tp); + +// Thin wrapper over the ReduceSum op +template +std::unique_ptr ReduceSum(const Tensor& input, const std::vector& input_shape_override, + const std::vector& 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 Diagonal(const Tensor& input, int64_t dim_1, int64_t dim_2, AllocatorPtr allocator); + +} // namespace EinsumOp + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/einsum_utils.cc b/onnxruntime/core/providers/cpu/math/einsum_utils.cc new file mode 100644 index 0000000000..8427d2eecb --- /dev/null +++ b/onnxruntime/core/providers/cpu/math/einsum_utils.cc @@ -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& 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 +static void FinalizeOutput(const Tensor& candidate_output, const std::vector& ordered_subscript_indices_in_candidate, + const std::vector& 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 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 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 +static std::unique_ptr PairwiseOperandProcess(const Tensor& left, + const TensorShape& left_shape_override, + const Tensor& right, + const TensorShape& right_shape_override, + const std::vector& 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(left_dims.size()); + int64_t right_rank = static_cast(right_dims.size()); + + std::unique_ptr current_left; + std::unique_ptr 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 lro; + std::vector lo; + std::vector 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(left, left_dims, {i}, allocator, tp); + } else if (has_right_dim) { + current_right = ReduceSum(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 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 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 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 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 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(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(*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& 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& EinsumComputePreprocessor::GetOutputDims() const { + return output_dims_; +} + +std::vector>& EinsumComputePreprocessor::GetPreprocessedInputTensors() { + return preprocessed_inputs_; +} + +const std::vector& EinsumComputePreprocessor::GetRawInputTensors() { + return inputs_; +} + +const std::vector& EinsumComputePreprocessor::GetHomogenizedInputDims() { + return homogenized_input_dims_; +} + +const std::vector& EinsumComputePreprocessor::GetMappedSubscriptIndicesToLastInputIndex() const { + return subscript_indices_to_last_input_; +} + +const std::vector& 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 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(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(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(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 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 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 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(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(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('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 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(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 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 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 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 permutation; + permutation.reserve(input_dims.size()); + for (auto& d : subscript_indices_to_input_index) { + if (d != -1) { + permutation.push_back(static_cast(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 +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 result; + + { + std::vector reduced_dims; + std::vector preserved_dims; // dims which were not reduced + std::vector 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(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(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 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(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(OpKernelContext* context, AllocatorPtr allocator, EinsumComputePreprocessor& einsum_compute_preprocessor); + +// int32_t +template Status EinsumTypedComputeProcessor(OpKernelContext* context, AllocatorPtr allocator, EinsumComputePreprocessor& einsum_compute_preprocessor); + +// double +template Status EinsumTypedComputeProcessor(OpKernelContext* context, AllocatorPtr allocator, EinsumComputePreprocessor& einsum_compute_preprocessor); + +// int64_t +template Status EinsumTypedComputeProcessor(OpKernelContext* context, AllocatorPtr allocator, EinsumComputePreprocessor& einsum_compute_preprocessor); + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/einsum_utils.h b/onnxruntime/core/providers/cpu/math/einsum_utils.h new file mode 100644 index 0000000000..d0e9a1df6c --- /dev/null +++ b/onnxruntime/core/providers/cpu/math/einsum_utils.h @@ -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 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& 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& 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>& GetPreprocessedInputTensors(); + + // Get raw inputs to the op + const std::vector& GetRawInputTensors(); + + // Get the "homogenized input dims" for each preprocessed/raw input + const std::vector& GetHomogenizedInputDims(); + + // For each subscript index, hold the last input the subscript index was seen in + const std::vector& GetMappedSubscriptIndicesToLastInputIndex() const; + + // For each subscript index, hold the index it corresponds to in the output's shape + const std::vector& 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& inputs_; + + // All preprocessed inputs + std::vector> preprocessed_inputs_; + + // Holds the preprocessed inputs' homogenized dims + std::vector 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 letter_to_count_; + + // Hold the assigned index corresponding to the letter seen + // `-1` means the corresponding letter wasn't seen at all + std::array 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 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 subscript_indices_to_dim_value_; + + // Holds the final calculated output dimensions + std::vector output_dims_; + + // All subscript indices in the equation for each input + std::vector> 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 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 +Status EinsumTypedComputeProcessor(OpKernelContext* context, AllocatorPtr allocator, + EinsumComputePreprocessor& einsum_compute_preprocessor); + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc b/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc index 92a976be82..f89e368e3b 100644 --- a/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc @@ -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 -bool PrepareForReduce(OpKernelContext* ctx, - FastAllocVector& transposedInputData, - Tensor** reducedTensor, +bool PrepareForReduce(const Tensor* input_tensor_ptr, + FastAllocVector& transposed_input_data, int64_t& block_size, int64_t& blocks, const std::vector& axes_, bool keepdims_, - bool check_no_transpose = false) { - const auto* input_tensor_ptr = ctx->Input(0); - ORT_ENFORCE(input_tensor_ptr != nullptr); - const Tensor& input = *input_tensor_ptr; + /*out*/ std::vector& 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(); } block_size = blocks = 1; - *reducedTensor = ctx->Output(0, input.Shape()); return true; } @@ -224,11 +231,11 @@ bool PrepareForReduce(OpKernelContext* ctx, std::vector 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(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(); - 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 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 Status ReduceL1::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; + std::vector reduced_dims; - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); + const Tensor* input = ctx->Input(0); + + bool no_transpose = PrepareForReduce(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(); if (no_transpose) { - const T* input_data = ctx->Input(0)->template Data(); + const T* input_data = input->template Data(); for (int64_t i = 0; i < block_size; ++i) { output_data[i] = ConstEigenVectorMap(input_data + (i * blocks), blocks).cwiseAbs().sum(); } } else { EigenVectorMap out_vec(output_data, block_size); - out_vec = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks).cwiseAbs().rowwise().sum(); + out_vec = ConstEigenMatrixMap(&transposed_input_data[0], block_size, blocks).cwiseAbs().rowwise().sum(); } return Status::OK(); @@ -381,24 +390,28 @@ Status ReduceL1::Compute(OpKernelContext* ctx) const { template Status ReduceL2::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; + std::vector reduced_dims; - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); + const Tensor* input = ctx->Input(0); + + bool no_transpose = PrepareForReduce(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(); if (no_transpose) { - const T* input_data = ctx->Input(0)->template Data(); + const T* input_data = input->template Data(); for (int64_t i = 0; i < block_size; ++i) { output_data[i] = ConstEigenVectorMap(input_data + (i * blocks), blocks).norm(); } } else { EigenVectorMap out_vec(output_data, block_size); - out_vec = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks).rowwise().norm(); + out_vec = ConstEigenMatrixMap(&transposed_input_data[0], block_size, blocks).rowwise().norm(); } return Status::OK(); @@ -406,24 +419,28 @@ Status ReduceL2::Compute(OpKernelContext* ctx) const { template Status ReduceLogSum::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; + std::vector reduced_dims; - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); + const Tensor* input = ctx->Input(0); + + bool no_transpose = PrepareForReduce(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(); if (no_transpose) { - const T* input_data = ctx->Input(0)->template Data(); + const T* input_data = input->template Data(); for (int64_t i = 0; i < block_size; ++i) { output_data[i] = ConstEigenVectorMap(input_data + (i * blocks), blocks).sum(); } } else { EigenVectorMap out_vec(output_data, block_size); - out_vec = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks).rowwise().sum(); + out_vec = ConstEigenMatrixMap(&transposed_input_data[0], block_size, blocks).rowwise().sum(); } for (int j = 0; j < block_size; ++j) { @@ -436,22 +453,26 @@ Status ReduceLogSum::Compute(OpKernelContext* ctx) const { template Status ReduceLogSumExp::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; - PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_); + std::vector reduced_dims; + const Tensor* input = ctx->Input(0); + + PrepareForReduce(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims); + + Tensor* reduced = ctx->Output(0, reduced_dims); T* output_data = reduced->template MutableData(); for (int j = 0; j < block_size; ++j) { T max_value = std::numeric_limits::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(std::exp(transposedInputData[i * block_size + j] - max_value)); + scaled_exp_sum += static_cast(std::exp(transposed_input_data[i * block_size + j] - max_value)); } *(output_data++) = static_cast(std::log(scaled_exp_sum) + max_value); } @@ -460,23 +481,27 @@ Status ReduceLogSumExp::Compute(OpKernelContext* ctx) const { template Status ReduceMax::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); + std::vector reduced_dims; + const Tensor* input = ctx->Input(0); + + bool no_transpose = PrepareForReduce(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(); if (no_transpose) { - const T* input_data = ctx->Input(0)->template Data(); + const T* input_data = input->template Data(); for (int64_t i = 0; i < block_size; ++i) { output_data[i] = ConstEigenVectorMap(input_data + (i * blocks), blocks).maxCoeff(); } } else { EigenVectorMap out_vec(output_data, block_size); - out_vec = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks).rowwise().maxCoeff(); + out_vec = ConstEigenMatrixMap(&transposed_input_data[0], block_size, blocks).rowwise().maxCoeff(); } return Status::OK(); @@ -484,11 +509,15 @@ Status ReduceMax::Compute(OpKernelContext* ctx) const { template Status ReduceMean::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); + std::vector reduced_dims; + const Tensor* input = ctx->Input(0); + + bool no_transpose = PrepareForReduce(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(); @@ -500,7 +529,7 @@ Status ReduceMean::Compute(OpKernelContext* ctx) const { concurrency::ThreadPool::TryBatchParallelFor(ctx->GetOperatorThreadPool(), block_size, lambda, 0); } else { EigenVectorMap out_vec(output_data, block_size); - out_vec = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks).rowwise().mean(); + out_vec = ConstEigenMatrixMap(&transposed_input_data[0], block_size, blocks).rowwise().mean(); } return Status::OK(); @@ -508,23 +537,27 @@ Status ReduceMean::Compute(OpKernelContext* ctx) const { template Status ReduceMin::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); + std::vector reduced_dims; + const Tensor* input = ctx->Input(0); + + bool no_transpose = PrepareForReduce(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(); if (no_transpose) { - const T* input_data = ctx->Input(0)->template Data(); + const T* input_data = input->template Data(); for (int64_t i = 0; i < block_size; ++i) { output_data[i] = ConstEigenVectorMap(input_data + (i * blocks), blocks).minCoeff(); } } else { EigenVectorMap out_vec(output_data, block_size); - out_vec = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks).rowwise().minCoeff(); + out_vec = ConstEigenMatrixMap(&transposed_input_data[0], block_size, blocks).rowwise().minCoeff(); } return Status::OK(); @@ -532,75 +565,110 @@ Status ReduceMin::Compute(OpKernelContext* ctx) const { template Status ReduceProd::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; + std::vector reduced_dims; + const Tensor* input = ctx->Input(0); - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); + bool no_transpose = PrepareForReduce(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(); if (no_transpose) { - const T* input_data = ctx->Input(0)->template Data(); + const T* input_data = input->template Data(); for (int64_t i = 0; i < block_size; ++i) { output_data[i] = ConstEigenVectorMap(input_data + (i * blocks), blocks).prod(); } } else { EigenVectorMap out_vec(output_data, block_size); - out_vec = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks).rowwise().prod(); + out_vec = ConstEigenMatrixMap(&transposed_input_data[0], block_size, blocks).rowwise().prod(); } return Status::OK(); } template -Status ReduceSum::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); - int64_t block_size; - int64_t blocks; - Tensor* reduced; - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); - - T* output_data = reduced->template MutableData(); - +static void ReduceSumCore(const T* input_data, T* output_data, bool no_transpose, + int64_t blocks, int64_t block_size, FastAllocVector& transposed_input_data, + concurrency::ThreadPool* tp) { if (no_transpose) { - const T* input_data = ctx->Input(0)->template Data(); 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>(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 out_vec(output_data, block_size); - out_vec = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks).rowwise().sum(); + out_vec = ConstEigenMatrixMap(&transposed_input_data[0], block_size, blocks).rowwise().sum(); } +} + +template +Tensor ReduceSum::Impl(const Tensor& input, const std::vector& reduce_axes, + AllocatorPtr allocator, concurrency::ThreadPool* tp, bool keep_dims, + const TensorShape* input_shape_override) { + FastAllocVector transposed_input_data(allocator); + int64_t block_size; + int64_t blocks; + std::vector reduced_dims; + + bool no_transpose = PrepareForReduce(&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(), output.template MutableData(), + no_transpose, blocks, block_size, transposed_input_data, tp); + + return output; +} + +template +Status ReduceSum::Compute(OpKernelContext* ctx) const { + FastAllocVector transposed_input_data(GetAllocator(*ctx)); + int64_t block_size; + int64_t blocks; + std::vector reduced_dims; + const Tensor* input = ctx->Input(0); + + bool no_transpose = PrepareForReduce(input, transposed_input_data, block_size, blocks, axes_, keepdims_, reduced_dims, true); + + auto* output = ctx->Output(0, reduced_dims); + + ReduceSumCore(input->template Data(), output->template MutableData(), + no_transpose, blocks, block_size, transposed_input_data, ctx->GetOperatorThreadPool()); return Status::OK(); } template Status ReduceSumSquare::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; + std::vector reduced_dims; + const Tensor* input = ctx->Input(0); - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); + bool no_transpose = PrepareForReduce(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(); if (no_transpose) { - const T* input_data = ctx->Input(0)->template Data(); + const T* input_data = input->template Data(); for (int64_t i = 0; i < block_size; ++i) { output_data[i] = ConstEigenVectorMap(input_data + (i * blocks), blocks).squaredNorm(); } } else { EigenVectorMap out_vec(output_data, block_size); - out_vec = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks).rowwise().squaredNorm(); + out_vec = ConstEigenMatrixMap(&transposed_input_data[0], block_size, blocks).rowwise().squaredNorm(); } return Status::OK(); @@ -608,13 +676,16 @@ Status ReduceSumSquare::Compute(OpKernelContext* ctx) const { template Status ArgMax::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); + std::vector reduced_dims; + const Tensor* input = ctx->Input(0); + bool no_transpose = PrepareForReduce(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(); Eigen::MatrixXf::Index maxIndex; @@ -642,7 +713,7 @@ Status ArgMax::Compute(OpKernelContext* ctx) const { } } } else { - auto matrixData = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks); + auto matrixData = ConstEigenMatrixMap(&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::Compute(OpKernelContext* ctx) const { template Status ArgMin::Compute(OpKernelContext* ctx) const { - FastAllocVector transposedInputData(GetAllocator(*ctx)); + FastAllocVector transposed_input_data(GetAllocator(*ctx)); int64_t block_size; int64_t blocks; - Tensor* reduced; - bool no_transpose = PrepareForReduce(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true); + std::vector reduced_dims; + const Tensor* input = ctx->Input(0); + bool no_transpose = PrepareForReduce(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(); Eigen::MatrixXf::Index minIndex; @@ -703,7 +777,7 @@ Status ArgMin::Compute(OpKernelContext* ctx) const { } } } else { - auto matrixData = ConstEigenMatrixMap(&transposedInputData[0], block_size, blocks); + auto matrixData = ConstEigenMatrixMap(&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::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; +template class ReduceSum; +template class ReduceSum; +template class ReduceSum; + } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/reduction/reduction_ops.h b/onnxruntime/core/providers/cpu/reduction/reduction_ops.h index 1afcb69399..ea8b17fc41 100644 --- a/onnxruntime/core/providers/cpu/reduction/reduction_ops.h +++ b/onnxruntime/core/providers/cpu/reduction/reduction_ops.h @@ -121,6 +121,12 @@ class ReduceSum final : public ReduceKernel { } 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& reduce_axes, + AllocatorPtr allocator, concurrency::ThreadPool* tp, bool keep_dims, + const TensorShape* input_shape_override = nullptr); }; template diff --git a/onnxruntime/core/providers/cpu/tensor/transpose.cc b/onnxruntime/core/providers/cpu/tensor/transpose.cc index 3978db4c8c..9a6466835c 100644 --- a/onnxruntime/core/providers/cpu/tensor/transpose.cc +++ b/onnxruntime/core/providers/cpu/tensor/transpose.cc @@ -174,8 +174,10 @@ static void DoTransposeEltWise(int64_t num_axes, const std::vector& tar } } -static Status DoUntypedTranspose(const std::vector& 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& 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& 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& 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& permutations, } } +// `input_shape_override` overrides the shape of `input` for compute purposes. static void SingleAxisTranspose(const std::vector& 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& permutations, size_t& return single_axis_moved; } -Status TransposeBase::DoTranspose(const std::vector& permutations, const Tensor& input, Tensor& output) { +//`input_shape_override` overrides the shape of `input` for compute purposes. +Status TransposeBase::DoTranspose(const std::vector& 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& 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); } } diff --git a/onnxruntime/core/providers/cpu/tensor/transpose.h b/onnxruntime/core/providers/cpu/tensor/transpose.h index 6fdd0bf1a3..3cb56f6ddc 100644 --- a/onnxruntime/core/providers/cpu/tensor/transpose.h +++ b/onnxruntime/core/providers/cpu/tensor/transpose.h @@ -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& permutations, const Tensor& input, Tensor& output); + static Status DoTranspose(const std::vector& permutations, const Tensor& input, Tensor& output, + const TensorShape* input_shape_override = nullptr); protected: TransposeBase(const OpKernelInfo& info) { diff --git a/onnxruntime/test/providers/cpu/math/einsum_test.cc b/onnxruntime/test/providers/cpu/math/einsum_test.cc new file mode 100644 index 0000000000..cf030e60f1 --- /dev/null +++ b/onnxruntime/test/providers/cpu/math/einsum_test.cc @@ -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("equation", "i->i"); + test.AddInput("x", {5}, {0.9f, 2.5f, 2.3f, 1.5f, -4.5f}); + test.AddOutput("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("equation", "i"); + test.AddInput("x", {5}, {0.9f, 2.5f, 2.3f, 1.5f, -4.5f}); + test.AddOutput("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("equation", "ji->ij"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "...i->i..."); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "...ji->...ij"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "ji"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "...ji"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "ij->i"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("y", {2}, {3.f, 7.f}); + test.Run(); +} + +TEST(Einsum, ExplicitEinsumAsReduceOp_2D_input_1) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "ij->j"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("y", {2}, {4.f, 6.f}); + test.Run(); +} +TEST(Einsum, ExplicitEinsumAsBatchedReduceOp_3D_input_0) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "...ji->...j"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "...ji->..."); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "i,j->ij"); + test.AddInput("x", {2}, {1.f, 2.f}); + test.AddInput("y", {2}, {3.f, 4.f}); + test.AddOutput("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("equation", "i,j->ji"); + test.AddInput("x", {2}, {1.f, 2.f}); + test.AddInput("y", {2}, {3.f, 4.f}); + test.AddOutput("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("equation", "i,j,k->jik"); + test.AddInput("x", {2}, {1.f, 2.f}); + test.AddInput("y", {2}, {3.f, 4.f}); + test.AddInput("z", {2}, {5.f, 6.f}); + test.AddOutput("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("equation", "i,j,k"); + test.AddInput("x", {2}, {1.f, 2.f}); + test.AddInput("y", {2}, {3.f, 4.f}); + test.AddInput("z", {2}, {5.f, 6.f}); + test.AddOutput("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("equation", "i,j,k"); + test.AddInput("x", {2}, {1.f, 2.f}); + test.AddInput("y", {2}, {3.f, 4.f}); + test.AddInput("z", {2}, {5.f, 6.f}); + test.AddOutput("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("equation", "ij,jk->ik"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "ij,jk,kl->li"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddInput("z", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {2, 2}, {37.f, 81.f, 54.f, 118.f}); + test.Run(); +} + +TEST(Einsum, ExplicitEinsumAsBatchedMatmul) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "bij,bjk->bik"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddInput("y", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "...ij,...jk->...ik"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddInput("y", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "...ij,bjk->...ik"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddInput("y", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "ij,jk->ki"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "ij,jk->ik"); + test.AddInput("x", {2, 1}, {2.f, 3.f}); + test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "ij,jk"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "ij,jk,kl"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddInput("z", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {2, 2}, {37.f, 54.f, 81.f, 118.f}); + test.Run(); +} +TEST(Einsum, ImplicitEinsumAsBatchedMatmul) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "bij,bjk"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddInput("y", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "...ij,...jk"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddInput("y", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "ij,jk"); + test.AddInput("x", {2, 1}, {2.f, 3.f}); + test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {2, 2}, {8.f, 12.f, 12.f, 18.f}); + test.Run(); +} + +TEST(Einsum, DiagonalWithMatmul) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "iij, jk"); + test.AddInput("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("y", {3, 3}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}); + test.AddOutput("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("equation", "ii->i"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {2}, {1.f, 4.f}); + test.Run(); +} + +TEST(Einsum, ExplicitEinsumAsDiagonalOp_1) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "iii->i"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {2}, {1.f, 4.f}); + test.Run(); +} + +TEST(Einsum, ExplicitEinsumAsDiagonalOpWithAxisReduced) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "iji->j"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {2}, {3.f, 7.f}); + test.Run(); +} + +TEST(Einsum, ExplicitEinsumAsDiagonalOpWithAxisPreserved) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "iji->ij"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {2, 2}, {1.f, 3.f, 2.f, 4.f}); + test.Run(); +} + +TEST(Einsum, ExplicitEinsumAsDiagonalOpWithTranspose) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "iji->ji"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", "iji->ji"); + test.AddInput("x", {2, 2, 2}, {1., 2., 3., 4., 1., 2., 3., 4.}); + test.AddOutput("o", {2, 2}, {1., 2., 3., 4.}); + test.Run(); +} + +TEST(Einsum, ExplicitEinsumAsDiagonalOpWithTranspose_int32) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "iji->ji"); + test.AddInput("x", {2, 2, 2}, {1, 2, 3, 4, 1, 2, 3, 4}); + test.AddOutput("o", {2, 2}, {1, 2, 3, 4}); + test.Run(); +} + +TEST(Einsum, ExplicitEinsumAsDiagonalOpWithTranspose_int64) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "iji->ji"); + test.AddInput("x", {2, 2, 2}, {1, 2, 3, 4, 1, 2, 3, 4}); + test.AddOutput("o", {2, 2}, {1, 2, 3, 4}); + test.Run(); +} +TEST(Einsum, ExplicitEinsumAsBatchedDiagonalOp) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "...ii->...i"); + test.AddInput("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("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("equation", "...iij->...j"); + test.AddInput("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("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("equation", "ii"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {}, {5.f}); + test.Run(); +} + +TEST(Einsum, ImplicitEinsumAsDiagonalOp_1) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "iii"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {}, {5.f}); + test.Run(); +} + +TEST(Einsum, ImplicitEinsumAsDiagonalOpWithAxisReduced) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "iji"); + test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {2}, {3.f, 7.f}); + test.Run(); +} + +TEST(Einsum, ImplicitEinsumAsBatchedDiagonalOp) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "...ii"); + test.AddInput("x", {2, 1, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {2, 1}, {5.f, 5.f}); + test.Run(); +} + +TEST(Einsum, ImplicitEinsumAsBatchedDiagonalOp_1) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "...iij"); + test.AddInput("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("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("equation", ",...i->...i"); + test.AddInput("x", {}, {10.f}); + test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", ",...i,->...i"); + test.AddInput("x", {}, {10.f}); + test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddInput("z", {}, {10.f}); + test.AddOutput("o", {2, 2}, {100.f, 200.f, 300.f, 400.f}); + test.Run(); +} +TEST(Einsum, ExplicitEinsumAsElementwiseMulOpWithAllScalars) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", ",->"); + test.AddInput("x", {}, {10.f}); + test.AddInput("y", {}, {2.f}); + test.AddOutput("o", {}, {20.f}); + test.Run(); +} + +TEST(Einsum, ExplicitEinsumReduceAxesInInputToScalarOutput) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "ij->"); + test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("o", {}, {10.f}); + test.Run(); +} + +// Implicit +TEST(Einsum, ImplicitEinsumAsElementwiseMulOpWithOneScalar) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", ",...i"); + test.AddInput("x", {}, {10.f}); + test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddOutput("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("equation", ",...i,,"); + test.AddInput("a", {}, {10.f}); + test.AddInput("b", {2, 2}, {1.f, 2.f, 3.f, 4.f}); + test.AddInput("c", {}, {10.f}); + test.AddInput("d", {}, {10.f}); + test.AddOutput("o", {2, 2}, {1000.f, 2000.f, 3000.f, 4000.f}); + test.Run(); +} +TEST(Einsum, ImplicitEinsumAsElementwiseMulOpWithAllScalars) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", ","); + test.AddInput("x", {}, {10.f}); + test.AddInput("y", {}, {2.f}); + test.AddOutput("o", {}, {20.f}); + test.Run(); +} + +// Tensor Contraction + +// Explicit +TEST(Einsum, ExplicitEinsumAsTensorContraction) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "abcd,ea->bcde"); + test.AddInput("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("y", {2, 2}, {1.f, 2.f, 1.f, 2.f}); + test.AddOutput("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("equation", "abcd,ea"); + test.AddInput("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("y", {2, 2}, {1.f, 2.f, 1.f, 2.f}); + test.AddOutput("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 diff --git a/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc b/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc index 3449828985..bcddd056d1 100644 --- a/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc +++ b/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc @@ -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",