Add a pod viewer tool for analyzing TPU pod performance.

tensorboard/plugins/profile/BUILD

@@ -90,6 +90,7 @@ py_binary(
         "profile_demo.google_chart_demo.json",
         "profile_demo.memory_viewer.json",
         "profile_demo.op_profile.json",
+        "profile_demo.pod_viewer.json",
     ],
     srcs_version = "PY2AND3",
     deps = [

tensorboard/plugins/profile/pod_viewer/details_card/BUILD

@@ -0,0 +1,19 @@
+package(default_visibility = ["//tensorboard:internal"])
+
+load("//tensorboard/defs:web.bzl", "tf_web_library")
+
+licenses(["notice"])  # Apache 2.0
+
+tf_web_library(
+    name = "details_card",
+    srcs = [
+        "details-card.html",
+        "details-card.ts",
+    ],
+    path = "/pod-viewer",
+    deps = [
+        "//tensorboard/components/tf_imports:polymer",
+        "//tensorboard/plugins/profile/pod_viewer/pod_viewer_common",
+        "@org_polymer_paper_card",
+    ],
+)

tensorboard/plugins/profile/pod_viewer/details_card/details-card.html

@@ -0,0 +1,100 @@
+<!--
+@license
+Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================
+-->
+
+<!--
+  details-card is a card that highlights detailed information of a selection.
+-->
+
+<link rel="import" href="../polymer/polymer.html">
+<link rel="import" href="../paper-card/paper-card.html">
+<link rel="import" href="pod-viewer-common.html">
+
+<dom-module id="details-card">
+  <style>
+paper-card {
+  max-height: 50vh;
+  overflow-y: auto;
+  --paper-card-header-color: white;
+  --paper-card-header: {
+    background-color: rgb(178,34,34);
+  }
+  width: 100%;
+}
+
+.card-content > div {
+  margin-bottom: 1em;
+  top: -5px;
+}
+
+.value {
+  font-weight: normal;
+  text-align: right;
+}
+
+.info {
+  font-size: 20px;
+  font-weight: bold;
+}
+
+.code-style {
+  font-size: 14px;
+  font-weight: normal;
+}
+  </style>
+  <template>
+    <paper-card id="card" heading="[[_name]]" hidden="[[!_name]]" elevation="2">
+      <template is="dom-repeat" items=[[nodes]] as="node">
+        <div class="card-content info">
+          <div hidden="[[!_isChannel(node)]]">
+            <p>Replica Id: <span class="value">[[node.replicaId]]</span></p>
+          </div>
+          <div hidden="[[_isStep(node)]]">
+            <p>Data Transferred: <span class="value">[[_sizeMiB(node.dataSize)]] MiB</span></p>
+            <p>Latency: <span class="value">[[_format(node.durationUs)]] µs</span></p>
+            <p>BW: <span class="value">
+                [[_bandwidth(node.dataSize, node.durationUs)]] GiB/s</span></p>
+          </div>
+          <div hidden="[[!_isChannel(node)]]">
+            <p>Send Delay: <span class="value">[[_format(node.sendDelayUs)]] µs</span></p>
+            <p>From: <span class="value">Chip[[_chipId(node.srcCoreId)]], Core[[_nodeId(node.srcCoreId)]]</span></p>
+            <p>To: <span class="value">Chip[[_chipId(node.dstCoreId)]], Core[[_nodeId(node.dstCoreId)]]</span></p>
+            <p>Hlo Names: </p>
+            <code class="code-style">
+              <template is="dom-repeat" items=[[node.hloNames]]>
+                "[[item]]"
+              </template>
+            </code>
+          </div>
+          <div hidden="[[!_hasReplicaGroups(node)]]">
+            <p>Replica Groups</p>
+            <code class="code-style">
+              <template is="dom-repeat" items=[[node.replicaGroups]]>
+                {[[item.replicaIds]]}<br>
+              </template>
+            </code>
+          </div>
+          <div hidden="[[!_isStep(node)]]">
+            <template is="dom-repeat" items=[[stepBreakdownLayers]]>
+              <p>[[item.label]]: <span class="value">
+                  [[_getStepBreakdownValue(node, item.key)]] µs
+                  [[_getStepBreakdownPct(node, item.key)]]</span></p>
+            </template>
+          </div>
+        </div>
+      </template>
+    </paper-card>
+  </template>
+  <script src="details-card.js"></script>
+</dom-module>

tensorboard/plugins/profile/pod_viewer/details_card/details-card.ts

@@ -0,0 +1,130 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+namespace pod_viewer_details_card {
+
+type DetailNode = podviewer.proto.ChannelInfo | podviewer.proto.PodStatsRecord
+                      | podviewer.proto.AllReduceOpInfo;
+
+Polymer({
+  is: 'details-card',
+  properties: {
+    nodes: {
+      type: Array,
+    },
+    _name: {
+      type: String,
+      computed: '_computeName(nodes)',
+    },
+    stepBreakdownLayers: {
+      type: Array,
+      value: () => [
+        {key: 'highFlopsComputeUs', label: 'High flops compute'},
+        {key: 'lowFlopsComputeUs', label: 'Low flops compute'},
+        {key: 'hostInfeedDurationUs', label: 'Infeed'},
+        {key: 'hostOutfeedDurationUs', label: 'Outfeed'},
+        {key: 'crsDurationUs', label: 'All reduce'},
+        {key: 'sendDurationUs', label: 'Send'},
+        {key: 'recvDurationUs', label: 'Recv'},
+      ],
+    },
+  },
+  _isAllReduce(node: DetailNode): node is podviewer.proto.AllReduceOpInfo {
+    return (<podviewer.proto.AllReduceOpInfo>node).replicaGroups != undefined;
+  },
+  _isChannel(node: DetailNode): node is podviewer.proto.ChannelInfo {
+    return (<podviewer.proto.ChannelInfo>node).channelId != undefined;
+  },
+  _isStep(node: DetailNode): node is podviewer.proto.PodStatsRecord {
+    return (<podviewer.proto.PodStatsRecord>node).hostName != undefined;
+  },
+  _hasReplicaGroups(node: podviewer.proto.AllReduceOpInfo): boolean {
+    return node.replicaGroups && node.replicaGroups.length > 0;
+  },
+  _computeName: function(nodes: Array<DetailNode>): string|undefined {
+    if (!nodes || nodes.length == 0) return;
+    const node = nodes[0];
+    if (this._isChannel(node)){
+      return 'Channel # ' + (<podviewer.proto.ChannelInfo>node).channelId;
+    } else if (this._isAllReduce(node)) {
+      return (<podviewer.proto.AllReduceOpInfo>node).name;
+    } else if (this._isStep(node)) {
+      return 'Step breakdown of chip '
+          + (<podviewer.proto.PodStatsRecord>node).chipId
+              + ', core ' + (<podviewer.proto.PodStatsRecord>node).nodeId;
+    }
+    return;
+  },
+  /**
+   * Converts from number of bytes to MiB.
+   */
+  _bytesToMiB: function(numBytes: number): number {
+    return numBytes / 1048576;
+  },
+  /**
+   * Return the formatted data size in MiB.
+   */
+  _sizeMiB: function(dataSize: undefined|number): string|undefined {
+    if (!dataSize) return;
+    return this._format(this._bytesToMiB(dataSize));
+  },
+  /**
+   * Return the formatted link bandwidth in GiB/s.
+   * The link bandwidth here is defined by the data size transferred over the
+   * duration between the start of the send operation to the end of the
+   * recv-done operation.
+   */
+  _bandwidth: function(
+      dataSize: undefined|number, duration: undefined|number):
+          string|undefined {
+    if (!dataSize || !duration) return;
+    return this._format(dataSize / duration / 1073.74);
+  },
+  /**
+   * Return the chip id given the global core id.
+   */
+  _chipId: function(coreId: number): number {
+    return Math.floor(coreId / 2);
+  },
+  /**
+   * Return the node ordinal given the global core id.
+   */
+  _nodeId: function(coreId: number): number {
+    return coreId & 1;
+  },
+  /**
+   * Format a number with two digits after the decimal point.
+   */
+  _format: function(number: undefined|number): string {
+    return number == null ? '' : number.toFixed(2);
+  },
+  /**
+   * Return a formatted value associated with a specific breakdown.
+   */
+  _getStepBreakdownValue:
+      function(node: undefined|podviewer.proto.PodStatsRecord,
+          key: undefined|string): string|undefined {
+    if (!key || !node) return;
+    return this._format(node[key]);
+  },
+  /**
+   * Return a the percentage of a specific breakdown.
+   */
+  _getStepBreakdownPct:
+      function(node: undefined|podviewer.proto.PodStatsRecord,
+          key: undefined|string): string|undefined {
+    if (!key || !node || !node.totalDurationUs) return;
+    return (node[key] / node.totalDurationUs * 100).toFixed(2) + '%';
+  },
+});
+
+} // namespace pod_viewer_details_card

tensorboard/plugins/profile/pod_viewer/pod_viewer_common/BUILD

@@ -0,0 +1,14 @@
+package(default_visibility = ["//tensorboard:internal"])
+
+load("//tensorboard/defs:web.bzl", "tf_web_library")
+
+licenses(["notice"])  # Apache 2.0
+
+tf_web_library(
+    name = "pod_viewer_common",
+    srcs = [
+        "pod-viewer-common.html",
+        "proto.ts",
+    ],
+    path = "/pod-viewer",
+)

tensorboard/plugins/profile/pod_viewer/pod_viewer_common/pod-viewer-common.html

@@ -0,0 +1,18 @@
+<!--
+@license
+Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+<script src="proto.js"></script>