Add print_pipeline_stage_assignment_summary()

Summary:
Add an API which prints how pipeline stages are assigned to stages.

TF2.4 Only
Fix T43335

Test Plan: CI

Reviewers: jackh, jakeh, #tensorflow, simonl, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved

Reviewed By: jackh, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved

Maniphest Tasks: T43335

Differential Revision: https://phabricator.sourcevertex.net/D49022

Author: georgepaw

tensorflow/python/ipu/keras/extensions/functional_extensions.py

@@ -23,6 +23,7 @@
 from tensorflow.python.ipu.keras.extensions import model_extensions
 from tensorflow.python.keras.engine import functional
 from tensorflow.python.training.tracking import base as trackable
+from tensorflow.python.util import nest
 
 
 class PipelineStage(object):
@@ -558,6 +559,53 @@ def reset_pipeline_stage_assignment(self):
     # Pipelining has changed therefore functions need to be recompiled.
     self._reset_ipu_extension()
 
+  def print_pipeline_stage_assignment_summary(self,
+                                              line_length=None,
+                                              print_fn=None):
+    """Prints a summary of the pipeline stage assignment of the model.
+
+    Arguments:
+        line_length: Total length of printed lines (e.g. set this to adapt the
+          display to different terminal window sizes).
+        print_fn: Print function to use. It will be called on each line of the
+          summary. You can set it to a custom function in order to capture the
+          string summary. It defaults to `print` (prints to stdout).
+    """
+    line_length = line_length or 89
+
+    def print_assignment_fn(assignment, print_row):
+      layer = assignment.layer
+      node_index = str(assignment.node_index)
+      inbound_layers = nest.flatten(assignment.inbound_layers)
+      pipeline_stage = str(assignment.pipeline_stage)
+
+      name = layer.name
+      input_layer_names = [l.name for l in inbound_layers]
+
+      cls_name = layer.__class__.__name__
+      if not input_layer_names:
+        first_input = ''
+      else:
+        first_input = input_layer_names[0]
+
+      fields = [
+          name + ' (' + cls_name + ') (' + node_index + ')', first_input,
+          pipeline_stage
+      ]
+      print_row(fields)
+
+      # Print other inputs on the new line.
+      if len(input_layer_names) > 1:
+        for i in range(1, len(input_layer_names)):
+          fields = ['', input_layer_names[i], '']
+          print_row(fields)
+
+    headers = ['Layer (type) (node index)', 'Input Layers', 'Pipeline Stage']
+    column_widths = [.4, .8, 1.]
+    self._print_pipeline_stage_assignment_summary_impl(print_assignment_fn,
+                                                       headers, column_widths,
+                                                       line_length, print_fn)
+
   @trackable.no_automatic_dependency_tracking
   def _get_pipeline_maximum_pipeline_stage(self):
     assert self._is_pipelined()

tensorflow/python/ipu/keras/extensions/model_extensions.py

@@ -1556,6 +1556,45 @@ def set_pipeline_stage_assignment(self, pipeline_stage_assignment):
   def reset_pipeline_stage_assignment(self):
     raise NotImplementedError
 
+  def print_pipeline_stage_assignment_summary(self,
+                                              line_length=None,
+                                              print_fn=None):
+    raise NotImplementedError
+
+  def _print_pipeline_stage_assignment_summary_impl(self, print_assignment_fn,
+                                                    headers, column_widths,
+                                                    line_length, print_fn):
+    """Implementation function for the print_pipeline_stage_assignment_summary.
+    """
+    print_fn = print_fn if print_fn else print
+    assignments = self.get_pipeline_stage_assignment()
+
+    positions = [int(line_length * p) for p in column_widths]
+
+    def print_row(fields):
+      assert len(fields) == len(positions)
+      line = ''
+      for i, field in enumerate(fields):
+        if i > 0:
+          line = line[:-1] + ' '
+        line += str(field)
+        line = line[:positions[i]]
+        line += ' ' * (positions[i] - len(line))
+      print_fn(line)
+
+    print_fn('Model: "{}"'.format(self.name))
+    print_fn('_' * line_length)
+    print_row(headers)
+    print_fn('=' * line_length)
+
+    for i, assignment in enumerate(assignments):
+      print_assignment_fn(assignment, print_row)
+
+      if i == len(assignments) - 1:
+        print_fn('=' * line_length)
+      else:
+        print_fn('_' * line_length)
+
   def _get_pipeline_maximum_pipeline_stage(self):
     """Returns the maximum pipeline stage assignment"""
     raise NotImplementedError

tensorflow/python/ipu/keras/extensions/sequential_extensions.py

@@ -459,6 +459,36 @@ def reset_pipeline_stage_assignment(self):
     # Pipelining has changed therefore functions need to be recompiled.
     self._reset_ipu_extension()
 
+  def print_pipeline_stage_assignment_summary(self,
+                                              line_length=None,
+                                              print_fn=None):
+    """Prints a summary of the pipeline stage assignment of the model.
+
+    Arguments:
+        line_length: Total length of printed lines (e.g. set this to adapt the
+          display to different terminal window sizes).
+        print_fn: Print function to use. It will be called on each line of the
+          summary. You can set it to a custom function in order to capture the
+          string summary. It defaults to `print` (prints to stdout).
+    """
+    line_length = line_length or 60
+
+    def print_assignment_fn(assignment, print_row):
+      layer = assignment.layer
+      pipeline_stage = str(assignment.pipeline_stage)
+
+      name = layer.name
+      cls_name = layer.__class__.__name__
+
+      fields = [name + ' (' + cls_name + ')', pipeline_stage]
+      print_row(fields)
+
+    headers = ['Layer (type)', 'Pipeline Stage']
+    column_widths = [.5, 1.]
+    self._print_pipeline_stage_assignment_summary_impl(print_assignment_fn,
+                                                       headers, column_widths,
+                                                       line_length, print_fn)
+
   @trackable.no_automatic_dependency_tracking
   def _get_pipeline_maximum_pipeline_stage(self):
     assert self._is_pipelined()

tensorflow/python/ipu/tests/keras/extensions/functional_pipeline_api_test.py

@@ -333,6 +333,64 @@ def testSetPipeliningOptions(self):
         self.assertEqual(m._pipelining_accumulate_outfeed, True)  # pylint: disable=protected-access
         self.assertEqual(m._experimental_pipelining_normalize_gradients, True)  # pylint: disable=protected-access
 
+  def testPrintPipelineStageSummary(self):
+    cfg = config.IPUConfig()
+
+    cfg.auto_select_ipus = 1
+    cfg.configure_ipu_system()
+
+    strategy = ipu_strategy.IPUStrategyV1()
+    with strategy.scope():
+      d1 = layers.Input(32, name="input_a")
+      d2 = layers.Input(32, name="input_b")
+      f1 = layers.Flatten()(d1)
+      f2 = layers.Flatten()(d2)
+      c1 = layers.Concatenate()([f1, f2])
+      x1 = layers.Dense(4)(c1)
+      m = training_module.Model((d1, d2), x1)
+
+      strings = []
+
+      def print_fn(x):
+        strings.append(x)
+
+      m.print_pipeline_stage_assignment_summary(line_length=85,
+                                                print_fn=print_fn)
+
+      # pylint: disable=line-too-long
+      self.assertEqual(strings[0], 'Model: "model"')
+      self.assertEqual(strings[1], '_' * 85)
+      self.assertEqual(
+          strings[2],
+          'Layer (type) (node index)         Input Layers                      Pipeline Stage   '
+      )
+      self.assertEqual(strings[3], '=' * 85)
+      self.assertEqual(
+          strings[4],
+          'flatten (Flatten) (0)             input_a                           None             '
+      )
+      self.assertEqual(strings[5], '_' * 85)
+      self.assertEqual(
+          strings[6],
+          'flatten_1 (Flatten) (0)           input_b                           None             '
+      )
+      self.assertEqual(strings[7], '_' * 85)
+      self.assertEqual(
+          strings[8],
+          'concatenate (Concatenate) (0)     flatten                           None             '
+      )
+      self.assertEqual(
+          strings[9],
+          '                                  flatten_1                                          '
+      )
+      self.assertEqual(strings[10], '_' * 85)
+      self.assertEqual(
+          strings[11],
+          'dense (Dense) (0)                 concatenate                       None             '
+      )
+      self.assertEqual(strings[12], '=' * 85)
+      # pylint: enable=line-too-long
+
 
 if __name__ == '__main__':
   test.main()

tensorflow/python/ipu/tests/keras/extensions/sequential_pipeline_api_test.py

@@ -246,6 +246,41 @@ def testSetPipeliningOptions(self):
         self.assertEqual(m._pipelining_accumulate_outfeed, True)  # pylint: disable=protected-access
         self.assertEqual(m._experimental_pipelining_normalize_gradients, True)  # pylint: disable=protected-access
 
+  def testPrintPipelineStageSummary(self):
+    cfg = config.IPUConfig()
+
+    cfg.auto_select_ipus = 1
+    cfg.configure_ipu_system()
+
+    strategy = ipu_strategy.IPUStrategyV1()
+    with strategy.scope():
+      m = sequential.Sequential([layers.Flatten(), layers.Dense(4)])
+
+      strings = []
+
+      def print_fn(x):
+        strings.append(x)
+
+      m.print_pipeline_stage_assignment_summary(line_length=65,
+                                                print_fn=print_fn)
+
+      # pylint: disable=line-too-long
+      self.assertEqual(strings[0], 'Model: "sequential"')
+      self.assertEqual(strings[1], '_' * 65)
+      self.assertEqual(
+          strings[2],
+          'Layer (type)                    Pipeline Stage                   ')
+      self.assertEqual(strings[3], '=' * 65)
+      self.assertEqual(
+          strings[4],
+          'flatten (Flatten)               None                             ')
+      self.assertEqual(strings[5], '_' * 65)
+      self.assertEqual(
+          strings[6],
+          'dense (Dense)                   None                             ')
+      self.assertEqual(strings[7], '=' * 65)
+      # pylint: enable=line-too-long
+
 
 if __name__ == '__main__':
   test.main()