Rework split_reshape simplification to handle multiple reshapes from same slice #2146

shivadbhavsar · 2023-08-31T22:28:17Z

This resolves an edge case found in multiple huggingface models where the following happens:

slice1 -> cont -> reshape1
       |
        -> cont -> reshape2

slice2 -> cont -> reshape1
       |
        -> cont -> reshape2

slice3 -> cont -> reshape1
       |
        -> cont -> reshape2

In some cases the find_split_reshape matcher will match with reshape2, but vec_rsp will consist of reshape1 dims causing a shape mismatch error. Solution is to include rsp when checking all reshapes are the same.

Added unit test reproduces this issue in the current build.

pfultz2 · 2023-08-31T22:55:16Z

It seems like this should be written as:

reshape1 -> slice1, slice2, slice3
reshape2 -> slice4,slice5,slice6

Rather than just stopping.

codecov · 2023-08-31T23:50:36Z

Codecov Report

Merging #2146 (4cc9076) into develop (c2e01b1) will increase coverage by 0.00%.
The diff coverage is 100.00%.

❗ Current head 4cc9076 differs from pull request most recent head 3400122. Consider uploading reports for the commit 3400122 to get more accurate results

@@           Coverage Diff            @@
##           develop    #2146   +/-   ##
========================================
  Coverage    91.48%   91.49%           
========================================
  Files          427      427           
  Lines        15938    15953   +15     
========================================
+ Hits         14581    14596   +15     
  Misses        1357     1357

Files Changed	Coverage Δ
src/simplify_algebra.cpp	`96.85% <100.00%> (+0.06%)`	⬆️

migraphx-bot · 2023-09-01T05:54:43Z

Test	Batch	Rate new 340012	Rate old c2e01b	Diff	Compare
torchvision-resnet50	64	2,282.79	2,284.87	-0.09%	✅
torchvision-resnet50_fp16	64	5,367.78	5,357.62	0.19%	✅
torchvision-densenet121	32	1,835.07	1,829.81	0.29%	✅
torchvision-densenet121_fp16	32	3,391.29	3,378.73	0.37%	✅
torchvision-inceptionv3	32	1,336.41	1,335.78	0.05%	✅
torchvision-inceptionv3_fp16	32	2,588.08	2,585.25	0.11%	✅
cadene-inceptionv4	16	678.73	679.39	-0.10%	✅
cadene-resnext64x4	16	589.72	590.19	-0.08%	✅
slim-mobilenet	64	7,213.88	7,216.76	-0.04%	✅
slim-nasnetalarge	64	237.05	236.57	0.20%	✅
slim-resnet50v2	64	2,528.69	2,530.71	-0.08%	✅
bert-mrpc-onnx	8	721.09	721.25	-0.02%	✅
bert-mrpc-tf	1	390.79	391.28	-0.12%	✅
pytorch-examples-wlang-gru	1	310.08	303.90	2.04%	✅
pytorch-examples-wlang-lstm	1	310.72	311.77	-0.34%	✅
torchvision-resnet50_1	1	556.07	555.70	0.07%	✅
torchvision-inceptionv3_1	1	307.56	306.97	0.19%	✅
cadene-dpn92_1	1	352.44	351.70	0.21%	✅
cadene-resnext101_1	1	220.50	220.52	-0.01%	✅
slim-vgg16_1	1	224.86	224.37	0.22%	✅
slim-mobilenet_1	1	1,483.04	1,468.80	0.97%	✅
slim-inceptionv4_1	1	221.72	221.26	0.21%	✅
onnx-taau-downsample	1	322.55	322.42	0.04%	✅
dlrm-criteoterabyte	1	21.68	21.67	0.02%	✅
dlrm-criteoterabyte_fp16	1	40.62	40.54	0.19%	✅
agentmodel	1	5,836.42	5,779.11	0.99%	✅
unet_fp16	2	55.08	55.07	0.02%	✅

This build is OK for merge ✅

migraphx-bot · 2023-09-01T05:54:45Z

:white_check_mark:bert-mrpc-onnx: PASSED: MIGraphX meets tolerance

:white_check_mark:bert-mrpc-tf: PASSED: MIGraphX meets tolerance

:white_check_mark:pytorch-examples-wlang-gru: PASSED: MIGraphX meets tolerance

:white_check_mark:pytorch-examples-wlang-lstm: PASSED: MIGraphX meets tolerance

:white_check_mark:torchvision-resnet50_1: PASSED: MIGraphX meets tolerance

