From 196b03bb55cbb72c75dc772cf1b6f88ea9cb55a3 Mon Sep 17 00:00:00 2001 From: Derek Murray Date: Thu, 30 Apr 2020 18:11:03 -0700 Subject: [PATCH] [ParseSequenceExampleOp] Optimize config construction. Several micro-optimizations that mostly benefit small input protos: 1. Reserve vector capacity before pushing new elements in a loop. 2. Use `vector::emplace_back()` instead of `vector::push_back()` to avoid move-constructing the temporary struct literals. 3. Construct the default value tensor using the optimized scalar constructors for `Tensor`, which avoid an unnecessary allocation. 4. Evaluate `Tensor::flat()` once for each of the key tensors, instead of re-evaluating it for each key, and similarly for the missing-assumed-empty tensor. PiperOrigin-RevId: 309337185 Change-Id: I3d546a3927a15379f811be80c44f81d4734d0016 --- .../core/kernels/example_parsing_ops.cc | 108 +++++++++++++----- 1 file changed, 80 insertions(+), 28 deletions(-) diff --git a/tensorflow/core/kernels/example_parsing_ops.cc b/tensorflow/core/kernels/example_parsing_ops.cc index 3412d00136e..483d7753e10 100644 --- a/tensorflow/core/kernels/example_parsing_ops.cc +++ b/tensorflow/core/kernels/example_parsing_ops.cc @@ -569,57 +569,109 @@ class ParseSequenceExampleOp : public OpKernel { const Tensor* dense_keys, const Tensor* sparse_keys, const Tensor* ragged_keys, const OpInputList& context_dense_defaults) const { + // Convert the tensors/attrs to ArraySlices once, instead of re-evaluating + // them in each loop iteration. + gtl::ArraySlice dense_keys_slice = + dense_keys + ? gtl::ArraySlice(dense_keys->flat().data(), + attrs_.num_context_dense) + : attrs_.context_dense_keys; + gtl::ArraySlice sparse_keys_slice = + sparse_keys + ? gtl::ArraySlice(sparse_keys->flat().data(), + attrs_.num_context_sparse) + : attrs_.context_sparse_keys; + gtl::ArraySlice ragged_keys_slice( + ragged_keys->flat().data(), attrs_.num_context_ragged); + example::FastParseExampleConfig config; + config.dense.reserve(attrs_.num_context_dense); for (int d = 0; d < attrs_.num_context_dense; ++d) { - const tstring& key = dense_keys ? dense_keys->flat()(d) - : attrs_.context_dense_keys[d]; - config.dense.push_back({key, attrs_.context_dense_types[d], - attrs_.context_dense_shapes[d], - context_dense_defaults[d], - false /* attrs_.context_variable_length[d] */, - 0 /*attrs_.context_elements_per_stride[d] */}); + const tstring& key = dense_keys_slice[d]; + config.dense.emplace_back(key, attrs_.context_dense_types[d], + attrs_.context_dense_shapes[d], + context_dense_defaults[d], + false /* attrs_.context_variable_length[d] */, + 0 /*attrs_.context_elements_per_stride[d] */); } + config.sparse.reserve(attrs_.num_context_sparse); for (int d = 0; d < attrs_.num_context_sparse; ++d) { - const tstring& key = sparse_keys ? sparse_keys->flat()(d) - : attrs_.context_sparse_keys[d]; - config.sparse.push_back({key, attrs_.context_sparse_types[d]}); + const tstring& key = sparse_keys_slice[d]; + config.sparse.emplace_back(key, attrs_.context_sparse_types[d]); } + config.ragged.reserve(attrs_.num_context_ragged); for (int d = 0; d < attrs_.num_context_ragged; ++d) { - config.ragged.push_back({ragged_keys->flat()(d), - attrs_.context_ragged_value_types[d], - attrs_.context_ragged_split_types[d]}); + config.ragged.emplace_back(ragged_keys_slice[d], + attrs_.context_ragged_value_types[d], + attrs_.context_ragged_split_types[d]); } return config; } + static Tensor ConstructDefaultScalar(DataType dtype) { + switch (dtype) { + case DT_INT64: + return Tensor(static_cast(0)); + case DT_FLOAT: + return Tensor(static_cast(0.0)); + case DT_STRING: + return Tensor(""); + default: + return Tensor(DT_INVALID); + } + } + example::FastParseExampleConfig MakeFeatureListConfig( const Tensor* dense_keys, const Tensor* sparse_keys, const Tensor* ragged_keys, const Tensor* feature_list_dense_missing_assumed_empty) const { + // Convert the tensors/attrs to ArraySlices once, instead of re-evaluating + // them in each loop iteration. + gtl::ArraySlice dense_keys_slice = + dense_keys + ? gtl::ArraySlice(dense_keys->flat().data(), + attrs_.num_feature_list_dense) + : attrs_.feature_list_dense_keys; + gtl::ArraySlice sparse_keys_slice = + sparse_keys + ? gtl::ArraySlice(sparse_keys->flat().data(), + attrs_.num_feature_list_sparse) + : attrs_.feature_list_sparse_keys; + gtl::ArraySlice ragged_keys_slice( + ragged_keys->flat().data(), attrs_.num_feature_list_ragged); + // Use an empty slice to indicate that the map in attrs_ should be used + // instead. + gtl::ArraySlice feature_list_dense_missing_assumed_empty_slice = + feature_list_dense_missing_assumed_empty + ? gtl::ArraySlice( + feature_list_dense_missing_assumed_empty->flat().data(), + attrs_.num_feature_list_dense) + : gtl::ArraySlice(nullptr, 0); + example::FastParseExampleConfig config; + config.dense.reserve(attrs_.num_feature_list_dense); for (int d = 0; d < attrs_.num_feature_list_dense; ++d) { - const tstring& key = dense_keys ? dense_keys->flat()(d) - : attrs_.feature_list_dense_keys[d]; + const tstring& key = dense_keys_slice[d]; bool missing_assumed_empty = - feature_list_dense_missing_assumed_empty - ? feature_list_dense_missing_assumed_empty->flat()(d) + !feature_list_dense_missing_assumed_empty_slice.empty() + ? feature_list_dense_missing_assumed_empty_slice[d] : attrs_.feature_list_dense_missing_assumed_empty.count(key) > 0; DataType dtype = attrs_.feature_list_dense_types[d]; - Tensor default_value = Tensor(dtype, TensorShape({})); - config.dense.push_back( - {key, dtype, attrs_.feature_list_dense_shapes[d], default_value, - missing_assumed_empty, - 0 /*attrs_.feature_list_elements_per_stride[d] */}); + config.dense.emplace_back( + key, dtype, attrs_.feature_list_dense_shapes[d], + ConstructDefaultScalar(dtype), missing_assumed_empty, + 0 /*attrs_.feature_list_elements_per_stride[d] */); } + config.sparse.reserve(attrs_.num_feature_list_sparse); for (int d = 0; d < attrs_.num_feature_list_sparse; ++d) { - const tstring& key = sparse_keys ? sparse_keys->flat()(d) - : attrs_.feature_list_sparse_keys[d]; - config.sparse.push_back({key, attrs_.feature_list_sparse_types[d]}); + const tstring& key = sparse_keys_slice[d]; + config.sparse.emplace_back(key, attrs_.feature_list_sparse_types[d]); } + config.ragged.reserve(attrs_.num_feature_list_ragged); for (int d = 0; d < attrs_.num_feature_list_ragged; ++d) { - config.ragged.push_back({ragged_keys->flat()(d), - attrs_.feature_list_ragged_value_types[d], - attrs_.feature_list_ragged_split_types[d]}); + config.ragged.emplace_back(ragged_keys_slice[d], + attrs_.feature_list_ragged_value_types[d], + attrs_.feature_list_ragged_split_types[d]); } return config; }