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:
TF_ManagedBuffer(void * data,size_t len,void (* deallocator)(void * data,size_t len,void * arg),void * deallocator_arg,bool owns_memory)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
~TF_ManagedBuffer()50 ~TF_ManagedBuffer() override {
51 (*deallocator_)(data(), len_, deallocator_arg_);
52 }
53
size()54 size_t size() const override { return len_; }
root_buffer()55 TensorBuffer* root_buffer() override { return this; }
FillAllocationDescription(tensorflow::AllocationDescription * proto)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
OwnsMemory()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:
Buffer(const Tensor & tensor)76 static TensorBuffer* Buffer(const Tensor& tensor) { return tensor.buf_; }
MakeTensor(TF_DataType type,const TensorShape & shape,TensorBuffer * 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:
TensorInterface()94 TensorInterface() {}
TensorInterface(tensorflow::Tensor t)95 explicit TensorInterface(tensorflow::Tensor t) : tensor_(std::move(t)) {}
~TensorInterface()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
Tensor()116 tensorflow::Tensor& Tensor() { return tensor_; }
117
118 private:
119 tensorflow::Tensor tensor_;
120 };
121
TensorFromInterface(AbstractTensorInterface * tensor)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 {
operatorTFTensorDeleter135 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