1 | /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. |
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | |
16 | #ifndef TENSORFLOW_C_TF_TENSOR_INTERNAL_H_ |
17 | #define TENSORFLOW_C_TF_TENSOR_INTERNAL_H_ |
18 | |
19 | #include <memory> |
20 | |
21 | #include "tensorflow/c/tensor_interface.h" |
22 | #include "tensorflow/c/tf_datatype.h" |
23 | #include "tensorflow/c/tf_tensor.h" |
24 | #include "tensorflow/core/framework/allocation_description.pb.h" |
25 | #include "tensorflow/core/framework/tensor.h" |
26 | #include "tensorflow/core/framework/tensor_shape.h" |
27 | #include "tensorflow/core/platform/casts.h" |
28 | |
29 | // Internal structures used by the C API. These are likely to change and should |
30 | // not be depended on. |
31 | |
32 | // This struct forms part of the C API's public interface. It must strictly be |
33 | // passed to or returned from C functions *by pointer*. Otherwise, changes to |
34 | // its internal structure will break the C API's binary interface. |
35 | typedef struct TF_Tensor { |
36 | tensorflow::AbstractTensorInterface* tensor; |
37 | } TF_Tensor; |
38 | |
39 | class TF_ManagedBuffer : public tensorflow::TensorBuffer { |
40 | public: |
41 | TF_ManagedBuffer(void* data, size_t len, |
42 | void (*deallocator)(void* data, size_t len, void* arg), |
43 | void* deallocator_arg, bool owns_memory) |
44 | : TensorBuffer(data), |
45 | len_(len), |
46 | deallocator_(deallocator), |
47 | deallocator_arg_(deallocator_arg), |
48 | owns_memory_(owns_memory) {} |
49 | |
50 | ~TF_ManagedBuffer() override { |
51 | (*deallocator_)(data(), len_, deallocator_arg_); |
52 | } |
53 | |
54 | size_t size() const override { return len_; } |
55 | TensorBuffer* root_buffer() override { return this; } |
56 | void FillAllocationDescription( |
57 | tensorflow::AllocationDescription* proto) const override { |
58 | int64_t rb = size(); |
59 | proto->set_requested_bytes(rb); |
60 | proto->set_allocator_name(tensorflow::cpu_allocator()->Name()); |
61 | } |
62 | |
63 | bool OwnsMemory() const override { return owns_memory_; } |
64 | |
65 | private: |
66 | const size_t len_; |
67 | void (*const deallocator_)(void* data, size_t len, void* arg); |
68 | void* const deallocator_arg_; |
69 | bool owns_memory_; |
70 | }; |
71 | |
72 | namespace tensorflow { |
73 | |
74 | class TensorCApi { |
75 | public: |
76 | static TensorBuffer* Buffer(const Tensor& tensor) { return tensor.buf_; } |
77 | static Tensor MakeTensor(TF_DataType type, const TensorShape& shape, |
78 | TensorBuffer* buf) { |
79 | return Tensor(static_cast<DataType>(type), shape, buf); |
80 | } |
81 | }; |
82 | |
83 | // Allocates tensor data buffer using specified allocator. |
84 | // `operation` is a name for this operation. |
85 | void* allocate_tensor(const char* operation, size_t len, Allocator* allocator); |
86 | |
87 | // Deallocates tensor data buffer. |
88 | // Defaults to deallocating using CPU allocator. You can pass pointer to |
89 | // a different Allocator as `arg`. |
90 | void deallocate_buffer(void* data, size_t len, void* arg); |
91 | |
92 | class TensorInterface : public AbstractTensorInterface { |
93 | public: |
94 | TensorInterface() {} |
95 | explicit TensorInterface(tensorflow::Tensor t) : tensor_(std::move(t)) {} |
96 | ~TensorInterface() override {} |
97 | |
98 | void Release() override; |
99 | |
100 | DataType Type() const override; |
101 | int NumDims() const override; |
102 | int64_t Dim(int dim_index) const override; |
103 | int64_t NumElements() const override; |
104 | size_t ByteSize() const override; |
105 | void* Data() const override; |
106 | bool IsAligned() const override; |
107 | bool CanMove() const override; |
108 | std::string SummarizeValue() const override; |
109 | |
110 | void SetShape(const int64_t* dims, int num_dims); |
111 | Status ToTensor(tensorflow::Tensor* dst) const; |
112 | Status BitcastFrom(const TensorInterface& from, DataType type, |
113 | const int64_t* new_dims, int num_new_dims); |
114 | Status FromProto(const tensorflow::TensorProto& from); |
115 | |
116 | tensorflow::Tensor& Tensor() { return tensor_; } |
117 | |
118 | private: |
119 | tensorflow::Tensor tensor_; |
120 | }; |
121 | |
122 | inline Tensor& TensorFromInterface(AbstractTensorInterface* tensor) { |
123 | return down_cast<TensorInterface*>(tensor)->Tensor(); |
124 | } |
125 | |
126 | Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst); |
127 | |
128 | TF_Tensor* TF_TensorFromTensor(const Tensor& src, Status* status); |
129 | |
130 | TF_Tensor* TF_TensorFromTensorShallow(const Tensor& src, Status* status); |
131 | |
132 | namespace internal { |
133 | |
134 | struct TFTensorDeleter { |
135 | void operator()(TF_Tensor* tf_tensor) const { TF_DeleteTensor(tf_tensor); } |
136 | }; |
137 | |
138 | } // namespace internal |
139 | |
140 | // Struct that wraps TF_Tensor to delete once out of scope. |
141 | using TF_TensorPtr = std::unique_ptr<TF_Tensor, internal::TFTensorDeleter>; |
142 | |
143 | } // namespace tensorflow |
144 | |
145 | #endif // TENSORFLOW_C_TF_TENSOR_INTERNAL_H_ |
146 | |