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

Enable type reduction for ConstantOfShape CPU kernel. (#6594)

Browse source 
* Enable type reduction for ConstantOfShape.
This commit is contained in:
Edward Chen 2021-02-12 00:27:25 -08:00 committed by GitHub
parent 950c941f11
commit 78e408dbe9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 225 additions and 190 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
onnxruntime/core/common/type_list.h
View file

@ -9,4 +9,9 @@ namespace onnxruntime {
template <typename... T>
struct TypeList {};
}
} // namespace onnxruntime
// type list type containing the given types
// Note: this is useful for passing TypeLists to macros which don't accept the
// comma-separated template arguments
#define ORT_TYPE_LIST(...) ::onnxruntime::TypeList<__VA_ARGS__>

155
onnxruntime/core/providers/cpu/generator/constant_of_shape.cc
View file

@ -1,130 +1,37 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/framework/tensorprotoutils.h"
#include "core/providers/cpu/generator/constant_of_shape.h"
#include "gsl/gsl"
#include "core/providers/cpu/generator/constant_of_shape_base.h"
#include "core/providers/op_kernel_type_control.h"
using namespace ::onnxruntime::common;
using namespace ONNX_NAMESPACE;
namespace onnxruntime {
ONNX_CPU_OPERATOR_KERNEL(
ConstantOfShape,
9,
KernelDefBuilder()
.TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("T2", std::vector<MLDataType>{
DataTypeImpl::GetTensorType<MLFloat16>(),
DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>(),
DataTypeImpl::GetTensorType<int8_t>(),
DataTypeImpl::GetTensorType<int16_t>(),
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>(),
DataTypeImpl::GetTensorType<uint8_t>(),
DataTypeImpl::GetTensorType<uint16_t>(),
DataTypeImpl::GetTensorType<uint32_t>(),
DataTypeImpl::GetTensorType<uint64_t>(),
DataTypeImpl::GetTensorType<bool>()}),
ConstantOfShape);
#define FETCH_VALUE_DATA(c_type) \
{ \
c_type val; \
auto unpack_status = UnpackTensor(t_proto, raw_data, raw_data_len, &val, 1); \
ORT_ENFORCE(unpack_status.IsOK(), "Value attribute unpacking failed:", unpack_status.ErrorMessage()); \
SetValue(sizeof(c_type), reinterpret_cast<void*>(&val)); \
}
void onnxruntime::ConstantOfShapeBase::SetValueFromTensorProto(const ONNX_NAMESPACE::TensorProto& t_proto) {
using namespace utils;
ORT_ENFORCE(utils::HasDataType(t_proto));
ORT_ENFORCE(TensorProto::DataType_IsValid(t_proto.data_type()));
ORT_ENFORCE(!utils::HasExternalData(t_proto), "Tensor proto with external data for value attribute is not supported.");
const auto tensor_type = static_cast<TensorProto_DataType>(t_proto.data_type());
const void* const raw_data = utils::HasRawData(t_proto) ? t_proto.raw_data().data() : nullptr;
const size_t raw_data_len = utils::HasRawData(t_proto) ? t_proto.raw_data().size() : 0;
switch (tensor_type) {
case TensorProto::BOOL:
FETCH_VALUE_DATA(bool);
break;
case TensorProto::FLOAT:
FETCH_VALUE_DATA(float);
break;
case TensorProto::FLOAT16:
FETCH_VALUE_DATA(MLFloat16);
break;
case TensorProto::DOUBLE:
FETCH_VALUE_DATA(double);
break;
case TensorProto::INT8:
FETCH_VALUE_DATA(int8_t);
break;
case TensorProto::INT16:
FETCH_VALUE_DATA(int16_t);
break;
case TensorProto::INT32:
FETCH_VALUE_DATA(int32_t);
break;
case TensorProto::INT64:
FETCH_VALUE_DATA(int64_t);
break;
case TensorProto::UINT8:
FETCH_VALUE_DATA(uint8_t);
break;
case TensorProto::UINT16:
FETCH_VALUE_DATA(uint16_t);
break;
case TensorProto::UINT32:
FETCH_VALUE_DATA(uint32_t);
break;
case TensorProto::UINT64:
FETCH_VALUE_DATA(uint64_t);
break;
default:
ORT_THROW("Unsupported value attribute datatype: ", tensor_type);
break;
}
namespace op_kernel_type_control {
ORT_SPECIFY_OP_KERNEL_ARG_SUPPORTED_TYPE_LIST_ALL_OPSETS(
kCpuExecutionProvider, kOnnxDomain, ConstantOfShape, Output, 0,
ConstantOfShapeDefaultOutputTypes);
}
#undef FETCH_VALUE_DATA
namespace {
using SupportedOutputTypes =
ORT_OP_KERNEL_ARG_SUPPORTED_TYPE_LIST_ALL_OPSETS(
kCpuExecutionProvider, kOnnxDomain, ConstantOfShape, Output, 0);
using EnabledOutputTypes =
ORT_OP_KERNEL_ARG_ENABLED_TYPE_LIST_ALL_OPSETS(
kCpuExecutionProvider, kOnnxDomain, ConstantOfShape, Output, 0);
class ConstantOfShape final : public ConstantOfShapeBase<EnabledOutputTypes>, public OpKernel {
public:
explicit ConstantOfShape(const OpKernelInfo& info) : ConstantOfShapeBase(info), OpKernel(info) {}
Status Compute(OpKernelContext* ctx) const override;
};
template <class T>
inline void FilloutOutput(T value, void* output_data, size_t size) {
auto out = gsl::make_span(reinterpret_cast<T*>(output_data), size);
std::fill(out.begin(), out.end(), value);
}
ConstantOfShapeBase::ConstantOfShapeBase(const OpKernelInfo& info) {
TensorProto t_proto;
if (info.GetAttr<TensorProto>("value", &t_proto).IsOK()) {
ORT_ENFORCE(t_proto.dims_size() == 1, "Must have a single dimension");
ORT_ENFORCE(t_proto.dims()[0] == 1, "Must have a single dimension of 1");
SetValueFromTensorProto(t_proto);
} else {
float f_value = 0.f;
SetValue(sizeof(float), reinterpret_cast<void*>(&f_value));
}
}
Status ConstantOfShapeBase::PrepareCompute(OpKernelContext* ctx, Tensor** output_tensor) const {
const auto shape_tensor = ctx->Input<Tensor>(0);
const auto& input_shape = shape_tensor->Shape();
// If empty the output is a scalar with empty shape
// TensorShape::Size() will still return 1 and we will output
// one value
std::vector<int64_t> output_dims;
ORT_ENFORCE(input_shape.NumDimensions() > 0, "Must have a valid input shape.");
const auto span = gsl::make_span(shape_tensor->Data<int64_t>(), input_shape.Size());
output_dims.insert(output_dims.end(), span.cbegin(), span.cend());
TensorShape output_shape(output_dims);
(*output_tensor) = ctx->Output(0, output_shape);
return Status::OK();
std::fill_n(reinterpret_cast<T*>(output_data), size, value);
}
Status ConstantOfShape::Compute(OpKernelContext* ctx) const {
@ -149,10 +56,22 @@ Status ConstantOfShape::Compute(OpKernelContext* ctx) const {
FilloutOutput(*(reinterpret_cast<const int64_t*>(value_ptr)), output_data, size);
break;
default:
ORT_THROW("Unsupported value attribute datatype with sizeof=: ", element_size);
break;
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unsupported output datatype with size: ", element_size);
}
return Status::OK();
}
} // namespace
ONNX_CPU_OPERATOR_KERNEL(
ConstantOfShape,
9,
KernelDefBuilder()
.TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("T2",
BuildKernelDefConstraintsFunctorFromTypeList<SupportedOutputTypes>{}(),
BuildKernelDefConstraintsFunctorFromTypeList<EnabledOutputTypes>{}()),
ConstantOfShape);
} // namespace onnxruntime

64
onnxruntime/core/providers/cpu/generator/constant_of_shape.h
View file

@ -1,64 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/common.h"
#include "core/framework/data_types.h"
#include "core/framework/op_kernel.h"
namespace onnxruntime {
class ConstantOfShapeBase {
protected:
ConstantOfShapeBase(const OpKernelInfo& info);
Status PrepareCompute(OpKernelContext* ctx, Tensor** output_tensor) const;
void* GetValuePtr() const { return p_value_; }
private:
union SizeBasedValue {
int8_t int8_;
int16_t int16_;
int32_t int32_;
int64_t int64_;
} s_value_;
void* p_value_;
void SetValue(size_t size, void* value) {
switch (size) {
case sizeof(int8_t):
s_value_.int8_ = *(reinterpret_cast<int8_t*>(value));
p_value_ = reinterpret_cast<void*>(&(s_value_.int8_));
break;
case sizeof(int16_t):
s_value_.int16_ = *(reinterpret_cast<int16_t*>(value));
p_value_ = reinterpret_cast<void*>(&(s_value_.int16_));
break;
case sizeof(int32_t):
s_value_.int32_ = *(reinterpret_cast<int32_t*>(value));
p_value_ = reinterpret_cast<void*>(&(s_value_.int32_));
break;
case sizeof(int64_t):
s_value_.int64_ = *(reinterpret_cast<int64_t*>(value));
p_value_ = reinterpret_cast<void*>(&(s_value_.int64_));
break;
default:
ORT_THROW("Unsupported value attribute datatype with sizeof=: ", size);
break;
}
}
void SetValueFromTensorProto(const ONNX_NAMESPACE::TensorProto&);
};
class ConstantOfShape final : public ConstantOfShapeBase, public OpKernel {
public:
explicit ConstantOfShape(const OpKernelInfo& info) : ConstantOfShapeBase(info), OpKernel(info) {};
Status Compute(OpKernelContext* ctx) const override;
};
} // namespace onnxruntime

136
onnxruntime/core/providers/cpu/generator/constant_of_shape_base.h Normal file
View file

@ -0,0 +1,136 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/common.h"
#include "core/common/type_list.h"
#include "core/framework/data_types.h"
#include "core/framework/data_types_internal.h"
#include "core/framework/op_kernel.h"
#include "core/framework/tensorprotoutils.h"
#include "core/providers/op_kernel_type_control_utils.h"
namespace onnxruntime {
using ConstantOfShapeDefaultOutputTypes =
TypeList<
MLFloat16,
float, double,
int8_t, int16_t, int32_t, int64_t,
uint8_t, uint16_t, uint32_t, uint64_t,
bool>;
template <typename EnabledOutputTypeList = ConstantOfShapeDefaultOutputTypes>
class ConstantOfShapeBase {
protected:
ConstantOfShapeBase(const OpKernelInfo& info) {
ONNX_NAMESPACE::TensorProto t_proto;
if (info.GetAttr<ONNX_NAMESPACE::TensorProto>("value", &t_proto).IsOK()) {
ORT_ENFORCE(t_proto.dims_size() == 1, "Must have a single dimension");
ORT_ENFORCE(t_proto.dims()[0] == 1, "Must have a single dimension of 1");
SetValueFromTensorProto(t_proto);
} else {
float f_value = 0.f;
SetValue(sizeof(float), reinterpret_cast<void*>(&f_value));
}
}
void* GetValuePtr() const { return p_value_; }
static Status PrepareCompute(OpKernelContext* ctx, Tensor** output_tensor) {
const auto shape_tensor = ctx->Input<Tensor>(0);
const auto& input_shape = shape_tensor->Shape();
// If empty the output is a scalar with empty shape
// TensorShape::Size() will still return 1 and we will output
// one value
ORT_RETURN_IF_NOT(input_shape.NumDimensions() > 0, "Must have a valid input shape.");
const auto span = shape_tensor->DataAsSpan<int64_t>();
TensorShape output_shape(span.begin(), span.size());
(*output_tensor) = ctx->Output(0, output_shape);
return Status::OK();
}
private:
union SizeBasedValue {
int8_t int8_;
int16_t int16_;
int32_t int32_;
int64_t int64_;
} s_value_;
void* p_value_;
void SetValue(size_t size, void* value) {
switch (size) {
case sizeof(int8_t):
s_value_.int8_ = *(reinterpret_cast<int8_t*>(value));
p_value_ = reinterpret_cast<void*>(&(s_value_.int8_));
break;
case sizeof(int16_t):
s_value_.int16_ = *(reinterpret_cast<int16_t*>(value));
p_value_ = reinterpret_cast<void*>(&(s_value_.int16_));
break;
case sizeof(int32_t):
s_value_.int32_ = *(reinterpret_cast<int32_t*>(value));
p_value_ = reinterpret_cast<void*>(&(s_value_.int32_));
break;
case sizeof(int64_t):
s_value_.int64_ = *(reinterpret_cast<int64_t*>(value));
p_value_ = reinterpret_cast<void*>(&(s_value_.int64_));
break;
default:
ORT_THROW("Unsupported value attribute datatype with size: ", size);
}
}
void SetValueFromTensorProto(const ONNX_NAMESPACE::TensorProto&);
};
#define CASE_FETCH_VALUE_DATA(c_type) \
case utils::ToTensorProtoElementType<c_type>(): { \
if (utils::HasType<EnabledOutputTypeList, c_type>()) { \
c_type val; \
ORT_THROW_IF_ERROR(utils::UnpackTensor(t_proto, raw_data, raw_data_len, &val, 1)); \
SetValue(sizeof(c_type), reinterpret_cast<void*>(&val)); \
handled = true; \
} \
break; \
}
template <typename EnabledOutputTypeList>
void ConstantOfShapeBase<EnabledOutputTypeList>::SetValueFromTensorProto(const ONNX_NAMESPACE::TensorProto& t_proto) {
ORT_ENFORCE(utils::HasDataType(t_proto));
ORT_ENFORCE(ONNX_NAMESPACE::TensorProto::DataType_IsValid(t_proto.data_type()));
ORT_ENFORCE(!utils::HasExternalData(t_proto),
"Tensor proto with external data for value attribute is not supported.");
const auto tensor_type = static_cast<ONNX_NAMESPACE::TensorProto_DataType>(t_proto.data_type());
const void* const raw_data = utils::HasRawData(t_proto) ? t_proto.raw_data().data() : nullptr;
const size_t raw_data_len = utils::HasRawData(t_proto) ? t_proto.raw_data().size() : 0;
bool handled = false;
switch (tensor_type) {
CASE_FETCH_VALUE_DATA(bool)
CASE_FETCH_VALUE_DATA(float)
CASE_FETCH_VALUE_DATA(MLFloat16)
CASE_FETCH_VALUE_DATA(double)
CASE_FETCH_VALUE_DATA(int8_t)
CASE_FETCH_VALUE_DATA(int16_t)
CASE_FETCH_VALUE_DATA(int32_t)
CASE_FETCH_VALUE_DATA(int64_t)
CASE_FETCH_VALUE_DATA(uint8_t)
CASE_FETCH_VALUE_DATA(uint16_t)
CASE_FETCH_VALUE_DATA(uint32_t)
CASE_FETCH_VALUE_DATA(uint64_t)
default:
ORT_THROW("Unsupported value attribute datatype: ", tensor_type);
}
ORT_ENFORCE(handled, "Unsupported value attribute datatype in this build: ", tensor_type);
}
#undef CASE_FETCH_VALUE_DATA
} // namespace onnxruntime

6
onnxruntime/core/providers/cuda/generator/constant_of_shape.h
View file

@ -6,15 +6,15 @@
#include "core/common/common.h"
#include "core/framework/data_types.h"
#include "core/providers/cuda/cuda_kernel.h"
#include "core/providers/cpu/generator/constant_of_shape.h"
#include "core/providers/cpu/generator/constant_of_shape_base.h"
#include "core/providers/cuda/shared_inc/cuda_utils.h"
namespace onnxruntime {
namespace cuda {
class ConstantOfShape final : public ConstantOfShapeBase, public CudaKernel {
class ConstantOfShape final : public ConstantOfShapeBase<>, public CudaKernel {
public:
explicit ConstantOfShape(const OpKernelInfo& info) : ConstantOfShapeBase(info), CudaKernel(info) {};
explicit ConstantOfShape(const OpKernelInfo& info) : ConstantOfShapeBase(info), CudaKernel(info) {}
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ConstantOfShape);

47
onnxruntime/core/providers/op_kernel_type_control.h
View file

@ -165,10 +165,10 @@ struct EnabledTypes {
* @param OpSet The opset that this set of supported types applies to.
* @param ArgDirection Direction of the given Op kernel argument - Input or Output.
* @param ArgIndex Index of the given Op kernel argument.
* @param ... The types.
* @param SupportedTypeList The types.
*/
#define ORT_SPECIFY_OP_KERNEL_ARG_SUPPORTED_TYPES( \
OpProvider, OpDomain, OpName, OpSet, ArgDirection, ArgIndex, ...) \
#define ORT_SPECIFY_OP_KERNEL_ARG_SUPPORTED_TYPE_LIST( \
OpProvider, OpDomain, OpName, OpSet, ArgDirection, ArgIndex, SupportedTypeList) \
class ORT_OP_KERNEL_TYPE_CTRL_INTERNAL_OP_TAG_CLASS_NAME(OpDomain, OpName); \
class ORT_OP_KERNEL_TYPE_CTRL_INTERNAL_PROVIDER_TAG_CLASS_NAME(OpProvider); \
template <> \
@ -177,9 +177,48 @@ struct EnabledTypes {
ORT_OP_KERNEL_TYPE_CTRL_INTERNAL_OP_KERNEL_ARG_TAG(OpDomain, OpName, ArgDirection, ArgIndex), \
ORT_OP_KERNEL_TYPE_CTRL_INTERNAL_PROVIDER_TAG_CLASS_NAME(OpProvider), \
OpSet>> { \
using types = ::onnxruntime::TypeList<__VA_ARGS__>; \
using types = SupportedTypeList; \
};
/**
* Specifies a supported set of types for a given Op kernel argument that is valid for all opsets.
* This should be specified with the Op kernel implementation.
*
* Note: This should be called from the onnxruntime::op_kernel_type_control namespace.
*
* @param OpProvider The Op provider.
* @param OpDomain The Op domain.
* @param OpName The Op name.
* @param ArgDirection Direction of the given Op kernel argument - Input or Output.
* @param ArgIndex Index of the given Op kernel argument.
* @param SupportedTypeList The types.
*/
#define ORT_SPECIFY_OP_KERNEL_ARG_SUPPORTED_TYPE_LIST_ALL_OPSETS( \
OpProvider, OpDomain, OpName, ArgDirection, ArgIndex, SupportedTypeList) \
ORT_SPECIFY_OP_KERNEL_ARG_SUPPORTED_TYPE_LIST(OpProvider, OpDomain, OpName, \
::onnxruntime::op_kernel_type_control::kAllOpSets, \
ArgDirection, ArgIndex, SupportedTypeList)
/**
* Specifies a supported set of types for a given Op kernel argument.
* This should be specified with the Op kernel implementation.
*
* Note: This should be called from the onnxruntime::op_kernel_type_control namespace.
*
* @param OpProvider The Op provider.
* @param OpDomain The Op domain.
* @param OpName The Op name.
* @param OpSet The opset that this set of supported types applies to.
* @param ArgDirection Direction of the given Op kernel argument - Input or Output.
* @param ArgIndex Index of the given Op kernel argument.
* @param ... The types.
*/
#define ORT_SPECIFY_OP_KERNEL_ARG_SUPPORTED_TYPES( \
OpProvider, OpDomain, OpName, OpSet, ArgDirection, ArgIndex, ...) \
ORT_SPECIFY_OP_KERNEL_ARG_SUPPORTED_TYPE_LIST( \
OpProvider, OpDomain, OpName, OpSet, ArgDirection, ArgIndex, \
ORT_TYPE_LIST(__VA_ARGS__))
/**
* Specifies a supported set of types for a given Op kernel argument that is valid for all opsets.
* This should be specified with the Op kernel implementation.