[XLA] Cache computations when creating reduces in algebraic simplifier or batchnorm...
authorBenjamin Kramer <kramerb@google.com>
Tue, 15 May 2018 17:23:27 +0000 (10:23 -0700)
committerTensorFlower Gardener <gardener@tensorflow.org>
Tue, 15 May 2018 17:26:09 +0000 (10:26 -0700)
commit2c83cddab0cd0de78f863e47d81b4427d6519eb7
tree900799da86a702a7962c34cbea909c81dd2d4f6c
parentc515b31314f0c4b551bded4961e096db98a998fc
[XLA] Cache computations when creating reduces in algebraic simplifier or batchnorm expander

Otherwise we create a lot of identical small computations. This shouldn't have
an effect except for cluttering the HLO, but turns out HloCSE doesn't look
inside of the computation of reduces, effectively never eliminating reduces
that were produced via this code path.

While there clean up some YAGNI, this only worked for F32 anyways, so just
hardcode it.

PiperOrigin-RevId: 196689316
tensorflow/compiler/xla/service/BUILD
tensorflow/compiler/xla/service/algebraic_simplifier.cc
tensorflow/compiler/xla/service/batchnorm_expander.cc