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 tensorflow::int64 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
108 Status ToTensor(tensorflow::Tensor* dst) const;
109 Status BitcastFrom(const TensorInterface& from, DataType type,
110 const int64_t* new_dims, int num_new_dims);
111
Tensor()112 tensorflow::Tensor& Tensor() { return tensor_; }
113
114 private:
115 tensorflow::Tensor tensor_;
116 };
117
TensorFromInterface(AbstractTensorInterface * tensor)118 inline Tensor& TensorFromInterface(AbstractTensorInterface* tensor) {
119 return down_cast<TensorInterface*>(tensor)->Tensor();
120 }
121
122 Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst);
123
124 TF_Tensor* TF_TensorFromTensor(const Tensor& src, Status* status);
125 } // namespace tensorflow
126
127 #endif // TENSORFLOW_C_TF_TENSOR_INTERNAL_H_
128