From 072fcb995a3fd658ee2461b59b159498c710513d Mon Sep 17 00:00:00 2001 From: Brennan Saeta Date: Tue, 9 Oct 2018 11:54:20 -0700 Subject: [tf.data] NUMA-aware MapAndBatch dataset. PiperOrigin-RevId: 216395709 --- ...pi_def_ExperimentalNumaMapAndBatchDataset.pbtxt | 58 + tensorflow/core/framework/model.h | 2 +- tensorflow/core/grappler/optimizers/data/BUILD | 35 + .../grappler/optimizers/data/graph_test_utils.cc | 16 + .../grappler/optimizers/data/graph_test_utils.h | 6 + .../data/map_and_batch_numa_aware_replacement.cc | 62 ++ .../data/map_and_batch_numa_aware_replacement.h | 48 + .../map_and_batch_numa_aware_replacement_test.cc | 112 ++ tensorflow/core/kernels/data/experimental/BUILD | 17 + .../experimental/numa_map_and_batch_dataset_op.cc | 1135 ++++++++++++++++++++ .../core/kernels/data/map_and_batch_dataset_op.cc | 38 +- tensorflow/core/ops/experimental_dataset_ops.cc | 26 + .../kernel_tests/map_and_batch_test.py | 280 +++-- .../experimental/kernel_tests/optimization/BUILD | 2 + .../optimization/model_dataset_op_test.py | 11 +- .../optimization/optimize_dataset_op_test.py | 16 + .../experimental/kernel_tests/serialization/BUILD | 15 + ...uma_map_and_batch_dataset_serialization_test.py | 95 ++ tensorflow/python/data/experimental/ops/BUILD | 1 + tensorflow/python/data/ops/dataset_ops.py | 7 +- .../api/golden/v1/tensorflow.data.-options.pbtxt | 4 + .../api/golden/v2/tensorflow.data.-options.pbtxt | 4 + 22 files changed, 1909 insertions(+), 81 deletions(-) create mode 100644 tensorflow/core/api_def/base_api/api_def_ExperimentalNumaMapAndBatchDataset.pbtxt create mode 100644 tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement.cc create mode 100644 tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement.h create mode 100644 tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement_test.cc create mode 100644 tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc create mode 100644 tensorflow/python/data/experimental/kernel_tests/serialization/numa_map_and_batch_dataset_serialization_test.py diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalNumaMapAndBatchDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalNumaMapAndBatchDataset.pbtxt new file mode 100644 index 0000000000..243922d969 --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalNumaMapAndBatchDataset.pbtxt @@ -0,0 +1,58 @@ +op { + graph_op_name: "ExperimentalNumaMapAndBatchDataset" + visibility: HIDDEN + in_arg { + name: "input_dataset" + description: <{}}}); } +NodeDef MakeMapAndBatchNode(StringPiece name, StringPiece input_node_name, + StringPiece batch_size_node_name, + StringPiece num_parallel_calls_node_name, + StringPiece drop_remainder_node_name, + StringPiece function_name) { + return test::function::NDef( + name, "MapAndBatchDatasetV2", + {string(input_node_name), "", string(batch_size_node_name), + string(num_parallel_calls_node_name), string(drop_remainder_node_name)}, + {{"predicate", FunctionDefHelper::FunctionRef(string(function_name))}, + {"Targuments", {}}, + {"output_shapes", gtl::ArraySlice{}}, + {"output_types", gtl::ArraySlice{}}}); +} + } // end namespace graph_tests_utils } // end namespace grappler } // end namespace tensorflow diff --git a/tensorflow/core/grappler/optimizers/data/graph_test_utils.h b/tensorflow/core/grappler/optimizers/data/graph_test_utils.h index ca0fde997d..f7891d5e1f 100644 --- a/tensorflow/core/grappler/optimizers/data/graph_test_utils.h +++ b/tensorflow/core/grappler/optimizers/data/graph_test_utils.h @@ -29,6 +29,12 @@ NodeDef MakeMapNode(StringPiece name, StringPiece input_node_name, NodeDef MakeFilterNode(StringPiece name, StringPiece input_node_name, StringPiece function_name = "IsZero"); +NodeDef MakeMapAndBatchNode(StringPiece name, StringPiece input_node_name, + StringPiece batch_size_node_name, + StringPiece num_parallel_calls_node_name, + StringPiece drop_remainder_node_name, + StringPiece function_name = "XTimesTwo"); + } // end namespace graph_tests_utils } // end namespace grappler } // end namespace tensorflow diff --git a/tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement.cc b/tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement.cc new file mode 100644 index 0000000000..452089eb67 --- /dev/null +++ b/tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement.cc @@ -0,0 +1,62 @@ +/* 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/grappler/optimizers/data/map_and_batch_numa_aware_replacement.h" + +#include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/grappler/clusters/cluster.h" +#include "tensorflow/core/grappler/grappler_item.h" +#include "tensorflow/core/grappler/mutable_graph_view.h" +#include "tensorflow/core/grappler/op_types.h" +#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h" +#include "tensorflow/core/grappler/optimizers/data/graph_utils.h" + +namespace tensorflow { +namespace grappler { +namespace { + +NodeDef MakeNumaAware(const NodeDef& node, MutableGraphView* graph) { + NodeDef numa_aware_node = node; + graph_utils::SetUniqueGraphNodeName("map_and_batch_numa_aware", + graph->GetGraph(), &numa_aware_node); + numa_aware_node.set_op("ExperimentalNumaMapAndBatchDataset"); + return numa_aware_node; +} + +} // namespace + +Status MapAndBatchNumaAwareReplacement::Optimize(Cluster* cluster, + const GrapplerItem& item, + GraphDef* output) { + *output = item.graph; + MutableGraphView graph(output); + std::set nodes_to_delete; + + for (const NodeDef& node : item.graph.node()) { + if (node.op() != "MapAndBatchDatasetV2") continue; + + auto* numa_node = graph.AddNode(MakeNumaAware(node, &graph)); + graph.ReplaceInput(node, *numa_node); + nodes_to_delete.insert(node.name()); + } + graph.DeleteNodes(nodes_to_delete); + return Status::OK(); +} + +REGISTER_GRAPH_OPTIMIZER_AS(MapAndBatchNumaAwareReplacement, + "map_and_batch_numa_aware_replacement"); + +} // namespace grappler +} // namespace tensorflow diff --git a/tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement.h b/tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement.h new file mode 100644 index 0000000000..3b2acd288b --- /dev/null +++ b/tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement.h @@ -0,0 +1,48 @@ +/* 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. +==============================================================================*/ + +#ifndef TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_MAP_AND_BATCH_NUMA_AWARE_REPLACEMENT_H_ +#define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_MAP_AND_BATCH_NUMA_AWARE_REPLACEMENT_H_ + +#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h" + +namespace tensorflow { +namespace grappler { + +class MapAndBatchNumaAwareReplacement : public CustomGraphOptimizer { + public: + MapAndBatchNumaAwareReplacement() = default; + ~MapAndBatchNumaAwareReplacement() override = default; + + string name() const override { + return "map_and_batch_numa_aware_replacement"; + } + + Status Init( + const tensorflow::RewriterConfig_CustomGraphOptimizer* config) override { + return Status::OK(); + } + + Status Optimize(Cluster* cluster, const GrapplerItem& item, + GraphDef* output) override; + + void Feedback(Cluster* cluster, const GrapplerItem& item, + const GraphDef& optimize_output, double result) override {} +}; + +} // namespace grappler +} // namespace tensorflow + +#endif // TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_MAP_AND_BATCH_NUMA_AWARE_REPLACEMENT_H_ diff --git a/tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement_test.cc b/tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement_test.cc new file mode 100644 index 0000000000..3c5c61d1c2 --- /dev/null +++ b/tensorflow/core/grappler/optimizers/data/map_and_batch_numa_aware_replacement_test.cc @@ -0,0 +1,112 @@ +/* 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/grappler/optimizers/data/map_and_batch_numa_aware_replacement.h" + +#include "tensorflow/core/framework/attr_value_util.h" +#include "tensorflow/core/framework/function_testlib.h" +#include "tensorflow/core/framework/tensor_testutil.h" +#include "tensorflow/core/grappler/grappler_item.h" +#include "tensorflow/core/grappler/optimizers/data/graph_test_utils.h" +#include "tensorflow/core/grappler/optimizers/data/graph_utils.h" + +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace grappler { +namespace { + +TEST(MapAndBatchNumaAwareReplacementTest, ReplaceSimple) { + using test::function::NDef; + GrapplerItem item; + item.graph = test::function::GDef( + { + NDef("start", "Const", {}, {{"value", 0}, {"dtype", DT_INT32}}), + NDef("stop", "Const", {}, {{"value", 10}, {"dtype", DT_INT32}}), + NDef("step", "Const", {}, {{"value", 1}, {"dtype", DT_INT32}}), + NDef("range", "RangeDataset", {"start", "stop", "step"}, {}), + NDef("batch_size", "Const", {}, {{"value", 3}, {"dtype", DT_INT32}}), + NDef("num_parallel_calls", "Const", {}, + {{"value", 5}, {"dtype", DT_INT32}}), + NDef("drop_remainder", "Const", {}, + {{"value", 0}, {"dtype", DT_BOOL}}), + graph_tests_utils::MakeMapAndBatchNode( + "map_and_batch", "range", "batch_size", "num_parallel_calls", + "drop_remainder"), + }, + // FunctionLib + { + test::function::XTimesTwo(), + }); + + MapAndBatchNumaAwareReplacement optimizer; + GraphDef output; + TF_ASSERT_OK(optimizer.Optimize(nullptr, item, &output)); + + EXPECT_FALSE(graph_utils::ContainsGraphNodeWithName("map_and_batch", output)); + EXPECT_FALSE(graph_utils::ContainsNodeWithOp("MapAndBatchDatasetV2", output)); + EXPECT_TRUE(graph_utils::ContainsNodeWithOp( + "ExperimentalNumaMapAndBatchDataset", output)); +} + +TEST(MapAndBatchNumaAawareReplacementTest, ReplaceWithExtraChild) { + using test::function::NDef; + GrapplerItem item; + item.graph = test::function::GDef( + { + NDef("start", "Const", {}, {{"value", 0}, {"dtype", DT_INT32}}), + NDef("stop", "Const", {}, {{"value", 10}, {"dtype", DT_INT32}}), + NDef("step", "Const", {}, {{"value", 1}, {"dtype", DT_INT32}}), + NDef("range", "RangeDataset", {"start", "stop", "step"}, {}), + NDef("batch_size", "Const", {}, {{"value", 3}, {"dtype", DT_INT32}}), + NDef("num_parallel_calls", "Const", {}, + {{"value", 5}, {"dtype", DT_INT32}}), + NDef("drop_remainder", "Const", {}, + {{"value", 0}, {"dtype", DT_BOOL}}), + graph_tests_utils::MakeMapAndBatchNode( + "map_and_batch", "range", "batch_size", "num_parallel_calls", + "drop_remainder"), + NDef("cache", "CacheDataset", {"map_and_batch"}, {}), + }, + // FunctionLib + { + test::function::XTimesTwo(), + }); + + MapAndBatchNumaAwareReplacement optimizer; + GraphDef output; + TF_ASSERT_OK(optimizer.Optimize(nullptr, item, &output)); + + EXPECT_FALSE(graph_utils::ContainsGraphNodeWithName("map_and_batch", output)); + EXPECT_FALSE(graph_utils::ContainsNodeWithOp("MapAndBatchDatasetV2", output)); + EXPECT_TRUE(graph_utils::ContainsNodeWithOp( + "ExperimentalNumaMapAndBatchDataset", output)); + EXPECT_TRUE(graph_utils::ContainsNodeWithOp("CacheDataset", output)); + + int numa_map_and_batch_component_id = graph_utils::FindGraphNodeWithOp( + "ExperimentalNumaMapAndBatchDataset", output); + auto& numa_map_and_batch_component = + output.node(numa_map_and_batch_component_id); + EXPECT_EQ(numa_map_and_batch_component.input(0), "range"); + + int cache_id = graph_utils::FindGraphNodeWithOp("CacheDataset", output); + auto& cache_node = output.node(cache_id); + EXPECT_EQ(cache_node.input(0), numa_map_and_batch_component.name()); +} + +} // namespace +} // namespace grappler +} // namespace tensorflow diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD index 43406db3ed..4cf5643bc0 100644 --- a/tensorflow/core/kernels/data/experimental/BUILD +++ b/tensorflow/core/kernels/data/experimental/BUILD @@ -102,6 +102,22 @@ tf_kernel_library( ], ) +tf_kernel_library( + name = "numa_map_and_batch_dataset_op", + srcs = ["numa_map_and_batch_dataset_op.cc"], + deps = [ + "//tensorflow/core:core_cpu_internal", + "//tensorflow/core:experimental_dataset_ops_op_lib", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core/kernels:inplace_ops", + "//tensorflow/core/kernels/data:captured_function", + "//tensorflow/core/kernels/data:dataset", + "@com_google_absl//absl/memory", + ], +) + tf_kernel_library( name = "unique_dataset_op", srcs = ["unique_dataset_op.cc"], @@ -132,6 +148,7 @@ tf_kernel_library( ":ignore_errors_dataset_op", ":indexed_dataset", ":lmdb_dataset_op", + ":numa_map_and_batch_dataset_op", ":prefetching_kernels", ":threadpool_dataset_op", ":unique_dataset_op", diff --git a/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc new file mode 100644 index 0000000000..d83edb9667 --- /dev/null +++ b/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc @@ -0,0 +1,1135 @@ +/* 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. +==============================================================================*/ +#define EIGEN_USE_THREADS + +#include +#include + +#include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/framework/dataset.h" +#include "tensorflow/core/framework/partial_tensor_shape.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/kernels/data/captured_function.h" +#include "tensorflow/core/kernels/data/dataset.h" +#include "tensorflow/core/kernels/inplace_ops_functor.h" +#include "tensorflow/core/lib/core/blocking_counter.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/gtl/cleanup.h" +#include "tensorflow/core/lib/random/random.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/cpu_info.h" +#include "tensorflow/core/platform/numa.h" +#include "tensorflow/core/platform/tracing.h" +#include "tensorflow/core/util/ptr_util.h" + +namespace tensorflow { +namespace data { +namespace { + +// kWindowSize is the fixed constant controlling the number of batch outputs +// each NumaWorkerBlock may be processing at a time. This is currently a +// constant and not user configurable to enable future performance optimizations +// in the implementation. +const int64 kWindowSize = 10; + +// Define a helper for more consistent logging. +#define WORKER_VLOG(verbose_level) \ + VLOG(verbose_level) << "WorkerThread (" << numa_node << ", " << thread_num \ + << "): " + +// See documentation in ../ops/dataset_ops.cc for a high-level +// description of the following op. + +class NumaMapAndBatchDatasetOp : public UnaryDatasetOpKernel { + public: + explicit NumaMapAndBatchDatasetOp(OpKernelConstruction* ctx) + : UnaryDatasetOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("f", &func_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_)); + } + + protected: + void MakeDataset(OpKernelContext* ctx, DatasetBase* input, + DatasetBase** output) override { + int64 batch_size; + OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "batch_size", &batch_size)); + OP_REQUIRES( + ctx, batch_size > 0, + errors::InvalidArgument("batch_size must be greater than zero.")); + + int64 num_parallel_calls; + OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "num_parallel_calls", + &num_parallel_calls)); + OP_REQUIRES(ctx, num_parallel_calls > 0 || num_parallel_calls == kAutoTune, + errors::InvalidArgument( + "num_parallel_calls must be greater than zero.")); + + bool drop_remainder; + OP_REQUIRES_OK(ctx, + ParseScalarArgument(ctx, "drop_remainder", &drop_remainder)); + + std::unique_ptr captured_func; + OP_REQUIRES_OK( + ctx, CapturedFunction::Create(func_, ctx, "other_arguments", + /* use_inter_op_parallelism = */ false, + &captured_func)); + + *output = new Dataset(ctx, input, batch_size, num_parallel_calls, + drop_remainder, output_types_, output_shapes_, func_, + std::move(captured_func)); + } + + private: + class Dataset : public DatasetBase { + public: + Dataset(OpKernelContext* ctx, const DatasetBase* input, int64 batch_size, + int64 num_parallel_calls, bool drop_remainder, + const DataTypeVector& output_types, + const std::vector& output_shapes, + const NameAttrList& func, + std::unique_ptr captured_func) + : DatasetBase(DatasetContext(ctx)), + input_(input), + batch_size_(batch_size), + num_parallel_calls_(num_parallel_calls), + drop_remainder_(drop_remainder), + output_types_(output_types), + output_shapes_(output_shapes), + func_(func), + captured_func_(std::move(captured_func)) { + input_->Ref(); + } + + ~Dataset() override { input_->Unref(); } + + std::unique_ptr MakeIteratorInternal( + const string& prefix) const override { + return std::unique_ptr( + new Iterator({this, strings::StrCat(prefix, "::NumaMapAndBatch")})); + } + + const DataTypeVector& output_dtypes() const override { + return output_types_; + } + + const std::vector& output_shapes() const override { + return output_shapes_; + } + + string DebugString() const override { + return "NumaMapAndBatchDatasetOp::Dataset"; + } + + protected: + Status AsGraphDefInternal(SerializationContext* ctx, + DatasetGraphDefBuilder* b, + Node** output) const override { + TF_RETURN_IF_ERROR(b->AddFunction(ctx, func_.name())); + Node* input_graph_node = nullptr; + TF_RETURN_IF_ERROR(b->AddInputDataset(ctx, input_, &input_graph_node)); + Node* batch_size_node; + TF_RETURN_IF_ERROR(b->AddScalar(batch_size_, &batch_size_node)); + Node* num_parallel_calls_node; + TF_RETURN_IF_ERROR( + b->AddScalar(num_parallel_calls_, &num_parallel_calls_node)); + Node* drop_remainder_node; + TF_RETURN_IF_ERROR(b->AddScalar(drop_remainder_, &drop_remainder_node)); + + DataTypeVector other_arguments_types; + other_arguments_types.reserve(captured_func_->captured_inputs().size()); + std::vector other_arguments; + other_arguments.reserve(captured_func_->captured_inputs().size()); + for (const Tensor& t : captured_func_->captured_inputs()) { + Node* node; + TF_RETURN_IF_ERROR(b->AddTensor(t, &node)); + other_arguments.emplace_back(node); + other_arguments_types.emplace_back(t.dtype()); + } + AttrValue f; + b->BuildAttrValue(func_, &f); + AttrValue other_arguments_types_attr; + b->BuildAttrValue(other_arguments_types, &other_arguments_types_attr); + + TF_RETURN_IF_ERROR(b->AddDataset( + this, + {std::make_pair(0, input_graph_node), + std::make_pair(2, batch_size_node), + std::make_pair(3, num_parallel_calls_node), + std::make_pair(4, drop_remainder_node)}, // Single tensor inputs. + {std::make_pair(1, other_arguments)}, // Tensor list inputs. + {std::make_pair("f", f), + std::make_pair("Targuments", other_arguments_types_attr)}, // Attrs + output)); + return Status::OK(); + } + + private: + class Iterator : public DatasetIterator { + public: + explicit Iterator(const Params& params) + : DatasetIterator(params), + mu_(std::make_shared()), + autotune_cond_var_(std::make_shared()), + num_parallel_calls_(std::make_shared( + params.dataset->num_parallel_calls_, mu_, autotune_cond_var_)) { + } + + ~Iterator() override { + mutex_lock l(*mu_); + cancelled_ = true; + VLOG(3) << "NumaMapAndBatchIterator::~Iterator: cancelling operations."; + for (size_t i = 0; i < workers_.size(); ++i) { + workers_[i]->manager.Cancel(); + } + VLOG(3) << "NumaMapAndBatchIterator::~Iterator: waiting for threads to " + "shut down."; + } + + Status Initialize(IteratorContext* ctx) override { + mutex_lock l(*mu_); + AddConstantParameter(ctx, "batch_size", dataset()->batch_size_); + if (num_parallel_calls_->value == kAutoTune) { + num_parallel_calls_->value = std::max(1, port::NUMANumNodes()); + AddTunableParameter(ctx, + /* name = */ "parallelism", + /* state = */ num_parallel_calls_, + /* min = */ num_parallel_calls_->value, + /* max = */ port::NumSchedulableCPUs()); + } else { + AddConstantParameter(ctx, "parallelism", num_parallel_calls_->value); + } + TF_RETURN_IF_ERROR( + dataset()->input_->MakeIterator(ctx, prefix(), &input_impl_)); + TF_RETURN_IF_ERROR(dataset()->captured_func_->Instantiate(ctx)); + return Status::OK(); + } + + Status GetNextInternal(IteratorContext* ctx, + std::vector* out_tensors, + bool* end_of_sequence) override { + auto cleanup = gtl::MakeCleanup( + [] { VLOG(3) << "GetNextInternal call returning."; }); + NumaWorkerBlock* worker = nullptr; + { + mutex_lock l(*mu_); + VLOG(3) << "GetNextInternal call; current block: " << cur_block_; + if (global_end_of_input_) { + *end_of_sequence = true; + return Status::OK(); + } + TF_RETURN_IF_ERROR(EnsureBackgroundThreadsStarted(ctx)); + worker = workers_[cur_block_].get(); + cur_block_ = (cur_block_ + 1) % workers_.size(); + } + TF_RETURN_IF_ERROR(worker->manager.GetBatch( + ctx, dataset()->drop_remainder_, &global_end_of_input_, out_tensors, + end_of_sequence)); + return Status::OK(); + } + + protected: + Status SaveInternal(IteratorStateWriter* writer) override { + mutex_lock l(*mu_); + for (size_t i = 0; i < workers_.size(); ++i) { + if (!workers_[i]->manager.Quiesce()) { + return errors::Cancelled( + "The iterator was deleted before it could reach a " + "checkpointable state."); + } + } + + TF_RETURN_IF_ERROR(SaveInput(writer, input_impl_)); + TF_RETURN_IF_ERROR( + writer->WriteScalar(full_name("num_workers"), workers_.size())); + + for (size_t i = 0; i < workers_.size(); ++i) { + size_t index = (cur_block_ + i) % workers_.size(); + TF_RETURN_IF_ERROR(workers_[index]->manager.Save(writer, this, i)); + } + return Status::OK(); + } + + Status RestoreInternal(IteratorContext* ctx, + IteratorStateReader* reader) override { + mutex_lock l(*mu_); + TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, input_impl_)); + int64 num_workers = -1; + TF_RETURN_IF_ERROR( + reader->ReadScalar(full_name("num_workers"), &num_workers)); + // Note: num_workers can be 0 if the iterator wasn't started when + // first checkpointed. + if (num_workers < 0) { + return errors::DataLoss( + "When restoring from checkpoint, we encountered a data " + "consistency error: num_workers has an invalid value: ", + num_workers); + } + if (port::NUMAEnabled()) { + int actual_numa_domains = port::NUMANumNodes(); + if (actual_numa_domains != num_workers && num_workers > 0) { + LOG(WARNING) << "# NUMA domains mismatch when restoring from " + "checkpoint: checkpoint has " + << num_workers + << " NUMA domains, while this host has: " + << actual_numa_domains << " NUMA domains."; + } + } + if (num_workers > 1 && !port::NUMAEnabled()) { + LOG(WARNING) << "NUMA is not enabled for this process, but restoring " + "a checkpoint that assumes " + << num_workers << " NUMA domains."; + } + workers_.resize(num_workers); + for (size_t i = 0; i < num_workers; ++i) { + workers_[i] = MakeUnique(this); + TF_RETURN_IF_ERROR( + workers_[i]->manager.Restore(ctx, reader, this, i)); + } + cur_block_ = 0; + return Status::OK(); + } + + private: + // NumaBlockManager manages all the state for a set of threads pinned to a + // single NUMA domain. + // + // The methods can be divided into 3 categories based on who should call + // them: + // + // (1) RunnerThread: WaitForInputSpace, PushInputs, SetEndOfInput. + // (2) WorkerThread: RetrieveInput, GetBatchTensors. + // RecordBatchEntryComplete + // (3) Client threads: GetBatch, Cancel, Save, Restore. + // + // Internally, we manage state in a circular buffer of size `kWindowSize`. + // There are 3 pointers into the circular buffer, and must maintain the + // following order: (1) next_input_batch_ (corresponding to the next input + // batch to be pulled from the input iterator), (2) next_input_ + // (corresponding to the batch the WorkerThreads should pull from for + // their next inputs), and (3) next_output_ corresponding to the next + // value to be consumed by the output iterator. + // + // Methods return errors::Cancelled if the iteration is cancelled before + // completing. + // + // NumaBlockManager is thread safe. + class NumaBlockManager { + public: + explicit NumaBlockManager(Iterator* itr) : itr_(itr) {} + + // WaitForInputSpace blocks until there is space in the circular buffer + // to begin processing a new batch of elements. + // + // Returns true when there is space, false if the Iterator is cancelled. + bool WaitForInputSpace(IteratorContext* ctx) { + mutex_lock l(mu_); + + size_t next = (next_input_batch_ + 1) % kWindowSize; + DCHECK(next < kWindowSize) << next; + + // Wait for space in the circular buffer. + while (!cancelled_ && batches_[next].state != BatchState::kEmpty) { + VLOG(3) << "Waiting for input space; next: " << next + << ", next_output_: " << next_output_ + << ", next_input_batch_: " << next_input_batch_; + itr_->RecordStop(ctx); + runner_cond_var_.wait(l); + itr_->RecordStart(ctx); + } + if (cancelled_) { + VLOG(3) << "WaitForInputSpace cancelled."; + return false; + } + + DCHECK(batches_[next].state == BatchState::kEmpty); + + next_input_batch_ = next; + return true; + } + + // PushInputs sets the inputs for the next batch as retrieved from the + // input iterator. + void PushInputs(const Status& status, + std::vector> inputs) { + mutex_lock l(mu_); + + DCHECK(next_input_ < kWindowSize) << next_input_; + DCHECK(batches_[next_input_batch_].state == BatchState::kEmpty); + DCHECK(batches_[next_input_batch_].next_input_to_process == 0) + << batches_[next_input_batch_].next_input_to_process; + DCHECK(batches_[next_input_batch_].status.ok()) + << batches_[next_input_batch_].status; + + batches_[next_input_batch_].inputs.swap(inputs); + batches_[next_input_batch_].state = BatchState::kInputsFilled; + batches_[next_input_batch_].status.Update(status); + if (batches_[next_input_batch_].status.ok()) { + worker_cond_var_.notify_all(); + } else { + client_cond_var_.notify_all(); + batches_[next_input_batch_].error_index = 0; + } + } + + // SetEndOfInput records the fact that we have reached the end of the + // input iterator, and that we should return end_of_sequence = true when + // we have exhaused all buffered batches. + void SetEndOfInput() { + mutex_lock l(mu_); + reached_eof_ = true; + worker_cond_var_.notify_all(); + client_cond_var_.notify_all(); + } + + // RetrieveInput gets the next input tuple to be mapped by a worker + // thread. + // + // Returns true if an input was retrieved, false if the iterator has + // been cancelled. + bool RetrieveInput(IteratorContext* ctx, std::vector* input, + uint64* index, size_t* sequence_number) { + mutex_lock l(mu_); + + // Wait for inputs to be ready. + while (!cancelled_ && + batches_[next_input_].state != BatchState::kInputsFilled) { + itr_->RecordStop(ctx); + worker_cond_var_.wait(l); + itr_->RecordStart(ctx); + } + + if (cancelled_) { + return false; + } + + DCHECK(batches_[next_input_].next_input_to_process < + batches_[next_input_].inputs.size()) + << "next_input_: " << next_input_ << ", next_input_to_process: " + << batches_[next_input_].next_input_to_process + << ", inputs.size(): " << batches_[next_input_].inputs.size() + << ", state: " << static_cast(batches_[next_input_].state) + << ", this: " << this; + *index = batches_[next_input_].next_input_to_process; + *sequence_number = next_input_; + input->swap(batches_[next_input_] + .inputs[batches_[next_input_].next_input_to_process]); + // Increment pointers. + batches_[next_input_].next_input_to_process++; + + if (batches_[next_input_].next_input_to_process == + batches_[next_input_].inputs.size()) { + batches_[next_input_].state = BatchState::kAllMapsStarted; + next_input_ = (next_input_ + 1) % kWindowSize; + } + return true; + } + + // GetBatchTensors returns a pointer to the output batch tensors for the + // worker thread to copy into. + // + // allocate_output is a function taking a batch size, and a pointer to + // the output tuple of Tensors to allocate them. The allocate_output + // function is called at most once per output batch. + std::vector* GetBatchTensors( + size_t sequence_number, + std::function*)> allocate_output) { + mutex_lock l(mu_); + DCHECK(sequence_number < kWindowSize) << sequence_number; + DCHECK(batches_[sequence_number].state == BatchState::kInputsFilled || + batches_[sequence_number].state == BatchState::kAllMapsStarted) + << sequence_number; + + if (batches_[sequence_number].outputs.empty()) { + allocate_output(batches_[sequence_number].inputs.size(), + &batches_[sequence_number].outputs); + } + return &batches_[sequence_number].outputs; + } + + // RecordBatchEntryComplete records an element of the batch has finished + // copying into the output tensors. + void RecordBatchEntryComplete(size_t sequence_number, uint64 index, + Status s) { + mutex_lock l(mu_); + DCHECK(sequence_number < kWindowSize) << sequence_number; + DCHECK(batches_[sequence_number].state == BatchState::kInputsFilled || + batches_[sequence_number].state == BatchState::kAllMapsStarted) + << sequence_number; + + batches_[sequence_number].num_outputs_complete++; + if (!s.ok() && batches_[sequence_number].error_index > index) { + batches_[sequence_number].status = s; + batches_[sequence_number].error_index = index; + } + + if (batches_[sequence_number].num_outputs_complete == + batches_[sequence_number].inputs.size()) { + DCHECK(batches_[sequence_number].state == + BatchState::kAllMapsStarted); + batches_[sequence_number].state = BatchState::kOutputsComplete; + batches_[sequence_number].inputs.clear(); // Eagerly save memory. + batches_[sequence_number].inputs.shrink_to_fit(); + client_cond_var_.notify_all(); + } + } + + // GetBatch retrieves the next output batch tensors. + Status GetBatch(IteratorContext* ctx, bool drop_remainder, + bool* global_eof, std::vector* out_tensor, + bool* end_of_sequence) { + mutex_lock l(mu_); + // Wait until one of 3 conditions occurs: + // (1) we're cancelled. + // (2) the state becomes kOutputsComplete + // (3) state is empty && reached_eof. + while (!cancelled_ && + batches_[next_output_].state != BatchState::kOutputsComplete && + !(reached_eof_ && + batches_[next_output_].state == BatchState::kEmpty)) { + VLOG(3) << "Waiting in GetBatch."; + itr_->RecordStop(ctx); + client_cond_var_.wait(l); + itr_->RecordStart(ctx); + } + + if (cancelled_) { + return errors::Cancelled( + "Cancelled in NumaMapAndBatch::GetNext call."); + } + + if (reached_eof_ && + batches_[next_output_].state == BatchState::kEmpty) { + VLOG(4) << "GetBatch returning end of sequence."; + *end_of_sequence = true; + *global_eof = true; + return Status::OK(); + } + + VLOG(3) << "Returning output index: " << next_output_ + << ", this: " << this; + + *end_of_sequence = false; + Status s = batches_[next_output_].status; + if (s.ok()) { + out_tensor->swap(batches_[next_output_].outputs); + } + // Handle early termination. + if (errors::IsOutOfRange(s)) { + *global_eof = true; + s = Status::OK(); + if (drop_remainder || batches_[next_output_].error_index == 0) { + *end_of_sequence = true; + } else { + std::vector true_outputs; + for (size_t i = 0; i < batches_[next_output_].outputs.size(); + ++i) { + TensorShape component_shape( + batches_[next_output_].outputs[i].shape()); + component_shape.set_dim(0, batches_[next_output_].error_index); + AllocatorAttributes attr; + attr.set_gpu_compatible(true); + Tensor component(ctx->allocator(attr), + batches_[next_output_].outputs[i].dtype(), + component_shape); + TF_RETURN_IF_ERROR(CopyPartialBatch( + &component, batches_[next_output_].outputs[i], + batches_[next_output_].error_index)); + true_outputs.emplace_back(std::move(component)); + } + out_tensor->swap(true_outputs); + } + } + + batches_[next_output_].Reset(); + next_output_ = (next_output_ + 1) % kWindowSize; + runner_cond_var_.notify_all(); + + return s; + } + + void Cancel() { + mutex_lock l(mu_); + VLOG(3) << "Cancelling NUMA block."; + cancelled_ = true; + runner_cond_var_.notify_all(); + worker_cond_var_.notify_all(); + client_cond_var_.notify_all(); + } + + // Waits until all the worker threads have completed their work and all + // internal state has reached a "safe-point" where we can safely + // checkpoint. + // + // Returns true if completed successfully, false if cancelled while + // waiting. + bool Quiesce() { + mutex_lock l(mu_); + VLOG(3) << "Waiting until the operations have quiesced."; + while (!cancelled_ && !AllMapOperationsFinished()) { + client_cond_var_.wait(l); + } + if (cancelled_) { + return false; + } + return true; + } + + Status Save(IteratorStateWriter* writer, Iterator* itr, size_t index) { + mutex_lock l(mu_); + string prefix = itr->full_name(strings::StrCat("numa_block_", index)); + if (reached_eof_) { + TF_RETURN_IF_ERROR(writer->WriteScalar( + strings::StrCat(prefix, "_end_of_input"), "")); + } + for (size_t i = 0; i < kWindowSize; ++i) { + size_t index = (next_output_ + i) % kWindowSize; + if (batches_[index].state == BatchState::kEmpty) { + break; + } + string batch_prefix = strings::StrCat(prefix, "_batch_", i); + TF_RETURN_IF_ERROR(writer->WriteScalar( + strings::StrCat(batch_prefix, "_code"), + static_cast(batches_[index].status.code()))); + if (!batches_[index].status.ok()) { + TF_RETURN_IF_ERROR( + writer->WriteScalar(strings::StrCat(batch_prefix, "_msg"), + batches_[index].status.error_message())); + TF_RETURN_IF_ERROR(writer->WriteScalar( + strings::StrCat(batch_prefix, "_error_index"), + batches_[index].error_index)); + } + + TF_RETURN_IF_ERROR(writer->WriteScalar( + strings::StrCat(batch_prefix, "_output_size"), + batches_[index].outputs.size())); + for (size_t j = 0; j < batches_[index].outputs.size(); ++j) { + string tensor_prefix = + strings::StrCat(batch_prefix, "_output_", j); + if (!batches_[index].status.ok()) { + DCHECK(batches_[index].error_index >= 0 && + batches_[index].error_index < + itr_->dataset()->batch_size_); + // If the batch is not full, we only store the first + // `error_index` values. The rest of the batch tensor might not + // be initialized, and accessing that will raise msan errors. + TF_RETURN_IF_ERROR(writer->WriteTensor( + tensor_prefix, batches_[index].outputs[j].Slice( + 0, batches_[index].error_index))); + } else { + TF_RETURN_IF_ERROR(writer->WriteTensor( + tensor_prefix, batches_[index].outputs[j])); + } + } + } + return Status::OK(); + } + + Status Restore(IteratorContext* ctx, IteratorStateReader* reader, + Iterator* itr, size_t index) { + mutex_lock l(mu_); + if (reached_eof_) { + return errors::FailedPrecondition( + "Already reached the end of the sequence."); + } + string prefix = itr->full_name(strings::StrCat("numa_block_", index)); + reached_eof_ = + reader->Contains(strings::StrCat(prefix, "_end_of_input")); + for (size_t i = 0; i < kWindowSize; ++i) { + string batch_prefix = strings::StrCat(prefix, "_batch_", i); + if (!reader->Contains(strings::StrCat(batch_prefix, "_code"))) { + break; + } + Batch batch; + batch.state = BatchState::kOutputsComplete; + int64 code_int; + TF_RETURN_IF_ERROR(reader->ReadScalar( + strings::StrCat(batch_prefix, "_code"), &code_int)); + error::Code code = static_cast(code_int); + if (code != error::Code::OK) { + string error_message; + TF_RETURN_IF_ERROR(reader->ReadScalar( + strings::StrCat(batch_prefix, "_msg"), &error_message)); + batch.status = Status(code, error_message); + int64 error_index_int = -1; + TF_RETURN_IF_ERROR(reader->ReadScalar( + strings::StrCat(batch_prefix, "_error_index"), + &error_index_int)); + if (error_index_int < 0 || + error_index_int > itr->dataset()->batch_size_) { + return errors::FailedPrecondition( + "Error index out of bounds when restoring from checkpoint; " + "error index: ", + error_index_int); + } + batch.error_index = static_cast(error_index_int); + } + int64 output_size = -1; + TF_RETURN_IF_ERROR(reader->ReadScalar( + strings::StrCat(batch_prefix, "_output_size"), &output_size)); + batch.outputs.reserve(output_size); + for (size_t j = 0; j < output_size; ++j) { + string tensor_name = strings::StrCat(batch_prefix, "_output_", j); + Tensor t; + TF_RETURN_IF_ERROR(reader->ReadTensor(tensor_name, &t)); + batch.outputs.emplace_back(std::move(t)); + } + batches_[i] = std::move(batch); + } + return Status::OK(); + } + + private: + bool AllMapOperationsFinished() EXCLUSIVE_LOCKS_REQUIRED(mu_) { + for (size_t i = 0; i < kWindowSize; ++i) { + if (batches_[i].state == BatchState::kInputsFilled || + batches_[i].state == BatchState::kAllMapsStarted) { + return false; + } + if (batches_[i].state != BatchState::kOutputsComplete && + !reached_eof_) { + return false; + } + } + return true; + } + + // Batches begin in the `kEmpty` state. Once the RunnerThread has + // filled the `inputs` to a `Batch`, it transitions to the + // `kInputsFilled` state. At this point, the Worker threads run the map + // function and copy the outputs appropriately. Once all worker threads + // have started, it transitions to `kAllMapsStarted`. After the outputs + // are complete, the GetNext call can consume the outputs, and return + // the batch to the kEmpty state. + enum class BatchState { + kEmpty, + kInputsFilled, + kAllMapsStarted, + kOutputsComplete, + }; + + // Batch captures all the state of an output batch as it progresses + // through the machinery. Once the RunnerThread fills inputs, it + // transitions to `kInputsFilled`. At this point, the worker threads can + // work on it, incrementing outputs_complete for every element of the + // input set that is copied into the output Tensors. Once all the input + // tuples have been processed (i.e. num_outputs_complete == + // inputs.size()), it transitions to the `kOutputsComplete` stage, where + // it is ready to be returned by a `GetBatch` call (called from + // `GetNextInternal`). + struct Batch { + BatchState state; + // Aggregates the Status of the input iterator's GetNext + // calls, in addition to the Status of the map function invocations. + // + // In the case where multiple non-OK statuses are encountered, we + // return the first one encountered. + Status status; + // In order to return the correct error status, we keep track of the + // error_index. + size_t error_index; + // The batch_size input tuples (or fewer in the case of the last + // batch). + // TODO(saeta): Avoid re-allocating vectors all the time! + std::vector> inputs; + std::vector outputs; + size_t next_input_to_process; + size_t num_outputs_complete; + + Batch() { Reset(); } + + // Resets the Batch state (e.g. after consuming the outputs). + void Reset() { + state = BatchState::kEmpty; + status = Status::OK(); + inputs.clear(); + inputs.shrink_to_fit(); + outputs.clear(); + outputs.shrink_to_fit(); + next_input_to_process = 0; + num_outputs_complete = 0; + error_index = -1; + } + }; + + Iterator* itr_; // Not owned. + mutex mu_; + Batch batches_[kWindowSize] GUARDED_BY(mu_); + size_t next_input_batch_ GUARDED_BY(mu_) = -1; + size_t next_input_ GUARDED_BY(mu_) = 0; + size_t next_output_ GUARDED_BY(mu_) = 0; + bool cancelled_ GUARDED_BY(mu_) = false; + bool reached_eof_ GUARDED_BY(mu_) = false; + + // The runner thread waits on this condition variable for space to be + // available. When the client thread takes a value out of the circular + // buffer, it notifies this condition variable that space is now + // available. + condition_variable runner_cond_var_ GUARDED_BY(mu_); + // The worker threads wait on this condition variable for available + // inputs. When the runner thread makes new inputs available, it + // notifies this condition variable. + condition_variable worker_cond_var_ GUARDED_BY(mu_); + // The client threads wait on this condition variable for avaiable + // batched outputs. When worker threads complete a batch, they notify + // this condition variable. + condition_variable client_cond_var_ GUARDED_BY(mu_); + }; + // Mark NumaBlockManager as a friend of Iterator in order to call + // protected Iterator methods during checkpointing. + friend NumaBlockManager; + + struct NumaWorkerBlock { + NumaBlockManager manager; + // TODO(saeta): Migrate to BackgroundWorker. + std::vector> threads; + + explicit NumaWorkerBlock(Iterator* itr) : manager(itr) {} + }; + + static void CustomNumaWorkerBlockDeleter(NumaWorkerBlock* ptr) { + ptr->~NumaWorkerBlock(); + port::NUMAFree(ptr, sizeof(NumaWorkerBlock)); + } + static void DefaultNumaWorkerBlockDeleter(NumaWorkerBlock* ptr) { + delete ptr; + } + + static Status CopyPartialBatch(Tensor* output, const Tensor& value, + int64 num_elements) { + switch (value.dtype()) { +#define HANDLE_TYPE(type) \ + case DataTypeToEnum::value: { \ + auto output_t = output->flat_outer_dims(); \ + auto value_t = value.flat_outer_dims(); \ + for (size_t i = 0; i < num_elements; i++) { \ + output_t.template chip<0>(i) = value_t.template chip<0>(i); \ + } \ + return Status::OK(); \ + } + TF_CALL_DATASET_TYPES(HANDLE_TYPE); +#undef HANDLE_TYPE + default: + return errors::InvalidArgument("Unsupported data type: ", + DataTypeString(value.dtype())); + } + return Status::OK(); + } + + Status EnsureBackgroundThreadsStarted(IteratorContext* ctx) + EXCLUSIVE_LOCKS_REQUIRED(*mu_) { + if (curr_num_parallel_calls_ >= num_parallel_calls_->value) { + // All necessary threads have been started. + curr_num_parallel_calls_ = num_parallel_calls_->value; + return Status::OK(); + } + + VLOG(4) << "Starting workers"; + bool numa_enabled = port::NUMAEnabled(); + + if (!numa_enabled) { + LOG(INFO) << "NUMA not enabled on this host."; + } + + int num_numa_nodes = port::NUMANumNodes(); + if (num_numa_nodes < 1) { + return errors::Internal("The number of NUMA nodes is invalid: ", + num_numa_nodes); + } + + // Only resize when empty to support restoring from checkpoints. + if (workers_.empty()) { + VLOG(3) << "# NUMA Nodes: " << num_numa_nodes + << ", # Parallel Calls: " << num_parallel_calls_->value; + workers_.resize(num_numa_nodes); + } else { + num_numa_nodes = workers_.size(); + } + + // Round up num_parallel_calls, with a minimum of 1. + const size_t num_threads_per_block = + std::max(1LL, (num_parallel_calls_->value + num_numa_nodes - 1) / + num_numa_nodes); + + VLOG(3) << "Starting " << num_threads_per_block * num_numa_nodes + << " worker threads, with " << num_threads_per_block + << " threads per block."; + + // Only allocate new_ctx if required. + std::shared_ptr new_ctx; + + for (int i = 0; i < num_numa_nodes; ++i) { + if (!workers_[i]) { + if (numa_enabled) { + // Allocate in appropriate NUMA domain. + // 4k page align. + void* ptr = port::NUMAMalloc(i, sizeof(NumaWorkerBlock), 0); + if (ptr != nullptr) { + NumaWorkerBlock* block = new (ptr) NumaWorkerBlock(this); + workers_[i] = + std::unique_ptr>( + block, CustomNumaWorkerBlockDeleter); + } else { + LOG(ERROR) << "Could not NUMA-allocate worker block: " << i; + } + } + // If the NUMA allocation fails, or NUMA is not enabled. + if (!workers_[i]) { + workers_[i] = + std::unique_ptr>( + new NumaWorkerBlock(this), DefaultNumaWorkerBlockDeleter); + } + } + // Be sure to start threads if num_parallel_calls_ has changed. + for (size_t j = workers_[i]->threads.size(); + j < num_threads_per_block; ++j) { + VLOG(3) << "Starting worker " << i << ", " << j; + if (!new_ctx) { + new_ctx = std::make_shared(*ctx); + } + workers_[i]->threads.emplace_back(ctx->env()->StartThread( + {}, + strings::StrCat("numa_map_and_batch_block_", i, "_thread_", j), + [this, new_ctx, i, j]() { WorkerThread(new_ctx, i, j); })); + VLOG(3) << "Worker " << i << ", " << j << " successfully started."; + } + } + if (!runner_thread_) { + if (!new_ctx) { + new_ctx = std::make_shared(*ctx); + } + runner_thread_.reset(ctx->env()->StartThread( + {}, "numa_map_runner_thread", + [this, new_ctx] { RunnerThread(new_ctx); })); + } + VLOG(3) << "All workers & runner thread started."; + return Status::OK(); + } + + void AllocateOutput(IteratorContext* ctx, size_t batch_size, + const std::vector& map_fn_outputs, + std::vector* batch_outputs) { + DCHECK(dataset()->output_dtypes().size() == + dataset()->output_shapes().size()); + DCHECK(map_fn_outputs.size() == dataset()->output_dtypes().size()); + for (size_t i = 0; i < dataset()->output_dtypes().size(); ++i) { + TensorShape component_shape({static_cast(batch_size)}); + component_shape.AppendShape(map_fn_outputs.at(i).shape()); + AllocatorAttributes attr; + attr.set_gpu_compatible(true); + Tensor component(ctx->allocator(attr), map_fn_outputs.at(i).dtype(), + component_shape); + batch_outputs->emplace_back(std::move(component)); + } + } + + void RunnerThread(std::shared_ptr ctx) + LOCKS_EXCLUDED(mu_) { + RecordStart(ctx.get()); + auto cleanup = gtl::MakeCleanup([this, &ctx] { + // Set end of input on all the managers in order to clean up in an + // orderly fashion. + VLOG(3) << "Setting End of Input on workers_[*]->manager"; + for (size_t i = 0; i < workers_.size(); ++i) { + workers_[i]->manager.SetEndOfInput(); + } + RecordStop(ctx.get()); + }); + + const size_t num_blocks = workers_.size(); + + while (true) { + for (size_t block = 0; block < num_blocks; ++block) { + VLOG(4) << "RunnerThread waiting for input space in block: " + << block; + if (TF_PREDICT_FALSE( + !workers_[block]->manager.WaitForInputSpace(ctx.get()))) { + VLOG(3) << "RunnerThread exiting due to cancellation."; + return; + } + VLOG(4) << "RunnerThread has space; pulling on upstream for block " + << block; + + Status s; + std::vector> inputs; + bool end_of_sequence = false; + for (size_t i = 0; i < dataset()->batch_size_; ++i) { + std::vector tuple; + s.Update( + input_impl_->GetNext(ctx.get(), &tuple, &end_of_sequence)); + if (!s.ok()) { + break; + } + if (end_of_sequence) { + VLOG(4) << "Runner thread encountered end of sequence."; + if (dataset()->drop_remainder_) { + return; + } + break; + } + inputs.push_back(std::move(tuple)); + } + + VLOG(4) << "Moving inputs to block " << block + << ", which has size: " << inputs.size(); + if (!s.ok() || !inputs.empty()) { + workers_[block]->manager.PushInputs(s, std::move(inputs)); + VLOG(4) << "Inputs moved into block " << block; + } + if (end_of_sequence) { + return; + } + } + } + } + + void WorkerThread(std::shared_ptr ctx, + const int numa_node, const int thread_num) { + RecordStart(ctx.get()); + WORKER_VLOG(3) << "started."; + auto stop_cleanup = + gtl::MakeCleanup([this, numa_node, thread_num, &ctx]() { + RecordStop(ctx.get()); + WORKER_VLOG(3) << "exiting."; + }); + + NumaWorkerBlock* block = workers_[numa_node].get(); + port::NUMASetThreadNodeAffinity(numa_node); + const int num_numa_nodes = port::NUMANumNodes(); + const int minimum_num_parallel_calls = thread_num * num_numa_nodes; + + while (true) { + // Put threads to sleep based on autotuner. + { + mutex_lock l(*mu_); + while (minimum_num_parallel_calls >= num_parallel_calls_->value && + !cancelled_) { + RecordStop(ctx.get()); + autotune_cond_var_->wait(l); + RecordStart(ctx.get()); + } + if (cancelled_) { + return; + } + } + + std::vector input; + uint64 index = 0; + size_t sequence_number = 0; + WORKER_VLOG(4) << "retrieving input."; + { + tracing::ScopedActivity trace( + "NumaMapAndBatch::Iterator::Worker::RetrieveInput"); + if (!block->manager.RetrieveInput(ctx.get(), &input, &index, + &sequence_number)) { + return; + } + } + + WORKER_VLOG(4) << "retrieved input; index: " << index + << ", sequence_number: " << sequence_number; + + std::vector return_values; + Status s; + { + tracing::ScopedActivity trace( + "NumaMapAndBatch::Iterator::Worker::FunctionExecution"); + s = dataset()->captured_func_->Run(ctx.get(), std::move(input), + &return_values); + } + WORKER_VLOG(4) << "ran function for index: " << index + << ", sequence_number: " << sequence_number; + + if (s.ok()) { + std::vector* output = block->manager.GetBatchTensors( + sequence_number, + [this, ctx, &return_values](size_t batch_size, + std::vector* output) { + AllocateOutput(ctx.get(), batch_size, return_values, output); + }); + WORKER_VLOG(4) << "copying tensors to batch output."; + { + tracing::ScopedActivity trace( + "NumaMapAndBatch::Iterator::Worker::BatchCopy"); + for (size_t i = 0; i < return_values.size() && s.ok(); ++i) { + Tensor& tensor = return_values.at(i); + Tensor* batch = &output->at(i); + if (tensor.NumElements() != + (batch->NumElements() / batch->dim_size(0))) { + s.Update(errors::InvalidArgument( + "Cannot add tensor to the batch: number of elements does " + "not match. Shapes are: [tensor]: ", + tensor.shape().DebugString(), + ", [batch]: ", batch->shape().DebugString())); + break; + } + s.Update(batch_util::CopyElementToSlice(std::move(tensor), + batch, index)); + } + } + } + + block->manager.RecordBatchEntryComplete(sequence_number, index, s); + WORKER_VLOG(4) << "finished index: " << index + << ", sequence_number: " << sequence_number; + } + } + + // mu_ protects shared internal state and is used to coordinate between + // the auto-tuner, client threads, worker threads, and the runner thread. + const std::shared_ptr mu_; + const std::shared_ptr autotune_cond_var_; + // The maximum number of parallel calls (can be auto-tuned). + const std::shared_ptr num_parallel_calls_; + + // Caches the last-seen value of num_parallel_calls_->value to + // short-circuit starting workers. + int64 curr_num_parallel_calls_ GUARDED_BY(*mu_) = 0; + + std::unique_ptr input_impl_; + int64 cur_block_ GUARDED_BY(*mu_) = 0; + bool global_end_of_input_ GUARDED_BY(*mu_) = false; + bool cancelled_ GUARDED_BY(*mu_) = false; + std::vector>> + workers_; // Const after initialization. + std::unique_ptr runner_thread_ GUARDED_BY(*mu_); + }; + + const DatasetBase* const input_; + const int64 batch_size_; + const int64 num_parallel_calls_; + const bool drop_remainder_; + const DataTypeVector output_types_; + const std::vector output_shapes_; + const NameAttrList func_; + const std::unique_ptr captured_func_; + }; + + DataTypeVector output_types_; + std::vector output_shapes_; + NameAttrList func_; +}; + +REGISTER_KERNEL_BUILDER( + Name("ExperimentalNumaMapAndBatchDataset").Device(DEVICE_CPU), + NumaMapAndBatchDatasetOp); + +} // namespace +} // namespace data +} // namespace tensorflow diff --git a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc index f45a239793..bae56828dc 100644 --- a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc +++ b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc @@ -324,6 +324,8 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { } private: + // BatchResult encapsulates the output batch, as well as anciliary + // metadata required to execute the fused map-and-batch operation. struct BatchResult { explicit BatchResult(int64 batch_size) { end_of_input = false; @@ -331,11 +333,23 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { num_elements = 0; output_allocated = false; status = Status::OK(); + status_offset = -1; } - void UpdateStatus(const Status& s) { - mutex_lock l(mu); - status.Update(s); + // UpdateStatus updates the batch's aggregate Status. + // + // In order to ensure that exactly the first non-OK status is returned + // (required to make the behavior is observably identical to a + // sequential execution of map followed by batch), we must also keep + // track of the offset into the batch that produced `s`. + void UpdateStatus(const Status& s, int64 offset) { + if (TF_PREDICT_FALSE(!s.ok())) { + mutex_lock l(mu); + if (status.ok() || offset < status_offset) { + status = s; + status_offset = offset; + } + } } mutex mu; @@ -344,6 +358,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { std::vector output; bool output_allocated GUARDED_BY(mu); Status status GUARDED_BY(mu); + int64 status_offset GUARDED_BY(mu); // Counts the number of outstanding calls for this batch. int64 num_calls; // access guarded by owner's mutex }; @@ -379,7 +394,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { std::shared_ptr> return_values = std::make_shared>(); auto done = [this, ctx, result, return_values, offset](Status status) { - result->UpdateStatus(status); + result->UpdateStatus(status, offset); if (status.ok()) { EnsureOutputAllocated(ctx, result, return_values); for (size_t i = 0; i < return_values->size(); ++i) { @@ -389,11 +404,14 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { (batch->NumElements() / batch->dim_size(0))) { TensorShape batch_shape = batch->shape(); batch_shape.RemoveDim(0); - result->UpdateStatus(errors::InvalidArgument( - "Cannot add tensor to the batch: number of elements does " - "not match. Shapes are: [tensor]: ", - tensor.shape().DebugString(), - ", [batch]: ", batch_shape.DebugString())); + result->UpdateStatus( + errors::InvalidArgument( + "Cannot add tensor to the batch: number of elements " + "does " + "not match. Shapes are: [tensor]: ", + tensor.shape().DebugString(), + ", [batch]: ", batch_shape.DebugString()), + offset); break; } // TODO(mrry): Add a version of DoParallelConcat that allows us to @@ -402,7 +420,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { Status copy_status = ::tensorflow::functor::DoParallelConcat( *dataset()->device_, tensor, offset, batch); if (!copy_status.ok()) { - result->UpdateStatus(copy_status); + result->UpdateStatus(copy_status, offset); break; } } diff --git a/tensorflow/core/ops/experimental_dataset_ops.cc b/tensorflow/core/ops/experimental_dataset_ops.cc index f6bd5dce26..bbbecc50f8 100644 --- a/tensorflow/core/ops/experimental_dataset_ops.cc +++ b/tensorflow/core/ops/experimental_dataset_ops.cc @@ -138,6 +138,32 @@ REGISTER_OP("ExperimentalAssertNextDataset") return shape_inference::ScalarShape(c); }); +REGISTER_OP("ExperimentalNumaMapAndBatchDataset") + .Input("input_dataset: variant") + .Input("other_arguments: Targuments") + .Input("batch_size: int64") + .Input("num_parallel_calls: int64") + .Input("drop_remainder: bool") + .Output("handle: variant") + .Attr("f: func") + .Attr("Targuments: list(type) >= 0") + .Attr("output_types: list(type) >= 1") + .Attr("output_shapes: list(shape) >= 1") + .SetShapeFn([](shape_inference::InferenceContext* c) { + // Use index from the end to retrieve the Input shapes, + // so that to avoid guessing the length of "other_arguments". + // batch_size, num_parallel_batches, and drop_remainder are 0-D scalars. + shape_inference::ShapeHandle unused; + TF_RETURN_IF_ERROR( + c->WithRank(c->input(c->num_inputs() - 3), 0, &unused)); + TF_RETURN_IF_ERROR( + c->WithRank(c->input(c->num_inputs() - 2), 0, &unused)); + TF_RETURN_IF_ERROR( + c->WithRank(c->input(c->num_inputs() - 1), 0, &unused)); + + return shape_inference::ScalarShape(c); + }); + REGISTER_OP("ExperimentalLMDBDataset") .Input("filenames: string") .Output("handle: variant") diff --git a/tensorflow/python/data/experimental/kernel_tests/map_and_batch_test.py b/tensorflow/python/data/experimental/kernel_tests/map_and_batch_test.py index d444c4082e..5ead6d1c75 100644 --- a/tensorflow/python/data/experimental/kernel_tests/map_and_batch_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/map_and_batch_test.py @@ -30,6 +30,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import sparse_tensor from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import script_ops from tensorflow.python.platform import test @@ -38,12 +39,17 @@ from tensorflow.python.platform import test class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): @parameterized.named_parameters( - ("Default", None, None), - ("SequentialCalls", 1, None), - ("ParallelCalls", 2, None), - ("ParallelBatches", None, 10), + ("Default", None, None, False), + ("SequentialCalls", 1, None, False), + ("ParallelCalls", 2, None, False), + ("ParallelBatches", None, 10, False), + ("DefaultNUMA", None, None, True), + ("SequentialCallsNUMA", 1, None, True), + ("ParallelCallsNUMA", 2, None, True), + ("ParallelBatchesNUMA", None, 10, True), ) - def testMapAndBatch(self, num_parallel_calls, num_parallel_batches): + def testMapAndBatch(self, num_parallel_calls, num_parallel_batches, + numa_aware): """Test a dataset that maps a TF function across its input elements.""" # The pipeline is TensorSliceDataset -> # RepeatDataset(count) -> MapAndBatchDataset(square_3, batch_size). @@ -57,14 +63,20 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): def _map_fn(x, y, z): return math_ops.square(x), math_ops.square(y), math_ops.square(z) - iterator = ( + dataset = ( dataset_ops.Dataset.from_tensor_slices(components).repeat(count).apply( batching.map_and_batch( map_func=_map_fn, batch_size=batch_size, num_parallel_calls=num_parallel_calls, - num_parallel_batches=num_parallel_batches)) - .make_initializable_iterator()) + num_parallel_batches=num_parallel_batches))) + + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + + iterator = dataset.make_initializable_iterator() init_op = iterator.initializer get_next = iterator.get_next() @@ -115,16 +127,25 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): sess.run(init_op, feed_dict={count: 14, batch_size: 0}) @parameterized.named_parameters( - ("Even", False), - ("Uneven", True), + ("Even", False, False), + ("Uneven", True, False), + ("EvenNUMA", False, True), + ("UnevenNUMA", True, True), ) - def testMapAndBatchPartialBatch(self, drop_remainder): - iterator = ( + def testMapAndBatchPartialBatch(self, drop_remainder, numa_aware): + dataset = ( dataset_ops.Dataset.range(10).apply( batching.map_and_batch( lambda x: array_ops.reshape(x * x, [1]), batch_size=4, - drop_remainder=drop_remainder)).make_one_shot_iterator()) + drop_remainder=drop_remainder))) + + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + iterator = dataset.make_one_shot_iterator() + if drop_remainder: self.assertEqual([4, 1], iterator.output_shapes.as_list()) else: @@ -138,11 +159,21 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(next_element) - def testMapAndBatchYieldsPartialBatch(self): - iterator = (dataset_ops.Dataset.range(10) - .apply(batching.map_and_batch( - lambda x: array_ops.reshape(x * x, [1]), 4)) - .make_one_shot_iterator()) + @parameterized.named_parameters( + ("Normal", False), + ("NUMA", True), + ) + def testMapAndBatchYieldsPartialBatch(self, numa_aware): + dataset = ( + dataset_ops.Dataset.range(10).apply( + batching.map_and_batch(lambda x: array_ops.reshape(x * x, [1]), 4))) + + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + + iterator = dataset.make_one_shot_iterator() self.assertEqual([None, 1], iterator.output_shapes.as_list()) next_element = iterator.get_next() with self.cached_session() as sess: @@ -152,10 +183,19 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(next_element) - def testMapAndBatchParallelGetNext(self): - iterator = (dataset_ops.Dataset.range(50000) - .apply(batching.map_and_batch(lambda x: x, batch_size=100)) - .make_one_shot_iterator()) + @parameterized.named_parameters( + ("Normal", False), + ("NUMA", True), + ) + def testMapAndBatchParallelGetNext(self, numa_aware): + dataset = dataset_ops.Dataset.range(50000).apply( + batching.map_and_batch(lambda x: x, batch_size=100)) + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + iterator = dataset.make_one_shot_iterator() + elements = [] for _ in range(100): elements.append(iterator.get_next()) @@ -165,17 +205,26 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): got.sort(key=lambda x: x[0]) expected = [] for j in range(100): - expected.append(range(i*10000+j*100, i*10000+(j+1)*100)) + expected.append(range(i * 10000 + j * 100, i * 10000 + (j + 1) * 100)) self.assertAllEqual(got, expected) with self.assertRaises(errors.OutOfRangeError): sess.run(elements) - def testMapAndBatchParallelGetNextDropRemainder(self): - iterator = ( - dataset_ops.Dataset.range(49999).apply( - batching.map_and_batch( - lambda x: x, batch_size=100, drop_remainder=True)) - .make_one_shot_iterator()) + @parameterized.named_parameters( + ("Normal", False), + ("NUMA", True), + ) + def testMapAndBatchParallelGetNextDropRemainder(self, numa_aware): + dataset = dataset_ops.Dataset.range(49999).apply( + batching.map_and_batch( + lambda x: x, batch_size=100, drop_remainder=True)) + + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + iterator = dataset.make_one_shot_iterator() + elements = [] for _ in range(100): elements.append(iterator.get_next()) @@ -185,19 +234,29 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): got.sort(key=lambda x: x[0]) expected = [] for j in range(100): - expected.append(range(i*10000+j*100, i*10000+(j+1)*100)) + expected.append(range(i * 10000 + j * 100, i * 10000 + (j + 1) * 100)) self.assertAllEqual(got, expected) with self.assertRaises(errors.OutOfRangeError): sess.run(elements) - def testMapAndBatchSparse(self): + @parameterized.named_parameters( + ("Normal", False), + ("NUMA", True), + ) + def testMapAndBatchSparse(self, numa_aware): def _sparse(i): return sparse_tensor.SparseTensorValue( indices=[[0]], values=(i * [1]), dense_shape=[1]) - iterator = dataset_ops.Dataset.range(10).apply( - batching.map_and_batch(_sparse, 5)).make_initializable_iterator() + dataset = dataset_ops.Dataset.range(10).apply( + batching.map_and_batch(_sparse, 5)) + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + iterator = dataset.make_initializable_iterator() + init_op = iterator.initializer get_next = iterator.get_next() @@ -214,21 +273,33 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - def testMapAndBatchFails(self): + @parameterized.named_parameters( + ("Normal", False), + ("NUMA", True), + ) + def testMapAndBatchFails(self, numa_aware): """Test a dataset that maps a TF function across its input elements.""" dataset = dataset_ops.Dataset.from_tensors( array_ops.check_numerics( constant_op.constant(1.0) / constant_op.constant(0.0), "oops")) batch_size = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = ( - dataset.apply(batching.map_and_batch(lambda x: x, batch_size)) - .make_initializable_iterator()) + dataset = dataset.apply(batching.map_and_batch(lambda x: x, batch_size)) + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + iterator = dataset.make_initializable_iterator() + init_op = iterator.initializer with self.cached_session() as sess: with self.assertRaisesRegexp(errors.InvalidArgumentError, "oops"): sess.run(init_op, feed_dict={batch_size: 14}) - def testMapAndBatchShapeMismatch(self): + @parameterized.named_parameters( + ("Normal", False), + ("NUMA", True), + ) + def testMapAndBatchShapeMismatch(self, numa_aware): """Test a dataset that maps a TF function across its input elements.""" def generator(): @@ -240,9 +311,13 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): dataset = dataset_ops.Dataset.from_generator( generator, output_types=dtypes.int32) batch_size = 4 - iterator = ( - dataset.apply(batching.map_and_batch(lambda x: x, batch_size)) - .make_initializable_iterator()) + dataset = dataset.apply(batching.map_and_batch(lambda x: x, batch_size)) + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + iterator = dataset.make_initializable_iterator() + init_op = iterator.initializer get_next = iterator.get_next() with self.cached_session() as sess: @@ -251,7 +326,11 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): "number of elements does not match"): sess.run(get_next) - def testMapAndBatchImplicitDispose(self): + @parameterized.named_parameters( + ("Normal", False), + ("NUMA", True), + ) + def testMapAndBatchImplicitDispose(self, numa_aware): # Tests whether a map and batch dataset will be cleaned up correctly when # the pipeline does not run it until exhaustion. # The pipeline is TensorSliceDataset -> RepeatDataset(1000) -> @@ -266,6 +345,10 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): dataset = dataset_ops.Dataset.from_tensor_slices(components).repeat( 1000).apply(batching.map_and_batch(_map_fn, batch_size=100)) dataset = dataset.prefetch(5) + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) iterator = dataset.make_one_shot_iterator() get_next = iterator.get_next() @@ -274,26 +357,38 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): sess.run(get_next) @parameterized.named_parameters( - ("1", 0), - ("2", 5), - ("3", 10), - ("4", 90), - ("5", 95), - ("6", 99), + ("1", 0, False), + ("2", 5, False), + ("3", 10, False), + ("4", 90, False), + ("5", 95, False), + ("6", 99, False), + ("1NUMA", 0, True), + ("2NUMA", 5, True), + ("3NUMA", 10, True), + ("4NUMA", 90, True), + ("5NUMA", 95, True), + ("6NUMA", 99, True), ) - def testMapAndBatchOutOfRangeError(self, threshold): + def testMapAndBatchOutOfRangeError(self, threshold, numa_aware): def raising_py_fn(i): - if i >= threshold: + if i == threshold: raise StopIteration() + elif i > threshold: + raise RuntimeError("Alternate error; you shouldn't see me! (i: %s)" % i) else: return i - iterator = ( - dataset_ops.Dataset.range(100).apply( - batching.map_and_batch( - lambda x: script_ops.py_func(raising_py_fn, [x], dtypes.int64), - batch_size=10)).make_one_shot_iterator()) + dataset = dataset_ops.Dataset.range(100).apply( + batching.map_and_batch( + lambda x: script_ops.py_func(raising_py_fn, [x], dtypes.int64), + batch_size=10)) + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + iterator = dataset.make_one_shot_iterator() get_next = iterator.get_next() with self.cached_session() as sess: @@ -307,25 +402,42 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): sess.run(get_next) @parameterized.named_parameters( - ("1", False, dtypes.bool), - ("2", -42, dtypes.int8), - ("3", -42, dtypes.int16), - ("4", -42, dtypes.int32), - ("5", -42, dtypes.int64), - ("6", 42, dtypes.uint8), - ("7", 42, dtypes.uint16), - ("8", 42.0, dtypes.float16), - ("9", 42.0, dtypes.float32), - ("10", 42.0, dtypes.float64), - ("11", b"hello", dtypes.string), + ("1", False, dtypes.bool, False), + ("2", -42, dtypes.int8, False), + ("3", -42, dtypes.int16, False), + ("4", -42, dtypes.int32, False), + ("5", -42, dtypes.int64, False), + ("6", 42, dtypes.uint8, False), + ("7", 42, dtypes.uint16, False), + ("8", 42.0, dtypes.float16, False), + ("9", 42.0, dtypes.float32, False), + ("10", 42.0, dtypes.float64, False), + ("11", b"hello", dtypes.string, False), + ("1NUMA", False, dtypes.bool, True), + ("2NUMA", -42, dtypes.int8, True), + ("3NUMA", -42, dtypes.int16, True), + ("4NUMA", -42, dtypes.int32, True), + ("5NUMA", -42, dtypes.int64, True), + ("6NUMA", 42, dtypes.uint8, True), + ("7NUMA", 42, dtypes.uint16, True), + ("8NUMA", 42.0, dtypes.float16, True), + ("9NUMA", 42.0, dtypes.float32, True), + ("10NUMA", 42.0, dtypes.float64, True), + ("11NUMA", b"hello", dtypes.string, True), ) - def testMapAndBatchTypes(self, element, dtype): + def testMapAndBatchTypes(self, element, dtype, numa_aware): + def gen(): yield element dataset = dataset_ops.Dataset.from_generator(gen, dtype).repeat(100).apply( batching.map_and_batch(lambda x: x, batch_size=10)) + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + get_next = dataset.make_one_shot_iterator().get_next() with self.cached_session() as sess: @@ -363,6 +475,40 @@ class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase): sess.run(iterator.initializer, feed_dict={captured_t: 42}) self.assertAllEqual([42] * 10, sess.run(get_next)) + @parameterized.named_parameters( + ("Normal", False), + ("NUMA", True), + ) + def testMapAndBatchControlFlow(self, numa_aware): + + def map_fn(x): + previous_cond_v2_value = control_flow_ops.ENABLE_COND_V2 + control_flow_ops.ENABLE_COND_V2 = True + return_value = control_flow_ops.cond(x < 50, lambda: x + 1, lambda: x * x) + control_flow_ops.ENABLE_COND_V2 = previous_cond_v2_value + return return_value + + dataset = dataset_ops.Dataset.range(100).apply( + batching.map_and_batch(map_fn, batch_size=10)) + if numa_aware: + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + iterator = dataset.make_one_shot_iterator() + get_next = iterator.get_next() + with self.cached_session() as sess: + for i in range(10): + print("Case %d" % i) + if i < 5: + self.assertAllEqual([i * 10 + j + 1 for j in range(10)], + sess.run(get_next)) + else: + self.assertAllEqual( + [((i * 10) + j) * ((i * 10) + j) for j in range(10)], + sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD b/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD index c92bb8b9bc..5a0a73fd83 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD @@ -161,6 +161,7 @@ py_test( "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", "//third_party/py/numpy", + "@absl_py//absl/testing:parameterized", ], ) @@ -199,6 +200,7 @@ py_test( deps = [ "//tensorflow/python:client_testlib", "//tensorflow/python:errors", + "//tensorflow/python/data/experimental/ops:batching", "//tensorflow/python/data/experimental/ops:optimization", "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/model_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/model_dataset_op_test.py index 82516356df..d38255a6ea 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/model_dataset_op_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/model_dataset_op_test.py @@ -19,6 +19,7 @@ from __future__ import print_function import time +from absl.testing import parameterized import numpy as np from tensorflow.python.data.experimental.ops import batching @@ -29,7 +30,7 @@ from tensorflow.python.ops import math_ops from tensorflow.python.platform import test -class ModelDatasetTest(test_base.DatasetTestBase): +class ModelDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): def testModelMap(self): k = 1024 * 1024 @@ -82,7 +83,11 @@ class ModelDatasetTest(test_base.DatasetTestBase): (np.median(deltas), np.mean(deltas), np.std(deltas), np.min(deltas), np.max(deltas))) - def testModelMapAndBatch(self): + @parameterized.named_parameters( + ("Default", False), + ("NUMA", True), + ) + def testModelMapAndBatch(self, numa_aware): batch_size = 16 k = 1024 * 1024 dataset = dataset_ops.Dataset.from_tensors((np.random.rand(1, 4 * k), @@ -95,6 +100,8 @@ class ModelDatasetTest(test_base.DatasetTestBase): batch_size=batch_size)) options = dataset_ops.Options() options.experimental_autotune = True + if numa_aware: + options.experimental_numa_aware = True iterator = dataset.with_options(options).make_one_shot_iterator() get_next = iterator.get_next() diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_op_test.py index 760cd8cc4e..2ef29796ab 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_op_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_op_test.py @@ -19,6 +19,7 @@ from __future__ import print_function import numpy as np +from tensorflow.python.data.experimental.ops import batching from tensorflow.python.data.experimental.ops import optimization from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops @@ -59,6 +60,21 @@ class OptimizeDatasetTest(test_base.DatasetTestBase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) + def testNumaAwareRewrite(self): + dataset = dataset_ops.Dataset.range(10).apply( + optimization.assert_next(["NumaMapAndBatch"])).apply( + batching.map_and_batch(lambda x: x * x, 10)) + options = dataset_ops.Options() + options.experimental_numa_aware = True + dataset = dataset.with_options(options) + iterator = dataset.make_one_shot_iterator() + get_next = iterator.get_next() + + with self.cached_session() as sess: + self.assertAllEqual([x * x for x in range(10)], sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + def testOptimizationStatefulFunction(self): dataset = dataset_ops.Dataset.range(10).map( lambda _: random_ops.random_uniform([])).batch(10) diff --git a/tensorflow/python/data/experimental/kernel_tests/serialization/BUILD b/tensorflow/python/data/experimental/kernel_tests/serialization/BUILD index e556b65b7c..a97cff9fbb 100644 --- a/tensorflow/python/data/experimental/kernel_tests/serialization/BUILD +++ b/tensorflow/python/data/experimental/kernel_tests/serialization/BUILD @@ -306,6 +306,21 @@ py_test( ], ) +py_test( + name = "numa_map_and_batch_dataset_serialization_test", + size = "medium", + srcs = ["numa_map_and_batch_dataset_serialization_test.py"], + srcs_version = "PY2AND3", + tags = ["no_pip"], + deps = [ + ":dataset_serialization_test_base", + "//tensorflow/python:client_testlib", + "//tensorflow/python:math_ops", + "//tensorflow/python/data/experimental/ops:batching", + "//tensorflow/python/data/ops:dataset_ops", + ], +) + py_test( name = "map_dataset_serialization_test", size = "medium", diff --git a/tensorflow/python/data/experimental/kernel_tests/serialization/numa_map_and_batch_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/numa_map_and_batch_dataset_serialization_test.py new file mode 100644 index 0000000000..04aab329cd --- /dev/null +++ b/tensorflow/python/data/experimental/kernel_tests/serialization/numa_map_and_batch_dataset_serialization_test.py @@ -0,0 +1,95 @@ +# 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. +# ============================================================================== +"""Tests for the MapAndBatchDataset serialization.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import math + +from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base +from tensorflow.python.data.experimental.ops import batching +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import test + + +class MapAndBatchDatasetSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): + + def testNumParallelBatches(self): + range_size = 11 + num_repeats = 2 + batch_size = 5 + total_outputs = range_size * num_repeats + num_outputs_drop_remainder = total_outputs // batch_size + num_outputs_keep_remainder = int(math.ceil(total_outputs / batch_size)) + num_parallel_batches = 2 + + def build_ds(range_start, drop_remainder=False): + + def _map_fn(x): + return math_ops.square(x) + + ds = dataset_ops.Dataset.range( + range_start, range_start + range_size).repeat(num_repeats).apply( + batching.map_and_batch( + map_func=_map_fn, + batch_size=batch_size, + num_parallel_batches=num_parallel_batches, + drop_remainder=drop_remainder)) + options = dataset_ops.Options() + options.experimental_numa_aware = True + return ds.with_options(options) + + self.run_core_tests(lambda: build_ds(10), lambda: build_ds(15), + num_outputs_keep_remainder) + self.run_core_tests(lambda: build_ds(10, True), lambda: build_ds(15, True), + num_outputs_drop_remainder) + + def testNumParallelCalls(self): + range_size = 11 + num_repeats = 2 + batch_size = 5 + total_outputs = range_size * num_repeats + num_outputs_drop_remainder = total_outputs // batch_size + num_outputs_keep_remainder = int(math.ceil(total_outputs / batch_size)) + num_parallel_calls = 7 + + def build_ds(range_start, drop_remainder=False): + + def _map_fn(x): + return math_ops.square(x) + + ds = dataset_ops.Dataset.range( + range_start, range_start + range_size).repeat(num_repeats).apply( + batching.map_and_batch( + map_func=_map_fn, + batch_size=batch_size, + num_parallel_calls=num_parallel_calls, + drop_remainder=drop_remainder)) + options = dataset_ops.Options() + options.experimental_numa_aware = True + return ds.with_options(options) + + self.run_core_tests(lambda: build_ds(10), lambda: build_ds(15), + num_outputs_keep_remainder) + self.run_core_tests(lambda: build_ds(10, True), lambda: build_ds(15, True), + num_outputs_drop_remainder) + + +if __name__ == "__main__": + test.main() + diff --git a/tensorflow/python/data/experimental/ops/BUILD b/tensorflow/python/data/experimental/ops/BUILD index 915d399f1b..46a9552b61 100644 --- a/tensorflow/python/data/experimental/ops/BUILD +++ b/tensorflow/python/data/experimental/ops/BUILD @@ -122,6 +122,7 @@ py_library( "//tensorflow/python:array_ops", "//tensorflow/python:dataset_ops_gen", "//tensorflow/python:dtypes", + "//tensorflow/python:experimental_dataset_ops_gen", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", "//tensorflow/python:tensor_shape", diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py index cf52f7529a..6195747671 100644 --- a/tensorflow/python/data/ops/dataset_ops.py +++ b/tensorflow/python/data/ops/dataset_ops.py @@ -1410,6 +1410,8 @@ class Options(object): "Whether to eliminate no-op transformations."), ("experimental_shuffle_and_repeat_fusion", bool, "Whether to fuse shuffle and repeat transformations."), + ("experimental_numa_aware", bool, + "Whether to use NUMA-aware operations."), ]: def _make_getter(name): # pylint: disable=no-self-argument @@ -1458,6 +1460,9 @@ class Options(object): for exp_opt in experimental_optimizations: if getattr(self, "experimental_" + exp_opt): result.append(exp_opt) + + if getattr(self, "experimental_numa_aware"): + result.append("map_and_batch_numa_aware_replacement") return result def merge(self, options): @@ -1485,7 +1490,7 @@ class Options(object): "experimental_map_and_filter_fusion", "experimental_map_fusion", "experimental_map_parallelization", "experimental_map_vectorization", "experimental_noop_elimination", - "experimental_shuffle_and_repeat_fusion" + "experimental_shuffle_and_repeat_fusion", "experimental_numa_aware", ]: this = getattr(result, name) that = getattr(other, name) diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.-options.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.-options.pbtxt index d15dccc173..22256996d3 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.data.-options.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.data.-options.pbtxt @@ -42,6 +42,10 @@ tf_class { name: "experimental_noop_elimination" mtype: "" } + member { + name: "experimental_numa_aware" + mtype: "" + } member { name: "experimental_shuffle_and_repeat_fusion" mtype: "" diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-options.pbtxt index d15dccc173..22256996d3 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.data.-options.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.data.-options.pbtxt @@ -42,6 +42,10 @@ tf_class { name: "experimental_noop_elimination" mtype: "" } + member { + name: "experimental_numa_aware" + mtype: "" + } member { name: "experimental_shuffle_and_repeat_fusion" mtype: "" -- cgit v1.2.3