Commit 90f257e0 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower Committed by TensorFlower Gardener
Browse files

Fix ReshapeMover bug with reshaped constants; add HloVerifiedTestBase.

An example of a bad ReshapeMover rewrite:

BEFORE
  %reshape.1 = f32[1,1,128] reshape(f32[1,128] %dot)
  %constant = f32[128] constant({...})
  %reshape.2 = f32[1,1,128] reshape(f32[128] %constant)
  %add = f32[1,1,128] add(f32[1,1,128] %reshape.1, f32[1,1,128] %reshape.2)

AFTER
  %constant = f32[128] constant({...})
  %add = f32[1,128] add(f32[1,128] %dot, f32[128] %constant)
  %reshape = f32[1,1,128] reshape(f32[1,128] %add)

The problem in AFTER is the add now contains an implicit broadcast. One way to
fix this is to re-shape the %constant to f32[1,128] before the %add.

Instead of that, the fix introduced in this CL is to simply prevent the
ReshapeMover from moving the reshapes in this case. A comment in
reshape_mover.cc describes the complexities that led to this choice.

Also added HloVerifiedTestBase, which keeps track of a default HloModule, and
automatically runs HloVerifier at the end of every test. This is useful for many
HLO tests; the tests of various passes can probably all use this. Three existing
issues in reshape_mover_test.cc were found and fixed as a result.

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