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

tensorflower-gardener · tensorflower-gardener · commit c4c19f129459 · 2016-07-06T12:47:38.000-07:00
Change: 126726018
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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);
diff --git 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";
diff --git 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";
diff --git 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"
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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";
diff --git 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 {
diff --git 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";
diff --git 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";
diff --git 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.
 //
diff --git 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;
