• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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