Skip to content

Add type reduction support to Slice and Transpose #6547

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
skottmckay merged 4 commits into from
Feb 5, 2021
Merged

Add type reduction support to Slice and Transpose #6547

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b0f2a07
Add type reduction support to Slice and Transpose
skottmckay Feb 3, 2021
b4f5d17
Address PR comments
skottmckay Feb 3, 2021
12d8f19
Address more PR comments
skottmckay Feb 4, 2021
4ab255a
Address PR comments
skottmckay Feb 4, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion include/onnxruntime/core/framework/data_types_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ constexpr ONNX_NAMESPACE::TensorProto_DataType ToTensorProtoElementType<BFloat16
function<int8_t>(__VA_ARGS__); \
break; \
case ONNX_NAMESPACE::TensorProto_DataType_UINT8: \
function<uint8_t>(__VA_ARGS__); \
function<uint8_t>(__VA_ARGS__); \
break; \
case ONNX_NAMESPACE::TensorProto_DataType_INT16: \
function<int16_t>(__VA_ARGS__); \
Expand Down
151 changes: 104 additions & 47 deletions onnxruntime/core/providers/cpu/tensor/slice.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,55 +4,75 @@
#include "core/providers/cpu/tensor/slice.h"
#include "core/providers/cpu/tensor/utils.h"
#include "core/providers/common.h"
#include "core/providers/op_kernel_type_control.h"
#include <unordered_map>
#include <limits>

using namespace ::onnxruntime::common;
using namespace std;

namespace onnxruntime {
namespace op_kernel_type_control {
// we're using one set of types for all opsets
ORT_SPECIFY_OP_KERNEL_ARG_SUPPORTED_TYPES_ALL_OPSETS(
kCpuExecutionProvider, kOnnxDomain, Slice, Input, 0,
ORT_OP_KERNEL_TYPE_CTRL_ALL_TENSOR_DATA_TYPES);

ORT_SPECIFY_OP_KERNEL_ARG_SUPPORTED_TYPES_ALL_OPSETS(
kCpuExecutionProvider, kOnnxDomain, Slice, Input, 1, int32_t, int64_t);
} // namespace op_kernel_type_control

namespace {
using EnabledDataTypes = ORT_OP_KERNEL_ARG_ENABLED_TYPE_LIST_ALL_OPSETS(kCpuExecutionProvider, kOnnxDomain,
Slice, Input, 0);
using EnabledIndiciesTypes = ORT_OP_KERNEL_ARG_ENABLED_TYPE_LIST_ALL_OPSETS(kCpuExecutionProvider, kOnnxDomain,
Comment thread
skottmckay marked this conversation as resolved.
Outdated
Slice, Input, 1);

const std::vector<MLDataType> dataTypeConstraints =
BuildKernelDefConstraintsFunctorFromTypeList<EnabledDataTypes>{}();

const std::vector<MLDataType> indicesTypeConstraints =
BuildKernelDefConstraintsFunctorFromTypeList<EnabledIndiciesTypes>{}();

// std::clamp doesn't exist until C++17 so create a local version
template <typename T>
const T& clamp(const T& v, const T& lo, const T& hi) {
if (v < lo) return lo;
if (v > hi) return hi;
return v;
}
} // namespace

ONNX_CPU_OPERATOR_VERSIONED_KERNEL(
Slice,
1, 9,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllTensorTypes()),
KernelDefBuilder().TypeConstraint("T", dataTypeConstraints),
Slice1);

