diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index 9f2898618d..c730b386f2 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -74,6 +74,8 @@ ORT_DEFINE_RELEASE(RunOptions); ORT_DEFINE_RELEASE(Session); ORT_DEFINE_RELEASE(SessionOptions); ORT_DEFINE_RELEASE(TensorTypeAndShapeInfo); +ORT_DEFINE_RELEASE(SequenceTypeInfo); +ORT_DEFINE_RELEASE(MapTypeInfo); ORT_DEFINE_RELEASE(TypeInfo); ORT_DEFINE_RELEASE(Value); ORT_DEFINE_RELEASE(ModelMetadata); @@ -285,11 +287,30 @@ struct TensorTypeAndShapeInfo : Base { std::vector GetShape() const; }; +struct SequenceTypeInfo : Base { + explicit SequenceTypeInfo(std::nullptr_t) {} + explicit SequenceTypeInfo(OrtSequenceTypeInfo* p) : Base{p} {} + + TypeInfo GetSequenceElementType() const; +}; + +struct MapTypeInfo : Base { + explicit MapTypeInfo(std::nullptr_t) {} + explicit MapTypeInfo(OrtMapTypeInfo* p) : Base{p} {} + + ONNXTensorElementDataType GetMapKeyType() const; + TypeInfo GetMapValueType() const; +}; + struct TypeInfo : Base { explicit TypeInfo(std::nullptr_t) {} explicit TypeInfo(OrtTypeInfo* p) : Base{p} {} Unowned GetTensorTypeAndShapeInfo() const; + Unowned GetSequenceTypeInfo() const; + Unowned GetMapTypeInfo() const; + + ONNXType GetONNXType() const; }; diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h index c41b569bc1..9888b9faba 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h @@ -622,6 +622,36 @@ inline Unowned TypeInfo::GetTensorTypeAndShapeInfo() con return Unowned(const_cast(out)); } +inline Unowned TypeInfo::GetSequenceTypeInfo() const { + const OrtSequenceTypeInfo* out; + ThrowOnError(GetApi().CastTypeInfoToSequenceTypeInfo(p_, &out)); + return Unowned{const_cast(out)}; +} + +inline TypeInfo SequenceTypeInfo::GetSequenceElementType() const { + OrtTypeInfo* output; + ThrowOnError(GetApi().GetSequenceElementType(p_, &output)); + return TypeInfo{output}; +} + +inline Unowned TypeInfo::GetMapTypeInfo() const { + const OrtMapTypeInfo* out; + ThrowOnError(GetApi().CastTypeInfoToMapTypeInfo(p_, &out)); + return Unowned{const_cast(out)}; +} + +inline ONNXTensorElementDataType MapTypeInfo::GetMapKeyType() const { + ONNXTensorElementDataType out; + ThrowOnError(GetApi().GetMapKeyType(p_, &out)); + return out; +} + +inline TypeInfo MapTypeInfo::GetMapValueType() const { + OrtTypeInfo* output; + ThrowOnError(GetApi().GetMapValueType(p_, &output)); + return TypeInfo{output}; +} + inline ONNXType TypeInfo::GetONNXType() const { ONNXType out; ThrowOnError(GetApi().GetOnnxTypeFromTypeInfo(p_, &out)); diff --git a/onnxruntime/test/shared_lib/test_nontensor_types.cc b/onnxruntime/test/shared_lib/test_nontensor_types.cc index 372699a7dd..5f68fa59ba 100644 --- a/onnxruntime/test/shared_lib/test_nontensor_types.cc +++ b/onnxruntime/test/shared_lib/test_nontensor_types.cc @@ -150,6 +150,41 @@ TEST(CApiTest, CreateGetVectorOfMapsStringFloat) { // support zipmap output typ } } +TEST(CApiTest, TypeInfoMap) { + // Creation + auto default_allocator = onnxruntime::make_unique(); + Ort::MemoryInfo info("Cpu", OrtDeviceAllocator, 0, OrtMemTypeDefault); + + const int64_t NUM_KV_PAIRS = 4; + std::vector keys{0, 1, 2, 3}; + std::vector dims = {NUM_KV_PAIRS}; + std::vector values{3.0f, 1.0f, 2.f, 0.f}; + // create key tensor + Ort::Value keys_tensor = Ort::Value::CreateTensor(info, keys.data(), keys.size() * sizeof(int64_t), + dims.data(), dims.size(), ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + // create value tensor + Ort::Value values_tensor = Ort::Value::CreateTensor(info, values.data(), values.size() * sizeof(float), + dims.data(), dims.size(), ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + + Ort::Value map_ort = Ort::Value::CreateMap(keys_tensor, values_tensor); + Ort::TypeInfo type_info = map_ort.GetTypeInfo(); + Ort::MapTypeInfo map_type_info = type_info.GetMapTypeInfo(); + + //Check key type + ASSERT_EQ(map_type_info.GetMapKeyType(), ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + + Ort::TypeInfo map_value_type_info = map_type_info.GetMapValueType(); + + //Check value type and shape + ASSERT_EQ(map_value_type_info.GetONNXType(), ONNX_TYPE_TENSOR); + // No shape present, as map values allow different shapes for each element + // ASSERT_EQ(map_value_type_info.GetTensorTypeAndShapeInfo().GetShape(), dims); + ASSERT_EQ(map_value_type_info.GetTensorTypeAndShapeInfo().GetElementType(), ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + + map_value_type_info.release(); + map_type_info.release(); +} + TEST(CApiTest, CreateGetSeqTensors) { // Creation auto default_allocator = onnxruntime::make_unique(); @@ -215,3 +250,31 @@ TEST(CApiTest, CreateGetSeqStringTensors) { } ASSERT_EQ(string_set, std::set(std::begin(string_input_data), std::end(string_input_data))); } + +TEST(CApiTest, TypeInfoSequence) { + // Creation + auto default_allocator = onnxruntime::make_unique(); + Ort::MemoryInfo info("Cpu", OrtDeviceAllocator, 0, OrtMemTypeDefault); + + std::vector in; + std::vector vals{3, 1, 2, 0}; + std::vector dims{1, 4}; + const int N = 2; + for (int i = 0; i < N; ++i) { + // create tensor + Ort::Value tensor = Ort::Value::CreateTensor(info, vals.data(), vals.size() * sizeof(int64_t), + dims.data(), dims.size(), ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + in.push_back(std::move(tensor)); + } + + Ort::Value seq_ort = Ort::Value::CreateSequence(in); + Ort::TypeInfo type_info = seq_ort.GetTypeInfo(); + Ort::SequenceTypeInfo seq_type_info = type_info.GetSequenceTypeInfo(); + + ASSERT_EQ(seq_type_info.GetSequenceElementType().GetONNXType(), ONNX_TYPE_TENSOR); + // No shape present, as sequence allows different shapes for each element + // ASSERT_EQ(seq_type_info.GetSequenceElementType().GetTensorTypeAndShapeInfo().GetShape(), dims); + ASSERT_EQ(seq_type_info.GetSequenceElementType().GetTensorTypeAndShapeInfo().GetElementType(), ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + + seq_type_info.release(); +} \ No newline at end of file