1 /* Copyright 2021 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_COMMON_GPU_MODEL_H_ 17 #define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_GPU_MODEL_H_ 18 19 #include <string> 20 #include <vector> 21 22 #include "absl/container/flat_hash_map.h" 23 #include "tensorflow/lite/delegates/gpu/common/gpu_model_generated.h" 24 #include "tensorflow/lite/delegates/gpu/common/model.h" 25 #include "tensorflow/lite/delegates/gpu/common/model_hints.h" 26 #include "tensorflow/lite/delegates/gpu/common/status.h" 27 #include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h" 28 29 namespace tflite { 30 namespace gpu { 31 32 struct GpuNode { 33 std::unique_ptr<GPUOperation> gpu_operation; 34 std::vector<ValueId> inputs; 35 std::vector<ValueId> outputs; 36 std::string name; 37 38 GpuNode() = default; 39 GpuNode(GpuNode&& node) = default; 40 GpuNode& operator=(GpuNode&& node) = default; 41 GpuNode(const GpuNode&) = delete; 42 GpuNode& operator=(const GpuNode&) = delete; 43 }; 44 45 struct CreateGpuModelInfo { 46 CalculationsPrecision precision; 47 TensorStorageType storage_type; 48 ModelHints hints; 49 50 // User can require specific layout for some tensors. 51 // This will guarantee that tensors with specific ids have exact specified 52 // layout. 53 // Some restrictions apply: 54 // 1) ValueId must be input or output id of GraphFloat32 55 // 2) data_type must be equal to DeduceDataTypeFromPrecision(precision); 56 // for example for precision F16, data_type must be FLOAT16 57 // 3) Layout must be without Batch dimension if tensor.shape.b == 1 58 // Layout must be with Batch dimension if tensor.shape.b != 1 59 // InitFromGraph will fail if gpu can not allocate tensor with requested 60 // tensor descriptor 61 // WARNING: This is an experimental API and subject to change. 62 // IMPORTANT: tensors ids from predefined / external_immutable_tensors / 63 // external_mutable_tensors should not intersect. 64 absl::flat_hash_map<ValueId, TensorDescriptor> predefined; 65 66 // User can provide immutable external tensors for inference context. 67 // Some restrictions apply: 68 // 1) ValueId must be input or output id of GraphFloat32 69 // 2) Provided ptrs must be valid during life of InferenceContext. 70 // 3) data_type must be equal to DeduceDataTypeFromPrecision(precision); 71 // for example for precision F16, data_type must be FLOAT16 72 // 4) Layout must be without Batch dimension if tensor.shape.b == 1 73 // Layout must be with Batch dimension if tensor.shape.b != 1 74 // InitFromGraph will fail if gpu can not allocate tensor with requested 75 // tensor descriptor 76 // WARNING: This is an experimental API and subject to change. 77 // IMPORTANT: tensors ids from predefined / external_immutable_tensors / 78 // external_mutable_tensors should not intersect. 79 absl::flat_hash_map<ValueId, GpuSpatialTensor*> external_immutable_tensors; 80 81 // User can provide mutable external tensors for inference context. 82 // HINT: Highly recommended to use other options if possible, this options 83 // will be with the worst performance. 84 // Some restrictions apply: 85 // 1) ValueId must be input or output id of GraphFloat32 86 // 2) data_type must be equal to DeduceDataTypeFromPrecision(precision); 87 // for example for precision F16, data_type must be FLOAT16 88 // 3) Layout must be without Batch dimension if tensor.shape.b == 1 89 // Layout must be with Batch dimension if tensor.shape.b != 1 90 // InitFromGraph will fail if gpu can not allocate tensor with requested 91 // tensor descriptor 92 // WARNING: This is an experimental API and subject to change. 93 // IMPORTANT: tensors ids from predefined / external_immutable_tensors / 94 // external_mutable_tensors should not intersect. 95 absl::flat_hash_map<ValueId, TensorDescriptor> external_mutable_tensors; 96 }; 97 98 struct GpuModel { 99 std::vector<std::pair<ValueId, ValueId>> input_ids_and_refs; 100 std::vector<std::pair<ValueId, ValueId>> variable_ids_and_refs; 101 std::vector<std::pair<ValueId, ValueId>> output_ids_and_refs; 102 std::vector<GpuNode> nodes; 103 absl::flat_hash_map<ValueId, TensorDescriptor> tensors; 104 absl::flat_hash_map<ValueId, TensorDescriptor> const_tensors; 105 }; 106 107 absl::Status GraphToGpuModel(const GraphFloat32& graph, 108 const CreateGpuModelInfo& create_info, 109 const GpuInfo& gpu_info, GpuModel* gpu_model); 110 111 flatbuffers::Offset<data::GpuModel> Encode( 112 const GpuModel& gpu_model, flatbuffers::FlatBufferBuilder* builder); 113 114 absl::Status Decode(const data::GpuModel* fb_gpu_model, GpuModel* gpu_model); 115 116 // This transformations MUST be applied to graph for correct work of GpuModel 117 // that will be created from graph 118 absl::Status RunGraphTransformsForGpuModel(GraphFloat32* graph); 119 120 } // namespace gpu 121 } // namespace tflite 122 123 #endif // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_GPU_MODEL_H_ 124