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:
		
							parent
							
								
									de43b973f8
								
							
						
					
					
						commit
						2b2ea4ca37
					
				| @ -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", | ||||
|     ], | ||||
| ) | ||||
|  | ||||
| @ -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
 | ||||
| @ -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_
 | ||||
| @ -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", | ||||
|  | ||||
| @ -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."; | ||||
|  | ||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user