Show layout analysis for convolution ops (tensorflow#484)

Author: sam-mccall

tensorboard/plugins/profile/tf_op_profile/tf-op-details.html

@@ -84,6 +84,31 @@
   font-style: italic;
   color: #666;
 }
+
+.layout {
+  display: table;
+  width: 50%;
+  border-spacing: 1em 0.3em;
+}
+.layout > * { display: table-row; }
+.layout > * > * { display: table-cell; }
+.layout .size, .layout .size-x { text-align: center; }
+.layout .size-x { font-weight: bold; color: #888; }
+.layout .semantics {
+  font-size: smaller;
+  font-weight: bold;
+  color: #48c;
+  text-transform: uppercase;
+  vertical-align: middle;
+}
+.layout .alignment {
+  /* Make the background extend outside the box */
+  position: absolute;
+  padding: 0.5em;
+  margin-top: -0.5em;
+  width: 50%;
+  box-sizing: border-box;
+}
   </style>
   <template>
     <paper-card id="card" heading="[[node.name]]" elevation="2">
@@ -103,6 +128,19 @@
         <div class="unavailable" hidden="[[!node.category]]">
           Select items within this category for performace details.
         </div>
+        <div hidden="[[!node.xla.layout]]">
+          <b>Layout: </b>
+          <div class="layout" hidden="[[!node.xla.layout]]">
+            <template is="dom-repeat" items="[[node.xla.layout.dimensions]]">
+              <div hidden="[[!index]]"><span class="size-x">×</span></div>
+              <div>
+                <span class="size">[[item.size]]</span>
+                <span class="semantics">[[item.semantics]]</span>
+                <span class="alignment" style$="background-color:[[_dimensionColor(item)]]">[[_dimensionHint(item)]]</span>
+              </div>
+            </template>
+          </div>
+        </div>
       </div>
     </paper-card>
   </template>
@@ -133,6 +171,20 @@
     return "Unknown";
   },
   _fused: function(node){ return node && node.xla && !node.metrics; },
+  _dimensionColor: function(dim) {
+    if (!dim || !dim.alignment) return null;
+    var ratio = dim.size / dim.alignment;
+    // Colors should grade harshly: 50% in a dimension is already very bad.
+    var harshCurve = (x) => 1 - Math.sqrt(1 - x);
+    return flameColor(ratio / Math.ceil(ratio), 1, 0.25, harshCurve);
+  },
+  _dimensionHint: function(dim) {
+    if (!dim || !dim.alignment) return null;
+    var mul = Math.ceil(dim.size / dim.alignment);
+    var mulSuffix = (mul == 1) ? "" : ": " + mul + " × " + dim.alignment;
+    if (dim.size % dim.alignment == 0) return "Exact fit" + mulSuffix;
+    return "Pad to " + (mul * dim.alignment) + mulSuffix;
+  },
 });
   </script>
 </dom-module>

tensorboard/plugins/profile/tf_op_profile/utils.ts

@@ -23,11 +23,13 @@ function rgba(red: number, green: number, blue: number, alpha: number) {
  * @param {number} fraction
  * @param {number=} brightness
  * @param {number=} opacity
+ * @param {Function=} curve mapping [0-1] to [0-1]
  * @return {string} An RGBA color.
  */
-export function flameColor(fraction: number, brightness = 1, opacity = 1) {
+export function flameColor(
+    fraction: number, brightness = 1, opacity = 1, curve = Math.sqrt) {
   if (isNaN(fraction)) return rgba(brightness, brightness, brightness, opacity);
-  fraction = Math.sqrt(fraction);  // Or everything is depressing and red.
+  fraction = curve(fraction); // Or everything is depressing and red.
   return (fraction < 0.5) ?
     rgba(brightness, 2 * fraction * brightness, 0, opacity) :
     rgba(2 * (1 - fraction) * brightness, brightness, 0, opacity);