[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
Loading
Please sign in to comment