1/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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.
35typedef struct TF_Tensor {
36 tensorflow::AbstractTensorInterface* tensor;
37} TF_Tensor;
38
39class 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
72namespace tensorflow {
73
74class 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.
85void* 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`.
90void deallocate_buffer(void* data, size_t len, void* arg);
91
92class 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
122inline Tensor& TensorFromInterface(AbstractTensorInterface* tensor) {
123 return down_cast<TensorInterface*>(tensor)->Tensor();
124}
125
126Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst);
127
128TF_Tensor* TF_TensorFromTensor(const Tensor& src, Status* status);
129
130TF_Tensor* TF_TensorFromTensorShallow(const Tensor& src, Status* status);
131
132namespace internal {
133
134struct 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.
141using TF_TensorPtr = std::unique_ptr<TF_Tensor, internal::TFTensorDeleter>;
142
143} // namespace tensorflow
144
145#endif // TENSORFLOW_C_TF_TENSOR_INTERNAL_H_
146