diff --git a/tensorflow/core/common_runtime/constant_folding.cc b/tensorflow/core/common_runtime/constant_folding.cc index e4427877c33cce..35e1f7e47be02c 100644 --- a/tensorflow/core/common_runtime/constant_folding.cc +++ b/tensorflow/core/common_runtime/constant_folding.cc @@ -22,6 +22,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "tensorflow/core/common_runtime/device_factory.h" #include "tensorflow/core/common_runtime/executor.h" @@ -72,13 +73,20 @@ bool GetShapeFromArgNode(const Node* node, TensorShapeProto* out_shape) { } return false; } +struct DynamicAncestorCacheEntry { + bool computed = false; + bool computing = false; + bool has_dynamic = false; + TensorShapeProto shape; +}; -bool GetShapeFromDirectDynamicSource(const Node* node, - TensorShapeProto* out_shape) { - return (IsShapeOp(node) || - node->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) && - GetShapeFromArgNode(node, out_shape); -} +using DynamicAncestorCache = + absl::flat_hash_map; + +bool GetShapeFromDynamicAncestor(const Node* node, + TensorShapeProto* out_shape); +bool GetShapeFromDynamicAncestor(const Node* node, TensorShapeProto* out_shape, + DynamicAncestorCache* cache); // For stateless RNGs ops, they are pure but device-dependent. Those ops are not // constant-foldable. @@ -272,9 +280,10 @@ bool IsConstantFoldable( shape_map, const std::function& consider, int64_t max_constant_size_in_bytes, - std::unordered_map>* shape_replacement_map) { + std::unordered_map>* shape_replacement_map, + DynamicAncestorCache* dynamic_ancestor_cache) { TensorShapeProto dynamic_shape; - if (GetShapeFromDirectDynamicSource(n, &dynamic_shape)) { + if (GetShapeFromDynamicAncestor(n, &dynamic_shape, dynamic_ancestor_cache)) { VLOG(1) << "Skipping constant folding for dynamic shape-derived node " << n->name() << " op=" << n->type_string() << " inferred_shape=" << dynamic_shape.DebugString(); @@ -367,10 +376,11 @@ void ConsiderConstantFoldableNode( Node* n, const ConstantFoldingOptions& opts, std::vector* nodes, std::unordered_map>* constant_control_deps, std::unordered_map>* shape_replacement_map, - bool* internal_node_inserted) { + bool* internal_node_inserted, + DynamicAncestorCache* dynamic_ancestor_cache) { if (!IsConstantFoldable(n, opts.shape_map, opts.consider, opts.max_constant_size_in_bytes, - shape_replacement_map)) { + shape_replacement_map, dynamic_ancestor_cache)) { return; } // A node is constant provided all of its non-control incoming Tensors come @@ -424,16 +434,19 @@ void FindConstantFoldableNodes( const Graph* graph, const ConstantFoldingOptions& opts, std::vector* nodes, std::unordered_map>* constant_control_deps, - std::unordered_map>* shape_replacement_map) { + std::unordered_map>* shape_replacement_map, + DynamicAncestorCache* dynamic_ancestor_cache) { bool internal_node_inserted = false; // Walk the nodes in data flow order. ReverseDFS( *graph, nullptr, [nodes, constant_control_deps, shape_replacement_map, - &internal_node_inserted, &opts](Node* n) { + &internal_node_inserted, &opts, + dynamic_ancestor_cache](Node* n) { ConsiderConstantFoldableNode(n, opts, nodes, constant_control_deps, shape_replacement_map, - &internal_node_inserted); + &internal_node_inserted, + dynamic_ancestor_cache); }, NodeComparatorName()); // If we have inserted just leaf level nodes, then there is nothing to fold. @@ -477,6 +490,37 @@ void AddNodeToConstantGraph( } } +bool GetShapeFromDynamicAncestor(const Node* node, TensorShapeProto* out_shape, + DynamicAncestorCache* cache) { + auto& entry = (*cache)[node]; + if (entry.computed) { + if (entry.has_dynamic) { + *out_shape = entry.shape; + } + return entry.has_dynamic; + } + if (entry.computing) { + return false; + } + + entry.computing = true; + if ((IsShapeOp(node) || + node->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) && + GetShapeFromArgNode(node, &entry.shape)) { + entry.has_dynamic = true; + } + entry.computing = false; + entry.computed = true; + if (entry.has_dynamic) { + *out_shape = entry.shape; + } + return entry.has_dynamic; +} + +bool GetShapeFromDynamicAncestor(const Node* node, TensorShapeProto* out_shape) { + DynamicAncestorCache cache; + return GetShapeFromDynamicAncestor(node, out_shape, &cache); +} // Replaces constant-foldable shape node n by a vector of constants in // constant_graph, which is being built up for subsequent evaluation of constant // propagation. node_map is the mapping of nodes in the original graph to nodes @@ -488,10 +532,17 @@ void AddShapeNodeToConstantGraph( const std::unordered_map>& shape_replacement_map, std::unordered_map>* node_map, - const ConstantFoldNameGenerator& generate_new_name, Graph* constant_graph) { + const ConstantFoldNameGenerator& generate_new_name, Graph* constant_graph, + DynamicAncestorCache* dynamic_ancestor_cache) { TensorShapeProto user_inferred_shape; const bool has_dynamic = - GetShapeFromDirectDynamicSource(n, &user_inferred_shape); + GetShapeFromDynamicAncestor(n, &user_inferred_shape, + dynamic_ancestor_cache); + if (has_dynamic) { + VLOG(1) << "AddShapeNodeToConstantGraph preserving dynamic metadata for " + << n->name() << " op=" << n->type_string() + << " inferred_shape=" << user_inferred_shape.DebugString(); + } std::vector& added = (*node_map)[n]; const string& node_name = n->name(); for (const Tensor& t : shape_replacement_map.at(n)) { @@ -524,7 +575,8 @@ Graph* GetConstantGraph( const std::unordered_map>& shape_replacement_map, std::map* tensors_to_fetch, - const ConstantFoldNameGenerator& generate_new_name) { + const ConstantFoldNameGenerator& generate_new_name, + DynamicAncestorCache* dynamic_ancestor_cache) { Graph* constant_graph = new Graph(orig_graph->op_registry()); std::unordered_map> node_map; node_map[orig_graph->source_node()] = {constant_graph->source_node()}; @@ -534,7 +586,8 @@ Graph* GetConstantGraph( AddNodeToConstantGraph(n, &node_map, constant_graph); } else { AddShapeNodeToConstantGraph(n, shape_replacement_map, &node_map, - generate_new_name, constant_graph); + generate_new_name, constant_graph, + dynamic_ancestor_cache); } } @@ -572,7 +625,8 @@ bool ReplaceTensorWithConstant( Graph* graph, const Device* partition_device, NodeAndOutput tensor, const Tensor& constant, const gtl::FlatSet& control_deps, int64_t max_constant_size_in_bytes, - const ConstantFoldNameGenerator& generate_new_name) { + const ConstantFoldNameGenerator& generate_new_name, + DynamicAncestorCache* dynamic_ancestor_cache) { // Be conservative when replacing a tensor with a constant, when not // running on CPU. // 1) Do not replace another constant. @@ -625,7 +679,13 @@ bool ReplaceTensorWithConstant( const string& node_name = n->name(); TensorShapeProto user_inferred_shape; const bool has_dynamic = - GetShapeFromDirectDynamicSource(tensor.first, &user_inferred_shape); + GetShapeFromDynamicAncestor(tensor.first, &user_inferred_shape, + dynamic_ancestor_cache); + if (has_dynamic) { + VLOG(1) << "ReplaceTensorWithConstant preserving dynamic metadata for " + << tensor.first->name() << " output=" << tensor.second + << " inferred_shape=" << user_inferred_shape.DebugString(); + } Node* constant_node; auto builder = NodeDefBuilder(generate_new_name(graph, node_name), "Const") .Attr("dtype", constant.dtype()) @@ -697,8 +757,10 @@ absl::Status ConstantFold(const ConstantFoldingOptions& opts, std::vector constant_foldable_nodes; std::unordered_map> constant_control_deps; std::unordered_map> shape_replacement_map; + DynamicAncestorCache dynamic_ancestor_cache; FindConstantFoldableNodes(graph, opts, &constant_foldable_nodes, - &constant_control_deps, &shape_replacement_map); + &constant_control_deps, &shape_replacement_map, + &dynamic_ancestor_cache); if (constant_foldable_nodes.empty()) { VLOG(1) << "No constant foldable nodes found"; *was_mutated = false; @@ -709,7 +771,8 @@ absl::Status ConstantFold(const ConstantFoldingOptions& opts, std::map tensors_to_fetch; std::unique_ptr constant_graph( GetConstantGraph(graph, constant_foldable_nodes, shape_replacement_map, - &tensors_to_fetch, generate_new_name)); + &tensors_to_fetch, generate_new_name, + &dynamic_ancestor_cache)); DumpGraph("Constant graph", constant_graph.get()); if (tensors_to_fetch.empty()) { @@ -765,7 +828,8 @@ absl::Status ConstantFold(const ConstantFoldingOptions& opts, constant_control_deps[tensors_to_replace[c].first]; if (ReplaceTensorWithConstant( graph, partition_device, tensors_to_replace[c], outputs[c], - control_deps, opts.max_constant_size_in_bytes, generate_new_name)) { + control_deps, opts.max_constant_size_in_bytes, generate_new_name, + &dynamic_ancestor_cache)) { ++num_nodes_replaced; } }