Parse tf.Example using arena allocation. Dense string benchmark shows 10-20% speedup.

Change: 126726018
2016-07-06 11:36:04 -08:00 · 2016-07-06 11:36:04 -08:00 · c4c19f1294
commit c4c19f1294
parent 72ae75bd0e
39 changed files with 46 additions and 41 deletions
--- a/tensorflow/core/example/example.proto
+++ b/tensorflow/core/example/example.proto
@ -3,7 +3,7 @@
 syntax = "proto3";

 import "tensorflow/core/example/feature.proto";
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "ExampleProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.example";
--- a/tensorflow/core/example/example_parser_configuration.proto
+++ b/tensorflow/core/example/example_parser_configuration.proto
@ -2,7 +2,7 @@

 syntax = "proto3";

-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "ExampleParserConfigurationProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.example";
--- a/tensorflow/core/example/feature.proto
+++ b/tensorflow/core/example/feature.proto
@ -54,7 +54,7 @@
 //

 syntax = "proto3";
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "FeatureProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.example";
--- a/tensorflow/core/framework/allocation_description.proto
+++ b/tensorflow/core/framework/allocation_description.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "AllocationDescriptionProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/attr_value.proto
+++ b/tensorflow/core/framework/attr_value.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "AttrValueProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/cost_graph.proto
+++ b/tensorflow/core/framework/cost_graph.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "CostGraphProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/device_attributes.proto
+++ b/tensorflow/core/framework/device_attributes.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "DeviceAttributesProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/function.proto
+++ b/tensorflow/core/framework/function.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "FunctionProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/graph.proto
+++ b/tensorflow/core/framework/graph.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "GraphProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/kernel_def.proto
+++ b/tensorflow/core/framework/kernel_def.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "KernelDefProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/log_memory.proto
+++ b/tensorflow/core/framework/log_memory.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "LogMemoryProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/op_def.proto
+++ b/tensorflow/core/framework/op_def.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "OpDefProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/step_stats.proto
+++ b/tensorflow/core/framework/step_stats.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "StepStatsProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/summary.proto
+++ b/tensorflow/core/framework/summary.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "SummaryProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/tensor.proto
+++ b/tensorflow/core/framework/tensor.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "TensorProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/tensor_description.proto
+++ b/tensorflow/core/framework/tensor_description.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "TensorDescriptionProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/tensor_shape.proto
+++ b/tensorflow/core/framework/tensor_shape.proto
@ -1,7 +1,7 @@
 // Protocol buffer representing the shape of tensors.

 syntax = "proto3";
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "TensorShapeProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/tensor_slice.proto
+++ b/tensorflow/core/framework/tensor_slice.proto
@ -1,7 +1,7 @@
 // Protocol buffer representing slices of a tensor

 syntax = "proto3";
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "TensorSliceProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/types.proto
+++ b/tensorflow/core/framework/types.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "TypesProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/variable.proto
+++ b/tensorflow/core/framework/variable.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "VariableProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/framework/versions.proto
+++ b/tensorflow/core/framework/versions.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "VersionsProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/kernels/example_parsing_ops.cc
+++ b/tensorflow/core/kernels/example_parsing_ops.cc
@ -148,8 +148,9 @@ class ExampleParserOp : public OpKernel {
      // Preallocate dense_values, since we know their sizes
      TensorShape out_shape;
      out_shape.AddDim(batch_size);
-      for (const int64 dim : dense_shapes_[d].dim_sizes())
+      for (const int64 dim : dense_shapes_[d].dim_sizes()) {
        out_shape.AddDim(dim);
+      }
      Tensor* out = nullptr;
      dense_values.allocate(d, out_shape, &out);

@ -193,11 +194,14 @@ class ExampleParserOp : public OpKernel {
    auto DoWork = [&ctx, &mu, &serialized_t, has_names, &names_t,
                   &fixed_len_features, &var_len_features, &output_dense_values,
                   &sparse_values_tmp](int64 start, int64 limit) {
-      Example ex;
      // Processing each Example in the batch starts here.
      for (std::size_t b = static_cast<size_t>(start);
           b < static_cast<size_t>(limit); ++b) {
-        bool parse_success = ParseProtoUnlimited(&ex, serialized_t(b));
+        // Benchmarks indicate that a tight Arena+Example is most performant.
+        protobuf::Arena arena;
+        // ex is owned by the arena.
+        Example* ex = protobuf::Arena::CreateMessage<Example>(&arena);
+        bool parse_success = ParseProtoUnlimited(ex, serialized_t(b));
        if (!TF_PREDICT_TRUE(parse_success)) {
          mutex_lock l(mu);
          ctx->CtxFailure(errors::InvalidArgument(
@ -206,7 +210,7 @@ class ExampleParserOp : public OpKernel {
        }
        const string& example_name = (has_names) ? names_t(b) : "<unknown>";
        Status s = SingleExampleProtoToTensors(
-            ex, example_name, b, fixed_len_features, var_len_features,
+            *ex, example_name, b, fixed_len_features, var_len_features,
            &output_dense_values, &sparse_values_tmp);
        if (!TF_PREDICT_TRUE(s.ok())) {
          mutex_lock l(mu);
--- a/tensorflow/core/kernels/reader_base.proto
+++ b/tensorflow/core/kernels/reader_base.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "ReaderBaseProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.kernels";
--- a/tensorflow/core/lib/core/error_codes.proto
+++ b/tensorflow/core/lib/core/error_codes.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow.error;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "ErrorCodesProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/platform/default/protobuf.h
+++ b/tensorflow/core/platform/default/protobuf.h
@ -19,6 +19,7 @@ limitations under the License.
 // IWYU pragma: private, include "third_party/tensorflow/core/platform/protobuf.h"
 // IWYU pragma: friend third_party/tensorflow/core/platform/protobuf.h

+#include "google/protobuf/arena.h"
 #include "google/protobuf/compiler/importer.h"
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/io/coded_stream.h"
--- a/tensorflow/core/protobuf/config.proto
+++ b/tensorflow/core/protobuf/config.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "ConfigProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/protobuf/master.proto
+++ b/tensorflow/core/protobuf/master.proto
@ -16,7 +16,7 @@ limitations under the License.
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "DistributedRuntimeProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.distruntime";
--- a/tensorflow/core/protobuf/meta_graph.proto
+++ b/tensorflow/core/protobuf/meta_graph.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "MetaGraphProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/protobuf/named_tensor.proto
+++ b/tensorflow/core/protobuf/named_tensor.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "NamedTensorProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/protobuf/queue_runner.proto
+++ b/tensorflow/core/protobuf/queue_runner.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "QueueRunnerProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
--- a/tensorflow/core/protobuf/saver.proto
+++ b/tensorflow/core/protobuf/saver.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "SaverProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.util";
--- a/tensorflow/core/protobuf/tensorflow_server.proto
+++ b/tensorflow/core/protobuf/tensorflow_server.proto
@ -18,7 +18,7 @@ syntax = "proto3";
 import "tensorflow/core/protobuf/config.proto";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "ServerProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.distruntime";
--- a/tensorflow/core/protobuf/worker.proto
+++ b/tensorflow/core/protobuf/worker.proto
@ -16,7 +16,7 @@ limitations under the License.
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "WorkerProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.distruntime";
--- a/tensorflow/core/util/event.proto
+++ b/tensorflow/core/util/event.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "EventProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.util";
--- a/tensorflow/core/util/memmapped_file_system.proto
+++ b/tensorflow/core/util/memmapped_file_system.proto
@ -15,7 +15,7 @@ limitations under the License.
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;

 // A message that describes one region of memmapped file.
 message MemmappedFileSystemDirectoryElement {
--- a/tensorflow/core/util/saved_tensor_slice.proto
+++ b/tensorflow/core/util/saved_tensor_slice.proto
@ -18,7 +18,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "SavedTensorSliceProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.util";
--- a/tensorflow/core/util/test_log.proto
+++ b/tensorflow/core/util/test_log.proto
@ -3,7 +3,7 @@ syntax = "proto3";

 import "google/protobuf/any.proto";

-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;
 option java_outer_classname = "TestLogProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.util.testlog";
--- a/tensorflow/python/training/checkpoint_state.proto
+++ b/tensorflow/python/training/checkpoint_state.proto
@ -1,7 +1,7 @@
 syntax = "proto3";

 package tensorflow;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;

 // Protocol buffer representing the checkpoint state.
 //
--- a/tensorflow/python/util/protobuf/compare_test.proto
+++ b/tensorflow/python/util/protobuf/compare_test.proto
@ -2,7 +2,7 @@
 syntax = "proto2";

 package compare_test;
-// option cc_enable_arenas = true;
+option cc_enable_arenas = true;

 enum Enum {
  A = 0;