diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index bebbc0ef53..c0c5cb8ef2 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -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); diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index df15d2d2ec..2ea0a0d238 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -252,7 +252,7 @@ struct AllocatorInfo : Base { struct CustomOpApi { CustomOpApi(const OrtCustomOpApi& api) : api_(api) {} - template // T is only implemented for float and int64_t + template // 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); diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h index 970155aeaa..ccb7af2616 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h @@ -393,6 +393,24 @@ inline int64_t CustomOpApi::KernelInfoGetAttribute(_In_ const OrtKernel return out; } +template <> +inline std::string CustomOpApi::KernelInfoGetAttribute(_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)); diff --git a/onnxruntime/core/session/custom_ops.cc b/onnxruntime/core/session/custom_ops.cc index 7d4218022a..3605850cfa 100644 --- a/onnxruntime/core/session/custom_ops.cc +++ b/onnxruntime/core/session/custom_ops.cc @@ -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(info)->GetAttr(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, diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 6c38782e43..c4ed8be620 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -11,6 +11,7 @@ #include #include #include +#include #include #include "core/common/logging/logging.h"