1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #ifndef TENSORFLOW_CORE_KERNELS_DATA_OPTIMIZE_DATASET_OP_H_ 16 #define TENSORFLOW_CORE_KERNELS_DATA_OPTIMIZE_DATASET_OP_H_ 17 18 #include "absl/container/flat_hash_set.h" 19 #include "tensorflow/core/framework/dataset.h" 20 #include "tensorflow/core/platform/platform.h" 21 22 // On mobile we do not provide optimize dataset op because not all of its 23 // dependencies are available there. The op is replaced with a no-op. 24 #if !defined(IS_MOBILE_PLATFORM) 25 namespace tensorflow { 26 namespace data { 27 28 class OptimizeDatasetOp : public UnaryDatasetOpKernel { 29 public: 30 static constexpr const char* const kDatasetType = "Optimize"; 31 static constexpr const char* const kInputDataset = "input_dataset"; 32 static constexpr const char* const kOptimizations = "optimizations"; 33 static constexpr const char* const kOptimizationsEnabled = 34 "optimizations_enabled"; 35 static constexpr const char* const kOptimizationsDisabled = 36 "optimizations_disabled"; 37 static constexpr const char* const kOptimizationsDefault = 38 "optimizations_default"; 39 static constexpr const char* const kOutputTypes = "output_types"; 40 static constexpr const char* const kOutputShapes = "output_shapes"; 41 static constexpr const char* const kOptimizationConfigs = 42 "optimization_configs"; 43 static constexpr const char* const kOptimizeDatasetV1 = "OptimizeDataset"; 44 static constexpr const char* const kOptimizeDatasetV2 = "OptimizeDatasetV2"; 45 46 // Creates and returns a OptimizeDatasetOp::Dataset in output, given the 47 // default optimizations and those that are enabled, disabled. This method is 48 // used to create the dataset without explicitly using the OptimizeDatasetOp. 49 static void MakeDatasetFromOptions( 50 OpKernelContext* ctx, DatasetBase* input, 51 const absl::flat_hash_set<tstring>& optimizations_enabled, 52 const absl::flat_hash_set<tstring>& optimizations_disabled, 53 const absl::flat_hash_set<tstring>& optimizations_default, 54 const absl::flat_hash_set<tstring>& optimization_configs, 55 DatasetBase** output); 56 57 explicit OptimizeDatasetOp(OpKernelConstruction* ctx); 58 59 protected: 60 void MakeDataset(OpKernelContext* ctx, DatasetBase* input, 61 DatasetBase** output) override; 62 63 private: 64 absl::flat_hash_set<tstring> optimization_configs_; 65 int op_version_ = 0; 66 }; 67 68 } // namespace data 69 } // namespace tensorflow 70 #else // !IS_MOBILE_PLATFORM 71 namespace tensorflow { 72 namespace data { 73 74 class OptimizeDatasetOp : public UnaryDatasetOpKernel { 75 public: 76 // Executes the logic of the OptimizeDatasetOp directly (as opposed to through 77 // executing the OptimizeDatasetOp op kernel). 78 static void MakeDatasetFromOptions( 79 OpKernelContext* ctx, DatasetBase* input, 80 const absl::flat_hash_set<tstring>& optimizations_enabled, 81 const absl::flat_hash_set<tstring>& optimizations_disabled, 82 const absl::flat_hash_set<tstring>& optimizations_default, 83 const absl::flat_hash_set<tstring>& optimization_configs, 84 DatasetBase** output); 85 86 explicit OptimizeDatasetOp(OpKernelConstruction* ctx); 87 88 protected: 89 void MakeDataset(OpKernelContext* ctx, DatasetBase* input, 90 DatasetBase** output) override; 91 }; 92 93 } // namespace data 94 } // namespace tensorflow 95 #endif // !IS_MOBILE_PLATFORM 96 97 #endif // TENSORFLOW_CORE_KERNELS_DATA_OPTIMIZE_DATASET_OP_H_ 98