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