microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
include/custom_op/tensor_api.h
588lines · modeblame
97ee9eb5Wenbing Li2 years ago | 1 | // Copyright (c) Microsoft Corporation. All rights reserved. |
| 2 | // Licensed under the MIT License. | |
| 3 | | |
64646279Wenbing Li2 years ago | 4 | #pragma once |
| 5 | #include <optional> | |
| 6 | #include <numeric> | |
| 7 | #include <type_traits> | |
311dd354Wenbing Li2 years ago | 8 | #include <assert.h> |
| 9 | | |
64646279Wenbing Li2 years ago | 10 | #include "onnxruntime_f16.h" |
| 11 | #include "kernel_context.h" | |
| 12 | | |
| 13 | namespace Ort { | |
| 14 | namespace Custom { | |
| 15 | | |
| 16 | template <typename T> | |
| 17 | struct Span { | |
| 18 | const T* data_ = {}; | |
| 19 | size_t size_ = {}; | |
| 20 | void Assign(const T* data, size_t size) { | |
| 21 | data_ = data; | |
| 22 | size_ = size; | |
| 23 | } | |
| 24 | size_t size() const { return size_; } | |
| 25 | T operator[](size_t indice) const { | |
| 26 | return data_[indice]; | |
| 27 | } | |
| 28 | const T* data() const { return data_; } | |
| 29 | }; | |
| 30 | | |
| 31 | | |
| 32 | #if ORT_API_VERSION >= 16 | |
| 33 | | |
| 34 | template <> | |
| 35 | struct Span<MFloat16> { | |
| 36 | const MFloat16* data_ = {}; | |
| 37 | size_t size_ = {}; | |
| 38 | void Assign(const MFloat16* data, size_t size) { | |
| 39 | data_ = data; | |
| 40 | size_ = size; | |
| 41 | } | |
| 42 | size_t size() const { return size_; } | |
| 43 | MFloat16 operator[](size_t indice) const { | |
| 44 | return data_[indice]; | |
| 45 | } | |
| 46 | const MFloat16* data() const { return data_; } | |
| 47 | }; | |
| 48 | | |
| 49 | template <> | |
| 50 | struct Span<BFloat16> { | |
| 51 | const BFloat16* data_ = {}; | |
| 52 | size_t size_ = {}; | |
| 53 | void Assign(const BFloat16* data, size_t size) { | |
| 54 | data_ = data; | |
| 55 | size_ = size; | |
| 56 | } | |
| 57 | size_t size() const { return size_; } | |
| 58 | BFloat16 operator[](size_t indice) const { | |
| 59 | return data_[indice]; | |
| 60 | } | |
| 61 | const BFloat16* data() const { return data_; } | |
| 62 | }; | |
| 63 | | |
| 64 | #endif | |
| 65 | | |
| 66 | class ITensorStorage{ | |
| 67 | public: | |
| 68 | virtual const std::vector<int64_t>& Shape() const = 0; | |
| 69 | virtual const void* DataRaw() const = 0; | |
| 70 | virtual bool IsInitialized() const = 0; | |
| 71 | virtual void* Initialize(const std::vector<int64_t>& shape, size_t element_size) = 0; | |
f0ef40d0Tang, Cheng2 years ago | 72 | virtual void* Release() = 0; |
beb9fbbaScott McKay2 years ago | 73 | virtual ~ITensorStorage() = default; |
64646279Wenbing Li2 years ago | 74 | }; |
| 75 | | |
| 76 | | |
| 77 | class IAllocator { | |
| 78 | public: | |
| 79 | virtual void* Alloc(size_t size) = 0; | |
| 80 | virtual void Free(void* p) = 0; | |
| 81 | }; | |
| 82 | | |
| 83 | | |
| 84 | class OrtEagerTensorStorage : public ITensorStorage { | |
| 85 | public: | |
| 86 | OrtEagerTensorStorage(const std::vector<int64_t>& shape, | |
| 87 | void* buffer) : buffer_(buffer), shape_(shape){ | |
| 88 | | |
| 89 | } | |
| 90 | | |
| 91 | OrtEagerTensorStorage(IAllocator* allocator) : allocator_(allocator){ | |
| 92 | } | |
| 93 | | |
beb9fbbaScott McKay2 years ago | 94 | ~OrtEagerTensorStorage() override{ |
64646279Wenbing Li2 years ago | 95 | if (allocator_ && buffer_) |
| 96 | allocator_->Free(buffer_); | |
| 97 | } | |
| 98 | | |
| 99 | const std::vector<int64_t>& Shape() const override { | |
| 100 | if (!IsInitialized()) | |
| 101 | ORTX_CXX_API_THROW("Tensor not initialized", ORT_RUNTIME_EXCEPTION); | |
| 102 | return *shape_; | |
| 103 | } | |
| 104 | | |
beb9fbbaScott McKay2 years ago | 105 | bool IsInitialized() const override { |
64646279Wenbing Li2 years ago | 106 | return shape_.has_value(); |
| 107 | } | |
| 108 | | |
| 109 | const void* DataRaw() const override { | |
| 110 | return buffer_; | |
| 111 | } | |
| 112 | | |
| 113 | void* Initialize(const std::vector<int64_t>& shape, size_t element_size) override { | |
| 114 | if (IsInitialized()) | |
| 115 | return buffer_; | |
| 116 | assert(allocator_); | |
| 117 | shape_ = shape; | |
| 118 | int64_t n_elem = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>()); | |
| 119 | auto buffer_size = n_elem * element_size; | |
| 120 | buffer_ = allocator_->Alloc(buffer_size); | |
| 121 | return buffer_; | |
| 122 | } | |
| 123 | | |
f0ef40d0Tang, Cheng2 years ago | 124 | void* Release() override { |
| 125 | void* tmp = buffer_; | |
| 126 | buffer_ = 0; | |
| 127 | shape_ = std::nullopt; | |
| 128 | return tmp; | |
| 129 | } | |
| 130 | | |
64646279Wenbing Li2 years ago | 131 | private: |
| 132 | void* buffer_ {}; | |
| 133 | std::optional<std::vector<int64_t>> shape_; | |
| 134 | // caller need to make sure the allocator is alive | |
f9290e8bWenbing Li2 years ago | 135 | IAllocator* allocator_{}; |
64646279Wenbing Li2 years ago | 136 | }; |
| 137 | | |
| 138 | template <typename TT> | |
| 139 | ONNXTensorElementDataType GetOrtDType(){ | |
| 140 | if constexpr (std::is_same<TT, bool>::value) | |
| 141 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL; | |
| 142 | else if constexpr (std::is_same<TT, float>::value) | |
| 143 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; | |
| 144 | else if constexpr (std::is_same<TT, double>::value) | |
| 145 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE; | |
| 146 | else if constexpr (std::is_same<TT, uint8_t>::value) | |
| 147 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8; | |
| 148 | else if constexpr (std::is_same<TT, int8_t>::value) | |
| 149 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8; | |
| 150 | else if constexpr (std::is_same<TT, uint16_t>::value) | |
| 151 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16; | |
| 152 | else if constexpr (std::is_same<TT, int16_t>::value) | |
| 153 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16; | |
| 154 | else if constexpr (std::is_same<TT, uint32_t>::value) | |
| 155 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32; | |
| 156 | else if constexpr (std::is_same<TT, int32_t>::value) | |
| 157 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; | |
| 158 | else if constexpr (std::is_same<TT, uint64_t>::value) | |
| 159 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64; | |
| 160 | else if constexpr (std::is_same<TT, int64_t>::value) | |
| 161 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; | |
| 162 | else if constexpr (std::is_same<TT, std::string>::value) | |
| 163 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING; | |
| 164 | ORTX_CXX_API_THROW("Unexpected type", ORT_RUNTIME_EXCEPTION); | |
| 165 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; | |
| 166 | } | |
| 167 | | |
| 168 | class TensorBase : public Arg { | |
| 169 | public: | |
beb9fbbaScott McKay2 years ago | 170 | virtual ~TensorBase() = default; |
64646279Wenbing Li2 years ago | 171 | |
| 172 | virtual ONNXTensorElementDataType Type() const = 0; | |
| 173 | virtual const std::vector<int64_t>& Shape() const = 0; | |
| 174 | virtual int64_t NumberOfElement() const = 0; | |
| 175 | virtual const void* DataRaw() const = 0; | |
| 176 | virtual size_t SizeInBytes() const = 0; | |
| 177 | }; | |
| 178 | | |
| 179 | template <typename T> | |
| 180 | class Tensor : public TensorBase { | |
| 181 | public: | |
| 182 | using TT = typename std::remove_reference<T>::type; | |
| 183 | Tensor(std::unique_ptr<ITensorStorage> tensor_storage) : storage_(std::move(tensor_storage)){ | |
| 184 | } | |
| 185 | | |
| 186 | Tensor(const std::vector<int64_t>& shape, void* buffer) : Tensor(std::make_unique<OrtEagerTensorStorage>(shape, buffer)) {} | |
| 187 | | |
| 188 | Tensor(IAllocator* allocator) : storage_(std::make_unique<OrtEagerTensorStorage>(allocator)){} | |
| 189 | | |
f0ef40d0Tang, Cheng2 years ago | 190 | Tensor(const Tensor& src) = delete; |
| 191 | | |
| 192 | Tensor& operator=(Tensor src) = delete; | |
| 193 | | |
| 194 | Tensor(Tensor&& other) : storage_(std::move(other.storage_)) { | |
| 195 | other.storage_ = nullptr; | |
| 196 | other.span_ = {}; | |
| 197 | } | |
| 198 | | |
| 199 | Tensor& operator=(Tensor&& other) | |
| 200 | { | |
| 201 | storage_ = std::move(other.storage_); | |
| 202 | other.span_ = {}; | |
| 203 | return *this; | |
| 204 | } | |
| 205 | | |
64646279Wenbing Li2 years ago | 206 | operator bool() const { |
f0ef40d0Tang, Cheng2 years ago | 207 | return storage_ && storage_->IsInitialized(); |
64646279Wenbing Li2 years ago | 208 | } |
| 209 | | |
| 210 | ONNXTensorElementDataType Type() const override { | |
| 211 | return GetOrtDType<T>(); | |
| 212 | } | |
| 213 | | |
| 214 | const std::vector<int64_t>& Shape() const override { | |
f0ef40d0Tang, Cheng2 years ago | 215 | if (!storage_) |
| 216 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 217 | return storage_->Shape(); |
| 218 | } | |
| 219 | | |
| 220 | int64_t NumberOfElement() const override { | |
f0ef40d0Tang, Cheng2 years ago | 221 | if (!storage_) |
| 222 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 223 | auto& shape = storage_->Shape(); |
| 224 | return std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>()); | |
| 225 | } | |
| 226 | | |
| 227 | std::string Shape2Str() const { | |
f0ef40d0Tang, Cheng2 years ago | 228 | if (!storage_) |
| 229 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
| 230 | if (storage_&& storage_->IsInitialized()) { | |
64646279Wenbing Li2 years ago | 231 | std::string shape_str; |
| 232 | auto& shape = storage_->Shape(); | |
| 233 | for (const auto& dim : shape) { | |
| 234 | shape_str.append(std::to_string(dim)); | |
| 235 | shape_str.append(", "); | |
| 236 | } | |
| 237 | return shape_str; | |
| 238 | } else { | |
| 239 | return "empty"; | |
| 240 | } | |
| 241 | } | |
f0ef40d0Tang, Cheng2 years ago | 242 | |
| 243 | void* Release() { | |
| 244 | if (!storage_) | |
| 245 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
| 246 | span_ = {}; | |
| 247 | return storage_->Release(); | |
| 248 | } | |
64646279Wenbing Li2 years ago | 249 | |
| 250 | const TT* Data() const { | |
f0ef40d0Tang, Cheng2 years ago | 251 | if (!storage_) |
| 252 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 253 | #if ORT_API_VERSION >= 16 |
| 254 | if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value) | |
| 255 | return reinterpret_cast<const TT*>(storage_->DataRaw()); | |
| 256 | else | |
| 257 | #endif | |
| 258 | return static_cast<const TT*>(storage_->DataRaw()); | |
| 259 | } | |
| 260 | | |
| 261 | const void* DataRaw() const override { | |
f0ef40d0Tang, Cheng2 years ago | 262 | if (!storage_) |
| 263 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 264 | return storage_->DataRaw(); |
| 265 | } | |
| 266 | | |
| 267 | size_t SizeInBytes() const override { | |
f0ef40d0Tang, Cheng2 years ago | 268 | if (!storage_) |
| 269 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 270 | return NumberOfElement() * sizeof(TT); |
| 271 | } | |
| 272 | | |
| 273 | TT* Allocate(const std::vector<int64_t>& shape) { | |
f0ef40d0Tang, Cheng2 years ago | 274 | if (!storage_) |
| 275 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 276 | // it should be OK to allocate multiple times |
| 277 | void* buffer = storage_->Initialize(shape, sizeof(TT)); | |
| 278 | #if ORT_API_VERSION >= 16 | |
| 279 | if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value) | |
| 280 | return reinterpret_cast<TT*>(buffer); | |
| 281 | else | |
| 282 | #endif | |
| 283 | return static_cast<TT*>(buffer); | |
| 284 | } | |
| 285 | | |
| 286 | const Span<T>& AsSpan() { | |
f0ef40d0Tang, Cheng2 years ago | 287 | if (!storage_) |
| 288 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 289 | #if ORT_API_VERSION >= 16 |
| 290 | if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value) { | |
| 291 | ORTX_CXX_API_THROW("AsSpan for MFloat16 / BFloat16 not implemented", ORT_RUNTIME_EXCEPTION); | |
| 292 | } | |
| 293 | else{ | |
| 294 | #endif | |
| 295 | auto& shape = storage_->Shape(); | |
| 296 | if (shape.size() != 1) { | |
| 297 | ORTX_CXX_API_THROW("to get a span, shape must be 1-D, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION); | |
| 298 | } | |
| 299 | span_.Assign(Data(), shape[0]); | |
| 300 | return span_; | |
| 301 | #if ORT_API_VERSION >= 16 | |
| 302 | } | |
| 303 | #endif | |
| 304 | } | |
| 305 | | |
| 306 | const T& AsScalar() { | |
f0ef40d0Tang, Cheng2 years ago | 307 | if (!storage_) |
| 308 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 309 | #if ORT_API_VERSION >= 16 |
| 310 | if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value) { | |
| 311 | ORTX_CXX_API_THROW("AsScalar for MFloat16 / BFloat16 not implemented", ORT_RUNTIME_EXCEPTION); | |
| 312 | } | |
| 313 | else{ | |
| 314 | #endif | |
| 315 | auto& shape = storage_->Shape(); | |
| 316 | if ((shape.size() == 1 && shape[0] != 1) || shape.size() > 1) { | |
| 317 | ORTX_CXX_API_THROW("to get a scalar, shape must be {1}, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION); | |
| 318 | } | |
| 319 | return *Data(); | |
| 320 | #if ORT_API_VERSION >= 16 | |
| 321 | } | |
| 322 | #endif | |
| 323 | } | |
| 324 | | |
| 325 | private: | |
| 326 | std::unique_ptr<ITensorStorage> storage_; | |
| 327 | Span<T> span_; | |
| 328 | }; | |
| 329 | | |
| 330 | template<typename T> | |
| 331 | class IStringTensorStorage{ | |
| 332 | public: | |
| 333 | using strings = std::vector<T>; | |
| 334 | virtual const std::vector<int64_t>& Shape() const = 0; | |
| 335 | virtual const void* DataRaw() const = 0; | |
| 336 | virtual const strings& Data() const = 0; | |
| 337 | virtual bool IsInitialized() const = 0; | |
| 338 | virtual void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) = 0; | |
| 339 | virtual void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) = 0; | |
beb9fbbaScott McKay2 years ago | 340 | virtual ~IStringTensorStorage() = default; |
64646279Wenbing Li2 years ago | 341 | }; |
| 342 | | |
| 343 | template<typename T> | |
| 344 | class EagerStringTensorStorage : public IStringTensorStorage<T>{ | |
| 345 | public: | |
| 346 | using strings = std::vector<T>; | |
| 347 | EagerStringTensorStorage(const strings& ss) : input_strings_(ss), shape_(std::vector<int64_t>{static_cast<int64_t>(ss.size())}){} | |
| 348 | | |
| 349 | EagerStringTensorStorage() {} | |
| 350 | | |
| 351 | const std::vector<int64_t>& Shape() const override { | |
| 352 | if (!IsInitialized()) | |
| 353 | ORTX_CXX_API_THROW("Tensor not initialized", ORT_RUNTIME_EXCEPTION); | |
| 354 | return *shape_; | |
| 355 | } | |
| 356 | | |
beb9fbbaScott McKay2 years ago | 357 | const void* DataRaw() const override { |
64646279Wenbing Li2 years ago | 358 | if (input_strings_.size() != 1) { |
| 359 | ORTX_CXX_API_THROW("DataRaw() only applies to string scalar", ORT_RUNTIME_EXCEPTION); | |
| 360 | } | |
| 361 | if constexpr (std::is_same<std::string_view, T>::value) | |
| 362 | return reinterpret_cast<const void*>(input_strings_[0].data()); | |
| 363 | else | |
| 364 | return reinterpret_cast<const void*>(input_strings_[0].c_str()); | |
| 365 | } | |
| 366 | | |
beb9fbbaScott McKay2 years ago | 367 | bool IsInitialized() const override { |
64646279Wenbing Li2 years ago | 368 | return shape_.has_value(); |
| 369 | } | |
| 370 | | |
beb9fbbaScott McKay2 years ago | 371 | void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) override { |
64646279Wenbing Li2 years ago | 372 | if constexpr (std::is_same<std::string_view, T>::value) |
| 373 | ORTX_CXX_API_THROW("Set output for string view tensor is not supported", ORT_RUNTIME_EXCEPTION); | |
| 374 | input_strings_.assign(ss.begin(), ss.end()); | |
| 375 | shape_ = dims; | |
| 376 | } | |
| 377 | | |
| 378 | const strings& Data() const override { | |
| 379 | return input_strings_; | |
| 380 | } | |
| 381 | | |
beb9fbbaScott McKay2 years ago | 382 | void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) override { |
64646279Wenbing Li2 years ago | 383 | if constexpr (std::is_same<std::string_view, T>::value) |
| 384 | ORTX_CXX_API_THROW("Set output for string view tensor is not supported", ORT_RUNTIME_EXCEPTION); | |
| 385 | | |
| 386 | for (const char* s : ss){ | |
| 387 | input_strings_.push_back(s); | |
| 388 | } | |
| 389 | shape_ = dims; | |
| 390 | } | |
| 391 | | |
| 392 | private: | |
| 393 | std::vector<T> input_strings_; | |
| 394 | std::optional<std::vector<int64_t>> shape_; | |
| 395 | }; | |
| 396 | | |
| 397 | template <> | |
| 398 | class Tensor<std::string> : public TensorBase { | |
| 399 | public: | |
| 400 | using strings = std::vector<std::string>; | |
| 401 | | |
| 402 | Tensor(std::unique_ptr<IStringTensorStorage<std::string>> storage) : storage_(std::move(storage)) {} | |
| 403 | | |
| 404 | Tensor(const strings& ss) : storage_(std::make_unique<EagerStringTensorStorage<std::string>>(ss)) {} | |
| 405 | | |
| 406 | Tensor() : storage_(std::make_unique<EagerStringTensorStorage<std::string>>()) {} | |
| 407 | | |
| 408 | ONNXTensorElementDataType Type() const override { | |
| 409 | return GetOrtDType<std::string>(); | |
| 410 | } | |
| 411 | | |
| 412 | const strings& Data() const { | |
| 413 | return storage_->Data(); | |
| 414 | } | |
| 415 | | |
| 416 | const std::vector<int64_t>& Shape() const override { | |
| 417 | return storage_->Shape(); | |
| 418 | } | |
| 419 | | |
| 420 | int64_t NumberOfElement() const override { | |
| 421 | auto& shape = storage_->Shape(); | |
| 422 | return std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>()); | |
| 423 | } | |
| 424 | | |
| 425 | std::string Shape2Str() const { | |
| 426 | if (storage_->IsInitialized()) { | |
| 427 | std::string shape_str; | |
| 428 | auto& shape = storage_->Shape(); | |
| 429 | for (const auto& dim : shape) { | |
| 430 | shape_str.append(std::to_string(dim)); | |
| 431 | shape_str.append(", "); | |
| 432 | } | |
| 433 | return shape_str; | |
| 434 | } else { | |
| 435 | return "empty"; | |
| 436 | } | |
| 437 | } | |
| 438 | | |
| 439 | const void* DataRaw() const override { | |
| 440 | return storage_->DataRaw(); | |
| 441 | } | |
| 442 | | |
| 443 | size_t SizeInBytes() const override { | |
| 444 | auto& ss = storage_->Data(); | |
| 445 | if (ss.size() != 1) { | |
| 446 | ORTX_CXX_API_THROW("SizeInBytes() only applies to string scalar", ORT_RUNTIME_EXCEPTION); | |
| 447 | } | |
| 448 | return ss[0].size(); | |
| 449 | } | |
| 450 | | |
| 451 | void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) { | |
| 452 | storage_->SetStringOutput(ss, dims); | |
| 453 | } | |
| 454 | void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) { | |
| 455 | storage_->SetStringOutput(ss, dims); | |
| 456 | } | |
| 457 | const Span<std::string>& AsSpan() { | |
| 458 | ORTX_CXX_API_THROW("span for TensorT of string not implemented", ORT_RUNTIME_EXCEPTION); | |
| 459 | } | |
| 460 | const std::string& AsScalar() { | |
| 461 | auto& ss = storage_->Data(); | |
| 462 | if (ss.size() != 1) { | |
| 463 | ORTX_CXX_API_THROW("to get a scalar, shape must be {1}, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION); | |
| 464 | } | |
| 465 | return ss[0]; | |
| 466 | } | |
| 467 | | |
| 468 | private: | |
| 469 | std::unique_ptr<IStringTensorStorage<std::string>> storage_; | |
| 470 | }; | |
| 471 | | |
| 472 | | |
| 473 | template <> | |
| 474 | class Tensor<std::string_view> : public TensorBase { | |
| 475 | public: | |
| 476 | using strings = std::vector<std::string_view>; | |
| 477 | | |
| 478 | Tensor(std::unique_ptr<IStringTensorStorage<std::string_view>> storage) : storage_(std::move(storage)) {} | |
| 479 | | |
| 480 | Tensor(const strings& ss) : storage_(std::make_unique<EagerStringTensorStorage<std::string_view>>(ss)) {} | |
| 481 | | |
| 482 | ONNXTensorElementDataType Type() const override { | |
| 483 | return GetOrtDType<std::string_view>(); | |
| 484 | } | |
| 485 | | |
| 486 | const strings& Data() const { | |
| 487 | return storage_->Data(); | |
| 488 | } | |
| 489 | | |
| 490 | const std::vector<int64_t>& Shape() const override { | |
| 491 | return storage_->Shape(); | |
| 492 | } | |
| 493 | | |
| 494 | int64_t NumberOfElement() const override { | |
| 495 | auto& shape = storage_->Shape(); | |
| 496 | return std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>()); | |
| 497 | } | |
| 498 | | |
| 499 | std::string Shape2Str() const { | |
| 500 | if (storage_->IsInitialized()) { | |
| 501 | std::string shape_str; | |
| 502 | auto& shape = storage_->Shape(); | |
| 503 | for (const auto& dim : shape) { | |
| 504 | shape_str.append(std::to_string(dim)); | |
| 505 | shape_str.append(", "); | |
| 506 | } | |
| 507 | return shape_str; | |
| 508 | } else { | |
| 509 | return "empty"; | |
| 510 | } | |
| 511 | } | |
| 512 | | |
| 513 | const void* DataRaw() const override { | |
| 514 | return storage_->DataRaw(); | |
| 515 | } | |
| 516 | | |
| 517 | size_t SizeInBytes() const override { | |
| 518 | auto& ss = storage_->Data(); | |
| 519 | if (ss.size() != 1) { | |
| 520 | ORTX_CXX_API_THROW("SizeInBytes() only applies to string scalar", ORT_RUNTIME_EXCEPTION); | |
| 521 | } | |
| 522 | return ss[0].size(); | |
| 523 | } | |
| 524 | | |
| 525 | void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) { | |
| 526 | storage_->SetStringOutput(ss, dims); | |
| 527 | } | |
| 528 | void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) { | |
| 529 | storage_->SetStringOutput(ss, dims); | |
| 530 | } | |
| 531 | const Span<std::string_view>& AsSpan() { | |
| 532 | ORTX_CXX_API_THROW("span for TensorT of string not implemented", ORT_RUNTIME_EXCEPTION); | |
| 533 | } | |
| 534 | const std::string_view& AsScalar() { | |
| 535 | auto& ss = storage_->Data(); | |
| 536 | if (ss.size() != 1) { | |
| 537 | ORTX_CXX_API_THROW("to get a scalar, shape must be {1}, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION); | |
| 538 | } | |
| 539 | return ss[0]; | |
| 540 | } | |
| 541 | | |
| 542 | private: | |
| 543 | std::unique_ptr<IStringTensorStorage<std::string_view>> storage_; | |
| 544 | }; | |
| 545 | | |
| 546 | | |
| 547 | template<typename ...Args> | |
| 548 | class NamedArgumentDict{ | |
| 549 | public: | |
| 550 | using ValueTuple = std::tuple<Args...>; | |
| 551 | | |
| 552 | NamedArgumentDict(const std::vector<const char*>& keys, const std::tuple<Args...>& args) : entries_(args) { | |
| 553 | for (const char* key : keys){ | |
| 554 | names_.push_back(key); | |
| 555 | } | |
| 556 | } | |
| 557 | | |
| 558 | template<typename T> | |
| 559 | T TryToGetAttributeWithDefault(const char* name, const T& default_value) const { | |
| 560 | return TryToGetAttributeWithDefaultInternal<0>(name, default_value); | |
| 561 | } | |
| 562 | | |
| 563 | private: | |
| 564 | template<size_t I, typename T> | |
| 565 | typename std::enable_if<I == sizeof...(Args), T>::type | |
| 566 | TryToGetAttributeWithDefaultInternal(const char* name, const T& default_value) const { | |
| 567 | return default_value; | |
| 568 | } | |
| 569 | | |
| 570 | template<size_t I, typename T> | |
| 571 | typename std::enable_if<I < sizeof...(Args), T>::type | |
| 572 | TryToGetAttributeWithDefaultInternal(const char* name, const T& default_value) const { | |
| 573 | if (names_[I] == name){ | |
| 574 | if constexpr (std::is_same<std::tuple_element_t<I, ValueTuple>, T>::value) | |
| 575 | return std::get<I>(entries_); | |
| 576 | else | |
| 577 | throw std::runtime_error("name matched but type is not"); | |
| 578 | } | |
| 579 | return TryToGetAttributeWithDefaultInternal<I+1>(name, default_value); | |
| 580 | } | |
| 581 | | |
| 582 | std::vector<std::string> names_; | |
| 583 | std::tuple<Args...> entries_; | |
| 584 | | |
| 585 | }; | |
| 586 | | |
| 587 | } | |
| 588 | } |