ONNX_CPU_OPERATOR_VERSIONED_KERNEL(
Slice,
10, 10,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::AllTensorTypes())
.TypeConstraint("Tind", {DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
.TypeConstraint("T", dataTypeConstraints)
.TypeConstraint("Tind", indicesTypeConstraints),
Slice10);

ONNX_CPU_OPERATOR_VERSIONED_KERNEL(
Slice,
11,
12,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::AllTensorTypes())
.TypeConstraint("Tind", {DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
.TypeConstraint("T", dataTypeConstraints)
.TypeConstraint("Tind", indicesTypeConstraints),
Slice10);

ONNX_CPU_OPERATOR_KERNEL(
Slice,
13,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::AllTensorTypes())
.TypeConstraint("Tind", {DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
.TypeConstraint("T", dataTypeConstraints)
.TypeConstraint("Tind", indicesTypeConstraints),
Slice10);
namespace {
// std::clamp doesn't exist until C++17 so create a local version
template <typename T>
const T& clamp(const T& v, const T& lo, const T& hi) {
if (v < lo) return lo;
if (v > hi) return hi;
return v;
}
} // namespace

// Check if it's possible to combine innermost dimensions so we copy larger blocks.
// Sets flattened_output_dims to nullptr if it is not.
Expand Down Expand Up @@ -218,19 +238,21 @@ Status SliceBase::PrepareForCompute(const std::vector<int64_t>& raw_starts,
}

// Slice V10 & DynamicSlice
void SliceBase::FillVectorsFromInput(const Tensor& start_tensor,
const Tensor& ends_tensor,
const Tensor* axes_tensor,
const Tensor* steps_tensor,
std::vector<int64_t>& input_starts,
std::vector<int64_t>& input_ends,
std::vector<int64_t>& input_axes,
std::vector<int64_t>& input_steps) {
ORT_ENFORCE(start_tensor.Shape().NumDimensions() == 1, "Starts must be a 1-D array");
ORT_ENFORCE(ends_tensor.Shape().NumDimensions() == 1, "Ends must be a 1-D array");
ORT_ENFORCE(start_tensor.Shape() == ends_tensor.Shape(), "Starts and ends shape mismatch");
ORT_ENFORCE(nullptr == axes_tensor || start_tensor.Shape() == axes_tensor->Shape(), "Starts and axes shape mismatch");
ORT_ENFORCE(nullptr == steps_tensor || start_tensor.Shape() == steps_tensor->Shape(), "Starts and steps shape mismatch");
Status SliceBase::FillVectorsFromInput(const Tensor& start_tensor,
Comment thread
skottmckay marked this conversation as resolved.
const Tensor& ends_tensor,
const Tensor* axes_tensor,
const Tensor* steps_tensor,
std::vector<int64_t>& input_starts,
std::vector<int64_t>& input_ends,
std::vector<int64_t>& input_axes,
std::vector<int64_t>& input_steps) {
ORT_RETURN_IF_NOT(start_tensor.Shape().NumDimensions() == 1, "Starts must be a 1-D array");
ORT_RETURN_IF_NOT(ends_tensor.Shape().NumDimensions() == 1, "Ends must be a 1-D array");
ORT_RETURN_IF_NOT(start_tensor.Shape() == ends_tensor.Shape(), "Starts and ends shape mismatch");
ORT_RETURN_IF_NOT(nullptr == axes_tensor || start_tensor.Shape() == axes_tensor->Shape(),
"Starts and axes shape mismatch");
ORT_RETURN_IF_NOT(nullptr == steps_tensor || start_tensor.Shape() == steps_tensor->Shape(),
"Starts and steps shape mismatch");

const auto& size = start_tensor.Shape().Size();
input_starts.resize(size);
Expand All @@ -241,7 +263,11 @@ void SliceBase::FillVectorsFromInput(const Tensor& start_tensor,
if (nullptr != steps_tensor)
input_steps.resize(size);

if (start_tensor.IsDataType<int32_t>()) {
// check for type reduction of supported indices types
constexpr bool int32_enabled = op_kernel_type_control::HasType<EnabledIndiciesTypes, int32_t>();
constexpr bool int64_enabled = op_kernel_type_control::HasType<EnabledIndiciesTypes, int64_t>();

if (int32_enabled && start_tensor.IsDataType<int32_t>()) {
std::copy(start_tensor.Data<int32_t>(), start_tensor.Data<int32_t>() + size, input_starts.begin());
std::copy(ends_tensor.Data<int32_t>(), ends_tensor.Data<int32_t>() + size, input_ends.begin());
if (nullptr != axes_tensor)
Expand All @@ -251,7 +277,7 @@ void SliceBase::FillVectorsFromInput(const Tensor& start_tensor,
std::copy(steps_tensor->Data<int32_t>(), steps_tensor->Data<int32_t>() + size, input_steps.begin());
}

else if (start_tensor.IsDataType<int64_t>()) {
else if (int64_enabled && start_tensor.IsDataType<int64_t>()) {
std::copy(start_tensor.Data<int64_t>(), start_tensor.Data<int64_t>() + size, input_starts.begin());
std::copy(ends_tensor.Data<int64_t>(), ends_tensor.Data<int64_t>() + size, input_ends.begin());
if (nullptr != axes_tensor)
Expand All @@ -261,10 +287,13 @@ void SliceBase::FillVectorsFromInput(const Tensor& start_tensor,
std::copy(steps_tensor->Data<int64_t>(), steps_tensor->Data<int64_t>() + size, input_steps.begin());
}

// should not reach this as no kernel is registered for this condition to be triggered - just an additional safety check
else {
ORT_THROW("Data type for starts and ends inputs' need to be int32_t or int64_t, but instead got ", start_tensor.DataType());
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
"Data type for starts and ends inputs' is not supported in this build. Got ",
start_tensor.DataType());
}

return Status::OK();
}

template <typename T>
Expand Down Expand Up @@ -304,22 +333,40 @@ static Status SliceImpl(OpKernelContext* ctx,
flattened_input_dims.back() = compute_metadata.p_flattened_output_dims_->back();
TensorShape input_shape(std::move(flattened_input_dims));

auto input_iterator = SliceIterator<T>(input_tensor, input_shape, compute_metadata.starts_, *compute_metadata.p_flattened_output_dims_, compute_metadata.steps_);
auto input_iterator = SliceIterator<T>(input_tensor, input_shape, compute_metadata.starts_,
*compute_metadata.p_flattened_output_dims_, compute_metadata.steps_);
create_output(input_iterator);
} else {
auto input_iterator = SliceIterator<T>(input_tensor, compute_metadata.starts_, compute_metadata.output_dims_, compute_metadata.steps_);
auto input_iterator = SliceIterator<T>(input_tensor, compute_metadata.starts_, compute_metadata.output_dims_,
compute_metadata.steps_);
create_output(input_iterator);
}

return Status::OK();
}

template <typename EnabledTypes, typename T>
static inline bool CallSliceImplIfEnabled(OpKernelContext* ctx,

@edgchen1 edgchen1 Feb 3, 2021 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CallSliceImplIfEnabled [](start = 19, length = 22)

wonder if this can be extended to call arbitrary enable-able code with a lambda... #Closed

@skottmckay skottmckay Feb 3, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

generic lambdas were c++14 and we're currently stuck at c++11, so I didn't think doing it as a lambda was possible.

In reply to: 569695618 [](ancestors = 569695618)

const Tensor& input_tensor,
SliceOp::PrepareForComputeMetadata& compute_metadata,
Status& status) {
constexpr bool enabled = op_kernel_type_control::HasTypeWithSameSize<EnabledTypes, T>();
if (enabled) {
status = SliceImpl<T>(ctx, input_tensor, compute_metadata);
}

return enabled;
}

Status SliceBase::Compute(OpKernelContext* ctx) const {
const auto* input_tensor_ptr = ctx->Input<Tensor>(0);
ORT_ENFORCE(input_tensor_ptr != nullptr, "Missing input tensor to be processed");
const auto& input_tensor = *input_tensor_ptr;
const auto& input_dimensions = input_tensor.Shape().GetDims();
if (input_dimensions.empty()) return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Cannot slice scalars");

if (input_dimensions.empty()) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Cannot slice scalars");
}

SliceOp::PrepareForComputeMetadata compute_metadata(input_dimensions);

// Slice V10 & DynamicSlice
Expand All @@ -340,28 +387,38 @@ Status SliceBase::Compute(OpKernelContext* ctx) const {

Status status = Status::OK();

bool supported = false;
if (input_tensor.IsDataTypeString()) {
status = SliceImpl<std::string>(ctx, input_tensor, compute_metadata);
if (op_kernel_type_control::HasType<EnabledDataTypes, std::string>()) {
supported = true;
status = SliceImpl<std::string>(ctx, input_tensor, compute_metadata);
}
} else {
const auto element_size = input_tensor.DataType()->Size();

// call SliceImpl
switch (element_size) {
case sizeof(uint32_t):
status = SliceImpl<uint32_t>(ctx, input_tensor, compute_metadata);
supported = CallSliceImplIfEnabled<EnabledDataTypes, uint32_t>(ctx, input_tensor, compute_metadata, status);
break;
case sizeof(uint64_t):
status = SliceImpl<uint64_t>(ctx, input_tensor, compute_metadata);
supported = CallSliceImplIfEnabled<EnabledDataTypes, uint64_t>(ctx, input_tensor, compute_metadata, status);
break;
case sizeof(uint16_t):
status = SliceImpl<uint16_t>(ctx, input_tensor, compute_metadata);
supported = CallSliceImplIfEnabled<EnabledDataTypes, uint16_t>(ctx, input_tensor, compute_metadata, status);
break;
case sizeof(uint8_t):
status = SliceImpl<uint8_t>(ctx, input_tensor, compute_metadata);
supported = CallSliceImplIfEnabled<EnabledDataTypes, uint8_t>(ctx, input_tensor, compute_metadata, status);
break;
default:
ORT_THROW("Unsupported input data type of ", input_tensor.DataType());
// leave 'supported' as false
break;
}
}

if (!supported) {
status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unsupported input data type of ", input_tensor.DataType());
}

return status;
}

Expand Down
16 changes: 8 additions & 8 deletions onnxruntime/core/providers/cpu/tensor/slice.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,14 +44,14 @@ class SliceBase {
SliceOp::PrepareForComputeMetadata& compute_metadata);

// Slice V10 & DynamicSlice
static void FillVectorsFromInput(const Tensor& start_tensor,
const Tensor& ends_tensor,
const Tensor* axes_tensor,
const Tensor* steps_tensor,
std::vector<int64_t>& input_starts,
std::vector<int64_t>& input_ends,
std::vector<int64_t>& input_axes,
std::vector<int64_t>& input_steps);
static Status FillVectorsFromInput(const Tensor& start_tensor,
const Tensor& ends_tensor,
const Tensor* axes_tensor,
const Tensor* steps_tensor,
std::vector<int64_t>& input_starts,
std::vector<int64_t>& input_ends,
std::vector<int64_t>& input_axes,
std::vector<int64_t>& input_steps);

protected:
SliceBase(const OpKernelInfo& info, bool dynamic = false)
Expand Down
Loading