diff --git a/onnxruntime/core/providers/cpu/tensor/concat.cc b/onnxruntime/core/providers/cpu/tensor/concat.cc index 102ee1acc8..3747df9f2f 100644 --- a/onnxruntime/core/providers/cpu/tensor/concat.cc +++ b/onnxruntime/core/providers/cpu/tensor/concat.cc @@ -2,9 +2,12 @@ // Licensed under the MIT License. #include "core/providers/cpu/tensor/concat.h" -#include "core/providers/common.h" + +#include "core/framework/element_type_lists.h" #include "core/framework/TensorSeq.h" +#include "core/providers/common.h" #include "core/providers/cpu/tensor/copy.h" +#include "core/providers/op_kernel_type_control.h" namespace onnxruntime { @@ -30,6 +33,24 @@ ONNX_CPU_OPERATOR_KERNEL( KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllTensorTypes()), Concat); +namespace op_kernel_type_control { +// we're using one set of types for all opsets +ORT_SPECIFY_OP_KERNEL_ARG_DEFAULT_TYPE_LIST_ALL_OPSETS( + kCpuExecutionProvider, kOnnxDomain, Concat, Input, 0, + element_type_lists::All); + +// Concat can be used with dimensions or indices so require int32_t and int64_t to be supported +ORT_SPECIFY_OP_KERNEL_ARG_REQUIRED_TYPES_ALL_OPSETS( + kCpuExecutionProvider, kOnnxDomain, Concat, Input, 0, int32_t, int64_t); +} // namespace op_kernel_type_control + +namespace { +using DataTypes = ORT_OP_KERNEL_ARG_DEFAULT_TYPE_LIST_ALL_OPSETS(kCpuExecutionProvider, kOnnxDomain, + Concat, Input, 0); +using EnabledDataTypes = ORT_OP_KERNEL_ARG_ENABLED_TYPE_LIST_ALL_OPSETS(kCpuExecutionProvider, kOnnxDomain, + Concat, Input, 0); +} // namespace + // this method will be shared between 'Concat' (CPU and GPU) and // 'ConcatFromSequence' ('concat' and 'stack' modes) to validate inputs Status ConcatBase::PrepareForCompute(OpKernelContext* ctx, @@ -244,13 +265,13 @@ Status ConcatBase::ComputeImpl(Prepare& p, OpKernelContext* ctx) const { continue; // parallel copy the data across - auto status = DispatchStridedCopy(ctx->GetOperatorThreadPool(), - *p.output_tensor, - initial_output_offset, - output_strides_for_copy, - prep.tensor->Shape(), - *prep.tensor, - StridesForTensor(*prep.tensor)); + auto status = DispatchStridedCopy(ctx->GetOperatorThreadPool(), + *p.output_tensor, + initial_output_offset, + output_strides_for_copy, + prep.tensor->Shape(), + *prep.tensor, + StridesForTensor(*prep.tensor)); ORT_RETURN_IF_ERROR(status); // advance along the axis that we are concatenating on (by the size of the axis of the tensor that we just copied) diff --git a/onnxruntime/core/providers/cpu/tensor/copy.cc b/onnxruntime/core/providers/cpu/tensor/copy.cc index 2eb1629554..ecfa96671a 100644 --- a/onnxruntime/core/providers/cpu/tensor/copy.cc +++ b/onnxruntime/core/providers/cpu/tensor/copy.cc @@ -20,142 +20,7 @@ std::vector StridesForTensor(const Tensor& tensor) { return strides; } -Status DispatchStridedCopy(concurrency::ThreadPool* thread_pool, - Tensor& dst, - std::ptrdiff_t dst_offset, - const std::vector dst_strides, - const TensorShape& copy_shape, - const Tensor& src, - const std::vector src_strides) { - ORT_ENFORCE(dst.DataType() == src.DataType(), "src and dst types must match"); - - // Manual dispatching: DispatchOnTensorType doesn't work here because we need to pass the type to the MutableData call -#define CALL_FOR_TYPE(T) \ - StridedCopy(thread_pool, dst.MutableData() + dst_offset, dst_strides, copy_shape, src.Data(), src_strides); \ - break - - auto tensor_type = dst.DataType()->AsPrimitiveDataType()->GetDataType(); - switch (tensor_type) { - case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: - CALL_FOR_TYPE(float); - case ONNX_NAMESPACE::TensorProto_DataType_BOOL: - CALL_FOR_TYPE(bool); - case ONNX_NAMESPACE::TensorProto_DataType_DOUBLE: - CALL_FOR_TYPE(double); - case ONNX_NAMESPACE::TensorProto_DataType_STRING: - CALL_FOR_TYPE(std::string); - case ONNX_NAMESPACE::TensorProto_DataType_INT8: - CALL_FOR_TYPE(int8_t); - case ONNX_NAMESPACE::TensorProto_DataType_UINT8: - CALL_FOR_TYPE(uint8_t); - case ONNX_NAMESPACE::TensorProto_DataType_INT16: - CALL_FOR_TYPE(int16_t); - case ONNX_NAMESPACE::TensorProto_DataType_UINT16: - CALL_FOR_TYPE(uint16_t); - case ONNX_NAMESPACE::TensorProto_DataType_INT32: - CALL_FOR_TYPE(int32_t); - case ONNX_NAMESPACE::TensorProto_DataType_UINT32: - CALL_FOR_TYPE(uint32_t); - case ONNX_NAMESPACE::TensorProto_DataType_INT64: - CALL_FOR_TYPE(int64_t); - case ONNX_NAMESPACE::TensorProto_DataType_UINT64: - CALL_FOR_TYPE(uint64_t); - case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: - CALL_FOR_TYPE(MLFloat16); - case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: - CALL_FOR_TYPE(BFloat16); - default: - ORT_ENFORCE(false, "Unknown tensor type of ", tensor_type); - } - return Status::OK(); -} - namespace { - -template -inline void Copy1DNonContiguous(T* dst, int64_t dst_stride, const T* src, int64_t src_stride, std::ptrdiff_t count) { - for (std::ptrdiff_t i = 0; i < count; i++) { - dst[0] = src[0]; - dst += dst_stride; - src += src_stride; - } -} - -template -inline void Copy1DContiguous(T* dst, const T* src, std::ptrdiff_t count) { - memcpy(dst, src, count * sizeof(T)); -} -template <> -inline void Copy1DContiguous(std::string* dst, const std::string* src, std::ptrdiff_t count) { - Copy1DNonContiguous(dst, 1, src, 1, count); -} - -template -inline void Copy1D(T* dst, int64_t dst_stride, const T* src, int64_t src_stride, std::ptrdiff_t count) { - if (dst_stride == 1 && src_stride == 1) { - Copy1DContiguous(dst, src, count); - } else { - Copy1DNonContiguous(dst, dst_stride, src, src_stride, count); - } -} - -template <> -inline void Copy1D(std::string* dst, int64_t dst_stride, const std::string* src, int64_t src_stride, std::ptrdiff_t count) { - // strings should always be copied using the for loop - Copy1DNonContiguous(dst, dst_stride, src, src_stride, count); -} - -struct NdCounter { - NdCounter(const std::vector& shape, std::ptrdiff_t first, std::ptrdiff_t last) - : dims(shape.size()), - last_dim_size(shape[dims - 1]), - current_offset(first), - last(last), - current_index(dims), - shape(shape) { - // compute the initial n-dimensional index - int64_t remaining_index = first; - // Iterate from dims to 1 so we don't roll over to positive on the bounds check - for (std::size_t dim = dims; dim > 0; dim--) { - auto shape_val = shape[dim - 1]; - current_index[dim - 1] = remaining_index % shape_val; - remaining_index /= shape_val; - } - } - - /* - Return the size of the largest step in the last dimension. - */ - std::ptrdiff_t NextStepSize() const { - auto elements_in_dimension = last_dim_size - current_index[dims - 1]; - std::ptrdiff_t span_end = std::min(last, current_offset + elements_in_dimension); - return span_end - current_offset; - } - - /* - Advance the counter by step_size elements. - */ - void Step(std::ptrdiff_t step_size) { - current_offset += step_size; - current_index[dims - 1] += step_size; - - // update the current_nd_idx if needed - std::size_t dim = dims - 1; - while (dim > 0 && current_index[dim] >= shape[dim]) { - current_index[dim] = 0; - dim--; - current_index[dim]++; - } - } - - const std::size_t dims; - const int64_t last_dim_size; - ptrdiff_t current_offset; - const ptrdiff_t last; - std::vector current_index; - const std::vector& shape; -}; - /* Check if we can coalesce dim with dim + 1. @@ -174,8 +39,8 @@ inline bool CanCoalesce( return true; } - for (const auto& strides_ : tensors_strides) { - std::vector& strides = strides_.get(); + for (const auto& cur_stride : tensors_strides) { + std::vector& strides = cur_stride.get(); if (shape_ndim * strides[ndim] != strides[dim]) { return false; } @@ -189,8 +54,8 @@ inline bool CanCoalesce( inline void CopyStride( std::initializer_list>>& tensors_strides, std::size_t dim, std::size_t ndim) { - for (const auto& strides_ : tensors_strides) { - std::vector& strides = strides_.get(); + for (const auto& cur_stride : tensors_strides) { + std::vector& strides = cur_stride.get(); strides[dim] = strides[ndim]; } } @@ -200,7 +65,8 @@ inline void CopyStride( /* Coalesce contiguous dimensions in the tensors. Operates inplace on the function arguments. */ -void CoalesceDimensions(std::initializer_list>>&& tensors_strides, std::vector& shape) { +void CoalesceDimensions(std::initializer_list>>&& tensors_strides, + std::vector& shape) { const std::size_t dims = shape.size(); // the current dimension is the one we are attempting to "coalesce onto" @@ -217,7 +83,7 @@ void CoalesceDimensions(std::initializer_list& strides = strides_.get(); + for (const auto& cur_stride : tensors_strides) { + std::vector& strides = cur_stride.get(); strides.resize(current_dim + 1); } } -template -void StridedCopy(concurrency::ThreadPool* thread_pool, - T* dst, - const std::vector& dst_strides_, - const TensorShape& copy_shape_, - const T* src, - const std::vector& src_strides_) { - // Coalesce dimensions - std::vector dst_strides = dst_strides_; - std::vector src_strides = src_strides_; - std::vector copy_shape(copy_shape_.GetDims()); - - CoalesceDimensions({dst_strides, src_strides}, copy_shape); - ORT_ENFORCE(dst_strides.size() == src_strides.size() && src_strides.size() == copy_shape.size(), "src and dst must have same shape"); - - const std::size_t dims = copy_shape.size(); - // We will iterate over the output dimensions - int64_t num_iterations = 1; - for (std::size_t dim = 0; dim < dims; dim++) { - num_iterations *= copy_shape[dim]; - } - - if (num_iterations <= 1) { - // scalar edge case - dst[0] = src[0]; - return; - } - - // TODOs for when we have strided tensors: - // - Reorder dimensions so that we iterate along the smallest strides first - - ORT_ENFORCE(dims > 0); - - if (dims <= 2 && src_strides[dims - 1] == 1 && dst_strides[dims - 1] == 1) { - // Fast path for 2D copies that skips the NdCounter required in the general case. - // This avoids overhead which becomes noticable at smaller iteration sizes. - // - // After coalescing, the case is actually quite common since all tensors in ORT are contiguous - - int64_t dst_stride = dims == 2 ? dst_strides[0] : 0; - int64_t src_stride = dims == 2 ? src_strides[0] : 0; - - // the size of contiguous spans that we can copy before having to advance the non-contiguous stride - int64_t contiguous_span_size = dims == 2 ? copy_shape[1] : copy_shape[0]; - - concurrency::ThreadPool::TryParallelFor( - thread_pool, num_iterations, - {static_cast(sizeof(T)), static_cast(sizeof(T)), 1.0F}, - [src_stride, dst_stride, dst, src, contiguous_span_size](std::ptrdiff_t first, std::ptrdiff_t last) { - // get the current inner and outer index - int64_t inner = first % contiguous_span_size; - int64_t outer = first / contiguous_span_size; - - std::ptrdiff_t dst_idx = outer * dst_stride + inner; - std::ptrdiff_t src_idx = outer * src_stride + inner; - - // Step 1: if there is anything left in the contiguous span that we are starting in, finish copying it - if (inner != 0) { - auto elements_to_copy = contiguous_span_size - inner; - // never copy more than what is in our partition - elements_to_copy = std::min(elements_to_copy, last - first); - Copy1DContiguous(dst + dst_idx, src + src_idx, elements_to_copy); - inner = 0; - outer++; - first += elements_to_copy; - - // reset the dst and src idx now that we are aligned to the start of a contiguous span - dst_idx = outer * dst_stride; - src_idx = outer * src_stride; - } - - // Step 2: copy contiguous span by contiguous span until we reach the penultimate span - while (first < last - contiguous_span_size) { - Copy1DContiguous(dst + dst_idx, src + src_idx, contiguous_span_size); - dst_idx += dst_stride; - src_idx += src_stride; - first += contiguous_span_size; - } - // Step 3: finish off the last (possibly partial) span manually, making sure that we don't go past the last - // element in our partition - ORT_ENFORCE(last >= first); - auto last_span_size = last - first; - Copy1DContiguous(dst + dst_idx, src + src_idx, last_span_size); - }); - } else { - concurrency::ThreadPool::TryParallelFor( - thread_pool, num_iterations, - {static_cast(sizeof(T)), static_cast(sizeof(T)), 1.0F}, - [copy_shape, dst_strides, dst, src, src_strides, dims](std::ptrdiff_t first, std::ptrdiff_t last) { - NdCounter counter(copy_shape, first, last); - - auto last_dst_stride = dst_strides[dims - 1]; - auto last_src_stride = src_strides[dims - 1]; - - auto iter_size = counter.NextStepSize(); - while (iter_size > 0) { - // Compute the src and dst addresses - std::ptrdiff_t dst_idx = 0; - std::ptrdiff_t src_idx = 0; - for (std::size_t dim = 0; dim < dims; dim++) { - dst_idx += counter.current_index[dim] * dst_strides[dim]; - src_idx += counter.current_index[dim] * src_strides[dim]; - } - // we can copy until the current dimension is done (or until we hit the last element we are trying to copy) - Copy1D(dst + dst_idx, last_dst_stride, src + src_idx, last_src_stride, iter_size); - - counter.Step(iter_size); - iter_size = counter.NextStepSize(); - } - ORT_ENFORCE(counter.current_offset == last); - }); - } -} } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/tensor/copy.h b/onnxruntime/core/providers/cpu/tensor/copy.h index fabbd36cb6..6eba9e3c31 100644 --- a/onnxruntime/core/providers/cpu/tensor/copy.h +++ b/onnxruntime/core/providers/cpu/tensor/copy.h @@ -1,5 +1,13 @@ -#include "core/providers/common.h" +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + #include "core/platform/threadpool.h" +#include "core/providers/common.h" +#include "core/providers/op_kernel_type_control.h" +#include "core/providers/op_kernel_type_control_utils.h" + #include namespace onnxruntime { @@ -7,21 +15,292 @@ namespace onnxruntime { void CoalesceDimensions( std::initializer_list>>&& tensors_strides, std::vector& shape); +std::vector StridesForTensor(const Tensor& tensor); + +namespace strided_copy_detail { + +template +void Copy1DNonContiguous(T* dst, int64_t dst_stride, const T* src, int64_t src_stride, std::ptrdiff_t count) { + for (std::ptrdiff_t i = 0; i < count; i++) { + dst[0] = src[0]; + dst += dst_stride; + src += src_stride; + } +} + +template +void Copy1DContiguous(T* dst, const T* src, std::ptrdiff_t count) { + if constexpr (std::is_same_v) { + Copy1DNonContiguous(dst, 1, src, 1, count); + } else { + memcpy(dst, src, count * sizeof(T)); + } +} + +template +void Copy1D(T* dst, int64_t dst_stride, const T* src, int64_t src_stride, std::ptrdiff_t count) { + if constexpr (std::is_same_v) { + // strings should always be copied using the for loop + Copy1DNonContiguous(dst, dst_stride, src, src_stride, count); + } else { + if (dst_stride == 1 && src_stride == 1) { + Copy1DContiguous(dst, src, count); + } else { + Copy1DNonContiguous(dst, dst_stride, src, src_stride, count); + } + } +} + +struct NdCounter { + NdCounter(const std::vector& shape, std::ptrdiff_t first, std::ptrdiff_t last) + : dims(shape.size()), + last_dim_size(shape[dims - 1]), + current_offset(first), + last(last), + current_index(dims), + shape(shape) { + // compute the initial n-dimensional index + int64_t remaining_index = first; + // Iterate from dims to 1 so we don't roll over to positive on the bounds check + for (std::size_t dim = dims; dim > 0; dim--) { + auto shape_val = shape[dim - 1]; + current_index[dim - 1] = remaining_index % shape_val; + remaining_index /= shape_val; + } + } + + /* + Return the size of the largest step in the last dimension. + */ + std::ptrdiff_t NextStepSize() const { + auto elements_in_dimension = last_dim_size - current_index[dims - 1]; + std::ptrdiff_t span_end = std::min(last, current_offset + elements_in_dimension); + return span_end - current_offset; + } + + /* + Advance the counter by step_size elements. + */ + void Step(std::ptrdiff_t step_size) { + current_offset += step_size; + current_index[dims - 1] += step_size; + + // update the current_nd_idx if needed + std::size_t dim = dims - 1; + while (dim > 0 && current_index[dim] >= shape[dim]) { + current_index[dim] = 0; + dim--; + current_index[dim]++; + } + } + + const std::size_t dims; + const int64_t last_dim_size; + ptrdiff_t current_offset; + const ptrdiff_t last; + std::vector current_index; + const std::vector& shape; +}; +} // namespace strided_copy_detail + +template +void StridedCopy(concurrency::ThreadPool* thread_pool, + T* dst, + const std::vector& dst_strides_in, + const TensorShape& copy_shape_in, + const T* src, + const std::vector& src_strides_in) { + // Coalesce dimensions + std::vector dst_strides = dst_strides_in; + std::vector src_strides = src_strides_in; + std::vector copy_shape(copy_shape_in.GetDims()); + + CoalesceDimensions({dst_strides, src_strides}, copy_shape); + ORT_ENFORCE(dst_strides.size() == src_strides.size() && + src_strides.size() == copy_shape.size() && + !copy_shape.empty(), + "src and dst must have same shape and not be rank 0."); + + const std::size_t dims = copy_shape.size(); + // We will iterate over the output dimensions + int64_t num_iterations = 1; + for (std::size_t dim = 0; dim < dims; dim++) { + num_iterations *= copy_shape[dim]; + } + + if (num_iterations <= 1) { + // scalar edge case + dst[0] = src[0]; + return; + } + + // TODOs for when we have strided tensors: + // - Reorder dimensions so that we iterate along the smallest strides first + + if (dims <= 2 && src_strides[dims - 1] == 1 && dst_strides[dims - 1] == 1) { + // Fast path for 2D copies that skips the NdCounter required in the general case. + // This avoids overhead which becomes noticeable at smaller iteration sizes. + // + // After coalescing, the case is actually quite common since all tensors in ORT are contiguous + + int64_t dst_stride = dims == 2 ? dst_strides[0] : 0; + int64_t src_stride = dims == 2 ? src_strides[0] : 0; + + // the size of contiguous spans that we can copy before having to advance the non-contiguous stride + int64_t contiguous_span_size = dims == 2 ? copy_shape[1] : copy_shape[0]; + + concurrency::ThreadPool::TryParallelFor( + thread_pool, num_iterations, + {static_cast(sizeof(T)), static_cast(sizeof(T)), 1.0F}, + [src_stride, dst_stride, dst, src, contiguous_span_size](std::ptrdiff_t first, std::ptrdiff_t last) { + // get the current inner and outer index + int64_t inner = first % contiguous_span_size; + int64_t outer = first / contiguous_span_size; + + std::ptrdiff_t dst_idx = outer * dst_stride + inner; + std::ptrdiff_t src_idx = outer * src_stride + inner; + + // Step 1: if there is anything left in the contiguous span that we are starting in, finish copying it + if (inner != 0) { + auto elements_to_copy = contiguous_span_size - inner; + // never copy more than what is in our partition + elements_to_copy = std::min(elements_to_copy, last - first); + strided_copy_detail::Copy1DContiguous(dst + dst_idx, src + src_idx, elements_to_copy); + inner = 0; + outer++; + first += elements_to_copy; + + // reset the dst and src idx now that we are aligned to the start of a contiguous span + dst_idx = outer * dst_stride; + src_idx = outer * src_stride; + } + + // Step 2: copy contiguous span by contiguous span until we reach the penultimate span + while (first < last - contiguous_span_size) { + strided_copy_detail::Copy1DContiguous(dst + dst_idx, src + src_idx, contiguous_span_size); + dst_idx += dst_stride; + src_idx += src_stride; + first += contiguous_span_size; + } + // Step 3: finish off the last (possibly partial) span manually, making sure that we don't go past the last + // element in our partition + ORT_ENFORCE(last >= first); + auto last_span_size = last - first; + strided_copy_detail::Copy1DContiguous(dst + dst_idx, src + src_idx, last_span_size); + }); + } else { + // enforce that the lambda doesn't change anything + const std::vector& const_dst_strides = dst_strides; + const std::vector& const_src_strides = src_strides; + const std::vector& const_copy_shape = copy_shape; + + concurrency::ThreadPool::TryParallelFor( + thread_pool, num_iterations, + {static_cast(sizeof(T)), static_cast(sizeof(T)), 1.0F}, + [&const_copy_shape, &const_dst_strides, dst, src, &const_src_strides, dims](std::ptrdiff_t first, + std::ptrdiff_t last) { + strided_copy_detail::NdCounter counter(const_copy_shape, first, last); + + auto last_dst_stride = const_dst_strides[dims - 1]; + auto last_src_stride = const_src_strides[dims - 1]; + + auto iter_size = counter.NextStepSize(); + while (iter_size > 0) { + // Compute the src and dst addresses + std::ptrdiff_t dst_idx = 0; + std::ptrdiff_t src_idx = 0; + for (std::size_t dim = 0; dim < dims; dim++) { + dst_idx += counter.current_index[dim] * const_dst_strides[dim]; + src_idx += counter.current_index[dim] * const_src_strides[dim]; + } + // we can copy until the current dimension is done (or until we hit the last element we are trying to copy) + strided_copy_detail::Copy1D(dst + dst_idx, last_dst_stride, src + src_idx, last_src_stride, iter_size); + + counter.Step(iter_size); + iter_size = counter.NextStepSize(); + } + ORT_ENFORCE(counter.current_offset == last); + }); + } +} + +// call StridedCopy if there is a type with the same size as T in the set of EnabledTypes +// e.g. if uint32_t is enabled all 4 byte types are supported +template +bool StridedCopyIfEnabled(concurrency::ThreadPool* thread_pool, + Tensor& dst, + std::ptrdiff_t dst_offset, + const std::vector& dst_strides, + const TensorShape& copy_shape, + const Tensor& src, + const std::vector& src_strides) { + constexpr bool enabled = utils::HasTypeWithSameSize(); + if (enabled) { + // T doesn't necessarily match the data type in src or dst so use reinterpret_cast. + // it will be a type with the same size though, which is all that matters given we're only copying bits. + StridedCopy(thread_pool, + reinterpret_cast(dst.MutableDataRaw()) + dst_offset, + dst_strides, copy_shape, + reinterpret_cast(src.DataRaw()), + src_strides); + } + + return enabled; +} + +// EnabledTypes is an onnxruntime::TypeList with the enabled types in this build. +// see "core/framework/element_type_lists.h" for default lists or the usage in +// onnxruntime/core/providers/cpu/tensor/concat.cc for +template Status DispatchStridedCopy(concurrency::ThreadPool* thread_pool, Tensor& dst, std::ptrdiff_t dst_offset, const std::vector dst_strides, const TensorShape& copy_shape, const Tensor& src, - const std::vector src_strides); + const std::vector src_strides) { + ORT_ENFORCE(dst.DataType() == src.DataType(), "src and dst types must match"); -template -void StridedCopy(concurrency::ThreadPool* thread_pool, - T* dst, - const std::vector& dst_strides, - const TensorShape& copy_shape, - const T* src, - const std::vector& src_strides); + bool supported = false; + if (src.IsDataTypeString()) { + if (utils::HasType()) { + supported = true; + StridedCopy(thread_pool, dst.MutableData() + dst_offset, dst_strides, copy_shape, + src.Data(), src_strides); + } + } else { + const auto element_size = src.DataType()->Size(); + switch (element_size) { + case sizeof(uint32_t): + supported = StridedCopyIfEnabled(thread_pool, dst, dst_offset, dst_strides, + copy_shape, src, src_strides); + break; + case sizeof(uint64_t): + supported = StridedCopyIfEnabled(thread_pool, dst, dst_offset, dst_strides, + copy_shape, src, src_strides); + break; + case sizeof(uint16_t): + supported = StridedCopyIfEnabled(thread_pool, dst, dst_offset, dst_strides, + copy_shape, src, src_strides); + break; + case sizeof(uint8_t): + static_assert(sizeof(bool) == sizeof(uint8_t), "Need to enable separate case for 'bool' on this platform."); + supported = StridedCopyIfEnabled(thread_pool, dst, dst_offset, dst_strides, + copy_shape, src, src_strides); + break; + // It's possible that bool is not 1 byte. static_assert above checks if we need to enable this on a platform. + //case sizeof(bool): + // supported = StridedCopyIfEnabled(thread_pool, dst, dst_offset, dst_strides, + // copy_shape, src, src_strides); + // break; + default: + // leave 'supported' as false + break; + } + } + + return !supported ? ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unsupported input data type of ", src.DataType()) + : Status::OK(); +} -std::vector StridesForTensor(const Tensor& tensor); } // namespace onnxruntime diff --git a/tools/python/util/ort_format_model/operator_type_usage_processors.py b/tools/python/util/ort_format_model/operator_type_usage_processors.py index 0d44c826d3..f3ea7aca18 100644 --- a/tools/python/util/ort_format_model/operator_type_usage_processors.py +++ b/tools/python/util/ort_format_model/operator_type_usage_processors.py @@ -344,6 +344,7 @@ def _create_operator_type_usage_processors(): # ops that are used to manipulate shapes or indices so require int32_t and int64_t to be available default_processor_onnx_ops_requiring_ints_for_input_0 = ['Add', + 'Concat', 'Div', 'Equal', 'Greater',