Commit 3e9bf087 authored by Nick Desaulniers's avatar Nick Desaulniers Committed by TensorFlower Gardener
Browse files

[XLA] optimize NearComparator#ExpectLiteralsNear()

While tracking down the issue of timeouts when running THE ISOLATOR, it was
observed that NearComparator#ExpectLiteralsNear() could be optimized in the
case of matching layouts to not compute multi indexes.

In the process of tracking down timeouts in THE ISOLATOR, I had assumed that
time spent was dominated by either generating input data, executing the input
data on various backends, or comparing the data. Never assume you know where
the time is spent in a program; the profiler may surprise you.

After making that optimization and then profiling the code before and after, I
was surprised by the profile. Image the shock, horror, and disgust I
experienced when discovering that runs of THE ISOLATOR were dominated (45%) by
calls to Literal#ToString() in NearComparator#ExpectLiteralsNear() for huge
(>120 million elements) literals that failed comparisons.  No wonder passing
shards of THE ISOLATOR were fast, and failing shards were slow.

Further, computing multi indexes many times is expensive enough (18%) to show
up in profiles, so avoid calculating it until it is necessary.

The optimizations in this patch:
* Don't call Literal#ToString() on huge literals that are going to get written
  to disk anyways. The utility of printing said literal to stdout is suspect.
* Initialize NearComparator#miscompares_ to false, only update miscompares_ and
  other stats when miscompare occurs.
* Split NearComparator#ExpectLiteralsNear() into two, since we only need to log
  and update stats if an actual miscompare occurs.
* Add fast path in NearComparator#ExpectLiteralsNear() for case of matching
  layouts, being careful not to compute multi index unless mismatch actually
  occurs.

This optimized NearComparator#ExpectLiteralsNear() for the case of many element
literals, with few miscompares. For many miscompares, we cannot avoid
calculating multi indexes, but can fast path for equal layouts. For zero
miscompares, we can at least fast path in the case of matching layouts.

Before this CL, a run of THE ISOLATOR for a single literal with >120 million
elements and a few miscompares took 379s (6.3m). With this CL, the same test
case now takes 44s.

Beautiful flame graphs omitted from public commit message, regrettably.

PiperOrigin-RevId: 183451138
parent b9494ce8
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment