mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Custom op shape inference API (#17737)
Add c/cxx API to allow custom ops do shape inference. --------- Co-authored-by: Randy Shuai <rashuai@microsoft.com>
This commit is contained in:
RandySheriffH
GitHub
parent
762703e037
commit
c6c3555d0e
11 changed files with 771 additions and 25 deletions
|
|
@ -299,6 +299,7 @@ ORT_RUNTIME_CLASS(DnnlProviderOptions);
|
|||
ORT_RUNTIME_CLASS(Op);
|
||||
ORT_RUNTIME_CLASS(OpAttr);
|
||||
ORT_RUNTIME_CLASS(Logger);
|
||||
ORT_RUNTIME_CLASS(ShapeInferContext);
|
||||
|
||||
#ifdef _WIN32
|
||||
typedef _Return_type_success_(return == 0) OrtStatus* OrtStatusPtr;
|
||||
|
|
@ -4444,6 +4445,73 @@ struct OrtApi {
|
|||
*/
|
||||
ORT_API2_STATUS(SetUserLoggingFunction, _Inout_ OrtSessionOptions* options,
|
||||
_In_ OrtLoggingFunction user_logging_function, _In_opt_ void* user_logging_param);
|
||||
|
||||
/**
|
||||
* Get number of input from OrtShapeInferContext
|
||||
*
|
||||
* \param[in] context
|
||||
* \param[out] out The number of inputs
|
||||
*
|
||||
* \since Version 1.17.
|
||||
*/
|
||||
ORT_API2_STATUS(ShapeInferContext_GetInputCount, _In_ const OrtShapeInferContext* context, _Out_ size_t* out);
|
||||
|
||||
/**
|
||||
* Get type and shape info of an input
|
||||
*
|
||||
* \param[in] context
|
||||
* \param[in] index The index of the input
|
||||
* \param[out] info Type shape info of the input
|
||||
*
|
||||
* \since Version 1.17.
|
||||
*/
|
||||
ORT_API2_STATUS(ShapeInferContext_GetInputTypeShape, _In_ const OrtShapeInferContext* context, _In_ size_t index, _Outptr_ OrtTensorTypeAndShapeInfo** info);
|
||||
|
||||
/**
|
||||
* Get attribute from OrtShapeInferContext. Note that OrtShapeInferContext is a per-node context, one could only read attribute from current node.
|
||||
*
|
||||
* \param[in] context
|
||||
* \param[in] attr_name Name of the attribute
|
||||
* \param[out] attr Handle of the attribute fetched
|
||||
*
|
||||
* \since Version 1.17.
|
||||
*/
|
||||
ORT_API2_STATUS(ShapeInferContext_GetAttribute, _In_ const OrtShapeInferContext* context, _In_ const char* attr_name, _Outptr_ const OrtOpAttr** attr);
|
||||
|
||||
/**
|
||||
* Set type and shape info of an ouput
|
||||
*
|
||||
* \param[in] context
|
||||
* \param[in] index The index of the ouput
|
||||
* \param[out] info Type shape info of the output
|
||||
*
|
||||
* \since Version 1.17.
|
||||
*/
|
||||
ORT_API2_STATUS(ShapeInferContext_SetOutputTypeShape, _In_ const OrtShapeInferContext* context, _In_ size_t index, _In_ const OrtTensorTypeAndShapeInfo* info);
|
||||
|
||||
/**
|
||||
* Set symbolic shape to type shape info
|
||||
*
|
||||
* \param[in] info Type shape info
|
||||
* \param[in] dim_params Symbolic strings
|
||||
* \param[in] dim_params_length Number of strings
|
||||
*
|
||||
* \since Version 1.17.
|
||||
*/
|
||||
ORT_API2_STATUS(SetSymbolicDimensions, _In_ OrtTensorTypeAndShapeInfo* info, _In_ const char* dim_params[], _In_ size_t dim_params_length);
|
||||
|
||||
/**
|
||||
* Read contents of an attribute to data
|
||||
*
|
||||
* \param[in] op_attr
|
||||
* \param[in] type Attribute type
|
||||
* \param[out] data Memory address to save raw content of the attribute
|
||||
* \param[in] len Number of bytes allowed to store in data
|
||||
* \param[out] out Number of bytes required to save the data when the call failed, or the real number of bytes saved to data on success
|
||||
*
|
||||
* \since Version 1.17.
|
||||
*/
|
||||
ORT_API2_STATUS(ReadOpAttr, _In_ const OrtOpAttr* op_attr, _In_ OrtOpAttrType type, _Inout_ void* data, _In_ size_t len, _Out_ size_t* out);
|
||||
};
|
||||
|
||||
/*
|
||||
|
|
@ -4535,6 +4603,8 @@ struct OrtCustomOp {
|
|||
|
||||
// Perform the computation step.
|
||||
OrtStatusPtr(ORT_API_CALL* KernelComputeV2)(_In_ void* op_kernel, _In_ OrtKernelContext* context);
|
||||
|
||||
OrtStatusPtr(ORT_API_CALL* InferOutputShapeFn)(_In_ const struct OrtCustomOp* op, _In_ OrtShapeInferContext*);
|
||||
};
|
||||
|
||||
/*
|
||||
|
|
|
|||
|
|
@ -2156,6 +2156,78 @@ struct Op : detail::Base<OrtOp> {
|
|||
size_t output_count);
|
||||
};
|
||||
|
||||
/// <summary>
|
||||
/// Provide access to per-node attributes and input shapes, so one could compute and set output shapes.
|
||||
/// </summary>
|
||||
struct ShapeInferContext {
|
||||
struct SymbolicInteger {
|
||||
SymbolicInteger(int64_t i) : i_(i), is_int_(true){};
|
||||
SymbolicInteger(const char* s) : s_(s), is_int_(false){};
|
||||
SymbolicInteger(const SymbolicInteger&) = default;
|
||||
SymbolicInteger(SymbolicInteger&&) = default;
|
||||
|
||||
SymbolicInteger& operator=(const SymbolicInteger&) = default;
|
||||
SymbolicInteger& operator=(SymbolicInteger&&) = default;
|
||||
|
||||
bool operator==(const SymbolicInteger& dim) const {
|
||||
if (is_int_ == dim.is_int_) {
|
||||
if (is_int_) {
|
||||
return i_ == dim.i_;
|
||||
} else {
|
||||
return std::string{s_} == std::string{dim.s_};
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool IsInt() const { return is_int_; }
|
||||
int64_t AsInt() const { return i_; }
|
||||
const char* AsSym() const { return s_; }
|
||||
|
||||
static constexpr int INVALID_INT_DIM = -2;
|
||||
|
||||
private:
|
||||
union {
|
||||
int64_t i_;
|
||||
const char* s_;
|
||||
};
|
||||
bool is_int_;
|
||||
};
|
||||
|
||||
using Shape = std::vector<SymbolicInteger>;
|
||||
|
||||
ShapeInferContext(const OrtApi* ort_api, OrtShapeInferContext* ctx);
|
||||
|
||||
const Shape& GetInputShape(size_t indice) const { return input_shapes_.at(indice); }
|
||||
|
||||
size_t GetInputCount() const { return input_shapes_.size(); }
|
||||
|
||||
Status SetOutputShape(size_t indice, const Shape& shape);
|
||||
|
||||
int64_t GetAttrInt(const char* attr_name);
|
||||
|
||||
using Ints = std::vector<int64_t>;
|
||||
Ints GetAttrInts(const char* attr_name);
|
||||
|
||||
float GetAttrFloat(const char* attr_name);
|
||||
|
||||
using Floats = std::vector<float>;
|
||||
Floats GetAttrFloats(const char* attr_name);
|
||||
|
||||
std::string GetAttrString(const char* attr_name);
|
||||
|
||||
using Strings = std::vector<std::string>;
|
||||
Strings GetAttrStrings(const char* attr_name);
|
||||
|
||||
private:
|
||||
const OrtOpAttr* GetAttrHdl(const char* attr_name) const;
|
||||
const OrtApi* ort_api_;
|
||||
OrtShapeInferContext* ctx_;
|
||||
std::vector<Shape> input_shapes_;
|
||||
};
|
||||
|
||||
using ShapeInferFn = Ort::Status (*)(Ort::ShapeInferContext&);
|
||||
|
||||
template <typename TOp, typename TKernel, bool WithStatus = false>
|
||||
struct CustomOpBase : OrtCustomOp {
|
||||
CustomOpBase() {
|
||||
|
|
@ -2206,6 +2278,8 @@ struct CustomOpBase : OrtCustomOp {
|
|||
static_cast<TKernel*>(op_kernel)->Compute(context);
|
||||
};
|
||||
}
|
||||
|
||||
SetShapeInferFn<TOp>(0);
|
||||
}
|
||||
|
||||
// Default implementation of GetExecutionProviderType that returns nullptr to default to the CPU provider
|
||||
|
|
@ -2257,6 +2331,20 @@ struct CustomOpBase : OrtCustomOp {
|
|||
return std::vector<std::string>{};
|
||||
}
|
||||
|
||||
template <typename C>
|
||||
decltype(&C::InferOutputShape) SetShapeInferFn(decltype(&C::InferOutputShape)) {
|
||||
OrtCustomOp::InferOutputShapeFn = [](const OrtCustomOp*, OrtShapeInferContext* ort_ctx) -> OrtStatusPtr {
|
||||
ShapeInferContext ctx(&GetApi(), ort_ctx);
|
||||
return C::InferOutputShape(ctx);
|
||||
};
|
||||
return {};
|
||||
}
|
||||
|
||||
template <typename C>
|
||||
void SetShapeInferFn(...) {
|
||||
OrtCustomOp::InferOutputShapeFn = {};
|
||||
}
|
||||
|
||||
protected:
|
||||
// Helper function that returns a map of session config entries specified by CustomOpBase::GetSessionConfigKeys.
|
||||
void GetSessionConfigs(std::unordered_map<std::string, std::string>& out, ConstSessionOptions options) const;
|
||||
|
|
|
|||
|
|
@ -8,6 +8,15 @@
|
|||
// the main C++ file with implementation details.
|
||||
|
||||
#include <cstring>
|
||||
#include <functional>
|
||||
|
||||
#define RETURN_ON_API_FAIL(expression) \
|
||||
{ \
|
||||
auto err = (expression); \
|
||||
if (err) { \
|
||||
return Status(err); \
|
||||
} \
|
||||
}
|
||||
|
||||
namespace Ort {
|
||||
|
||||
|
|
@ -1883,4 +1892,154 @@ void CustomOpBase<TOp, TKernel, WithStatus>::GetSessionConfigs(std::unordered_ma
|
|||
}
|
||||
}
|
||||
|
||||
inline ShapeInferContext::ShapeInferContext(const OrtApi* ort_api,
|
||||
OrtShapeInferContext* ctx) : ort_api_(ort_api), ctx_(ctx) {
|
||||
size_t input_count = 0;
|
||||
Ort::ThrowOnError(ort_api_->ShapeInferContext_GetInputCount(ctx_, &input_count));
|
||||
for (size_t ith_input = 0; ith_input < input_count; ++ith_input) {
|
||||
OrtTensorTypeAndShapeInfo* info{};
|
||||
Ort::ThrowOnError(ort_api_->ShapeInferContext_GetInputTypeShape(ctx, ith_input, &info));
|
||||
TensorTypeAndShapeInfo type_shape_info(info);
|
||||
auto integer_shape = type_shape_info.GetShape();
|
||||
std::vector<const char*> symbolic_shape(integer_shape.size(), {});
|
||||
type_shape_info.GetSymbolicDimensions(&symbolic_shape[0], integer_shape.size());
|
||||
Shape shape;
|
||||
for (size_t ith = 0; ith < integer_shape.size(); ++ith) {
|
||||
if (symbolic_shape[ith] && std::string{symbolic_shape[ith]}.size() > 0) {
|
||||
shape.emplace_back(symbolic_shape[ith]);
|
||||
} else {
|
||||
shape.emplace_back(integer_shape[ith]);
|
||||
}
|
||||
}
|
||||
input_shapes_.push_back(std::move(shape));
|
||||
type_shape_info.release();
|
||||
}
|
||||
}
|
||||
|
||||
inline Status ShapeInferContext::SetOutputShape(size_t indice, const Shape& shape) {
|
||||
OrtTensorTypeAndShapeInfo* info = {};
|
||||
RETURN_ON_API_FAIL(ort_api_->CreateTensorTypeAndShapeInfo(&info));
|
||||
|
||||
using InfoPtr = std::unique_ptr<OrtTensorTypeAndShapeInfo, std::function<void(OrtTensorTypeAndShapeInfo*)>>;
|
||||
|
||||
InfoPtr info_ptr(info, [this](OrtTensorTypeAndShapeInfo* obj) {
|
||||
ort_api_->ReleaseTensorTypeAndShapeInfo(obj);
|
||||
});
|
||||
|
||||
std::vector<int64_t> integer_dims;
|
||||
std::vector<const char*> symbolic_dims;
|
||||
|
||||
for (const auto dim : shape) {
|
||||
if (dim.IsInt()) {
|
||||
integer_dims.push_back(dim.IsInt());
|
||||
symbolic_dims.push_back("");
|
||||
} else {
|
||||
if (!dim.AsSym() || std::string{dim.AsSym()}.empty()) {
|
||||
ORT_CXX_API_THROW("Symbolic dim must not be an empty string", ORT_INVALID_ARGUMENT);
|
||||
}
|
||||
integer_dims.push_back(SymbolicInteger::INVALID_INT_DIM);
|
||||
symbolic_dims.push_back(dim.AsSym());
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_ON_API_FAIL(ort_api_->SetDimensions(info, integer_dims.data(), integer_dims.size()));
|
||||
RETURN_ON_API_FAIL(ort_api_->SetSymbolicDimensions(info, symbolic_dims.data(), symbolic_dims.size()));
|
||||
RETURN_ON_API_FAIL(ort_api_->ShapeInferContext_SetOutputTypeShape(ctx_, indice, info));
|
||||
return Status{nullptr};
|
||||
}
|
||||
|
||||
inline int64_t ShapeInferContext::GetAttrInt(const char* attr_name) {
|
||||
const auto* attr = GetAttrHdl(attr_name);
|
||||
int64_t i = {};
|
||||
size_t out = {};
|
||||
Ort::ThrowOnError(ort_api_->ReadOpAttr(attr, ORT_OP_ATTR_INT, &i, sizeof(i), &out));
|
||||
return i;
|
||||
}
|
||||
|
||||
inline ShapeInferContext::Ints ShapeInferContext::GetAttrInts(const char* attr_name) {
|
||||
const auto* attr = GetAttrHdl(attr_name);
|
||||
int64_t i = {};
|
||||
size_t out = {};
|
||||
// first call to get the bytes needed
|
||||
auto status = ort_api_->ReadOpAttr(attr, ORT_OP_ATTR_INTS, &i, sizeof(i), &out);
|
||||
if (status) {
|
||||
size_t num_i = out / sizeof(int64_t);
|
||||
ShapeInferContext::Ints ints(num_i, 0);
|
||||
Ort::ThrowOnError(ort_api_->ReadOpAttr(attr, ORT_OP_ATTR_INTS, ints.data(), out, &out));
|
||||
return ints;
|
||||
} else {
|
||||
return {i};
|
||||
}
|
||||
}
|
||||
|
||||
inline float ShapeInferContext::GetAttrFloat(const char* attr_name) {
|
||||
const auto* attr = GetAttrHdl(attr_name);
|
||||
float f = {};
|
||||
size_t out = {};
|
||||
Ort::ThrowOnError(ort_api_->ReadOpAttr(attr, ORT_OP_ATTR_FLOAT, &f, sizeof(f), &out));
|
||||
return f;
|
||||
}
|
||||
|
||||
inline ShapeInferContext::Floats ShapeInferContext::GetAttrFloats(const char* attr_name) {
|
||||
const auto* attr = GetAttrHdl(attr_name);
|
||||
float f = {};
|
||||
size_t out = {};
|
||||
// first call to get the bytes needed
|
||||
auto status = ort_api_->ReadOpAttr(attr, ORT_OP_ATTR_FLOATS, &f, sizeof(f), &out);
|
||||
if (status) {
|
||||
size_t num_f = out / sizeof(float);
|
||||
ShapeInferContext::Floats floats(num_f, 0);
|
||||
Ort::ThrowOnError(ort_api_->ReadOpAttr(attr, ORT_OP_ATTR_FLOATS, floats.data(), out, &out));
|
||||
return floats;
|
||||
} else {
|
||||
return {f};
|
||||
}
|
||||
}
|
||||
|
||||
inline std::string ShapeInferContext::GetAttrString(const char* attr_name) {
|
||||
const auto* attr = GetAttrHdl(attr_name);
|
||||
char c = {};
|
||||
size_t out = {};
|
||||
// first call to get the bytes needed
|
||||
auto status = ort_api_->ReadOpAttr(attr, ORT_OP_ATTR_STRING, &c, sizeof(char), &out);
|
||||
if (status) {
|
||||
std::vector<char> chars(out, '\0');
|
||||
Ort::ThrowOnError(ort_api_->ReadOpAttr(attr, ORT_OP_ATTR_STRING, chars.data(), out, &out));
|
||||
return {chars.data()};
|
||||
} else {
|
||||
return {c};
|
||||
}
|
||||
}
|
||||
|
||||
inline ShapeInferContext::Strings ShapeInferContext::GetAttrStrings(const char* attr_name) {
|
||||
const auto* attr = GetAttrHdl(attr_name);
|
||||
char c = {};
|
||||
size_t out = {};
|
||||
// first call to get the bytes needed
|
||||
auto status = ort_api_->ReadOpAttr(attr, ORT_OP_ATTR_STRINGS, &c, sizeof(char), &out);
|
||||
if (status) {
|
||||
std::vector<char> chars(out, '\0');
|
||||
Ort::ThrowOnError(ort_api_->ReadOpAttr(attr, ORT_OP_ATTR_STRINGS, chars.data(), out, &out));
|
||||
ShapeInferContext::Strings strings;
|
||||
char* char_st = chars.data();
|
||||
char* char_ed = char_st + out;
|
||||
while (char_st < char_ed) {
|
||||
strings.emplace_back(char_st);
|
||||
while (*char_st != '\0') {
|
||||
char_st++;
|
||||
}
|
||||
char_st++;
|
||||
}
|
||||
return strings;
|
||||
} else {
|
||||
return {std::string{c}};
|
||||
}
|
||||
}
|
||||
|
||||
inline const OrtOpAttr* ShapeInferContext::GetAttrHdl(const char* attr_name) const {
|
||||
const OrtOpAttr* attr_hdl = {};
|
||||
Ort::ThrowOnError(ort_api_->ShapeInferContext_GetAttribute(ctx_, attr_name, &attr_hdl));
|
||||
return attr_hdl;
|
||||
}
|
||||
|
||||
} // namespace Ort
|
||||
|
|
|
|||
|
|
@ -835,6 +835,8 @@ struct OrtLiteCustomOp : public OrtCustomOp {
|
|||
OrtCustomOp::CreateKernelV2 = {};
|
||||
OrtCustomOp::KernelComputeV2 = {};
|
||||
OrtCustomOp::KernelCompute = {};
|
||||
|
||||
OrtCustomOp::InferOutputShapeFn = {};
|
||||
}
|
||||
|
||||
const std::string op_name_;
|
||||
|
|
@ -870,8 +872,10 @@ struct OrtLiteCustomFunc : public OrtLiteCustomOp {
|
|||
|
||||
OrtLiteCustomFunc(const char* op_name,
|
||||
const char* execution_provider,
|
||||
ComputeFn compute_fn) : OrtLiteCustomOp(op_name, execution_provider),
|
||||
compute_fn_(compute_fn) {
|
||||
ComputeFn compute_fn,
|
||||
ShapeInferFn shape_infer_fn = {}) : OrtLiteCustomOp(op_name, execution_provider),
|
||||
compute_fn_(compute_fn),
|
||||
shape_infer_fn_(shape_infer_fn) {
|
||||
ParseArgs<Args...>(input_types_, output_types_);
|
||||
|
||||
OrtCustomOp::KernelCompute = [](void* op_kernel, OrtKernelContext* context) {
|
||||
|
|
@ -894,12 +898,22 @@ struct OrtLiteCustomFunc : public OrtLiteCustomOp {
|
|||
OrtCustomOp::KernelDestroy = [](void* op_kernel) {
|
||||
delete reinterpret_cast<Kernel*>(op_kernel);
|
||||
};
|
||||
|
||||
if (shape_infer_fn_) {
|
||||
OrtCustomOp::InferOutputShapeFn = [](const OrtCustomOp* op, OrtShapeInferContext* ort_ctx) -> OrtStatusPtr {
|
||||
auto shape_info_fn = static_cast<const MyType*>(op)->shape_infer_fn_;
|
||||
ShapeInferContext ctx(&GetApi(), ort_ctx);
|
||||
return shape_info_fn(ctx);
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
OrtLiteCustomFunc(const char* op_name,
|
||||
const char* execution_provider,
|
||||
ComputeFnReturnStatus compute_fn_return_status) : OrtLiteCustomOp(op_name, execution_provider),
|
||||
compute_fn_return_status_(compute_fn_return_status) {
|
||||
ComputeFnReturnStatus compute_fn_return_status,
|
||||
ShapeInferFn shape_infer_fn = {}) : OrtLiteCustomOp(op_name, execution_provider),
|
||||
compute_fn_return_status_(compute_fn_return_status),
|
||||
shape_infer_fn_(shape_infer_fn) {
|
||||
ParseArgs<Args...>(input_types_, output_types_);
|
||||
|
||||
OrtCustomOp::KernelComputeV2 = [](void* op_kernel, OrtKernelContext* context) -> OrtStatusPtr {
|
||||
|
|
@ -922,10 +936,19 @@ struct OrtLiteCustomFunc : public OrtLiteCustomOp {
|
|||
OrtCustomOp::KernelDestroy = [](void* op_kernel) {
|
||||
delete reinterpret_cast<Kernel*>(op_kernel);
|
||||
};
|
||||
|
||||
if (shape_infer_fn_) {
|
||||
OrtCustomOp::InferOutputShapeFn = [](const OrtCustomOp* op, OrtShapeInferContext* ort_ctx) -> OrtStatusPtr {
|
||||
auto shape_info_fn = static_cast<const MyType*>(op)->shape_infer_fn_;
|
||||
ShapeInferContext ctx(&GetApi(), ort_ctx);
|
||||
return shape_info_fn(ctx);
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
ComputeFn compute_fn_ = {};
|
||||
ComputeFnReturnStatus compute_fn_return_status_ = {};
|
||||
ShapeInferFn shape_infer_fn_ = {};
|
||||
}; // struct OrtLiteCustomFunc
|
||||
|
||||
/////////////////////////// OrtLiteCustomStruct ///////////////////////////
|
||||
|
|
@ -964,7 +987,7 @@ struct OrtLiteCustomStruct : public OrtLiteCustomOp {
|
|||
OrtLiteCustomStruct(const char* op_name,
|
||||
const char* execution_provider) : OrtLiteCustomOp(op_name,
|
||||
execution_provider) {
|
||||
init(&CustomOp::Compute);
|
||||
SetCompute(&CustomOp::Compute);
|
||||
|
||||
OrtCustomOp::CreateKernel = [](const OrtCustomOp* this_, const OrtApi* ort_api, const OrtKernelInfo* info) {
|
||||
auto kernel = std::make_unique<Kernel>();
|
||||
|
|
@ -979,10 +1002,12 @@ struct OrtLiteCustomStruct : public OrtLiteCustomOp {
|
|||
OrtCustomOp::KernelDestroy = [](void* op_kernel) {
|
||||
delete reinterpret_cast<Kernel*>(op_kernel);
|
||||
};
|
||||
|
||||
SetShapeInfer<CustomOp>(0);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
void init(CustomComputeFn<Args...>) {
|
||||
void SetCompute(CustomComputeFn<Args...>) {
|
||||
ParseArgs<Args...>(input_types_, output_types_);
|
||||
OrtCustomOp::KernelCompute = [](void* op_kernel, OrtKernelContext* context) {
|
||||
auto kernel = reinterpret_cast<Kernel*>(op_kernel);
|
||||
|
|
@ -993,7 +1018,7 @@ struct OrtLiteCustomStruct : public OrtLiteCustomOp {
|
|||
}
|
||||
|
||||
template <typename... Args>
|
||||
void init(CustomComputeFnReturnStatus<Args...>) {
|
||||
void SetCompute(CustomComputeFnReturnStatus<Args...>) {
|
||||
ParseArgs<Args...>(input_types_, output_types_);
|
||||
OrtCustomOp::KernelComputeV2 = [](void* op_kernel, OrtKernelContext* context) -> OrtStatusPtr {
|
||||
auto kernel = reinterpret_cast<Kernel*>(op_kernel);
|
||||
|
|
@ -1002,6 +1027,20 @@ struct OrtLiteCustomStruct : public OrtLiteCustomOp {
|
|||
return std::apply([kernel](Args const&... t_args) { Status status = kernel->custom_op_->Compute(t_args...); return status.release(); }, t);
|
||||
};
|
||||
}
|
||||
|
||||
template <typename C>
|
||||
decltype(&C::InferOutputShape) SetShapeInfer(decltype(&C::InferOutputShape)) {
|
||||
OrtCustomOp::InferOutputShapeFn = [](const OrtCustomOp*, OrtShapeInferContext* ort_ctx) -> OrtStatusPtr {
|
||||
ShapeInferContext ctx(&GetApi(), ort_ctx);
|
||||
return C::InferOutputShape(ctx);
|
||||
};
|
||||
return {};
|
||||
}
|
||||
|
||||
template <typename C>
|
||||
void SetShapeInfer(...) {
|
||||
OrtCustomOp::InferOutputShapeFn = {};
|
||||
}
|
||||
}; // struct OrtLiteCustomStruct
|
||||
|
||||
/////////////////////////// CreateLiteCustomOp ////////////////////////////
|
||||
|
|
@ -1009,17 +1048,19 @@ struct OrtLiteCustomStruct : public OrtLiteCustomOp {
|
|||
template <typename... Args>
|
||||
OrtLiteCustomOp* CreateLiteCustomOp(const char* op_name,
|
||||
const char* execution_provider,
|
||||
void (*custom_compute_fn)(Args...)) {
|
||||
void (*custom_compute_fn)(Args...),
|
||||
Status (*shape_infer_fn)(ShapeInferContext&) = {}) {
|
||||
using LiteOp = OrtLiteCustomFunc<Args...>;
|
||||
return std::make_unique<LiteOp>(op_name, execution_provider, custom_compute_fn).release();
|
||||
return std::make_unique<LiteOp>(op_name, execution_provider, custom_compute_fn, shape_infer_fn).release();
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
OrtLiteCustomOp* CreateLiteCustomOp(const char* op_name,
|
||||
const char* execution_provider,
|
||||
Status (*custom_compute_fn_v2)(Args...)) {
|
||||
Status (*custom_compute_fn_v2)(Args...),
|
||||
Status (*shape_infer_fn)(ShapeInferContext&) = {}) {
|
||||
using LiteOp = OrtLiteCustomFunc<Args...>;
|
||||
return std::make_unique<LiteOp>(op_name, execution_provider, custom_compute_fn_v2).release();
|
||||
return std::make_unique<LiteOp>(op_name, execution_provider, custom_compute_fn_v2, shape_infer_fn).release();
|
||||
}
|
||||
|
||||
template <typename CustomOp>
|
||||
|
|
|
|||
|
|
@ -85,6 +85,16 @@ ORT_API_STATUS_IMPL(OrtApis::GetSymbolicDimensions,
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
ORT_API_STATUS_IMPL(OrtApis::SetSymbolicDimensions,
|
||||
_In_ struct OrtTensorTypeAndShapeInfo* info,
|
||||
_In_ const char** names, _In_ size_t dim_params_length) {
|
||||
info->dim_params.clear();
|
||||
for (size_t idx = 0; idx < dim_params_length; ++idx) {
|
||||
info->dim_params.push_back(names[idx]);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
ORT_API_STATUS_IMPL(OrtApis::GetTensorShapeElementCount,
|
||||
_In_ const OrtTensorTypeAndShapeInfo* this_ptr, _Out_ size_t* out) {
|
||||
API_IMPL_BEGIN
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ static constexpr uint32_t min_ort_version_with_variadic_io_support = 14;
|
|||
|
||||
#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
|
||||
static constexpr uint32_t min_ort_version_with_compute_v2_support = 17;
|
||||
static constexpr uint32_t min_ort_version_with_shape_inference = 17;
|
||||
#endif
|
||||
|
||||
#if !defined(DISABLE_FLOAT8_TYPES)
|
||||
|
|
@ -37,6 +38,231 @@ static constexpr uint32_t min_ort_version_with_compute_v2_support = 17;
|
|||
#define SUPPORTED_TENSOR_TYPES DataTypeImpl::AllTensorTypesIRv4()
|
||||
#endif
|
||||
|
||||
#if defined(ORT_MINIMAL_BUILD)
|
||||
struct OrtShapeInferContext {
|
||||
size_t GetInputCount() const { return 0; }
|
||||
OrtTensorTypeAndShapeInfo* GetInputTypeShape(size_t) const { return {}; }
|
||||
onnxruntime::Status SetOutputTypeShape(size_t, const OrtTensorTypeAndShapeInfo*) const {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "OrtShapeInferContext::SetOutputTypeShape not implemented for minimal build");
|
||||
}
|
||||
const ONNX_NAMESPACE::AttributeProto* GetAttr(const char*) const { return {}; }
|
||||
};
|
||||
#else
|
||||
struct OrtShapeInferContext {
|
||||
OrtShapeInferContext(ONNX_NAMESPACE::InferenceContext& ctx) : ctx_(ctx) {
|
||||
auto num_inputs = ctx_.getNumInputs();
|
||||
for (size_t ith_input = 0; ith_input < num_inputs; ++ith_input) {
|
||||
const auto* input_type = ctx_.getInputType(ith_input);
|
||||
const auto& value_case = input_type->value_case();
|
||||
ORT_ENFORCE(value_case == ONNX_NAMESPACE::TypeProto::kTensorType, "shape inference not yet supported for non-tensor types");
|
||||
const auto& shape_proto = input_type->tensor_type().shape();
|
||||
const auto& type_proto = input_type->tensor_type();
|
||||
auto elem_type = ::onnxruntime::utils::CApiElementTypeFromProtoType(type_proto.elem_type());
|
||||
auto tensor_shape = ::onnxruntime::utils::GetTensorShapeFromTensorShapeProto(shape_proto);
|
||||
auto symbolic_dims = GetSymbolicDims(shape_proto);
|
||||
input_type_shapes_.emplace_back(OrtTensorTypeAndShapeInfo::GetTensorShapeAndTypeHelper(elem_type, tensor_shape, &symbolic_dims).release());
|
||||
}
|
||||
}
|
||||
|
||||
~OrtShapeInferContext() = default;
|
||||
size_t GetInputCount() const { return input_type_shapes_.size(); }
|
||||
|
||||
OrtTensorTypeAndShapeInfo* GetInputTypeShape(size_t idx) const {
|
||||
return input_type_shapes_.at(idx).get();
|
||||
}
|
||||
|
||||
onnxruntime::Status SetOutputTypeShape(size_t index, const OrtTensorTypeAndShapeInfo* info) const {
|
||||
ORT_RETURN_IF_NOT(info, "Invalid shape info");
|
||||
ONNX_NAMESPACE::TensorShapeProto shape_proto;
|
||||
const auto& symbolic_dims = info->dim_params;
|
||||
const auto& integer_dims = info->shape.GetDims();
|
||||
ORT_RETURN_IF_NOT(symbolic_dims.size() == integer_dims.size(), "symbolic and integer dims mismatch!");
|
||||
for (size_t ith = 0; ith < symbolic_dims.size(); ith++) {
|
||||
auto* dim_proto = shape_proto.add_dim();
|
||||
if (symbolic_dims[ith].size() > 0) {
|
||||
dim_proto->set_dim_param(symbolic_dims[ith]);
|
||||
} else {
|
||||
dim_proto->set_dim_value(integer_dims[ith]);
|
||||
}
|
||||
}
|
||||
ONNX_NAMESPACE::updateOutputShape(ctx_, index, shape_proto);
|
||||
return onnxruntime::Status::OK();
|
||||
}
|
||||
|
||||
const ONNX_NAMESPACE::AttributeProto* GetAttr(const char* attr_name) const {
|
||||
return ctx_.getAttribute(attr_name);
|
||||
}
|
||||
|
||||
private:
|
||||
static std::vector<std::string> GetSymbolicDims(const ONNX_NAMESPACE::TensorShapeProto& shape_proto) {
|
||||
std::vector<std::string> symblic_dims;
|
||||
for (int ith = 0; ith < shape_proto.dim_size(); ith++) {
|
||||
const auto& dim = shape_proto.dim(ith);
|
||||
if (::onnxruntime::utils::HasDimValue(dim)) {
|
||||
symblic_dims.emplace_back();
|
||||
} else {
|
||||
symblic_dims.emplace_back(dim.dim_param());
|
||||
}
|
||||
}
|
||||
return symblic_dims;
|
||||
}
|
||||
ONNX_NAMESPACE::InferenceContext& ctx_;
|
||||
using TypeShapePtr = std::unique_ptr<OrtTensorTypeAndShapeInfo>;
|
||||
onnxruntime::InlinedVector<TypeShapePtr> input_type_shapes_;
|
||||
};
|
||||
#endif
|
||||
|
||||
ORT_API_STATUS_IMPL(OrtApis::ShapeInferContext_GetInputCount, _In_ const OrtShapeInferContext* context, _Out_ size_t* out) {
|
||||
API_IMPL_BEGIN
|
||||
*out = context->GetInputCount();
|
||||
return nullptr;
|
||||
API_IMPL_END
|
||||
}
|
||||
|
||||
ORT_API_STATUS_IMPL(OrtApis::ShapeInferContext_GetInputTypeShape, _In_ const OrtShapeInferContext* context, _In_ size_t index, _Outptr_ OrtTensorTypeAndShapeInfo** info) {
|
||||
API_IMPL_BEGIN
|
||||
*info = context->GetInputTypeShape(index);
|
||||
if (*info) {
|
||||
return nullptr;
|
||||
} else {
|
||||
return OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Failed to fetch type shape info for the index.");
|
||||
}
|
||||
API_IMPL_END
|
||||
}
|
||||
|
||||
ORT_API_STATUS_IMPL(OrtApis::ShapeInferContext_GetAttribute, _In_ const OrtShapeInferContext* context, _In_ const char* attr_name, _Outptr_ const OrtOpAttr** attr) {
|
||||
API_IMPL_BEGIN
|
||||
*attr = reinterpret_cast<const OrtOpAttr*>(context->GetAttr(attr_name));
|
||||
if (*attr) {
|
||||
return nullptr;
|
||||
} else {
|
||||
return OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Attribute does not exist.");
|
||||
}
|
||||
API_IMPL_END
|
||||
}
|
||||
|
||||
ORT_API_STATUS_IMPL(OrtApis::ReadOpAttr,
|
||||
_In_ const OrtOpAttr* op_attr,
|
||||
_In_ OrtOpAttrType type,
|
||||
_Inout_ void* data,
|
||||
_In_ size_t len,
|
||||
_Out_ size_t* out) {
|
||||
API_IMPL_BEGIN
|
||||
|
||||
if (!op_attr) {
|
||||
return OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Invalid attribute.");
|
||||
}
|
||||
|
||||
auto attr = reinterpret_cast<const ONNX_NAMESPACE::AttributeProto*>(op_attr);
|
||||
OrtStatusPtr ret = nullptr;
|
||||
*out = 0;
|
||||
|
||||
if (type == OrtOpAttrType::ORT_OP_ATTR_FLOAT) {
|
||||
if (len < sizeof(float)) {
|
||||
ret = OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Size of data not large enough to hold a float.");
|
||||
} else {
|
||||
if (attr->has_f()) {
|
||||
auto output_f = reinterpret_cast<float*>(data);
|
||||
*output_f = attr->f();
|
||||
} else {
|
||||
ret = OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Attribute has no float value.");
|
||||
}
|
||||
}
|
||||
*out = sizeof(float);
|
||||
|
||||
} else if (type == OrtOpAttrType::ORT_OP_ATTR_FLOATS) {
|
||||
const auto& floats = attr->floats();
|
||||
auto num_floats = floats.size();
|
||||
|
||||
if (len < sizeof(float) * num_floats) {
|
||||
ret = OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Size of data not large enough to hold the array of floats.");
|
||||
} else {
|
||||
auto output_f = reinterpret_cast<float*>(data);
|
||||
for (auto f : floats) {
|
||||
*output_f = f;
|
||||
output_f++;
|
||||
}
|
||||
}
|
||||
*out = num_floats * sizeof(float);
|
||||
|
||||
} else if (type == OrtOpAttrType::ORT_OP_ATTR_INT) {
|
||||
if (len < sizeof(int)) {
|
||||
ret = OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Size of data not large enough to hold an int64.");
|
||||
} else {
|
||||
if (attr->has_i()) {
|
||||
auto output_i = reinterpret_cast<int64_t*>(data);
|
||||
*output_i = attr->i();
|
||||
} else {
|
||||
ret = OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Attribute has no int64 value.");
|
||||
}
|
||||
}
|
||||
*out = sizeof(int64_t);
|
||||
|
||||
} else if (type == OrtOpAttrType::ORT_OP_ATTR_INTS) {
|
||||
const auto& ints = attr->ints();
|
||||
auto num_ints = ints.size();
|
||||
|
||||
if (len < sizeof(int64_t) * num_ints) {
|
||||
ret = OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Size of data not large enough to hold the array of int64.");
|
||||
} else {
|
||||
auto output_i = reinterpret_cast<int64_t*>(data);
|
||||
for (auto i : ints) {
|
||||
*output_i = i;
|
||||
output_i++;
|
||||
}
|
||||
}
|
||||
*out = num_ints * sizeof(int64_t);
|
||||
|
||||
} else if (type == OrtOpAttrType::ORT_OP_ATTR_STRING) {
|
||||
const auto& s = attr->s();
|
||||
if (len < s.size() + 1) {
|
||||
ret = OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Size of data not large enough to hold the string.");
|
||||
} else {
|
||||
char* output_c = reinterpret_cast<char*>(data);
|
||||
for (char c : s) {
|
||||
*output_c++ = c;
|
||||
}
|
||||
*output_c = '\0';
|
||||
}
|
||||
*out = s.size() + 1;
|
||||
|
||||
} else if (type == OrtOpAttrType::ORT_OP_ATTR_STRINGS) {
|
||||
const auto& ss = attr->strings();
|
||||
size_t num_bytes = 0;
|
||||
for_each(ss.begin(), ss.end(), [&num_bytes](const std::string& s) { num_bytes += s.size() + 1; });
|
||||
|
||||
if (len < num_bytes) {
|
||||
ret = OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Size of data not large enough to hold the array of strings.");
|
||||
} else {
|
||||
char* output_c = reinterpret_cast<char*>(data);
|
||||
for (const auto& s : ss) {
|
||||
for (char c : s) {
|
||||
*output_c++ = c;
|
||||
}
|
||||
*output_c++ = '\0';
|
||||
}
|
||||
}
|
||||
*out = num_bytes;
|
||||
|
||||
} else {
|
||||
ret = OrtApis::CreateStatus(OrtErrorCode::ORT_INVALID_ARGUMENT, "Unknown attribute type.");
|
||||
}
|
||||
|
||||
return ret;
|
||||
API_IMPL_END
|
||||
}
|
||||
|
||||
ORT_API_STATUS_IMPL(OrtApis::ShapeInferContext_SetOutputTypeShape, _In_ const OrtShapeInferContext* context, _In_ size_t index, _In_ const OrtTensorTypeAndShapeInfo* info) {
|
||||
API_IMPL_BEGIN
|
||||
auto status = context->SetOutputTypeShape(index, info);
|
||||
if (status.IsOK()) {
|
||||
return nullptr;
|
||||
} else {
|
||||
return OrtApis::CreateStatus(static_cast<OrtErrorCode>(status.Code()), status.ErrorMessage().c_str());
|
||||
}
|
||||
API_IMPL_END
|
||||
}
|
||||
|
||||
ORT_API_STATUS_IMPL(OrtApis::KernelInfoGetAttribute_float, _In_ const OrtKernelInfo* info, _In_ const char* name, _Out_ float* out) {
|
||||
API_IMPL_BEGIN
|
||||
auto status = reinterpret_cast<const onnxruntime::OpKernelInfo*>(info)->GetAttr<float>(name, out);
|
||||
|
|
@ -596,6 +822,13 @@ ONNX_NAMESPACE::OpSchema CreateSchema(const std::string& domain, const OrtCustom
|
|||
schema.SetDomain(domain);
|
||||
schema.SinceVersion(1);
|
||||
schema.AllowUncheckedAttributes();
|
||||
|
||||
if (op->version >= min_ort_version_with_shape_inference && op->InferOutputShapeFn) {
|
||||
schema.TypeAndShapeInferenceFunction([op](ONNX_NAMESPACE::InferenceContext& infer_ctx) {
|
||||
OrtShapeInferContext ctx(infer_ctx);
|
||||
op->InferOutputShapeFn(op, &ctx);
|
||||
});
|
||||
}
|
||||
return schema;
|
||||
}
|
||||
|
||||
|
|
@ -764,10 +997,14 @@ common::Status CreateCustomRegistry(gsl::span<OrtCustomOpDomain* const> op_domai
|
|||
for (auto schema_iter : schema_map) {
|
||||
schemas.push_back(schema_iter.second);
|
||||
InlinedVector<const KernelDef*> kernel_defs = std::move(kernel_def_map[schema_iter.first]);
|
||||
ONNX_NAMESPACE::InferenceFunction infer_fn = [kernel_defs](ONNX_NAMESPACE::InferenceContext& infer_ctx) {
|
||||
auto infer_fn = schemas.back().GetTypeAndShapeInferenceFunction();
|
||||
ONNX_NAMESPACE::InferenceFunction extended_infer_fn = [infer_fn, kernel_defs](ONNX_NAMESPACE::InferenceContext& infer_ctx) {
|
||||
InferOutputTypes(kernel_defs, infer_ctx);
|
||||
if (infer_fn) {
|
||||
infer_fn(infer_ctx);
|
||||
}
|
||||
};
|
||||
schemas.back().TypeAndShapeInferenceFunction(infer_fn);
|
||||
schemas.back().TypeAndShapeInferenceFunction(extended_infer_fn);
|
||||
}
|
||||
|
||||
ORT_RETURN_IF_ERROR(output->RegisterOpSet(schemas,
|
||||
|
|
|
|||
|
|
@ -2715,6 +2715,12 @@ static constexpr OrtApi ort_api_1_to_17 = {
|
|||
// End of Version 16 - DO NOT MODIFY ABOVE (see above text for more information)
|
||||
|
||||
&OrtApis::SetUserLoggingFunction,
|
||||
&OrtApis::ShapeInferContext_GetInputCount,
|
||||
&OrtApis::ShapeInferContext_GetInputTypeShape,
|
||||
&OrtApis::ShapeInferContext_GetAttribute,
|
||||
&OrtApis::ShapeInferContext_SetOutputTypeShape,
|
||||
&OrtApis::SetSymbolicDimensions,
|
||||
&OrtApis::ReadOpAttr,
|
||||
};
|
||||
|
||||
// OrtApiBase can never change as there is no way to know what version of OrtApiBase is returned by OrtGetApiBase.
|
||||
|
|
|
|||
|
|
@ -491,6 +491,14 @@ ORT_API_STATUS_IMPL(GetTensorRTProviderOptionsByName, _In_ const OrtTensorRTProv
|
|||
ORT_API_STATUS_IMPL(UpdateCUDAProviderOptionsWithValue, _Inout_ OrtCUDAProviderOptionsV2* cuda_options, _In_ const char* key, _In_ void* value);
|
||||
ORT_API_STATUS_IMPL(GetCUDAProviderOptionsByName, _In_ const OrtCUDAProviderOptionsV2* cuda_options, _In_ const char* key, _Outptr_ void** ptr);
|
||||
ORT_API_STATUS_IMPL(KernelContext_GetResource, _In_ const OrtKernelContext* context, _In_ int resource_version, _In_ int resource_id, _Outptr_ void** stream);
|
||||
|
||||
ORT_API_STATUS_IMPL(SetUserLoggingFunction, _Inout_ OrtSessionOptions* options,
|
||||
_In_ OrtLoggingFunction user_logging_function, _In_opt_ void* user_logging_param);
|
||||
ORT_API_STATUS_IMPL(ShapeInferContext_GetInputCount, _In_ const OrtShapeInferContext* context, _Out_ size_t* out);
|
||||
ORT_API_STATUS_IMPL(ShapeInferContext_GetInputTypeShape, _In_ const OrtShapeInferContext* context, _In_ size_t index, _Outptr_ OrtTensorTypeAndShapeInfo** info);
|
||||
ORT_API_STATUS_IMPL(ShapeInferContext_GetAttribute, _In_ const OrtShapeInferContext* context, _In_ const char* attr_name, _Outptr_ const OrtOpAttr** attr);
|
||||
ORT_API_STATUS_IMPL(ShapeInferContext_SetOutputTypeShape, _In_ const OrtShapeInferContext* context, _In_ size_t index, _In_ const OrtTensorTypeAndShapeInfo* info);
|
||||
ORT_API_STATUS_IMPL(SetSymbolicDimensions, _In_ OrtTensorTypeAndShapeInfo* info, _In_ const char* dim_params[], _In_ size_t dim_params_length);
|
||||
ORT_API_STATUS_IMPL(ReadOpAttr, _In_ const OrtOpAttr* op_attr, _In_ OrtOpAttrType type, _Inout_ void* data, _In_ size_t len, _Out_ size_t* out);
|
||||
|
||||
} // namespace OrtApis
|
||||
|
|
|
|||
|
|
@ -183,6 +183,7 @@ static constexpr PATH_TYPE SEQUENCE_MODEL_URI = TSTR("testdata/sequence_length.o
|
|||
static constexpr PATH_TYPE SEQUENCE_MODEL_URI_2 = TSTR("testdata/optional_sequence_tensor.onnx");
|
||||
#endif
|
||||
static constexpr PATH_TYPE CUSTOM_OP_MODEL_URI = TSTR("testdata/foo_1.onnx");
|
||||
static constexpr PATH_TYPE CUSTOM_OP_LIBRARY_ATTR_TESTER_URI = TSTR("testdata/custom_op_library/attr_tester.onnx");
|
||||
static constexpr PATH_TYPE CUSTOM_OP_LIBRARY_TEST_MODEL_URI = TSTR("testdata/custom_op_library/custom_op_test.onnx");
|
||||
static constexpr PATH_TYPE CUSTOM_OP_LIBRARY_COPY_TENSOR_ARRAY_2 = TSTR("testdata/custom_op_library/copy_2_inputs_2_outputs.onnx");
|
||||
static constexpr PATH_TYPE CUSTOM_OP_LIBRARY_COPY_TENSOR_ARRAY_3 = TSTR("testdata/custom_op_library/copy_3_inputs_3_outputs.onnx");
|
||||
|
|
@ -1408,6 +1409,36 @@ TEST(CApiTest, test_custom_op_library) {
|
|||
#endif
|
||||
}
|
||||
|
||||
#if defined(__ANDROID__)
|
||||
TEST(CApiTest, DISABLED_test_custom_op_shape_infer_attr) {
|
||||
// To accomodate a reduced op build pipeline
|
||||
#elif defined(REDUCED_OPS_BUILD) && defined(USE_CUDA)
|
||||
TEST(CApiTest, DISABLED_test_custom_op_shape_infer_attr) {
|
||||
#else
|
||||
TEST(CApiTest, test_custom_op_shape_infer_attr) {
|
||||
#endif
|
||||
std::vector<Input> inputs(1);
|
||||
inputs[0].name = "input_0";
|
||||
inputs[0].dims = {5};
|
||||
inputs[0].values = {1.f, 2.f, 3.f, 4.f, 5.f};
|
||||
|
||||
// prepare expected inputs and outputs
|
||||
std::vector<int64_t> expected_dims_y = {5};
|
||||
std::vector<float> expected_values_y = {6.f, 12.f, 18.f, 24.f, 30.f};
|
||||
|
||||
onnxruntime::PathString lib_name;
|
||||
#if defined(_WIN32)
|
||||
lib_name = ORT_TSTR("custom_op_library.dll");
|
||||
#elif defined(__APPLE__)
|
||||
lib_name = ORT_TSTR("libcustom_op_library.dylib");
|
||||
#else
|
||||
lib_name = ORT_TSTR("./libcustom_op_library.so");
|
||||
#endif
|
||||
|
||||
TestInference<float>(*ort_env, CUSTOM_OP_LIBRARY_ATTR_TESTER_URI, inputs, "output_0", expected_dims_y,
|
||||
expected_values_y, 0, nullptr, lib_name.c_str());
|
||||
}
|
||||
|
||||
// It has memory leak. The OrtCustomOpDomain created in custom_op_library.cc:RegisterCustomOps function was not freed
|
||||
#if defined(__ANDROID__)
|
||||
TEST(CApiTest, test_custom_op_library_copy_variadic) {
|
||||
|
|
@ -3201,6 +3232,15 @@ struct Merge {
|
|||
strings_out->SetStringOutput(string_pool, {static_cast<int64_t>(string_pool.size())});
|
||||
return Ort::Status(nullptr);
|
||||
}
|
||||
static Ort::Status InferOutputShape(Ort::ShapeInferContext& ctx) {
|
||||
auto input_count = ctx.GetInputCount();
|
||||
if (input_count != 2) {
|
||||
return Ort::Status("input count should be 2", OrtErrorCode::ORT_INVALID_ARGUMENT);
|
||||
}
|
||||
Ort::ShapeInferContext::Shape shape_1 = {{-1}};
|
||||
ctx.SetOutputShape(0, shape_1);
|
||||
return Ort::Status(nullptr);
|
||||
}
|
||||
bool reverse_ = false;
|
||||
};
|
||||
|
||||
|
|
|
|||
BIN
onnxruntime/test/testdata/custom_op_library/attr_tester.onnx
vendored
Normal file
BIN
onnxruntime/test/testdata/custom_op_library/attr_tester.onnx
vendored
Normal file
Binary file not shown.
|
|
@ -7,25 +7,47 @@
|
|||
|
||||
#include "onnxruntime_lite_custom_op.h"
|
||||
|
||||
#define CUSTOM_ENFORCE(cond, msg) \
|
||||
if (!(cond)) { \
|
||||
ORT_CXX_API_THROW(msg, OrtErrorCode::ORT_RUNTIME_EXCEPTION); \
|
||||
}
|
||||
|
||||
using namespace Ort::Custom;
|
||||
|
||||
namespace Cpu {
|
||||
|
||||
Ort::Status KernelOne(const Ort::Custom::Tensor<float>& X,
|
||||
struct KernelOne {
|
||||
KernelOne(const OrtApi*, const OrtKernelInfo*) {}
|
||||
|
||||
Ort::Status Compute(const Ort::Custom::Tensor<float>& X,
|
||||
const Ort::Custom::Tensor<float>& Y,
|
||||
Ort::Custom::Tensor<float>& Z) {
|
||||
if (X.NumberOfElement() != Y.NumberOfElement()) {
|
||||
return Ort::Status("x and y has different number of elements", OrtErrorCode::ORT_INVALID_ARGUMENT);
|
||||
if (X.NumberOfElement() != Y.NumberOfElement()) {
|
||||
return Ort::Status("x and y has different number of elements", OrtErrorCode::ORT_INVALID_ARGUMENT);
|
||||
}
|
||||
auto x_shape = X.Shape();
|
||||
auto x_raw = X.Data();
|
||||
auto y_raw = Y.Data();
|
||||
auto z_raw = Z.Allocate(x_shape);
|
||||
for (int64_t i = 0; i < Z.NumberOfElement(); ++i) {
|
||||
z_raw[i] = x_raw[i] + y_raw[i];
|
||||
}
|
||||
return Ort::Status{nullptr};
|
||||
}
|
||||
auto x_shape = X.Shape();
|
||||
auto x_raw = X.Data();
|
||||
auto y_raw = Y.Data();
|
||||
auto z_raw = Z.Allocate(x_shape);
|
||||
for (int64_t i = 0; i < Z.NumberOfElement(); ++i) {
|
||||
z_raw[i] = x_raw[i] + y_raw[i];
|
||||
|
||||
static Ort::Status InferOutputShape(Ort::ShapeInferContext& ctx) {
|
||||
auto input_count = ctx.GetInputCount();
|
||||
if (input_count != 2) {
|
||||
return Ort::Status("input count should be 2", OrtErrorCode::ORT_INVALID_ARGUMENT);
|
||||
}
|
||||
Ort::ShapeInferContext::Shape shape_3_5 = {{3}, {5}};
|
||||
if (ctx.GetInputShape(0) != shape_3_5 ||
|
||||
ctx.GetInputShape(1) != shape_3_5) {
|
||||
return Ort::Status("input shape mismatch", OrtErrorCode::ORT_INVALID_ARGUMENT);
|
||||
}
|
||||
return Ort::Status{nullptr};
|
||||
}
|
||||
return Ort::Status{nullptr};
|
||||
}
|
||||
};
|
||||
|
||||
// lite custom op as a function
|
||||
void KernelTwo(const Ort::Custom::Tensor<float>& X,
|
||||
|
|
@ -206,8 +228,68 @@ Ort::Status CopyTensorArrayCombined(const Ort::Custom::Tensor<float>& first_inpu
|
|||
return CopyTensorArrayAllVariadic<T>(other_inputs, other_outputs);
|
||||
}
|
||||
|
||||
Ort::Status AttrTesterIntFloatCompute(const Ort::Custom::Tensor<float>& X, Ort::Custom::Tensor<float>& Z) {
|
||||
auto x_shape = X.Shape();
|
||||
auto x_raw = X.Data();
|
||||
auto z_raw = Z.Allocate(x_shape);
|
||||
for (int64_t i = 0; i < X.NumberOfElement(); ++i) {
|
||||
z_raw[i] = x_raw[i] * 2;
|
||||
}
|
||||
return Ort::Status{nullptr};
|
||||
}
|
||||
|
||||
Ort::Status AttrTesterIntFloatShapeInfer(Ort::ShapeInferContext& ctx) {
|
||||
CUSTOM_ENFORCE(ctx.GetAttrInt("a_int") == 1, "int attr mismatch");
|
||||
CUSTOM_ENFORCE(ctx.GetAttrFloat("a_float") == 2.f, "float attr mismatch");
|
||||
std::vector<int64_t> ints{3, 4, 5};
|
||||
CUSTOM_ENFORCE(ctx.GetAttrInts("ints") == ints, "ints attr mismatch");
|
||||
std::vector<float> floats{6, 7, 8};
|
||||
CUSTOM_ENFORCE(ctx.GetAttrFloats("floats") == floats, "floats attr mismatch");
|
||||
auto input_shape = ctx.GetInputShape(0);
|
||||
CUSTOM_ENFORCE(input_shape.size() == 1 &&
|
||||
!input_shape[0].IsInt() &&
|
||||
std::string{input_shape[0].AsSym()} == "d",
|
||||
"input dim is not symbolic");
|
||||
ctx.SetOutputShape(0, input_shape);
|
||||
return Ort::Status{nullptr};
|
||||
}
|
||||
|
||||
struct AttrTesterStringKernel {
|
||||
void Compute(OrtKernelContext* context) {
|
||||
Ort::KernelContext ctx(context);
|
||||
auto input_X = ctx.GetInput(0);
|
||||
const auto* X = input_X.GetTensorData<float>();
|
||||
auto dimensions = input_X.GetTensorTypeAndShapeInfo().GetShape();
|
||||
auto output = ctx.GetOutput(0, dimensions);
|
||||
auto* out = output.GetTensorMutableData<float>();
|
||||
const size_t size = output.GetTensorTypeAndShapeInfo().GetElementCount();
|
||||
for (size_t i = 0; i < size; i++) {
|
||||
out[i] = X[i] * 3;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct AttrTesterStringOp : Ort::CustomOpBase<AttrTesterStringOp, AttrTesterStringKernel> {
|
||||
void* CreateKernel(const OrtApi&, const OrtKernelInfo*) const {
|
||||
return std::make_unique<AttrTesterStringKernel>().release();
|
||||
};
|
||||
const char* GetName() const { return "AttrTesterString"; };
|
||||
const char* GetExecutionProviderType() const { return "CPUExecutionProvider"; };
|
||||
size_t GetInputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetInputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
size_t GetOutputTypeCount() const { return 1; };
|
||||
ONNXTensorElementDataType GetOutputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
|
||||
|
||||
static Ort::Status InferOutputShape(Ort::ShapeInferContext& ctx) {
|
||||
CUSTOM_ENFORCE(ctx.GetAttrString("a_string") == "iamastring", "string attr mismatch");
|
||||
std::vector<std::string> strings{"more", "strings"};
|
||||
CUSTOM_ENFORCE(ctx.GetAttrStrings("strings") == strings, "strings attr mismatch");
|
||||
return Ort::Status{nullptr};
|
||||
}
|
||||
};
|
||||
|
||||
void RegisterOps(Ort::CustomOpDomain& domain) {
|
||||
static const std::unique_ptr<OrtLiteCustomOp> c_CustomOpOne{Ort::Custom::CreateLiteCustomOp("CustomOpOne", "CPUExecutionProvider", KernelOne)};
|
||||
static const std::unique_ptr<OrtLiteCustomOp> c_CustomOpOne{Ort::Custom::CreateLiteCustomOp<KernelOne>("CustomOpOne", "CPUExecutionProvider")};
|
||||
static const std::unique_ptr<OrtLiteCustomOp> c_CustomOpTwo{Ort::Custom::CreateLiteCustomOp("CustomOpTwo", "CPUExecutionProvider", KernelTwo)};
|
||||
static const std::unique_ptr<OrtLiteCustomOp> c_MulTopOpFloat{Ort::Custom::CreateLiteCustomOp("MulTop", "CPUExecutionProvider", MulTop<float>)};
|
||||
static const std::unique_ptr<OrtLiteCustomOp> c_MulTopOpInt32{Ort::Custom::CreateLiteCustomOp("MulTop", "CPUExecutionProvider", MulTop<int32_t>)};
|
||||
|
|
@ -218,6 +300,9 @@ void RegisterOps(Ort::CustomOpDomain& domain) {
|
|||
static const std::unique_ptr<OrtLiteCustomOp> c_CopyTensorArrayAllVariadic{Ort::Custom::CreateLiteCustomOp("CopyTensorArrayAllVariadic", "CPUExecutionProvider", CopyTensorArrayAllVariadic<float>)};
|
||||
static const std::unique_ptr<OrtLiteCustomOp> c_CopyTensorArrayCombined{Ort::Custom::CreateLiteCustomOp("CopyTensorArrayCombined", "CPUExecutionProvider", CopyTensorArrayCombined<float>)};
|
||||
|
||||
static const std::unique_ptr<OrtLiteCustomOp> c_AtterTesterIntFloat{Ort::Custom::CreateLiteCustomOp("AttrTesterIntFloat", "CPUExecutionProvider", AttrTesterIntFloatCompute, AttrTesterIntFloatShapeInfer)};
|
||||
static const AttrTesterStringOp c_AtterTesterString;
|
||||
|
||||
#if !defined(DISABLE_FLOAT8_TYPES)
|
||||
static const CustomOpOneFloat8 c_CustomOpOneFloat8;
|
||||
static const std::unique_ptr<OrtLiteCustomOp> c_FilterFloat8{Ort::Custom::CreateLiteCustomOp("FilterFloat8", "CPUExecutionProvider", FilterFloat8)};
|
||||
|
|
@ -233,6 +318,8 @@ void RegisterOps(Ort::CustomOpDomain& domain) {
|
|||
domain.Add(c_Box.get());
|
||||
domain.Add(c_CopyTensorArrayAllVariadic.get());
|
||||
domain.Add(c_CopyTensorArrayCombined.get());
|
||||
domain.Add(c_AtterTesterIntFloat.get());
|
||||
domain.Add(&c_AtterTesterString);
|
||||
|
||||
#if !defined(DISABLE_FLOAT8_TYPES)
|
||||
domain.Add(&c_CustomOpOneFloat8);
|
||||
|
|
|
|||
Loading…
Reference in a new issue