🔴torchvision-inceptionv3_1: FAILED: MIGraphX is not within tolerance - check verbose output

🔴cadene-dpn92_1: FAILED: MIGraphX is not within tolerance - check verbose output

:white_check_mark:cadene-resnext101_1: PASSED: MIGraphX meets tolerance

:white_check_mark:slim-vgg16_1: PASSED: MIGraphX meets tolerance

:white_check_mark:slim-mobilenet_1: PASSED: MIGraphX meets tolerance

🔴slim-inceptionv4_1: FAILED: MIGraphX is not within tolerance - check verbose output

:white_check_mark:dlrm-criteoterabyte: PASSED: MIGraphX meets tolerance

:white_check_mark:agentmodel: PASSED: MIGraphX meets tolerance

:white_check_mark:unet: PASSED: MIGraphX meets tolerance

umangyadav · 2023-09-05T15:45:02Z

In some cases the find_split_reshape matcher will match with reshape2, but vec_rsp will consist of reshape1 dims causing a shape mismatch error

Looks like Matcher was assuming slice only has one output but is not enforcing it.

        // Only want to apply this optimization if each split output is followed by
        // a contiguous op and a reshape
        if(std::any_of(split_outputs.begin(), split_outputs.end(), [](auto i) {
               if(i->outputs().size() == 1)
               {
                   auto cont = i->outputs().front();
                   return cont->outputs().size() != 1;
               }
               return false;
           }))
        {
            return;
        }

This part if outputs are more than 1 then it is returning false, but it should return true i think.

shivadbhavsar · 2023-09-05T17:18:08Z

In some cases the find_split_reshape matcher will match with reshape2, but vec_rsp will consist of reshape1 dims causing a shape mismatch error

Looks like Matcher was assuming slice only has one output but is not enforcing it.
        // Only want to apply this optimization if each split output is followed by
        // a contiguous op and a reshape
        if(std::any_of(split_outputs.begin(), split_outputs.end(), [](auto i) {
               if(i->outputs().size() == 1)
               {
                   auto cont = i->outputs().front();
                   return cont->outputs().size() != 1;
               }
               return false;
           }))
        {
            return;
        }
This part if outputs are more than 1 then it is returning false, but it should return true i think.

That was my initial fix actually. But I figured it was coded this way because we wanted to allow the splits to have other outputs but only simpilfy the instructions leading to the reshapes. I guess its not really doing that though because we just assume single output (and that its cont --> reshape) for each slice when we create vec_rsp.

In any case, this would solve the bug but it wont do what @pfultz2 is suggesting. For that we'd have to add more checks for what going into vec_rsp.

…slice

shivadbhavsar · 2023-09-06T00:51:17Z

This turned into a larger update than intended but here is a summary of changes:

The initial bug was in part due to a set of (slice -> cont -> reshape)s where the reshapes did not use the entire input shape. Eg. Input: [4, 6, 30], Slices: [4, 2, 30], [4, 2, 30], Reshapes: [4, 60], [4, 60]. Here, our current logic would try to create a new reshape (from input [4, 6, 60]) of size [4, 120], which is invalid (should be [4, 180]). Allowing this kind of simplification required modifying how the new slice indices are computed.
It no longer assumes that the slices have only a single (contiguous -> reshape) attached. We do not modify the existing slice -> contiguous so any nodes using these as inputs should still be valid. If no nodes use these as inputs, dead code elim will take care of it.

src/simplify_algebra.cpp

umangyadav · 2023-09-06T15:36:26Z

Unet accuracy check has failed on compilation. Can you check ?

shivadbhavsar · 2023-09-07T23:02:42Z

Unet accuracy check has failed on compilation. Can you check ?

This should be fixed in the new run

pfultz2 · 2023-09-07T23:09:09Z

src/simplify_algebra.cpp

