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

Ryanunderhill/c api string arg (#1436)

Browse source 
* Add string attribute interface for C API.

* Add string attribute interface for C++ API accordingly.

* Update comment to say that string is also valid
This commit is contained in:
Ryan Hill 2019-07-19 19:53:37 -07:00 committed by GitHub
parent f0b9a814e5
commit 9e2fa69785
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 39 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
include/onnxruntime/core/session/onnxruntime_c_api.h
View file

@ -548,6 +548,7 @@ struct OrtCustomOpApi {
*/
OrtStatus*(ORT_API_CALL* KernelInfoGetAttribute_float)(_In_ const OrtKernelInfo* info, _In_ const char* name, _Out_ float* out);
OrtStatus*(ORT_API_CALL* KernelInfoGetAttribute_int64)(_In_ const OrtKernelInfo* info, _In_ const char* name, _Out_ int64_t* out);
OrtStatus*(ORT_API_CALL* KernelInfoGetAttribute_string)(_In_ const OrtKernelInfo* info, _In_ const char* name, _Out_ char* out, _Inout_ size_t* size);
OrtStatus*(ORT_API_CALL* GetTensorTypeAndShape)(_In_ const OrtValue* value, _Out_ OrtTensorTypeAndShapeInfo** out);

2
include/onnxruntime/core/session/onnxruntime_cxx_api.h
View file

@ -252,7 +252,7 @@ struct AllocatorInfo : Base<OrtAllocatorInfo> {
struct CustomOpApi {
CustomOpApi(const OrtCustomOpApi& api) : api_(api) {}
template <typename T> // T is only implemented for float and int64_t
template <typename T> // T is only implemented for float, int64_t, and string
T KernelInfoGetAttribute(_In_ const OrtKernelInfo* info, _In_ const char* name);
OrtTensorTypeAndShapeInfo* GetTensorTypeAndShape(_In_ const OrtValue* value);

18
include/onnxruntime/core/session/onnxruntime_cxx_inline.h
View file

@ -393,6 +393,24 @@ inline int64_t CustomOpApi::KernelInfoGetAttribute<int64_t>(_In_ const OrtKernel
return out;
}
template <>
inline std::string CustomOpApi::KernelInfoGetAttribute<std::string>(_In_ const OrtKernelInfo* info, _In_ const char* name) {
size_t size = 0;
std::string out;
OrtStatus* status = api_.KernelInfoGetAttribute_string(info, name, nullptr, &size);
// The status should be ORT_INVALID_ARGUMENT because the size is insufficient to hold the string
if (OrtGetErrorCode(status) == ORT_INVALID_ARGUMENT) {
OrtReleaseStatus(status);
out.resize(size);
ORT_THROW_ON_ERROR(api_.KernelInfoGetAttribute_string(info, name, &out[0], &size));
out.resize(size - 1); // remove the terminating character '\0'
} else {
ORT_THROW_ON_ERROR(status);
}
return out;
}
inline OrtTensorTypeAndShapeInfo* CustomOpApi::GetTensorTypeAndShape(_In_ const OrtValue* value) {
OrtTensorTypeAndShapeInfo* out;
ORT_THROW_ON_ERROR(api_.GetTensorTypeAndShape(value, &out));

18
onnxruntime/core/session/custom_ops.cc
View file

@ -50,9 +50,27 @@ ORT_API_STATUS_IMPL(OrtKernelContext_GetOutput, OrtKernelContext* context, _In_
return nullptr;
};
ORT_API_STATUS_IMPL(OrtKernelInfoGetAttribute_string, _In_ const OrtKernelInfo* info, _In_ const char* name, _Out_ char* out, _Inout_ size_t *size) {
std::string value;
auto status = reinterpret_cast<const onnxruntime::OpKernelInfo*>(info)->GetAttr<std::string>(name, &value);
if (status.IsOK()) {
if (*size >= value.size() + 1) {
std::memcpy(out, value.data(), value.size());
out[value.size()] = '\0';
*size = value.size();
return nullptr;
} else {
*size = value.size() + 1;
return OrtCreateStatus(ORT_INVALID_ARGUMENT, "Result buffer is not large enough");
}
}
return onnxruntime::ToOrtStatus(status);
}
constexpr OrtCustomOpApi g_custom_op_api = {
&OrtKernelInfoGetAttribute_float,
&OrtKernelInfoGetAttribute_int64,
&OrtKernelInfoGetAttribute_string,
&OrtGetTensorTypeAndShape,

1
onnxruntime/core/session/inference_session.cc
View file

@ -11,6 +11,7 @@
#include <sstream>
#include <unordered_set>
#include <list>
#include <string>
#include <thread>
#include "core/common/logging/logging.h"