[tf.data] Cleaning up mobile build target.

PiperOrigin-RevId: 333597626
Change-Id: I63be2d4b736d7d1423c93999a0ac2184257eadb4

tensorflow/core/kernels/data/BUILD

@@ -173,6 +173,7 @@ tf_kernel_library(
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:dataset_ops_op_lib",
         "//tensorflow/core:framework",
+        "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/grappler:graph_topology_view",
         "//tensorflow/core/grappler/utils:traversal",
@@ -540,6 +541,7 @@ tf_cc_test(
 tf_kernel_library(
     name = "model_dataset_op",
     srcs = ["model_dataset_op.cc"],
+    hdrs = ["model_dataset_op.h"],
     deps = [
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:dataset_ops_op_lib",
@@ -1444,15 +1446,11 @@ filegroup(
             "*.h",
         ],
         exclude = [
-            "dataset_ops*",  # includes grappler dependency, which isn't supported on mobile.
+            "dataset_test_base.*",
-            "optimize_dataset_op.*",  # includes grappler dependency, which isn't supported on mobile.
-            "model_dataset_op.*",  # not supported on mobile.
-            "rewrite_utils*",  # includes grappler dependency, which isn't supported on mobile.
             "*test.cc",
-            "*test.h",
-            "*_test_*",
         ],
     ),
+    visibility = ["//tensorflow:__subpackages__"],
 )
 
 tf_kernel_library(

tensorflow/core/kernels/data/dataset_ops.cc

@@ -12,9 +12,11 @@ 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/core/kernels/data/dataset_ops.h"
 
+// On mobile we do not provide this functionality because not all of its
+// dependencies are available there.
+#if !defined(IS_MOBILE_PLATFORM)
 #include "tensorflow/core/common_runtime/graph_constructor.h"
 #include "tensorflow/core/common_runtime/graph_runner.h"
 #include "tensorflow/core/common_runtime/process_function_library_runtime.h"
@@ -168,3 +170,4 @@ REGISTER_KERNEL_BUILDER(Name("DatasetFromGraph").Device(DEVICE_CPU),
 
 }  // namespace data
 }  // namespace tensorflow
+#endif  // !IS_MOBILE_PLATFORM

tensorflow/core/kernels/data/dataset_ops.h

@@ -12,10 +12,14 @@ 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_CORE_KERNELS_DATA_DATASET_OPS_H_
 #define TENSORFLOW_CORE_KERNELS_DATA_DATASET_OPS_H_
 
+#include "tensorflow/core/platform/platform.h"
+
+// On mobile we do not provide this functionality because not all of its
+// dependencies are available there.
+#if !defined(IS_MOBILE_PLATFORM)
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/op_kernel.h"
 
@@ -61,5 +65,6 @@ class DatasetFromGraphOp : public OpKernel {
 
 }  // namespace data
 }  // namespace tensorflow
+#endif  // !IS_MOBILE_PLATFORM
 
 #endif  // TENSORFLOW_CORE_KERNELS_DATA_DATASET_OPS_H_

tensorflow/core/kernels/data/model_dataset_op.cc

@@ -12,7 +12,11 @@ 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/core/kernels/data/model_dataset_op.h"
 
+// On mobile we do not provide model dataset op because not all of its
+// dependencies are available there. The op is replaced with a no-op.
+#if !defined(IS_MOBILE_PLATFORM)
 #include "absl/memory/memory.h"
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/metrics.h"
@@ -32,42 +36,13 @@ constexpr int64 kOptimizationPeriodThresholdMs = 60 * EnvTime::kSecondsToMillis;
 // Default share of available RAM that can be used by model's internal buffers.
 constexpr double kRamBudgetShare = 0.5;
 
-class ModelDatasetOp : public UnaryDatasetOpKernel {
+}  // namespace
- public:
-  static constexpr const char* const kAlgorithm = "algorithm";
-  static constexpr const char* const kCpuBudget = "cpu_budget";
-  static constexpr const char* const kRamBudget = "ram_budget";
 
-  explicit ModelDatasetOp(OpKernelConstruction* ctx)
+/* static */ constexpr const char* const ModelDatasetOp::kAlgorithm;
-      : UnaryDatasetOpKernel(ctx) {
+/* static */ constexpr const char* const ModelDatasetOp::kCpuBudget;
-    if (ctx->HasAttr(kAlgorithm)) {
+/* static */ constexpr const char* const ModelDatasetOp::kRamBudget;
-      int64 algorithm;
-      OP_REQUIRES_OK(ctx, ctx->GetAttr(kAlgorithm, &algorithm));
-      algorithm_ = model::AutotuneAlgorithm(algorithm);
-    } else {
-      algorithm_ = model::AutotuneAlgorithm::HILL_CLIMB;
-    }
-    OP_REQUIRES_OK(ctx, ctx->GetAttr(kCpuBudget, &cpu_budget_));
-    OP_REQUIRES(ctx, cpu_budget_ >= 0,
-                errors::InvalidArgument("CPU budget must be positive but is ",
-                                        cpu_budget_, "."));
-    if (ctx->HasAttr(kRamBudget)) {
-      OP_REQUIRES_OK(ctx, ctx->GetAttr(kRamBudget, &ram_budget_));
-    } else {
-      ram_budget_ = 0;
-    }
-    OP_REQUIRES(ctx, ram_budget_ >= 0,
-                errors::InvalidArgument("RAM budget must be positive but is ",
-                                        ram_budget_, "."));
-  }
 
-  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
+class ModelDatasetOp::Dataset : public DatasetBase {
-                   DatasetBase** output) override {
-    *output = new Dataset(ctx, input, algorithm_, cpu_budget_, ram_budget_);
-  }
-
- private:
-  class Dataset : public DatasetBase {
  public:
   Dataset(OpKernelContext* ctx, const DatasetBase* input,
           model::AutotuneAlgorithm algorithm, int64 cpu_budget,
@@ -99,8 +74,7 @@ class ModelDatasetOp : public UnaryDatasetOpKernel {
 
   int64 Cardinality() const override { return input_->Cardinality(); }
 
-    Status InputDatasets(
+  Status InputDatasets(std::vector<const DatasetBase*>* inputs) const override {
-        std::vector<const DatasetBase*>* inputs) const override {
     inputs->push_back(input_);
     return Status::OK();
   }
@@ -156,8 +130,8 @@ class ModelDatasetOp : public UnaryDatasetOpKernel {
     Status Initialize(IteratorContext* ctx) override {
       IteratorContext::Params params(ctx);
       params.model = model_;
-        return dataset()->input_->MakeIterator(
+      return dataset()->input_->MakeIterator(IteratorContext(std::move(params)),
-            IteratorContext(std::move(params)), this, prefix(), &input_impl_);
+                                             this, prefix(), &input_impl_);
     }
 
     Status GetNextInternal(IteratorContext* ctx,
@@ -219,11 +193,10 @@ class ModelDatasetOp : public UnaryDatasetOpKernel {
       while (true) {
         {
           mutex_lock l(mu_);
-            while (!cancelled_ &&
+          while (!cancelled_ && last_optimization_ms + optimization_period_ms >
-                   last_optimization_ms + optimization_period_ms >
                                     current_time_ms) {
-              auto wait_ms = last_optimization_ms + optimization_period_ms -
+            auto wait_ms =
-                             current_time_ms;
+                last_optimization_ms + optimization_period_ms - current_time_ms;
             VLOG(2) << "Waiting for " << wait_ms << " ms.";
             cond_var_.wait_for(l, std::chrono::milliseconds(wait_ms));
             current_time_ms = EnvTime::NowMicros() / EnvTime::kMillisToMicros;
@@ -288,13 +261,58 @@ class ModelDatasetOp : public UnaryDatasetOpKernel {
   const int64 ram_budget_;
 };
 
-  model::AutotuneAlgorithm algorithm_;
+ModelDatasetOp::ModelDatasetOp(OpKernelConstruction* ctx)
-  int64 cpu_budget_;
+    : UnaryDatasetOpKernel(ctx) {
-  int64 ram_budget_;
+  if (ctx->HasAttr(kAlgorithm)) {
-};
+    int64 algorithm;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr(kAlgorithm, &algorithm));
+    algorithm_ = model::AutotuneAlgorithm(algorithm);
+  } else {
+    algorithm_ = model::AutotuneAlgorithm::HILL_CLIMB;
+  }
+  OP_REQUIRES_OK(ctx, ctx->GetAttr(kCpuBudget, &cpu_budget_));
+  OP_REQUIRES(ctx, cpu_budget_ >= 0,
+              errors::InvalidArgument("CPU budget must be positive but is ",
+                                      cpu_budget_, "."));
+  if (ctx->HasAttr(kRamBudget)) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr(kRamBudget, &ram_budget_));
+  } else {
+    ram_budget_ = 0;
+  }
+  OP_REQUIRES(ctx, ram_budget_ >= 0,
+              errors::InvalidArgument("RAM budget must be positive but is ",
+                                      ram_budget_, "."));
+}
 
+void ModelDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
+                                 DatasetBase** output) {
+  *output = new ModelDatasetOp::Dataset(ctx, input, algorithm_, cpu_budget_,
+                                        ram_budget_);
+}
+
+namespace {
 REGISTER_KERNEL_BUILDER(Name("ModelDataset").Device(DEVICE_CPU),
                         ModelDatasetOp);
 }  // namespace
 }  // namespace data
 }  // namespace tensorflow
+#else  // !IS_MOBILE_PLATFORM
+namespace tensorflow {
+namespace data {
+
+ModelDatasetOp::ModelDatasetOp(OpKernelConstruction* ctx)
+    : UnaryDatasetOpKernel(ctx) {}
+
+void ModelDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
+                                 DatasetBase** output) {
+  input->Ref();
+  *output = input;
+}
+
+namespace {
+REGISTER_KERNEL_BUILDER(Name("ModelDataset").Device(DEVICE_CPU),
+                        ModelDatasetOp);
+}  // namespace
+}  // namespace data
+}  // namespace tensorflow
+#endif  // !IS_MOBILE_PLATFORM

tensorflow/core/kernels/data/model_dataset_op.h

@@ -0,0 +1,70 @@
+/* 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_CORE_KERNELS_DATA_MODEL_DATASET_OP_H_
+#define TENSORFLOW_CORE_KERNELS_DATA_MODEL_DATASET_OP_H_
+
+#include "tensorflow/core/platform/platform.h"
+
+// On mobile we do not provide model dataset op because not all of its
+// dependencies are available there. The op is replaced with a no-op.
+#if !defined(IS_MOBILE_PLATFORM)
+#include "tensorflow/core/framework/dataset.h"
+#include "tensorflow/core/framework/model.h"
+
+namespace tensorflow {
+namespace data {
+
+class ModelDatasetOp : public UnaryDatasetOpKernel {
+ public:
+  static constexpr const char* const kAlgorithm = "algorithm";
+  static constexpr const char* const kCpuBudget = "cpu_budget";
+  static constexpr const char* const kRamBudget = "ram_budget";
+
+  explicit ModelDatasetOp(OpKernelConstruction* ctx);
+
+ protected:
+  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
+                   DatasetBase** output) override;
+
+ private:
+  class Dataset;
+
+  model::AutotuneAlgorithm algorithm_;
+  int64 cpu_budget_;
+  int64 ram_budget_;
+};
+
+}  // namespace data
+}  // namespace tensorflow
+#else  // !IS_MOBILE_PLATFORM
+#include "tensorflow/core/framework/dataset.h"
+
+namespace tensorflow {
+namespace data {
+
+class ModelDatasetOp : public UnaryDatasetOpKernel {
+ public:
+  explicit ModelDatasetOp(OpKernelConstruction* ctx);
+
+ protected:
+  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
+                   DatasetBase** output) override;
+};
+
+}  // namespace data
+}  // namespace tensorflow
+#endif  // !IS_MOBILE_PLATFORM
+
+#endif  // TENSORFLOW_CORE_KERNELS_DATA_MODEL_DATASET_OP_H_

tensorflow/core/kernels/data/model_dataset_op_mobile.cc

@@ -1,37 +0,0 @@
-/* Copyright 2018 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.
-==============================================================================*/
-#include "tensorflow/core/framework/dataset.h"
-
-namespace tensorflow {
-namespace data {
-
-class ModelDatasetOp : public UnaryDatasetOpKernel {
- public:
-  explicit ModelDatasetOp(OpKernelConstruction* ctx)
-      : UnaryDatasetOpKernel(ctx) {}
-
-  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
-                   DatasetBase** output) {
-    input->Ref();
-    *output = input;
-  }
-};
-
-namespace {
-REGISTER_KERNEL_BUILDER(Name("ModelDataset").Device(DEVICE_CPU),
-                        ModelDatasetOp);
-}  // namespace
-}  // namespace data
-}  // namespace tensorflow

tensorflow/core/kernels/data/optimize_dataset_op.cc

@@ -14,6 +14,9 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/core/kernels/data/optimize_dataset_op.h"
 
+// On mobile we do not provide optimize dataset op because not all of its
+// dependencies are available there. The op is replaced with a no-op.
+#if !defined(IS_MOBILE_PLATFORM)
 #include <map>
 
 #include "tensorflow/core/framework/partial_tensor_shape.h"
@@ -27,9 +30,6 @@ limitations under the License.
 namespace tensorflow {
 namespace data {
 
-// See documentation in ../../ops/dataset_ops.cc for a high-level
-// description of the following op.
-
 /* static */ constexpr const char* const OptimizeDatasetOp::kDatasetType;
 /* static */ constexpr const char* const OptimizeDatasetOp::kInputDataset;
 /* static */ constexpr const char* const OptimizeDatasetOp::kOptimizations;
@@ -178,3 +178,25 @@ REGISTER_KERNEL_BUILDER(Name("OptimizeDatasetV2").Device(DEVICE_CPU),
 }  // namespace
 }  // namespace data
 }  // namespace tensorflow
+#else  // !IS_MOBILE_PLATFORM
+namespace tensorflow {
+namespace data {
+
+OptimizeDatasetOp::OptimizeDatasetOp(OpKernelConstruction* ctx)
+    : UnaryDatasetOpKernel(ctx) {}
+
+void OptimizeDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
+                                    DatasetBase** output) {
+  input->Ref();
+  *output = input;
+}
+
+namespace {
+REGISTER_KERNEL_BUILDER(Name("OptimizeDataset").Device(DEVICE_CPU),
+                        OptimizeDatasetOp);
+REGISTER_KERNEL_BUILDER(Name("OptimizeDatasetV2").Device(DEVICE_CPU),
+                        OptimizeDatasetOp);
+}  // namespace
+}  // namespace data
+}  // namespace tensorflow
+#endif  // !IS_MOBILE_PLATFORM

tensorflow/core/kernels/data/optimize_dataset_op.h

@@ -15,6 +15,11 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_KERNELS_DATA_OPTIMIZE_DATASET_OP_H_
 #define TENSORFLOW_CORE_KERNELS_DATA_OPTIMIZE_DATASET_OP_H_
 
+#include "tensorflow/core/platform/platform.h"
+
+// On mobile we do not provide optimize dataset op because not all of its
+// dependencies are available there. The op is replaced with a no-op.
+#if !defined(IS_MOBILE_PLATFORM)
 #include "tensorflow/core/framework/dataset.h"
 
 namespace tensorflow {
@@ -54,5 +59,23 @@ class OptimizeDatasetOp : public UnaryDatasetOpKernel {
 
 }  // namespace data
 }  // namespace tensorflow
+#else  // !IS_MOBILE_PLATFORM
+#include "tensorflow/core/framework/dataset.h"
+
+namespace tensorflow {
+namespace data {
+
+class OptimizeDatasetOp : public UnaryDatasetOpKernel {
+ public:
+  explicit OptimizeDatasetOp(OpKernelConstruction* ctx);
+
+ protected:
+  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
+                   DatasetBase** output) override;
+};
+
+}  // namespace data
+}  // namespace tensorflow
+#endif  // !IS_MOBILE_PLATFORM
 
 #endif  // TENSORFLOW_CORE_KERNELS_DATA_OPTIMIZE_DATASET_OP_H_

tensorflow/core/kernels/data/optimize_dataset_op_mobile.cc

@@ -1,39 +0,0 @@
-/* Copyright 2018 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.
-==============================================================================*/
-#include "tensorflow/core/framework/dataset.h"
-
-namespace tensorflow {
-namespace data {
-
-class OptimizeDatasetOp : public UnaryDatasetOpKernel {
- public:
-  explicit OptimizeDatasetOp(OpKernelConstruction* ctx)
-      : UnaryDatasetOpKernel(ctx) {}
-
-  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
-                   DatasetBase** output) {
-    input->Ref();
-    *output = input;
-  }
-};
-
-namespace {
-REGISTER_KERNEL_BUILDER(Name("OptimizeDataset").Device(DEVICE_CPU),
-                        OptimizeDatasetOp);
-REGISTER_KERNEL_BUILDER(Name("OptimizeDatasetV2").Device(DEVICE_CPU),
-                        OptimizeDatasetOp);
-}  // namespace
-}  // namespace data
-}  // namespace tensorflow

