add gpu target to avg_pool

PiperOrigin-RevId: 302578305
Change-Id: I22c469b634feb30c4b094cb8704031c14e827ea1

tensorflow/compiler/mlir/lite/BUILD

@@ -206,6 +206,7 @@ cc_library(
     name = "tensorflow_lite",
     srcs = [
         "experimental/estimators/estimator.h",
+        "experimental/estimators/gpu_estimator.h.inc",
         "ir/tfl_ops.cc",
         "ir/tfl_ops.cc.inc",
         "ir/tfl_ops.h.inc",

tensorflow/compiler/mlir/lite/experimental/estimators/gpu_estimator.h.inc

@@ -0,0 +1,31 @@
+/* Copyright 2020 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_COMPILER_MLIR_LITE_EXPERIMENTAL_ESTIMATORS_GPU_ESTIMATOR_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_ESTIMATORS_GPU_ESTIMATOR_H_
+
+template <>
+class TFLiteCostEstimator<AveragePool2DOp, hardware::GPU> {
+ public:
+  static double GetCost(mlir::Operation* op) {
+    llvm::errs() << "No defined cost function for op: "
+                 << op->getName().getStringRef().str();
+    return 0.0;
+  }
+
+  static bool IsSupported(mlir::Operation* op) { return true; }
+};
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_ESTIMATORS_GPU_ESTIMATOR_H_

tensorflow/compiler/mlir/lite/ir/tfl_ops.h

@@ -53,6 +53,8 @@ class TensorFlowLiteDialect : public Dialect {
 #include "tensorflow/compiler/mlir/lite/ir/tfl_ops_interface.h.inc"
 #define GET_OP_CLASSES
 #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h.inc"
+// Include all specializes estimators below this line
+#include "tensorflow/compiler/mlir/lite/experimental/estimators/gpu_estimator.h.inc"
 
 }  // end namespace TFL
 }  // end namespace mlir

tensorflow/compiler/mlir/lite/ir/tfl_ops.td

@@ -480,7 +480,10 @@ Note this is a custom op that is not supported in the standard runtime.
 }
 
 def TFL_AveragePool2DOp:
-    TFL_Op<"average_pool_2d", [NoSideEffect, SameOperandsAndResultsScale]> {
+    TFL_Op<"average_pool_2d",
+           [NoSideEffect,
+            SameOperandsAndResultsScale,
+            TFL_GpuTargetOp]> {
   let summary = "Average_pool_2d operator";
 
   let description = [{