chromium/third_party/tflite/src/tensorflow/lite/delegates/utils.h

/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_LITE_DELEGATES_UTILS_H_
#define TENSORFLOW_LITE_DELEGATES_UTILS_H_

// Utility functions and classes for implementing delegates.

#include <functional>
#include <limits>
#include <set>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/util.h"

namespace tflite {
namespace delegates {

// Creates a new Read/Write tensor having the same shape as the original, but
// with a different type. Note that this might void existing references to
// tensors.
TfLiteStatus CreateNewTensorWithDifferentType(TfLiteContext* context,
                                              const int original_tensor_index,
                                              TfLiteType new_type,
                                              TfLiteTensor** new_tensor,
                                              int* new_tensor_index);

// Retrieves the corresponding TfLiteContext of a subgraph given a subgraph
// index and switches to the delegate context for this subgraph. If an invalid
// subgraph index is given, returns kTfLiteError.
// NOTE: This function is expected to be paired with ReleaseSubgraphContext()
// once the delegate preparation is done and/or the delegate context functions
// are no longer needed.
TfLiteStatus AcquireSubgraphContext(const TfLiteContext* context,
                                    int subgraph_index,
                                    TfLiteContext** acquired_context);

// Releases the subgraph context by switching back to the TFLite kernel
// context for this specified subgraph.
// NOTE: This function is expected to be used after AcquireSubgraphContext()
// once the delegate preparation is done and/or the delegate context functions
// are no longer needed.
TfLiteStatus ReleaseSubgraphContext(const TfLiteContext* context,
                                    int subgraph_index);

// Marks the subgraph with the given index as delegation-skippable. Returns
// kTfLiteOk if the given subgraph index is valid and is successfully marked
// as delegation-skippable, and an error status if the subgraph index is
// invalid.
// If a subgraph is delegation-skippable, then the subgraph will be handled by
// a TfLiteDelegate (and that the delegate is supposed to be already aware of
// this state), and therefore, TfLiteInterpreter can skip invoking
// `ModifyGraphWithDelegate` on this subgraph.
// NOTE: This function is expected to be called only when the subgraph that
// `subgraph_index` is pointing to should be skipped by
// interpreter::ModifyGraphWithDelegate (e.g. the subgraph is part of the list
// of callee subgraphs of the same control flow node, and all of those callees
// are supported by the same delegate at once).
//
// For example, this function can be used when the delegate is handling
// control flow ops like while op. E.g. A while op has condition subgraph
// indexed at `i` and body subgraph indexed at `j`. The op can be delegated
// when the following condition satisfied:
//   1. The delegate supports while op
//   2. Both condition subgraph `i` and body subgraph `j` can be fully
//   delegated by the delegate.
// Then if the delegate decides to support the while node along with both body
// and condition subgraphs, it should mark subgraphs `i` and `j` skippable so
// those two subgraphs won't be delegated separately again after being
// absorbed by the parent subgraph.
// WARNING: It is the delegate's responsibility to define when to skip
// subgraph->ModifyGraphWithDelegate, to check any edge cases (i.e. multiple
// references to the subgraph that `subgraph_index` is pointing to), and to mark
// that subgraph as skippable using this function.
// NOTE: Entry point for C node plugin API.
TfLiteStatus MarkSubgraphAsDelegationSkippable(const TfLiteContext* context,
                                               int subgraph_index);

IsNodeSupportedFn;

// A utility class to help model graph parition.
// Note the class *needs* to be used in TfLiteDelegate::Prepare.
class GraphPartitionHelper { … };

// Specialized partitioner for graphs that possibly contain fp16 tensors.
//
// From nodes that accept fp16 inputs, this delegates the following:
// 1. All nodes (except DEQUANTIZE) that are supported with constant fp16 inputs
// by the delegate (in the TFLite graph, these nodes take in dequantized FP32
// outputs).
// 2. All fp16 DEQUANTIZE nodes that have *all* their consumers in the *first*
// delegated partition. This is because TFLite's partitioning algorithm
// greedily puts all such nodes in the first partition.
class FP16GraphPartitionHelper : public GraphPartitionHelper { … };

}  // namespace delegates
}  // namespace tflite

#endif  // TENSORFLOW_LITE_DELEGATES_UTILS_H_