microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
include/custom_op/tensor_api.h
601lines · 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; | |
cbed8fd5Wenbing Li2 years ago | 177 | |
| 178 | virtual std::byte* AllocateRaw(const std::vector<int64_t>& shape) = 0; | |
64646279Wenbing Li2 years ago | 179 | }; |
| 180 | | |
| 181 | template <typename T> | |
| 182 | class Tensor : public TensorBase { | |
| 183 | public: | |
| 184 | using TT = typename std::remove_reference<T>::type; | |
| 185 | Tensor(std::unique_ptr<ITensorStorage> tensor_storage) : storage_(std::move(tensor_storage)){ | |
| 186 | } | |
| 187 | | |
| 188 | Tensor(const std::vector<int64_t>& shape, void* buffer) : Tensor(std::make_unique<OrtEagerTensorStorage>(shape, buffer)) {} | |
| 189 | | |
| 190 | Tensor(IAllocator* allocator) : storage_(std::make_unique<OrtEagerTensorStorage>(allocator)){} | |
| 191 | | |
f0ef40d0Tang, Cheng2 years ago | 192 | Tensor(const Tensor& src) = delete; |
| 193 | | |
| 194 | Tensor& operator=(Tensor src) = delete; | |
| 195 | | |
| 196 | Tensor(Tensor&& other) : storage_(std::move(other.storage_)) { | |
| 197 | other.storage_ = nullptr; | |
| 198 | other.span_ = {}; | |
| 199 | } | |
| 200 | | |
| 201 | Tensor& operator=(Tensor&& other) | |
| 202 | { | |
| 203 | storage_ = std::move(other.storage_); | |
| 204 | other.span_ = {}; | |
| 205 | return *this; | |
| 206 | } | |
| 207 | | |
64646279Wenbing Li2 years ago | 208 | operator bool() const { |
f0ef40d0Tang, Cheng2 years ago | 209 | return storage_ && storage_->IsInitialized(); |
64646279Wenbing Li2 years ago | 210 | } |
| 211 | | |
| 212 | ONNXTensorElementDataType Type() const override { | |
| 213 | return GetOrtDType<T>(); | |
| 214 | } | |
| 215 | | |
| 216 | const std::vector<int64_t>& Shape() const override { | |
f0ef40d0Tang, Cheng2 years ago | 217 | if (!storage_) |
| 218 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 219 | return storage_->Shape(); |
| 220 | } | |
| 221 | | |
| 222 | int64_t NumberOfElement() const override { | |
f0ef40d0Tang, Cheng2 years ago | 223 | if (!storage_) |
| 224 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 225 | auto& shape = storage_->Shape(); |
| 226 | return std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>()); | |
| 227 | } | |
| 228 | | |
| 229 | std::string Shape2Str() const { | |
f0ef40d0Tang, Cheng2 years ago | 230 | if (!storage_) |
| 231 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
| 232 | if (storage_&& storage_->IsInitialized()) { | |
64646279Wenbing Li2 years ago | 233 | std::string shape_str; |
| 234 | auto& shape = storage_->Shape(); | |
| 235 | for (const auto& dim : shape) { | |
| 236 | shape_str.append(std::to_string(dim)); | |
| 237 | shape_str.append(", "); | |
| 238 | } | |
| 239 | return shape_str; | |
| 240 | } else { | |
| 241 | return "empty"; | |
| 242 | } | |
| 243 | } | |
f0ef40d0Tang, Cheng2 years ago | 244 | |
| 245 | void* Release() { | |
| 246 | if (!storage_) | |
| 247 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
| 248 | span_ = {}; | |
| 249 | return storage_->Release(); | |
| 250 | } | |
64646279Wenbing Li2 years ago | 251 | |
| 252 | const TT* Data() const { | |
f0ef40d0Tang, Cheng2 years ago | 253 | if (!storage_) |
| 254 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 255 | #if ORT_API_VERSION >= 16 |
| 256 | if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value) | |
| 257 | return reinterpret_cast<const TT*>(storage_->DataRaw()); | |
| 258 | else | |
| 259 | #endif | |
| 260 | return static_cast<const TT*>(storage_->DataRaw()); | |
| 261 | } | |
| 262 | | |
| 263 | const void* DataRaw() const override { | |
f0ef40d0Tang, Cheng2 years ago | 264 | if (!storage_) |
| 265 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 266 | return storage_->DataRaw(); |
| 267 | } | |
| 268 | | |
| 269 | size_t SizeInBytes() const override { | |
f0ef40d0Tang, Cheng2 years ago | 270 | if (!storage_) |
| 271 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 272 | return NumberOfElement() * sizeof(TT); |
| 273 | } | |
| 274 | | |
| 275 | TT* Allocate(const std::vector<int64_t>& shape) { | |
f0ef40d0Tang, Cheng2 years ago | 276 | if (!storage_) |
| 277 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 278 | // it should be OK to allocate multiple times |
| 279 | void* buffer = storage_->Initialize(shape, sizeof(TT)); | |
| 280 | #if ORT_API_VERSION >= 16 | |
| 281 | if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value) | |
| 282 | return reinterpret_cast<TT*>(buffer); | |
| 283 | else | |
| 284 | #endif | |
| 285 | return static_cast<TT*>(buffer); | |
| 286 | } | |
| 287 | | |
cbed8fd5Wenbing Li2 years ago | 288 | std::byte* AllocateRaw(const std::vector<int64_t>& shape) override { |
| 289 | return reinterpret_cast<std::byte*>(Allocate(shape)); | |
| 290 | } | |
| 291 | | |
64646279Wenbing Li2 years ago | 292 | const Span<T>& AsSpan() { |
f0ef40d0Tang, Cheng2 years ago | 293 | if (!storage_) |
| 294 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 295 | #if ORT_API_VERSION >= 16 |
| 296 | if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value) { | |
| 297 | ORTX_CXX_API_THROW("AsSpan for MFloat16 / BFloat16 not implemented", ORT_RUNTIME_EXCEPTION); | |
| 298 | } | |
| 299 | else{ | |
| 300 | #endif | |
| 301 | auto& shape = storage_->Shape(); | |
| 302 | if (shape.size() != 1) { | |
| 303 | ORTX_CXX_API_THROW("to get a span, shape must be 1-D, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION); | |
| 304 | } | |
| 305 | span_.Assign(Data(), shape[0]); | |
| 306 | return span_; | |
| 307 | #if ORT_API_VERSION >= 16 | |
| 308 | } | |
| 309 | #endif | |
| 310 | } | |
| 311 | | |
| 312 | const T& AsScalar() { | |
f0ef40d0Tang, Cheng2 years ago | 313 | if (!storage_) |
| 314 | ORTX_CXX_API_THROW("tensor not initialized.", ORT_RUNTIME_EXCEPTION); | |
64646279Wenbing Li2 years ago | 315 | #if ORT_API_VERSION >= 16 |
| 316 | if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value) { | |
| 317 | ORTX_CXX_API_THROW("AsScalar for MFloat16 / BFloat16 not implemented", ORT_RUNTIME_EXCEPTION); | |
| 318 | } | |
| 319 | else{ | |
| 320 | #endif | |
| 321 | auto& shape = storage_->Shape(); | |
| 322 | if ((shape.size() == 1 && shape[0] != 1) || shape.size() > 1) { | |
| 323 | ORTX_CXX_API_THROW("to get a scalar, shape must be {1}, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION); | |
| 324 | } | |
| 325 | return *Data(); | |
| 326 | #if ORT_API_VERSION >= 16 | |
| 327 | } | |
| 328 | #endif | |
| 329 | } | |
| 330 | | |
| 331 | private: | |
| 332 | std::unique_ptr<ITensorStorage> storage_; | |
| 333 | Span<T> span_; | |
| 334 | }; | |
| 335 | | |
| 336 | template<typename T> | |
| 337 | class IStringTensorStorage{ | |
| 338 | public: | |
| 339 | using strings = std::vector<T>; | |
| 340 | virtual const std::vector<int64_t>& Shape() const = 0; | |
| 341 | virtual const void* DataRaw() const = 0; | |
| 342 | virtual const strings& Data() const = 0; | |
| 343 | virtual bool IsInitialized() const = 0; | |
| 344 | virtual void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) = 0; | |
| 345 | virtual void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) = 0; | |
beb9fbbaScott McKay2 years ago | 346 | virtual ~IStringTensorStorage() = default; |
64646279Wenbing Li2 years ago | 347 | }; |
| 348 | | |
| 349 | template<typename T> | |
| 350 | class EagerStringTensorStorage : public IStringTensorStorage<T>{ | |
| 351 | public: | |
| 352 | using strings = std::vector<T>; | |
| 353 | EagerStringTensorStorage(const strings& ss) : input_strings_(ss), shape_(std::vector<int64_t>{static_cast<int64_t>(ss.size())}){} | |
| 354 | | |
| 355 | EagerStringTensorStorage() {} | |
| 356 | | |
| 357 | const std::vector<int64_t>& Shape() const override { | |
| 358 | if (!IsInitialized()) | |
| 359 | ORTX_CXX_API_THROW("Tensor not initialized", ORT_RUNTIME_EXCEPTION); | |
| 360 | return *shape_; | |
| 361 | } | |
| 362 | | |
beb9fbbaScott McKay2 years ago | 363 | const void* DataRaw() const override { |
64646279Wenbing Li2 years ago | 364 | if (input_strings_.size() != 1) { |
| 365 | ORTX_CXX_API_THROW("DataRaw() only applies to string scalar", ORT_RUNTIME_EXCEPTION); | |
| 366 | } | |
| 367 | if constexpr (std::is_same<std::string_view, T>::value) | |
| 368 | return reinterpret_cast<const void*>(input_strings_[0].data()); | |
| 369 | else | |
| 370 | return reinterpret_cast<const void*>(input_strings_[0].c_str()); | |
| 371 | } | |
| 372 | | |
beb9fbbaScott McKay2 years ago | 373 | bool IsInitialized() const override { |
64646279Wenbing Li2 years ago | 374 | return shape_.has_value(); |
| 375 | } | |
| 376 | | |
beb9fbbaScott McKay2 years ago | 377 | void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) override { |
64646279Wenbing Li2 years ago | 378 | if constexpr (std::is_same<std::string_view, T>::value) |
| 379 | ORTX_CXX_API_THROW("Set output for string view tensor is not supported", ORT_RUNTIME_EXCEPTION); | |
| 380 | input_strings_.assign(ss.begin(), ss.end()); | |
| 381 | shape_ = dims; | |
| 382 | } | |
| 383 | | |
| 384 | const strings& Data() const override { | |
| 385 | return input_strings_; | |
| 386 | } | |
| 387 | | |
beb9fbbaScott McKay2 years ago | 388 | void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) override { |
64646279Wenbing Li2 years ago | 389 | if constexpr (std::is_same<std::string_view, T>::value) |
| 390 | ORTX_CXX_API_THROW("Set output for string view tensor is not supported", ORT_RUNTIME_EXCEPTION); | |
| 391 | | |
| 392 | for (const char* s : ss){ | |
| 393 | input_strings_.push_back(s); | |
| 394 | } | |
| 395 | shape_ = dims; | |
| 396 | } | |
| 397 | | |
| 398 | private: | |
| 399 | std::vector<T> input_strings_; | |
| 400 | std::optional<std::vector<int64_t>> shape_; | |
| 401 | }; | |
| 402 | | |
| 403 | template <> | |
| 404 | class Tensor<std::string> : public TensorBase { | |
| 405 | public: | |
| 406 | using strings = std::vector<std::string>; | |
| 407 | | |
| 408 | Tensor(std::unique_ptr<IStringTensorStorage<std::string>> storage) : storage_(std::move(storage)) {} | |
| 409 | | |
| 410 | Tensor(const strings& ss) : storage_(std::make_unique<EagerStringTensorStorage<std::string>>(ss)) {} | |
| 411 | | |
| 412 | Tensor() : storage_(std::make_unique<EagerStringTensorStorage<std::string>>()) {} | |
| 413 | | |
| 414 | ONNXTensorElementDataType Type() const override { | |
| 415 | return GetOrtDType<std::string>(); | |
| 416 | } | |
| 417 | | |
| 418 | const strings& Data() const { | |
| 419 | return storage_->Data(); | |
| 420 | } | |
| 421 | | |
| 422 | const std::vector<int64_t>& Shape() const override { | |
| 423 | return storage_->Shape(); | |
| 424 | } | |
| 425 | | |
| 426 | int64_t NumberOfElement() const override { | |
| 427 | auto& shape = storage_->Shape(); | |
| 428 | return std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>()); | |
| 429 | } | |
| 430 | | |
| 431 | std::string Shape2Str() const { | |
| 432 | if (storage_->IsInitialized()) { | |
| 433 | std::string shape_str; | |
| 434 | auto& shape = storage_->Shape(); | |
| 435 | for (const auto& dim : shape) { | |
| 436 | shape_str.append(std::to_string(dim)); | |
| 437 | shape_str.append(", "); | |
| 438 | } | |
| 439 | return shape_str; | |
| 440 | } else { | |
| 441 | return "empty"; | |
| 442 | } | |
| 443 | } | |
| 444 | | |
| 445 | const void* DataRaw() const override { | |
| 446 | return storage_->DataRaw(); | |
| 447 | } | |
| 448 | | |
| 449 | size_t SizeInBytes() const override { | |
| 450 | auto& ss = storage_->Data(); | |
| 451 | if (ss.size() != 1) { | |
| 452 | ORTX_CXX_API_THROW("SizeInBytes() only applies to string scalar", ORT_RUNTIME_EXCEPTION); | |
| 453 | } | |
| 454 | return ss[0].size(); | |
| 455 | } | |
| 456 | | |
cbed8fd5Wenbing Li2 years ago | 457 | std::byte* AllocateRaw(const std::vector<int64_t>& shape) override { |
| 458 | ORTX_CXX_API_THROW("AllocateRaw() not supported for string tensor", ORT_RUNTIME_EXCEPTION); | |
| 459 | } | |
| 460 | | |
64646279Wenbing Li2 years ago | 461 | void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) { |
| 462 | storage_->SetStringOutput(ss, dims); | |
| 463 | } | |
| 464 | void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) { | |
| 465 | storage_->SetStringOutput(ss, dims); | |
| 466 | } | |
176c1d01Wenbing Li1 years ago | 467 | const Span<std::string>& AsSpan() const { |
64646279Wenbing Li2 years ago | 468 | ORTX_CXX_API_THROW("span for TensorT of string not implemented", ORT_RUNTIME_EXCEPTION); |
| 469 | } | |
176c1d01Wenbing Li1 years ago | 470 | const std::string& AsScalar() const { |
64646279Wenbing Li2 years ago | 471 | auto& ss = storage_->Data(); |
| 472 | if (ss.size() != 1) { | |
| 473 | ORTX_CXX_API_THROW("to get a scalar, shape must be {1}, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION); | |
| 474 | } | |
| 475 | return ss[0]; | |
| 476 | } | |
| 477 | | |
| 478 | private: | |
| 479 | std::unique_ptr<IStringTensorStorage<std::string>> storage_; | |
| 480 | }; | |
| 481 | | |
| 482 | | |
| 483 | template <> | |
| 484 | class Tensor<std::string_view> : public TensorBase { | |
| 485 | public: | |
| 486 | using strings = std::vector<std::string_view>; | |
| 487 | | |
| 488 | Tensor(std::unique_ptr<IStringTensorStorage<std::string_view>> storage) : storage_(std::move(storage)) {} | |
| 489 | | |
| 490 | Tensor(const strings& ss) : storage_(std::make_unique<EagerStringTensorStorage<std::string_view>>(ss)) {} | |
| 491 | | |
| 492 | ONNXTensorElementDataType Type() const override { | |
| 493 | return GetOrtDType<std::string_view>(); | |
| 494 | } | |
| 495 | | |
| 496 | const strings& Data() const { | |
| 497 | return storage_->Data(); | |
| 498 | } | |
| 499 | | |
| 500 | const std::vector<int64_t>& Shape() const override { | |
| 501 | return storage_->Shape(); | |
| 502 | } | |
| 503 | | |
| 504 | int64_t NumberOfElement() const override { | |
| 505 | auto& shape = storage_->Shape(); | |
| 506 | return std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>()); | |
| 507 | } | |
| 508 | | |
| 509 | std::string Shape2Str() const { | |
| 510 | if (storage_->IsInitialized()) { | |
| 511 | std::string shape_str; | |
| 512 | auto& shape = storage_->Shape(); | |
| 513 | for (const auto& dim : shape) { | |
| 514 | shape_str.append(std::to_string(dim)); | |
| 515 | shape_str.append(", "); | |
| 516 | } | |
| 517 | return shape_str; | |
| 518 | } else { | |
| 519 | return "empty"; | |
| 520 | } | |
| 521 | } | |
| 522 | | |
| 523 | const void* DataRaw() const override { | |
| 524 | return storage_->DataRaw(); | |
| 525 | } | |
| 526 | | |
| 527 | size_t SizeInBytes() const override { | |
| 528 | auto& ss = storage_->Data(); | |
| 529 | if (ss.size() != 1) { | |
| 530 | ORTX_CXX_API_THROW("SizeInBytes() only applies to string scalar", ORT_RUNTIME_EXCEPTION); | |
| 531 | } | |
| 532 | return ss[0].size(); | |
| 533 | } | |
| 534 | | |
cbed8fd5Wenbing Li2 years ago | 535 | std::byte* AllocateRaw(const std::vector<int64_t>& shape) override { |
| 536 | ORTX_CXX_API_THROW("AllocateRaw() not supported for string tensor", ORT_RUNTIME_EXCEPTION); | |
| 537 | } | |
| 538 | | |
64646279Wenbing Li2 years ago | 539 | void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) { |
| 540 | storage_->SetStringOutput(ss, dims); | |
| 541 | } | |
| 542 | void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) { | |
| 543 | storage_->SetStringOutput(ss, dims); | |
| 544 | } | |
176c1d01Wenbing Li1 years ago | 545 | const Span<std::string_view>& AsSpan() const { |
64646279Wenbing Li2 years ago | 546 | ORTX_CXX_API_THROW("span for TensorT of string not implemented", ORT_RUNTIME_EXCEPTION); |
| 547 | } | |
176c1d01Wenbing Li1 years ago | 548 | const std::string_view& AsScalar() const { |
64646279Wenbing Li2 years ago | 549 | auto& ss = storage_->Data(); |
| 550 | if (ss.size() != 1) { | |
| 551 | ORTX_CXX_API_THROW("to get a scalar, shape must be {1}, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION); | |
| 552 | } | |
| 553 | return ss[0]; | |
| 554 | } | |
| 555 | | |
| 556 | private: | |
| 557 | std::unique_ptr<IStringTensorStorage<std::string_view>> storage_; | |
| 558 | }; | |
| 559 | | |
| 560 | template<typename ...Args> | |
| 561 | class NamedArgumentDict{ | |
| 562 | public: | |
| 563 | using ValueTuple = std::tuple<Args...>; | |
| 564 | | |
| 565 | NamedArgumentDict(const std::vector<const char*>& keys, const std::tuple<Args...>& args) : entries_(args) { | |
| 566 | for (const char* key : keys){ | |
| 567 | names_.push_back(key); | |
| 568 | } | |
| 569 | } | |
| 570 | | |
| 571 | template<typename T> | |
| 572 | T TryToGetAttributeWithDefault(const char* name, const T& default_value) const { | |
| 573 | return TryToGetAttributeWithDefaultInternal<0>(name, default_value); | |
| 574 | } | |
| 575 | | |
| 576 | private: | |
| 577 | template<size_t I, typename T> | |
| 578 | typename std::enable_if<I == sizeof...(Args), T>::type | |
| 579 | TryToGetAttributeWithDefaultInternal(const char* name, const T& default_value) const { | |
| 580 | return default_value; | |
| 581 | } | |
| 582 | | |
| 583 | template<size_t I, typename T> | |
| 584 | typename std::enable_if<I < sizeof...(Args), T>::type | |
| 585 | TryToGetAttributeWithDefaultInternal(const char* name, const T& default_value) const { | |
| 586 | if (names_[I] == name){ | |
| 587 | if constexpr (std::is_same<std::tuple_element_t<I, ValueTuple>, T>::value) | |
| 588 | return std::get<I>(entries_); | |
| 589 | else | |
| 590 | throw std::runtime_error("name matched but type is not"); | |
| 591 | } | |
| 592 | return TryToGetAttributeWithDefaultInternal<I+1>(name, default_value); | |
| 593 | } | |
| 594 | | |
| 595 | std::vector<std::string> names_; | |
| 596 | std::tuple<Args...> entries_; | |
| 597 | | |
| 598 | }; | |
| 599 | | |
| 600 | } | |
| 601 | } |