Move flex ops whitelist to a library outside toco

This enables its reuse from other places outside toco.

PiperOrigin-RevId: 249776820
This commit is contained in:
A. Unique TensorFlower 2019-05-23 22:34:01 -07:00 committed by TensorFlower Gardener
parent de43b973f8
commit 2b2ea4ca37
5 changed files with 23 additions and 13 deletions

View File

@ -2,7 +2,8 @@
# This is a TF Lite delegate that is powered by TensorFlow's Eager.
#
package(default_visibility = [
"//visibility:private",
"//tensorflow/compiler/mlir/lite:__subpackages__",
"//tensorflow/lite/toco/tflite:__subpackages__",
])
licenses(["notice"]) # Apache 2.0
@ -241,3 +242,13 @@ tf_cc_test(
"@com_google_googletest//:gtest",
],
)
cc_library(
name = "whitelisted_flex_ops_lib",
srcs = [
"whitelisted_flex_ops.cc",
],
hdrs = [
"whitelisted_flex_ops.h",
],
)

View File

@ -12,12 +12,12 @@ 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.
==============================================================================*/
#include "tensorflow/lite/toco/tflite/whitelisted_flex_ops.h"
#include "tensorflow/lite/delegates/flex/whitelisted_flex_ops.h"
#include <set>
namespace toco {
namespace tflite {
namespace flex {
bool IsWhitelistedFlexOp(const std::string& tensorflow_op_name) {
static const std::set<std::string>* whitelisted_flex_ops =
@ -473,5 +473,5 @@ bool IsWhitelistedFlexOp(const std::string& tensorflow_op_name) {
whitelisted_flex_ops->end();
}
} // namespace flex
} // namespace tflite
} // namespace toco

View File

@ -12,13 +12,13 @@ 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_TOCO_TFLITE_WHITELISTED_FLEX_OPS_H_
#define TENSORFLOW_LITE_TOCO_TFLITE_WHITELISTED_FLEX_OPS_H_
#ifndef TENSORFLOW_LITE_DELEGATES_FLEX_WHITELISTED_FLEX_OPS_H_
#define TENSORFLOW_LITE_DELEGATES_FLEX_WHITELISTED_FLEX_OPS_H_
#include <string>
namespace toco {
namespace tflite {
namespace flex {
// Whether the given op has been statically whitelisted for flex export.
//
@ -29,7 +29,7 @@ namespace tflite {
// TODO(b/118389105): Automate generation of the whitelisted flex ops.
bool IsWhitelistedFlexOp(const std::string& tensorflow_op_name);
} // namespace flex
} // namespace tflite
} // namespace toco
#endif // TENSORFLOW_LITE_TOCO_TFLITE_WHITELISTED_FLEX_OPS_H_
#endif // TENSORFLOW_LITE_DELEGATES_FLEX_WHITELISTED_FLEX_OPS_H_

View File

@ -14,8 +14,6 @@ cc_library(
name = "operator",
srcs = [
"operator.cc",
"whitelisted_flex_ops.cc",
"whitelisted_flex_ops.h",
],
hdrs = [
"builtin_operator.h",
@ -31,6 +29,7 @@ cc_library(
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:ptr_util",
"//tensorflow/lite/delegates/flex:whitelisted_flex_ops_lib",
"//tensorflow/lite/schema:schema_fbs",
"//tensorflow/lite/toco:graph_transformations",
"//tensorflow/lite/toco:model",

View File

@ -22,6 +22,7 @@ limitations under the License.
// TODO(ycling): Consider refactoring to extract the LSTM definition out of
// graph_transformation module.
#include "tensorflow/lite/delegates/flex/whitelisted_flex_ops.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/toco/graph_transformations/lstm_utils.h"
#include "tensorflow/lite/toco/model.h"
@ -29,7 +30,6 @@ limitations under the License.
#include "tensorflow/lite/toco/tflite/custom_operator.h"
#include "tensorflow/lite/toco/tflite/simple_operator.h"
#include "tensorflow/lite/toco/tflite/types.h"
#include "tensorflow/lite/toco/tflite/whitelisted_flex_ops.h"
namespace toco {
@ -2646,7 +2646,7 @@ bool ShouldExportAsFlexOp(bool enable_select_tf_ops,
return false;
}
if (!IsWhitelistedFlexOp(tensorflow_op_name)) {
if (!::tflite::flex::IsWhitelistedFlexOp(tensorflow_op_name)) {
LOG(WARNING) << "Op " << tensorflow_op_name
<< " is a valid TensorFlow op but has not been whitelisted for"
" the TensorFlow Lite flex op set.";