xref: /aosp_15_r20/external/tensorflow/tensorflow/lite/delegates/gpu/metal/metal_spatial_tensor.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948) 
1 /* Copyright 2020 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 #include "tensorflow/lite/delegates/gpu/metal/metal_spatial_tensor.h"
17 
18 #include <cstring>
19 #include <memory>
20 #include <utility>
21 #include <vector>
22 
23 #include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
24 
25 namespace tflite {
26 namespace gpu {
27 namespace metal {
28 namespace {
29 
CreateTextureBuffer(id<MTLBuffer> buffer,uint64_t buffer_offset,const TensorDescriptor & descriptor,id<MTLTexture> * texture)30 absl::Status CreateTextureBuffer(id<MTLBuffer> buffer, uint64_t buffer_offset,
31                                  const TensorDescriptor& descriptor,
32                                  id<MTLTexture>* texture) {
33   std::vector<uint64_t> storage_dims = descriptor.GetStorageDims();
34   if (@available(macOS 10.14, iOS 12.0, tvOS 12.0, *)) {
35     const size_t data_size = storage_dims[0] * descriptor.GetElementSize() *
36                              SizeOf(descriptor.GetDataType());
37     MTLTextureDescriptor* texture_desc = [[MTLTextureDescriptor alloc] init];
38     texture_desc.width = storage_dims[0];
39     texture_desc.pixelFormat =
40         DataTypeToRGBAPixelFormat(descriptor.GetDataType(), false);
41     texture_desc.textureType = MTLTextureTypeTextureBuffer;
42     texture_desc.usage = MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite;
43     texture_desc.storageMode = buffer.storageMode;
44     *texture = [buffer newTextureWithDescriptor:texture_desc
45                                          offset:buffer_offset
46                                     bytesPerRow:data_size];
47     if (!*texture) {
48       return absl::UnknownError("Failed to allocate id<MTLTexture>");
49     }
50   } else {
51     return absl::UnknownError(
52         "TensorStorageType::IMAGE_BUFFER available only in iOS 12/tvOS "
53         "12/macOS 10.14 and higher.");
54   }
55   return absl::OkStatus();
56 }
57 
AllocateTensorMemory(id<MTLDevice> device,const TensorDescriptor & descriptor,id<MTLBuffer> * buffer,id<MTLTexture> * texture)58 absl::Status AllocateTensorMemory(id<MTLDevice> device,
59                                   const TensorDescriptor& descriptor,
60                                   id<MTLBuffer>* buffer,
61                                   id<MTLTexture>* texture) {
62   std::vector<uint64_t> storage_dims = descriptor.GetStorageDims();
63   const void* data_ptr =
64       descriptor.GetData().empty() ? nullptr : descriptor.GetData().data();
65   switch (descriptor.GetStorageType()) {
66     case TensorStorageType::BUFFER:
67     case TensorStorageType::IMAGE_BUFFER: {
68       const size_t data_size = storage_dims[0] * descriptor.GetElementSize() *
69                                SizeOf(descriptor.GetDataType());
70       if (data_ptr) {
71         *buffer = [device newBufferWithBytes:data_ptr
72                                       length:data_size
73                                      options:MTLResourceStorageModeShared];
74       } else {
75         *buffer = [device newBufferWithLength:data_size
76                                       options:MTLResourceStorageModeShared];
77       }
78       if (!*buffer) {
79         return absl::UnknownError("Failed to allocate id<MTLBuffer>");
80       }
81       if (descriptor.GetStorageType() == TensorStorageType::IMAGE_BUFFER) {
82         RETURN_IF_ERROR(CreateTextureBuffer(*buffer, 0, descriptor, texture));
83       }
84       return absl::OkStatus();
85     }
86     case TensorStorageType::TEXTURE_2D: {
87       MTLTextureDescriptor* texture_desc = [MTLTextureDescriptor
88           texture2DDescriptorWithPixelFormat:DataTypeToRGBAPixelFormat(
89                                                  descriptor.GetDataType(),
90                                                  false)
91                                        width:storage_dims[0]
92                                       height:storage_dims[1]
93                                    mipmapped:NO];
94       texture_desc.textureType = MTLTextureType2D;
95       texture_desc.usage =
96           MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite;
97       texture_desc.storageMode = MTLStorageModePrivate;
98 
99       *texture = [device newTextureWithDescriptor:texture_desc];
100       if (!*texture) {
101         return absl::UnknownError("Failed to allocate id<MTLTexture>");
102       }
103       if (data_ptr) {
104         WriteDataToTexture2D(*texture, device, data_ptr);
105       }
106       return absl::OkStatus();
107     }
108     case TensorStorageType::TEXTURE_3D: {
109       MTLTextureDescriptor* texture_desc = [[MTLTextureDescriptor alloc] init];
110       texture_desc.width = storage_dims[0];
111       texture_desc.height = storage_dims[1];
112       texture_desc.depth = storage_dims[2];
113       texture_desc.pixelFormat =
114           DataTypeToRGBAPixelFormat(descriptor.GetDataType(), false);
115       texture_desc.textureType = MTLTextureType3D;
116       texture_desc.usage =
117           MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite;
118       texture_desc.storageMode = MTLStorageModePrivate;
119 
120       *texture = [device newTextureWithDescriptor:texture_desc];
121       if (!*texture) {
122         return absl::UnknownError("Failed to allocate id<MTLTexture>");
123       }
124       if (data_ptr) {
125         WriteDataToTexture3D(*texture, device, data_ptr);
126       }
127       return absl::OkStatus();
128     }
129     case TensorStorageType::TEXTURE_ARRAY: {
130       MTLTextureDescriptor* texture_desc = [[MTLTextureDescriptor alloc] init];
131       texture_desc.width = storage_dims[0];
132       texture_desc.height = storage_dims[1];
133       texture_desc.arrayLength = storage_dims[2];
134       texture_desc.pixelFormat =
135           DataTypeToRGBAPixelFormat(descriptor.GetDataType(), false);
136       texture_desc.textureType = MTLTextureType2DArray;
137       texture_desc.usage =
138           MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite;
139       texture_desc.storageMode = MTLStorageModePrivate;
140 
141       *texture = [device newTextureWithDescriptor:texture_desc];
142       if (!*texture) {
143         return absl::UnknownError("Failed to allocate id<MTLTexture>");
144       }
145       if (data_ptr) {
146         WriteDataToTexture2DArray(*texture, device, data_ptr);
147       }
148       return absl::OkStatus();
149     }
150     case TensorStorageType::SINGLE_TEXTURE_2D:
151     default:
152       return absl::InternalError("Unsupported tensor storage type");
153   }
154 }
155 }  // namespace
156 
MetalSpatialTensor(id<MTLBuffer> buffer,id<MTLTexture> texture,bool memory_owner,bool texture_mem_owner,const TensorDescriptor & descriptor)157 MetalSpatialTensor::MetalSpatialTensor(id<MTLBuffer> buffer,
158                                        id<MTLTexture> texture,
159                                        bool memory_owner,
160                                        bool texture_mem_owner,
161                                        const TensorDescriptor& descriptor)
162     : memory_(buffer),
163       texture_mem_(texture),
164       memory_owner_(memory_owner),
165       texture_mem_owner_(texture_mem_owner),
166       descriptor_(descriptor) {}
167 
MetalSpatialTensor(MetalSpatialTensor && tensor)168 MetalSpatialTensor::MetalSpatialTensor(MetalSpatialTensor&& tensor)
169     : memory_(tensor.memory_),
170       texture_mem_(tensor.texture_mem_),
171       memory_owner_(tensor.memory_owner_),
172       texture_mem_owner_(tensor.texture_mem_owner_),
173       descriptor_(std::move(tensor.descriptor_)),
174       aligned_texture_width_(tensor.aligned_texture_width_),
175       buffer_offset_(tensor.buffer_offset_) {
176   tensor.memory_ = nullptr;
177 }
178 
operator =(MetalSpatialTensor && tensor)179 MetalSpatialTensor& MetalSpatialTensor::operator=(MetalSpatialTensor&& tensor) {
180   if (this != &tensor) {
181     Release();
182     std::swap(memory_, tensor.memory_);
183     std::swap(texture_mem_, tensor.texture_mem_);
184     std::swap(memory_owner_, tensor.memory_owner_);
185     std::swap(texture_mem_owner_, tensor.texture_mem_owner_);
186     descriptor_ = std::move(tensor.descriptor_);
187     std::swap(aligned_texture_width_, tensor.aligned_texture_width_);
188     std::swap(buffer_offset_, tensor.buffer_offset_);
189   }
190   return *this;
191 }
192 
Release()193 void MetalSpatialTensor::Release() {
194   if (memory_owner_ && memory_) {
195     memory_ = nullptr;
196   }
197   if (texture_mem_owner_ && texture_mem_) {
198     texture_mem_ = nullptr;
199   }
200 }
201 
GetGPUResources(const GPUObjectDescriptor * obj_ptr,GPUResourcesWithValue * resources) const202 absl::Status MetalSpatialTensor::GetGPUResources(
203     const GPUObjectDescriptor* obj_ptr,
204     GPUResourcesWithValue* resources) const {
205   const auto* buffer_desc = dynamic_cast<const BufferDescriptor*>(obj_ptr);
206   if (buffer_desc) {
207     if (descriptor_.GetStorageType() != TensorStorageType::BUFFER) {
208       return absl::InvalidArgumentError(
209           "Tensor can be used with BufferDescriptor only wtih "
210           "TensorStorageType::BUFFER.");
211     }
212     resources->buffers.push_back({"buffer", {memory_, buffer_offset_}});
213     return absl::OkStatus();
214   }
215   const auto* tensor_desc = dynamic_cast<const TensorDescriptor*>(obj_ptr);
216   if (!tensor_desc) {
217     return absl::InvalidArgumentError("Expected TensorDescriptor on input.");
218   }
219   tensor_desc->GetGpuResources(descriptor_.GetBHWDCShape(),
220                                &resources->generic);
221 
222   if (descriptor_.GetStorageType() == TensorStorageType::BUFFER) {
223     resources->buffers.push_back({"buffer", {memory_, buffer_offset_}});
224   } else if (descriptor_.GetStorageType() == TensorStorageType::TEXTURE_2D) {
225     if (obj_ptr->GetAccess() == AccessType::WRITE &&
226         tensor_desc->GetUseBufferForWriteOnlyTexture2d()) {
227       resources->AddInt("aligned_texture_width", aligned_texture_width_);
228       resources->buffers.push_back({"buffer", {memory_, buffer_offset_}});
229     } else {
230       resources->images2d.push_back({"image2d", texture_mem_});
231     }
232   } else if (descriptor_.GetStorageType() == TensorStorageType::TEXTURE_3D) {
233     resources->images3d.push_back({"image3d", texture_mem_});
234   } else if (descriptor_.GetStorageType() == TensorStorageType::TEXTURE_ARRAY) {
235     resources->image2d_arrays.push_back({"image2d_array", texture_mem_});
236   } else if (descriptor_.GetStorageType() == TensorStorageType::IMAGE_BUFFER) {
237     if (obj_ptr->GetAccess() == AccessType::WRITE &&
238         tensor_desc->GetUseBufferForWriteOnlyImageBuffer()) {
239       resources->buffers.push_back({"buffer", {memory_, buffer_offset_}});
240     } else {
241       resources->image_buffers.push_back({"image_buffer", texture_mem_});
242     }
243   }
244 
245   return absl::OkStatus();
246 }
247 
CreateFromDescriptor(const TensorDescriptor & desc,id<MTLDevice> device)248 absl::Status MetalSpatialTensor::CreateFromDescriptor(
249     const TensorDescriptor& desc, id<MTLDevice> device) {
250   desc.CopyWithoutData(&descriptor_);
251   memory_owner_ = true;
252   id<MTLBuffer> buffer;
253   id<MTLTexture> texture;
254   RETURN_IF_ERROR(AllocateTensorMemory(device, desc, &buffer, &texture));
255   memory_ = buffer;
256   texture_mem_ = texture;
257   return absl::OkStatus();
258 }
259 
UploadDescriptorData(const TensorDescriptor & desc,id<MTLDevice> device)260 absl::Status MetalSpatialTensor::UploadDescriptorData(
261     const TensorDescriptor& desc, id<MTLDevice> device) {
262   return WriteData(device, desc.GetData().data());
263 }
264 
ToDescriptor(TensorDescriptor * desc,id<MTLDevice> device) const265 absl::Status MetalSpatialTensor::ToDescriptor(TensorDescriptor* desc,
266                                               id<MTLDevice> device) const {
267   *desc = descriptor_;
268   std::vector<uint8_t> data(GetMemorySizeInBytes());
269   RETURN_IF_ERROR(ReadData(device, data.data()));
270   desc->SetData(std::move(data));
271   return absl::OkStatus();
272 }
273 
WriteData(id<MTLDevice> device,const void * ptr)274 absl::Status MetalSpatialTensor::WriteData(id<MTLDevice> device,
275                                            const void* ptr) {
276   switch (descriptor_.GetStorageType()) {
277     case TensorStorageType::BUFFER:
278     case TensorStorageType::IMAGE_BUFFER:
279       std::memcpy(
280           reinterpret_cast<uint8_t*>([memory_ contents]) + buffer_offset_, ptr,
281           GetMemorySizeInBytes());
282       break;
283     case TensorStorageType::TEXTURE_2D:
284       WriteDataToTexture2D(texture_mem_, device, ptr);
285       break;
286     case TensorStorageType::TEXTURE_3D:
287       WriteDataToTexture3D(texture_mem_, device, ptr);
288       break;
289     case TensorStorageType::TEXTURE_ARRAY:
290       WriteDataToTexture2DArray(texture_mem_, device, ptr);
291       break;
292     case TensorStorageType::SINGLE_TEXTURE_2D:
293     default:
294       return absl::InternalError("Unsupported tensor storage type");
295   }
296   return absl::OkStatus();
297 }
298 
ReadData(id<MTLDevice> device,void * ptr) const299 absl::Status MetalSpatialTensor::ReadData(id<MTLDevice> device,
300                                           void* ptr) const {
301   switch (descriptor_.GetStorageType()) {
302     case TensorStorageType::BUFFER:
303     case TensorStorageType::IMAGE_BUFFER:
304       std::memcpy(
305           ptr, reinterpret_cast<uint8_t*>([memory_ contents]) + buffer_offset_,
306           GetMemorySizeInBytes());
307       break;
308     case TensorStorageType::TEXTURE_2D:
309       ReadDataFromTexture2D(texture_mem_, device, ptr);
310       break;
311     case TensorStorageType::TEXTURE_3D:
312       ReadDataFromTexture3D(texture_mem_, device, ptr);
313       break;
314     case TensorStorageType::TEXTURE_ARRAY:
315       ReadDataFromTexture2DArray(texture_mem_, device, ptr);
316       break;
317     case TensorStorageType::SINGLE_TEXTURE_2D:
318     default:
319       return absl::InternalError("Unsupported tensor storage type");
320   }
321   return absl::OkStatus();
322 }
323 
SetBufferHandle(id<MTLBuffer> buffer)324 absl::Status MetalSpatialTensor::SetBufferHandle(id<MTLBuffer> buffer) {
325   if (memory_owner_) {
326     return absl::InvalidArgumentError(
327         "SetBufferHandle can be used only with shared "
328         "Tensors(CreateSharedBufferTensor).");
329   }
330   if (memory_ == buffer) {
331     return absl::OkStatus();
332   }
333   memory_ = buffer;
334   if (descriptor_.GetStorageType() == TensorStorageType::IMAGE_BUFFER) {
335     id<MTLTexture> texture_buffer = nullptr;
336     RETURN_IF_ERROR(
337         CreateTextureBuffer(memory_, 0, descriptor_, &texture_buffer));
338     texture_mem_ = texture_buffer;
339   }
340   return absl::OkStatus();
341 }
342 
GetBufferHandle() const343 id<MTLBuffer> MetalSpatialTensor::GetBufferHandle() const { return memory_; }
344 
CreateTensor(id<MTLDevice> device,const TensorDescriptor & descriptor,MetalSpatialTensor * result)345 absl::Status CreateTensor(id<MTLDevice> device,
346                           const TensorDescriptor& descriptor,
347                           MetalSpatialTensor* result) {
348   id<MTLBuffer> buffer;
349   id<MTLTexture> texture;
350   RETURN_IF_ERROR(AllocateTensorMemory(device, descriptor, &buffer, &texture));
351   *result = MetalSpatialTensor(buffer, texture, true, true, descriptor);
352   return absl::OkStatus();
353 }
354 
CreateTensorSharedBuffer(id<MTLBuffer> buffer,const TensorDescriptor & descriptor,MetalSpatialTensor * result,uint64_t buffer_offset)355 absl::Status CreateTensorSharedBuffer(id<MTLBuffer> buffer,
356                                       const TensorDescriptor& descriptor,
357                                       MetalSpatialTensor* result,
358                                       uint64_t buffer_offset) {
359   id<MTLTexture> texture_buffer = nullptr;
360   if (buffer &&
361       descriptor.GetStorageType() == TensorStorageType::IMAGE_BUFFER) {
362     RETURN_IF_ERROR(CreateTextureBuffer(buffer, buffer_offset, descriptor,
363                                         &texture_buffer));
364   }
365   *result = MetalSpatialTensor(buffer, texture_buffer, false, true, descriptor);
366   result->buffer_offset_ = buffer_offset;
367   return absl::OkStatus();
368 }
369 
CreateTensorSharedImage2DBuffer(id<MTLBuffer> buffer,const TensorDescriptor & descriptor,int row_bytes_alignment,MetalSpatialTensor * result,uint64_t buffer_offset)370 absl::Status CreateTensorSharedImage2DBuffer(id<MTLBuffer> buffer,
371                                              const TensorDescriptor& descriptor,
372                                              int row_bytes_alignment,
373                                              MetalSpatialTensor* result,
374                                              uint64_t buffer_offset) {
375   std::vector<uint64_t> storage_dims = descriptor.GetStorageDims();
376   const int width = storage_dims[0];
377   const int height = storage_dims[1];
378   const int channels = descriptor.GetElementSize();
379   MTLTextureDescriptor* texture_desc = [[MTLTextureDescriptor alloc] init];
380   texture_desc.width = width;
381   texture_desc.height = height;
382   texture_desc.depth = 1;
383   texture_desc.textureType = MTLTextureType2D;
384   texture_desc.arrayLength = 1;
385   texture_desc.mipmapLevelCount = 1;
386   texture_desc.sampleCount = 1;
387   texture_desc.pixelFormat =
388       DataTypeToRGBAPixelFormat(descriptor.GetDataType(), false);
389   texture_desc.usage = MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite;
390   texture_desc.storageMode = buffer.storageMode;
391   const size_t pixel_size = channels * SizeOf(descriptor.GetDataType());
392   const size_t bytes_per_row = width * pixel_size;
393   const size_t bytes_per_row_aligned =
394       AlignByN(bytes_per_row, row_bytes_alignment);
395   id<MTLTexture> texture_buffer =
396       [buffer newTextureWithDescriptor:texture_desc
397                                 offset:buffer_offset
398                            bytesPerRow:bytes_per_row_aligned];
399   if (!texture_buffer) {
400     return absl::UnknownError("Failed to allocate id<MTLTexture>.");
401   }
402   if (bytes_per_row_aligned % pixel_size != 0) {
403     return absl::UnknownError("Alignment mismatch.");
404   }
405   *result = MetalSpatialTensor(buffer, texture_buffer, false, true, descriptor);
406   result->aligned_texture_width_ = bytes_per_row_aligned / pixel_size;
407   result->buffer_offset_ = buffer_offset;
408   return absl::OkStatus();
409 }
410 
GetFastestStorageType(const GpuInfo & gpu_info)411 TensorStorageType GetFastestStorageType(const GpuInfo& gpu_info) {
412   const bool a7_or_a8 =
413       gpu_info.IsApple() && (gpu_info.apple_info.IsA7GenerationGpu() ||
414                              gpu_info.apple_info.IsA8GenerationGpu());
415   if (a7_or_a8) {
416     return TensorStorageType::TEXTURE_2D;
417   } else {
418     return TensorStorageType::BUFFER;
419   }
420 }
421 
422 }  // namespace metal
423 }  // namespace gpu
424 }  // namespace tflite
425