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 #ifndef TENSORFLOW_LITE_EXPERIMENTAL_DELEGATES_COREML_COREML_DELEGATE_KERNEL_H_ 16 #define TENSORFLOW_LITE_EXPERIMENTAL_DELEGATES_COREML_COREML_DELEGATE_KERNEL_H_ 17 18 #include "tensorflow/lite/c/common.h" 19 #include "tensorflow/lite/delegates/coreml/builders/op_builder.h" 20 #import "tensorflow/lite/delegates/coreml/coreml_executor.h" 21 22 namespace tflite { 23 namespace delegates { 24 namespace coreml { 25 26 // Represents a subgraph in TFLite that will be delegated to CoreML. 27 // It is abstracted as a single kernel node in the main TFLite graph and 28 // implements Init/Prepare/Invoke as TFLite kernel nodes. 29 class CoreMlDelegateKernel { 30 public: CoreMlDelegateKernel(int coreml_version)31 explicit CoreMlDelegateKernel(int coreml_version) 32 : coreml_version_(coreml_version) {} 33 // Initialize the delegated graph and add required nodes. 34 TfLiteStatus Init(TfLiteContext* context, const TfLiteDelegateParams* params); 35 36 // Any preparation work needed for the delegated graph. 37 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node); 38 39 // Allocates delegated tensordefs for graph I/O & execute it. 40 TfLiteStatus Invoke(TfLiteContext* context, TfLiteNode* node); 41 42 ~CoreMlDelegateKernel(); 43 44 private: 45 // Builds the ML Model protocol buffer 46 TfLiteStatus BuildModel(TfLiteContext* context, 47 const TfLiteDelegateParams* params); 48 49 // Adds the output tensors to the model generated. 50 void AddOutputTensors(const TfLiteIntArray* output_tensors, 51 TfLiteContext* context); 52 53 // Adds the input tensors to the model generated. 54 void AddInputTensors(const TfLiteIntArray* output_tensors, 55 TfLiteContext* context); 56 57 std::unique_ptr<delegates::coreml::GraphBuilder> builder_; 58 std::unique_ptr<CoreML::Specification::Model> model_; 59 ::CoreMlExecutor* executor_; 60 int coreml_version_; 61 62 std::vector<int> input_tensor_ids_; 63 std::vector<TensorData> inputs_; 64 std::vector<TensorData> outputs_; 65 }; 66 67 } // namespace coreml 68 } // namespace delegates 69 } // namespace tflite 70 71 #endif // TENSORFLOW_LITE_EXPERIMENTAL_DELEGATES_COREML_COREML_DELEGATE_KERNEL_H_ 72