xref: /aosp_15_r20/external/tensorflow/tensorflow/lite/delegates/gpu/cl/tensor.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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_LITE_DELEGATES_GPU_CL_TENSOR_H_
17 #define TENSORFLOW_LITE_DELEGATES_GPU_CL_TENSOR_H_
18 
19 #include <cstdint>
20 #include <memory>
21 
22 #include "tensorflow/lite/delegates/gpu/cl/cl_command_queue.h"
23 #include "tensorflow/lite/delegates/gpu/cl/cl_context.h"
24 #include "tensorflow/lite/delegates/gpu/cl/cl_device.h"
25 #include "tensorflow/lite/delegates/gpu/cl/cl_memory.h"
26 #include "tensorflow/lite/delegates/gpu/cl/gpu_object.h"
27 #include "tensorflow/lite/delegates/gpu/cl/util.h"
28 #include "tensorflow/lite/delegates/gpu/common/data_type.h"
29 #include "tensorflow/lite/delegates/gpu/common/shape.h"
30 #include "tensorflow/lite/delegates/gpu/common/status.h"
31 #include "tensorflow/lite/delegates/gpu/common/task/gpu_tensor.h"
32 #include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
33 #include "tensorflow/lite/delegates/gpu/common/tensor.h"
34 #include "tensorflow/lite/delegates/gpu/common/types.h"
35 
36 namespace tflite {
37 namespace gpu {
38 namespace cl {
39 
40 class Tensor : public GPUObject, public GpuSpatialTensor {
41  public:
Tensor()42   Tensor()
43       : memory_(nullptr), image_buffer_memory_(nullptr), memory_owner_(true) {}
44   Tensor(cl_mem memory, bool memory_owner, const TensorDescriptor& descriptor);
45   Tensor(cl_mem memory, bool memory_owner, cl_mem image_buffer_memory,
46          const TensorDescriptor& descriptor);
47 
48   // Move only
49   Tensor(Tensor&& tensor);
50   Tensor& operator=(Tensor&& tensor);
51   Tensor(const Tensor&) = delete;
52   Tensor& operator=(const Tensor&) = delete;
53 
~Tensor()54   ~Tensor() override { Release(); }
55 
56   absl::Status GetGPUResources(const GPUObjectDescriptor* obj_ptr,
57                                GPUResourcesWithValue* resources) const override;
58 
Width()59   int Width() const override { return descriptor_.GetBHWDCShape().w; }
Height()60   int Height() const override { return descriptor_.GetBHWDCShape().h; }
Depth()61   int Depth() const override { return descriptor_.GetBHWDCShape().d; }
Channels()62   int Channels() const override { return descriptor_.GetBHWDCShape().c; }
Slices()63   int Slices() const override {
64     return DivideRoundUp(descriptor_.GetBHWDCShape().c, 4);
65   }
Batch()66   int Batch() const override { return descriptor_.GetBHWDCShape().b; }
67 
GetDescriptor()68   TensorDescriptor GetDescriptor() const override { return descriptor_; }
GetDataType()69   DataType GetDataType() const { return descriptor_.GetDataType(); }
GetStorageType()70   TensorStorageType GetStorageType() const {
71     return descriptor_.GetStorageType();
72   }
GetMemorySizeInBytes()73   uint64_t GetMemorySizeInBytes() const {
74     return descriptor_.GetMemorySizeInBytes();
75   }
76 
77   cl_mem GetMemoryPtr() const;
78 
79   // This function returns buffer memory ptr for IMAGE_BUFFER instead of image
80   // memory ptr.
81   cl_mem GetMemoryPtrForWriting() const;
82 
83   absl::Status CreateFromDescriptor(const TensorDescriptor& desc,
84                                     CLContext* context);
85   absl::Status UploadDescriptorData(const TensorDescriptor& desc,
86                                     CLCommandQueue* queue);
87   absl::Status ToDescriptor(TensorDescriptor* desc,
88                             CLCommandQueue* queue) const;
89 
90  private:
91   friend absl::Status CreateTensorSharedImage2DBuffer(
92       const CLContext& context, cl_mem memory,
93       const TensorDescriptor& descriptor, int width_pixel_alignment,
94       Tensor* result);
95 
96   absl::Status WriteData(const void* ptr, CLCommandQueue* queue);
97   absl::Status ReadData(void* ptr, CLCommandQueue* queue) const;
98 
99   void Release();
100 
101   cl_mem memory_;
102   cl_mem image_buffer_memory_;  // for IMAGE_BUFFER/TEXTURE_2D/SINGLE_TEXTURE_2D
103   bool memory_owner_;
104   bool buffer_based_ = false;
105   TensorDescriptor descriptor_;
106   // for use with TEXTURE_2D and when texture created from buffer.
107   int aligned_texture_width_;
108 };
109 
110 using TensorPtr = std::shared_ptr<Tensor>;
111 
112 absl::Status AllocateTensorMemory(const CLContext& context,
113                                   const TensorDescriptor& descriptor,
114                                   CLMemory* result);
115 
116 absl::Status CreateTensor(const CLContext& context,
117                           const TensorDescriptor& descriptor, Tensor* result);
118 
119 absl::Status CreateTensorShared(const CLContext& context, cl_mem memory,
120                                 const TensorDescriptor& descriptor,
121                                 Tensor* result);
122 
123 absl::Status CreateTensorSharedImage2DBuffer(const CLContext& context,
124                                              cl_mem memory,
125                                              const TensorDescriptor& descriptor,
126                                              int width_pixel_alignment,
127                                              Tensor* result);
128 
129 
130 }  // namespace cl
131 }  // namespace gpu
132 }  // namespace tflite
133 
134 #endif  // TENSORFLOW_LITE_DELEGATES_GPU_CL_TENSOR_H_
135