+        // cannot create a valid reshape for simplification
+        if(input_size % rsp_fixed_size != 0)
+        {
+            return;


Is it possible to add a unit test for this case?

causten · 2023-09-16T15:32:50Z

But I figured it was coded this way

Is there enough comments in the code to help understand why something was done a certain way? If you need to make a code change for any other reason feel free to improve the code comments.

src/simplify_algebra.cpp

shivadbhavsar · 2023-09-19T01:08:13Z

For more clarity heres an example walking through the code (this simplification wouldnt work with the original implementation):

Original Instuctions:
Input: {128, 96}
Slices: 
s0 slice(axes={0}, starts={0}, ends={8})(inp) -> {8, 96}
s1 slice(axes={0}, starts={8}, ends={16})(inp) -> {8, 96}
s2 slice(axes={0}, starts={16}, ends={24})(inp) -> {8, 96}
s3 slice(axes={0}, starts={24}, ends={128})(inp) -> {104, 96}
Reshapes:
reshape(dims={2, 4, 96})(s0)
reshape(dims={2, 4, 96})(s1)
reshape(dims={2, 4, 96})(s2)

Intermediate Variables:
axis = 0
slc_lens = {8, 96}
slc_axis_len = input[axis] = 128
slc_dim_size = 8*96 = 768

rsp_lens = {2, 4, 96}, strides: {96*4*2, 96*4, 96, 1} = {768, 384, 96, 1} (with added stride)
rsp_axis = <index where (rsp_stride == slc_dim_size)> = 0

rsp_out_lens (initial) = {1, 4, 96}
rsp_fixed_size = 1*4*96 = 384
rsp_axis_len = 128 * 96 / 384 = 32 
rsp_out_lens (final) = {32, 4, 96}

new_starts = {0*32/128,  8*32/128, 16*32/128} = {0, 2, 4}
new_ends   = {8*32/128, 16*32/128, 24*32/128} = {2, 4, 6}


New Instructions:
reshape(dims={32, 4, 96})
slice(axes={0}, starts={0}, ends={2}) -> {2, 4, 96}
slice(axes={0}, starts={2}, ends={4}) -> {2, 4, 96}
slice(axes={0}, starts={4}, ends={6}) -> {2, 4, 96}

src/simplify_algebra.cpp

umangyadav · 2023-09-19T13:04:32Z

For more clarity heres an example walking through the code (this simplification wouldnt work with the original implementation):

Original Instuctions:
Input: {128, 96}
Slices: 
s0 slice(axes={0}, starts={0}, ends={8})(inp) -> {8, 96}
s1 slice(axes={0}, starts={8}, ends={16})(inp) -> {8, 96}
s2 slice(axes={0}, starts={16}, ends={24})(inp) -> {8, 96}
s3 slice(axes={0}, starts={24}, ends={128})(inp) -> {104, 96}
Reshapes:
reshape(dims={2, 4, 96})(s0)
reshape(dims={2, 4, 96})(s1)
reshape(dims={2, 4, 96})(s2)

Intermediate Variables:
axis = 0
slc_lens = {8, 96}
slc_axis_len = input[axis] = 128
slc_dim_size = 8*96 = 768

rsp_lens = {2, 4, 96}, strides: {96*4*2, 96*4, 96, 1} = {768, 384, 96, 1} (with added stride)
rsp_axis = <index where (rsp_stride == slc_dim_size)> = 0

rsp_out_lens (initial) = {1, 4, 96}
rsp_fixed_size = 1*4*96 = 384
rsp_axis_len = 128 * 96 / 384 = 32 
rsp_out_lens (final) = {32, 4, 96}

new_starts = {0*32/128,  8*32/128, 16*32/128} = {0, 2, 4}
new_ends   = {8*32/128, 16*32/128, 24*32/128} = {2, 4, 6}


New Instructions:
reshape(dims={32, 4, 96})
slice(axes={0}, starts={0}, ends={2}) -> {2, 4, 96}
slice(axes={0}, starts={2}, ends={4}) -> {2, 4, 96}
slice(axes={0}, starts={4}, ends={6}) -> {2, 4, 96}

It would be good to add this a test instead of one with reorder_reshape_slice_uneven_slice which would have worked without your changes as well. Better Keep both.

umangyadav · 2023-09-19T13:43:27Z

@shivadbhavsar

onnx-taau-downsample	1	247.81	322.80	-23.23%	🔴

Can you check this one ?

split_reshape bug fix

7b68324

shivadbhavsar requested review from pfultz2 and umangyadav August 31, 2023 22:28

Merge branch 'develop' into fix_split_reshape

47ebed5

shivadbhavsar mentioned this pull request Aug 31, 2023

Dynamo Benchmarks - Simplify algebra reshape size error #2136

Closed

allow slice reshape simplification for multiple reshapes from single …

d103134

…slice

shivadbhavsar changed the title ~~Include matched reshape instruction in split_reshape in same_ops check~~ Rework split_reshape simplification to handle multiple reshapes from same slice Sep 6, 2023

pfultz2 reviewed Sep 6, 2023

View reviewed changes

src/simplify_algebra.cpp Outdated Show resolved Hide resolved

causten assigned shivadbhavsar Sep 6, 2023

shivadbhavsar and others added 4 commits September 6, 2023 19:46

use elements() to get arg size

4b286e5

Merge branch 'develop' into fix_split_reshape

f3ccae3

Merge branch 'develop' into fix_split_reshape

6f5ade4

handle edge case where dimension is not a multiple of the split sizes

a349cf0

pfultz2 reviewed Sep 7, 2023

View reviewed changes

add test case for invalid reshape axis len

dfe9465

shivadbhavsar requested a review from pfultz2 September 8, 2023 19:28

Merge branch 'develop' into fix_split_reshape

50cccbe

shivadbhavsar linked an issue Sep 12, 2023 that may be closed by this pull request

Dynamo Benchmarks - Simplify algebra reshape size error #2136

Closed

causten requested a review from TedThemistokleous September 15, 2023 17:21

TedThemistokleous added roadmap Tasks to finish for a release enhancement New feature or request labels Sep 15, 2023

TedThemistokleous added TorchMIGraphX high priority A PR with high priority for review and merging. labels Sep 15, 2023

TedThemistokleous approved these changes Sep 15, 2023

View reviewed changes

umangyadav reviewed Sep 18, 2023

View reviewed changes

src/simplify_algebra.cpp Show resolved Hide resolved

umangyadav reviewed Sep 18, 2023

View reviewed changes

src/simplify_algebra.cpp Show resolved Hide resolved

shivadbhavsar added 2 commits September 18, 2023 19:22

add check for single axis when performing split reshape simplification

0849f48

add comments with examples for clarity

2ebb711

shivadbhavsar requested a review from umangyadav September 19, 2023 00:56

umangyadav reviewed Sep 19, 2023

View reviewed changes

src/simplify_algebra.cpp Show resolved Hide resolved

umangyadav approved these changes Sep 19, 2023

View reviewed changes

shivadbhavsar and others added 2 commits September 19, 2023 17:16

add additional test case for slice_reshape

01f35e5

Merge branch 'develop' into fix_split_reshape

3400122

causten merged commit 0bb8508 into develop Sep 19, 2023
12 checks passed

causten deleted the fix_split_reshape branch September 19, 2023 23:47

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Rework split_reshape simplification to handle multiple reshapes from same slice #2146

Rework split_reshape simplification to handle multiple reshapes from same slice #2146

shivadbhavsar commented Aug 31, 2023

pfultz2 commented Aug 31, 2023

codecov bot commented Aug 31, 2023 •

edited

Loading

migraphx-bot commented Sep 1, 2023 •

edited

Loading

migraphx-bot commented Sep 1, 2023 •

edited

Loading

umangyadav commented Sep 5, 2023

shivadbhavsar commented Sep 5, 2023

shivadbhavsar commented Sep 6, 2023

umangyadav commented Sep 6, 2023

shivadbhavsar commented Sep 7, 2023

pfultz2 Sep 7, 2023

shivadbhavsar Sep 8, 2023

causten commented Sep 16, 2023

shivadbhavsar commented Sep 19, 2023

umangyadav commented Sep 19, 2023 •

edited

Loading

umangyadav commented Sep 19, 2023

Rework split_reshape simplification to handle multiple reshapes from same slice #2146

Rework split_reshape simplification to handle multiple reshapes from same slice #2146

Conversation

shivadbhavsar commented Aug 31, 2023

pfultz2 commented Aug 31, 2023

codecov bot commented Aug 31, 2023 • edited Loading

Codecov Report

migraphx-bot commented Sep 1, 2023 • edited Loading

migraphx-bot commented Sep 1, 2023 • edited Loading

umangyadav commented Sep 5, 2023

shivadbhavsar commented Sep 5, 2023

shivadbhavsar commented Sep 6, 2023

umangyadav commented Sep 6, 2023

shivadbhavsar commented Sep 7, 2023

pfultz2 Sep 7, 2023

Choose a reason for hiding this comment

shivadbhavsar Sep 8, 2023

Choose a reason for hiding this comment

causten commented Sep 16, 2023

shivadbhavsar commented Sep 19, 2023

umangyadav commented Sep 19, 2023 • edited Loading

umangyadav commented Sep 19, 2023

codecov bot commented Aug 31, 2023 •

edited

Loading

migraphx-bot commented Sep 1, 2023 •

edited

Loading

migraphx-bot commented Sep 1, 2023 •

edited

Loading

umangyadav commented Sep 19, 2023 •

edited

Loading