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/core/framework/allocation_description.pb.h"
24 #include "tensorflow/core/framework/tensor.h"
25 #include "tensorflow/core/framework/tensor_shape.h"
26 #include "tensorflow/core/platform/casts.h"
27
28 // Internal structures used by the C API. These are likely to change and should
29 // not be depended on.
30
31 // This struct forms part of the C API's public interface. It must strictly be
32 // passed to or returned from C functions *by pointer*. Otherwise, changes to
33 // its internal structure will break the C API's binary interface.
34 typedef struct TF_Tensor {
35 tensorflow::AbstractTensorInterface* tensor;
36 } TF_Tensor;
37
38 class TF_ManagedBuffer : public tensorflow::TensorBuffer {
39 public:
TF_ManagedBuffer(void * data,size_t len,void (* deallocator)(void * data,size_t len,void * arg),void * deallocator_arg,bool owns_memory)40 TF_ManagedBuffer(void* data, size_t len,
41 void (*deallocator)(void* data, size_t len, void* arg),
42 void* deallocator_arg, bool owns_memory)
43 : TensorBuffer(data),
44 len_(len),
45 deallocator_(deallocator),
46 deallocator_arg_(deallocator_arg),
47 owns_memory_(owns_memory) {}
48
~TF_ManagedBuffer()49 ~TF_ManagedBuffer() override {
50 (*deallocator_)(data(), len_, deallocator_arg_);
51 }
52
size()53 size_t size() const override { return len_; }
root_buffer()54 TensorBuffer* root_buffer() override { return this; }
FillAllocationDescription(tensorflow::AllocationDescription * proto)55 void FillAllocationDescription(
56 tensorflow::AllocationDescription* proto) const override {
57 int64_t rb = size();
58 proto->set_requested_bytes(rb);
59 proto->set_allocator_name(tensorflow::cpu_allocator()->Name());
60 }
61
OwnsMemory()62 bool OwnsMemory() const override { return owns_memory_; }
63
64 private:
65 const size_t len_;
66 void (*const deallocator_)(void* data, size_t len, void* arg);
67 void* const deallocator_arg_;
68 bool owns_memory_;
69 };
70
71 namespace tensorflow {
72
73 class TensorCApi {
74 public:
Buffer(const Tensor & tensor)75 static TensorBuffer* Buffer(const Tensor& tensor) { return tensor.buf_; }
MakeTensor(TF_DataType type,const TensorShape & shape,TensorBuffer * buf)76 static Tensor MakeTensor(TF_DataType type, const TensorShape& shape,
77 TensorBuffer* buf) {
78 return Tensor(static_cast<DataType>(type), shape, buf);
79 }
80 };
81
82 // Allocates tensor data buffer using specified allocator.
83 // `operation` is a name for this operation.
84 void* allocate_tensor(const char* operation, size_t len, Allocator* allocator);
85
86 // Deallocates tensor data buffer.
87 // Defaults to deallocating using CPU allocator. You can pass pointer to
88 // a different Allocator as `arg`.
89 void deallocate_buffer(void* data, size_t len, void* arg);
90
91 class TensorInterface : public AbstractTensorInterface {
92 public:
TensorInterface()93 TensorInterface() {}
TensorInterface(tensorflow::Tensor t)94 explicit TensorInterface(tensorflow::Tensor t) : tensor_(std::move(t)) {}
~TensorInterface()95 ~TensorInterface() override {}
96
97 void Release() override;
98
99 DataType Type() const override;
100 int NumDims() const override;
101 int64_t Dim(int dim_index) const override;
102 int64_t NumElements() const override;
103 size_t ByteSize() const override;
104 void* Data() const override;
105 bool IsAligned() const override;
106 bool CanMove() const override;
107 std::string SummarizeValue() const override;
108
109 Status ToTensor(tensorflow::Tensor* dst) const;
110 Status BitcastFrom(const TensorInterface& from, DataType type,
111 const int64_t* new_dims, int num_new_dims);
112 Status FromProto(const tensorflow::TensorProto& from);
113
Tensor()114 tensorflow::Tensor& Tensor() { return tensor_; }
115
116 private:
117 tensorflow::Tensor tensor_;
118 };
119
TensorFromInterface(AbstractTensorInterface * tensor)120 inline Tensor& TensorFromInterface(AbstractTensorInterface* tensor) {
121 return down_cast<TensorInterface*>(tensor)->Tensor();
122 }
123
124 Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst);
125
126 TF_Tensor* TF_TensorFromTensor(const Tensor& src, Status* status);
127 } // namespace tensorflow
128
129 #endif // TENSORFLOW_C_TF_TENSOR_INTERNAL_H_
130