Add tensor output shapes to CostGraphDef.

Update StepStatsCollector::BuildCostModel() to populate these fields.
Misc. cleanup and optimizations in step_stats_collector.cc and costmodel.cc.
Change: 135541791

1 files changed, 13 insertions, 4 deletions

diff --git a/tensorflow/core/graph/costmodel.h b/tensorflow/core/graph/costmodel.h
index e9bb6e669b..6eee2e3e66 100644
--- a/tensorflow/core/graph/costmodel.h
+++ b/tensorflow/core/graph/costmodel.h
@@ -20,13 +20,16 @@ limitations under the License.
 #include <vector>
 
 #include "tensorflow/core/framework/cost_graph.pb.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/graph/types.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
 #include "tensorflow/core/lib/gtl/array_slice.h"
 #include "tensorflow/core/platform/macros.h"
 
 namespace tensorflow {
-typedef std::unordered_map<string, int32> NodeNameToCostIdMap;
+typedef std::unordered_map<StringPiece, int32, StringPiece::Hasher>
+    NodeNameToCostIdMap;
 
 class StepStats;
 
@@ -104,14 +107,19 @@ class CostModel {
   // Check that an estimate is available for every OP node in graph.
   void CheckInitialized(const Graph& graph) const;
 
-  // Records the maximum size in bytes of the tensor generated by "output_slot"
-  // of "node".
-  void RecordMaxMemorySize(const Node* node, int output_slot, Bytes bytes);
+  // Records the maximum size in bytes and optionally the corresponding shape of
+  // the tensor generated by "output_slot" of "node". If
+  void RecordMaxMemorySize(const Node* node, int output_slot, Bytes bytes,
+                           const TensorShapeProto& tensor_shape);
 
   // Returns the maximum size in bytes of the tensor generated by "output_slot"
   // of "node".
   Bytes MaxMemorySize(const Node* node, int output_slot) const;
 
+  // Returns the shape corresponding to the largest memory size of the tensor
+  // generated by "output_slot" of "node".
+  TensorShapeProto MaxMemoryShape(const Node* node, int output_slot) const;
+
   // Returns the size in bytes of temporary memory consumed by "node".
   Bytes TempMemorySize(const Node* node) const;
 
@@ -169,6 +177,7 @@ class CostModel {
   struct MemUsage {
     Bytes temp_memory_size;
     gtl::InlinedVector<Bytes, 2> output_port_mem;
+    gtl::InlinedVector<TensorShapeProto, 2> output_port_shape;
   };
   std::vector<MemUsage> max_mem_usage_;