tensorflow/core/kernels/data/rewrite_utils.cc

@@ -12,9 +12,11 @@ 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/core/kernels/data/rewrite_utils.h"
 
+// On mobile we do not provide this functionality because not all of its
+// dependencies are available there.
+#if !defined(IS_MOBILE_PLATFORM)
 #include "tensorflow/core/common_runtime/graph_constructor.h"
 #include "tensorflow/core/common_runtime/graph_runner.h"
 #include "tensorflow/core/common_runtime/metrics.h"
@@ -222,3 +224,4 @@ Status RewriteDataset(OpKernelContext* ctx, const DatasetBase* input,
 
 }  // namespace data
 }  // namespace tensorflow
+#endif  // !IS_MOBILE_PLATFORM

tensorflow/core/kernels/data/rewrite_utils.h

@@ -15,6 +15,11 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_KERNELS_DATA_REWRITE_UTILS_H_
 #define TENSORFLOW_CORE_KERNELS_DATA_REWRITE_UTILS_H_
 
+#include "tensorflow/core/platform/platform.h"
+
+// On mobile we do not provide this functionality because not all of its
+// dependencies are available there.
+#if !defined(IS_MOBILE_PLATFORM)
 #include "tensorflow/core/common_runtime/function.h"
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/function.h"
@@ -31,5 +36,6 @@ Status RewriteDataset(OpKernelContext* ctx, const DatasetBase* input,
 
 }  // namespace data
 }  // namespace tensorflow
+#endif  // !IS_MOBILE_PLATFORM
 
 #endif  // TENSORFLOW_CORE_KERNELS_DATA_REWRITE_UTILS_H_