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 #ifndef TENSORFLOW_LITE_DELEGATES_GPU_METAL_METAL_ARGUMENTS_H_ 16 #define TENSORFLOW_LITE_DELEGATES_GPU_METAL_METAL_ARGUMENTS_H_ 17 18 #import <Metal/Metal.h> 19 20 #include <map> 21 #include <string> 22 #include <vector> 23 24 #include "tensorflow/lite/delegates/gpu/common/status.h" 25 #include "tensorflow/lite/delegates/gpu/common/task/arguments.h" 26 #include "tensorflow/lite/delegates/gpu/common/task/gpu_object_desc.h" 27 #include "tensorflow/lite/delegates/gpu/metal/gpu_object.h" 28 #include "tensorflow/lite/delegates/gpu/metal/metal_device.h" 29 30 namespace tflite { 31 namespace gpu { 32 namespace metal { 33 34 class MetalArguments : public ArgumentsBinder { 35 public: 36 MetalArguments() = default; 37 38 absl::Status Init(const std::map<std::string, std::string>& linkables, 39 MetalDevice* device, Arguments* args, std::string* code); 40 41 // Move only 42 MetalArguments(MetalArguments&& args) = default; 43 MetalArguments& operator=(MetalArguments&& args) = default; 44 MetalArguments(const MetalArguments&) = delete; 45 MetalArguments& operator=(const MetalArguments&) = delete; 46 47 absl::Status SetInt(const std::string& name, int value) override; 48 absl::Status SetFloat(const std::string& name, float value) override; 49 absl::Status SetHalf(const std::string& name, half value) override; 50 absl::Status SetObjectRef(const std::string& name, const GPUObject& object); 51 52 void Encode(id<MTLComputeCommandEncoder> encoder, int buffer_offset, 53 int texture_offset = 0) const; 54 55 private: 56 // creates structure with layout: 57 // struct uniforms_buffer { 58 // int val_0; 59 // int val_1; 60 // float val_2; 61 // int dummy; // for alignment 62 // }; 63 std::string ScalarArgumentsToStructWithScalarFields(Arguments* args, 64 std::string* code); 65 66 // creates structure with layout: 67 // struct uniforms_buffer { 68 // int4 val_0_val_1_dummy_dummy; 69 // float4 val_2_dummy_dummy_dummy; 70 // }; 71 std::string ScalarArgumentsToStructWithVec4Fields(Arguments* args, 72 std::string* code); 73 74 absl::Status AllocateObjects(const Arguments& args, id<MTLDevice> device); 75 absl::Status AddObjectArgs(Arguments* args); 76 77 void AddGPUResources(const std::string& name, const GPUResources& resources, 78 Arguments* args); 79 80 std::string GetListOfArgs(int buffer_offset, int textures_offset = 0); 81 82 absl::Status SetGPUResources(const std::string& name, 83 const GPUResourcesWithValue& resources); 84 85 void AddBuffer(const std::string& name, const GPUBufferDescriptor& desc); 86 void AddImage2D(const std::string& name, const GPUImage2DDescriptor& desc); 87 void AddImage2DArray(const std::string& name, 88 const GPUImage2DArrayDescriptor& desc); 89 void AddImage3D(const std::string& name, const GPUImage3DDescriptor& desc); 90 void AddImageBuffer(const std::string& name, 91 const GPUImageBufferDescriptor& desc); 92 93 absl::Status SetBuffer(const std::string& name, id<MTLBuffer> handle); 94 absl::Status SetImage2D(const std::string& name, id<MTLTexture> handle); 95 absl::Status SetImage2DArray(const std::string& name, id<MTLTexture> handle); 96 absl::Status SetImage3D(const std::string& name, id<MTLTexture> handle); 97 absl::Status SetImageBuffer(const std::string& name, id<MTLTexture> handle); 98 99 absl::Status SetObjectsResources(const Arguments& args); 100 101 absl::Status ResolveSelectorsPass( 102 const GpuInfo& gpu_info, const Arguments& args, 103 const std::map<std::string, std::string>& linkables, std::string* code); 104 105 absl::Status ResolveSelector( 106 const GpuInfo& gpu_info, const Arguments& args, 107 const std::map<std::string, std::string>& linkables, 108 const std::string& object_name, const std::string& selector, 109 const std::vector<std::string>& function_args, 110 const std::vector<std::string>& template_args, std::string* result); 111 112 void ResolveObjectNames(const std::string& object_name, 113 const std::vector<std::string>& member_names, 114 std::string* code); 115 116 void ResolveArgsPass(std::string* code); 117 118 static constexpr char kArgsPrefix[] = "args."; 119 struct IntValue { 120 int value; 121 122 // many arguments generated automatically and not used 123 // to reduce amount of data transferred we adding this optimization 124 bool active = false; 125 126 // offset to shared storage. 127 uint32_t bytes_offset = -1; 128 }; 129 std::map<std::string, IntValue> int_values_; 130 131 struct FloatValue { 132 float value; 133 134 // many arguments generated automatically and not used 135 // to reduce amount of data transferred we adding this optimization 136 bool active = false; 137 138 // offset to shared storage. 139 uint32_t bytes_offset = -1; 140 }; 141 std::map<std::string, FloatValue> float_values_; 142 std::vector<uint8_t> const_data_; 143 144 struct MetalBufferDescriptor { 145 GPUBufferDescriptor desc; 146 id<MTLBuffer> handle; 147 }; 148 struct MetalImage2DDescriptor { 149 GPUImage2DDescriptor desc; 150 id<MTLTexture> handle; 151 }; 152 struct MetalImage2DArrayDescriptor { 153 GPUImage2DArrayDescriptor desc; 154 id<MTLTexture> handle; 155 }; 156 struct MetalImage3DDescriptor { 157 GPUImage3DDescriptor desc; 158 id<MTLTexture> handle; 159 }; 160 struct MetalImageBufferDescriptor { 161 GPUImageBufferDescriptor desc; 162 id<MTLTexture> handle; 163 }; 164 165 std::map<std::string, MetalBufferDescriptor> buffers_; 166 std::map<std::string, MetalImage2DDescriptor> images2d_; 167 std::map<std::string, MetalImage2DArrayDescriptor> image2d_arrays_; 168 std::map<std::string, MetalImage3DDescriptor> images3d_; 169 std::map<std::string, MetalImageBufferDescriptor> image_buffers_; 170 171 std::map<std::string, GPUObjectDescriptorPtr> object_refs_; 172 std::vector<GPUObjectPtr> objects_; 173 }; 174 175 } // namespace metal 176 } // namespace gpu 177 } // namespace tflite 178 179 #endif // TENSORFLOW_LITE_DELEGATES_GPU_METAL_METAL_ARGUMENTS_H_ 180