115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 315907659888a3e36e8de3d5a95de8d3327cb7c46Derek MurrayLicensed under the Apache License, Version 2.0 (the "License"); 415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murrayyou may not use this file except in compliance with the License. 515907659888a3e36e8de3d5a95de8d3327cb7c46Derek MurrayYou may obtain a copy of the License at 615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray http://www.apache.org/licenses/LICENSE-2.0 815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 915907659888a3e36e8de3d5a95de8d3327cb7c46Derek MurrayUnless required by applicable law or agreed to in writing, software 1015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murraydistributed under the License is distributed on an "AS IS" BASIS, 1115907659888a3e36e8de3d5a95de8d3327cb7c46Derek MurrayWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 1215907659888a3e36e8de3d5a95de8d3327cb7c46Derek MurraySee the License for the specific language governing permissions and 1315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murraylimitations under the License. 1415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray==============================================================================*/ 1515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 1615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray#include "tensorflow/core/framework/partial_tensor_shape.h" 1715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray#include "tensorflow/core/framework/tensor.h" 18a5b2a0c9a3335d10c4dd3dfdff96149f74a4d120Jiri Simsa#include "tensorflow/core/kernels/data/dataset.h" 19a5b2a0c9a3335d10c4dd3dfdff96149f74a4d120Jiri Simsa#include "tensorflow/core/kernels/data/stats_aggregator.h" 2015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray#include "tensorflow/core/lib/random/random.h" 2115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 2215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murraynamespace tensorflow { 2315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murraynamespace { 2415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 2515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// This op defines a `Dataset` that passes through its input elements and 2615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// records the latency of producing each element in the context's 2715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// `StatsAggregator`. 2815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// 2915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// TODO(mrry): It is likely that many *StatsDatasetOp kernels will have the 3015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// same or similar structure. We should abstract the common boilerplate into 3115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// a base case and/or investigate how to make general-purpose *StatsDatasetOp 3215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// kernels that use TensorFlow functions to represent their logic. For example, 3315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// if the performance were adequate, we might replace this kernel with an 3415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// implementation that executes functions before and after the `GetNext()` call 3515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// on the input, each executing an op that gets the current time and performing 3615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray// the subtraction. 3715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murrayclass LatencyStatsDatasetOp : public UnaryDatasetOpKernel { 3815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray public: 3915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray explicit LatencyStatsDatasetOp(OpKernelConstruction* ctx) 4015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray : UnaryDatasetOpKernel(ctx) {} 4115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 4215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray void MakeDataset(OpKernelContext* ctx, DatasetBase* input, 4315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray DatasetBase** output) override { 4415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray string tag; 4515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "tag", &tag)); 46ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal *output = new Dataset(ctx, input, std::move(tag)); 4715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 4815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 4915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray private: 50ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal class Dataset : public GraphDatasetBase { 5115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray public: 52ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal explicit Dataset(OpKernelContext* ctx, const DatasetBase* input, string tag) 53ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal : GraphDatasetBase(ctx), input_(input), tag_(std::move(tag)) { 5415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray input_->Ref(); 5515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 5615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 5715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray ~Dataset() override { input_->Unref(); } 5815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 5915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray std::unique_ptr<IteratorBase> MakeIterator( 6015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray const string& prefix) const override { 6115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray return std::unique_ptr<IteratorBase>( 6215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray new Iterator({this, strings::StrCat(prefix, "::LatencyStats")})); 6315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 6415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 6515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray const DataTypeVector& output_dtypes() const override { 6615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray return input_->output_dtypes(); 6715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 6815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray const std::vector<PartialTensorShape>& output_shapes() const override { 6915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray return input_->output_shapes(); 7015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 7115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 7215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray string DebugString() override { return "LatencyStatsDatasetOp::Dataset"; } 7315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 74ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal protected: 75ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal Status AsGraphDefInternal(OpKernelContext* ctx, DatasetGraphDefBuilder* b, 76ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal Node** output) const override { 77ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal Node* input_node; 78ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal TF_RETURN_IF_ERROR(b->AddParentDataset(ctx, input_, &input_node)); 79ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal Node* tag_node; 80ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal TF_RETURN_IF_ERROR(b->AddScalar(tag_, &tag_node)); 81ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal TF_RETURN_IF_ERROR(b->AddDataset(this, {input_node, tag_node}, output)); 82ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal return Status::OK(); 83ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal } 84ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal 8515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray private: 8615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray class Iterator : public DatasetIterator<Dataset> { 8715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray public: 8815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray explicit Iterator(const Params& params) 8915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray : DatasetIterator<Dataset>(params), 9015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray input_impl_(params.dataset->input_->MakeIterator(params.prefix)) {} 9115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 9215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray Status GetNextInternal(IteratorContext* ctx, 9315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray std::vector<Tensor>* out_tensors, 9415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray bool* end_of_sequence) override { 95ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal tf_shared_lock l(mu_); 9615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray uint64 start = ctx->env()->NowMicros(); 9715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray Status s = input_impl_->GetNext(ctx, out_tensors, end_of_sequence); 9815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray uint64 end = ctx->env()->NowMicros(); 9915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray auto stats_aggregator = ctx->stats_aggregator(); 10015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray if (stats_aggregator && !*end_of_sequence) { 10115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray ctx->stats_aggregator()->AddToHistogram( 10215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray dataset()->tag_, {static_cast<double>(end - start)}); 10315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 10415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray return s; 10515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 10615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 107ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal protected: 108ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal Status SaveInternal(IteratorStateWriter* writer) override { 109ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal mutex_lock l(mu_); 110ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal TF_RETURN_IF_ERROR(SaveParent(writer, input_impl_)); 111ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal return Status::OK(); 112ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal } 113ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal 114dda43c93e0b1cbee0b3215a2ed3fa21afd87e702Derek Murray Status RestoreInternal(IteratorContext* ctx, 115ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal IteratorStateReader* reader) override { 116ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal mutex_lock l(mu_); 117ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal TF_RETURN_IF_ERROR(RestoreParent(ctx, reader, input_impl_)); 118ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal return Status::OK(); 119ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal } 120ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal 12115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray private: 122ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal mutex mu_; 123ed24130f90c2c45db0473df3e9158d4895ce326bShivani Agrawal std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_); 12415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray }; 12515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 12615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray const DatasetBase* const input_; 12715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray const string tag_; 12815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray }; 12915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray}; 13015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 13115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murrayclass BytesProducedStatsDatasetOp : public UnaryDatasetOpKernel { 13215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray public: 13315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray explicit BytesProducedStatsDatasetOp(OpKernelConstruction* ctx) 13415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray : UnaryDatasetOpKernel(ctx) {} 13515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 13615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray void MakeDataset(OpKernelContext* ctx, DatasetBase* input, 13715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray DatasetBase** output) override { 13815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray string tag; 13915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "tag", &tag)); 14022fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal *output = new Dataset(ctx, input, std::move(tag)); 14115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 14215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 14315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray private: 14422fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal class Dataset : public GraphDatasetBase { 14515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray public: 14622fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal explicit Dataset(OpKernelContext* ctx, const DatasetBase* input, string tag) 14722fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal : GraphDatasetBase(ctx), input_(input), tag_(std::move(tag)) { 14815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray input_->Ref(); 14915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 15015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 15115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray ~Dataset() override { input_->Unref(); } 15215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 15315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray std::unique_ptr<IteratorBase> MakeIterator( 15415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray const string& prefix) const override { 15515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray return std::unique_ptr<IteratorBase>(new Iterator( 15615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray {this, strings::StrCat(prefix, "::BytesProducedStats")})); 15715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 15815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 15915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray const DataTypeVector& output_dtypes() const override { 16015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray return input_->output_dtypes(); 16115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 16215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray const std::vector<PartialTensorShape>& output_shapes() const override { 16315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray return input_->output_shapes(); 16415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 16515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 16615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray string DebugString() override { 16715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray return "BytesProducedStatsDatasetOp::Dataset"; 16815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 16915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 17022fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal protected: 17122fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal Status AsGraphDefInternal(OpKernelContext* ctx, DatasetGraphDefBuilder* b, 17222fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal Node** output) const override { 17322fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal Node* input_node; 17422fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal TF_RETURN_IF_ERROR(b->AddParentDataset(ctx, input_, &input_node)); 17522fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal Node* tag_node; 17622fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal TF_RETURN_IF_ERROR(b->AddScalar(tag_, &tag_node)); 17722fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal TF_RETURN_IF_ERROR(b->AddDataset(this, {input_node, tag_node}, output)); 17822fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal return Status::OK(); 17922fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal } 18022fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal 18115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray private: 18215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray class Iterator : public DatasetIterator<Dataset> { 18315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray public: 18415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray explicit Iterator(const Params& params) 18515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray : DatasetIterator<Dataset>(params), 18615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray input_impl_(params.dataset->input_->MakeIterator(params.prefix)) {} 18715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 18815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray Status GetNextInternal(IteratorContext* ctx, 18915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray std::vector<Tensor>* out_tensors, 19015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray bool* end_of_sequence) override { 19122fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal tf_shared_lock l(mu_); 19215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray Status s = input_impl_->GetNext(ctx, out_tensors, end_of_sequence); 19315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray auto stats_aggregator = ctx->stats_aggregator(); 19415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray if (stats_aggregator && s.ok() && !*end_of_sequence) { 19515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray size_t total_bytes = 0; 19615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray for (const Tensor& t : *out_tensors) { 19715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray total_bytes += t.TotalBytes(); 19815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 19915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray ctx->stats_aggregator()->AddToHistogram( 20015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray dataset()->tag_, {static_cast<double>(total_bytes)}); 20115907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 20215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray return s; 20315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray } 20415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 20522fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal protected: 20622fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal Status SaveInternal(IteratorStateWriter* writer) override { 20722fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal mutex_lock l(mu_); 20822fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal TF_RETURN_IF_ERROR(SaveParent(writer, input_impl_)); 20922fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal return Status::OK(); 21022fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal } 21122fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal 212dda43c93e0b1cbee0b3215a2ed3fa21afd87e702Derek Murray Status RestoreInternal(IteratorContext* ctx, 21322fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal IteratorStateReader* reader) override { 21422fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal mutex_lock l(mu_); 21522fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal TF_RETURN_IF_ERROR(RestoreParent(ctx, reader, input_impl_)); 21622fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal return Status::OK(); 21722fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal } 21822fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal 21915907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray private: 22022fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal mutex mu_; 22122fe6558a958c6cc81d16d371031c06e262b1c83Shivani Agrawal std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_); 22215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray }; 22315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 22415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray const DatasetBase* const input_; 22515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray const string tag_; 22615907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray }; 22715907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray}; 22815907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 22915907659888a3e36e8de3d5a95de8d3327cb7c46Derek MurrayREGISTER_KERNEL_BUILDER(Name("LatencyStatsDataset").Device(DEVICE_CPU), 23015907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray LatencyStatsDatasetOp); 23115907659888a3e36e8de3d5a95de8d3327cb7c46Derek MurrayREGISTER_KERNEL_BUILDER(Name("BytesProducedStatsDataset").Device(DEVICE_CPU), 23215907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray BytesProducedStatsDatasetOp); 23315907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray 23415907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray} // namespace 23515907659888a3e36e8de3d5a95de8d3327cb7c46Derek Murray} // namespace tensorflow 236