experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Predicate Reduce(Dot(....)) under enable_dot_strength_reduction.

Browse Source 
PiperOrigin-RevId: 295896500
Change-Id: I07fda5d17b160f8ea1492c71dee9b6d58204d50b
This commit is contained in:
Blake Hechtman 2020-02-18 22:05:49 -08:00 committed by TensorFlower Gardener
parent eebf50dd9e
commit 7c1bc443fa
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/compiler/xla/service/algebraic_simplifier.cc
View File

@ -3727,7 +3727,8 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* hlo) {
// Convert Reduce(Dot(X,Y)) to Dot(X,Y) if any of the dimensions reduced were
// batch dimensions of the dot. The transformation supports reducing other
// dimensions as well.
if (Match(arg, m::Dot(&dot, m::Op(&lhs), m::Op(&rhs)).WithOneUser()) &&
if (options_.enable_dot_strength_reduction() &&
Match(arg, m::Dot(&dot, m::Op(&lhs), m::Op(&rhs)).WithOneUser()) &&
Match(reduce->to_apply()->root_instruction(),
m::Add(m::Parameter(), m::Parameter())) &&
absl::c_any_of(reduce->dimensions(), [&](int64 dim) {