[Q] about reduce(reshape(x)) swap in algebraic_simplifier.cc HandleReduce

[Alexander Pivovarov]

Hello

I have a question about AlgebraicSimplifierVisitor::HandleReduce.Handle

reduce(reshape(x)) case in particular

Lets look at the following hlo example:

Input: 

%transpose.0 = bf16[20,256,8,160]{3,2,1,0} transpose(%reshape.0), dimensions={0,2,1,3} 

%reshape.1 = bf16[20,256,1280]{2,1,0} reshape(%transpose.0) 

%reduce.0 = bf16[1280]{0} reduce(%reshape.1, %constant.0), dimensions={0,1}, to_apply=%AddComputation.7

Output (after AlgebraicSimplifierVisitor::HandleReduce):

%transpose.0 = bf16[20,256,8,160]{3,2,1,0} transpose(%reshape.0), dimensions={0,2,1,3} 

%reduce.0 = bf16[8,160]{1,0} reduce(%transpose.0, %constant.0), dimensions={0,1}, to_apply=%AddComputation.7 

%reshape.1 = bf16[1280]{0} reshape(%reduce.0)

The code comments says:
Case2: This is beneficial because the reduce will now operate on less data.

The statement above is not 100% clear to me.

Why 

bf16[1280] reduce(bf16[20,256,1280])          (1)

operate on less data than 

bf16[8,160] reduce(bf16[20,256,8,160])         (2)

???

In the first and the second cases input size is 6,553,600 elems and output size is 1,280 elems - no difference in the amount of data.

The only diff is in input/output shapes and tensor ranks:

- In the first case input is 3d and output is 1d

- In the second case input is 4d and output is 2d

I'd say that the first case (without this optimization) looks a bit "simpler" because reduce deals with lower rank input/output tensors.

Docs and links:

AlgebraicSimplifierVisitor::HandleReduce Code comments

https://github.com/openxla/xla/blob/main/xla/service/algebraic_simplifier.cc#L6985 

// Handle two cases of reduce(reshape(x)). 

// 

// 1. The reshape collapses/expands only dimensions that are being reduced. 

// In this case we can just reduce those dimensions and skip the reshape. 

// 2. The reshape collapses/expands only dimensions that are *not* being 

// reduced. In this case we can do the reshape after the reduce. This is 

// beneficial because the reduce will now operate on less data.

https://github.com/openxla/xla/blob/main/xla/service/algebraic_simplifier.cc#L7041 

// Case 2: Check whether the reshape only modifies non-reduction dimensions. 

// Equivalently, the reduction dimensions are all preserved by the reshape.

Handle reduce(reshape(x)) Case 2 was added on Mar 29, 2023

https://github.com/openxla/xla/commit/4195a90a414564d9ec260042475c6945078c7e4a

Thank you

Alex

[Rahul Joshi]

I suspect it's a typo in the comment, it should be: "This is beneficial because the reshape will now operate on less data".

--
You received this message because you are subscribed to the Google Groups "OpenXLA Discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to openxla-discu...@openxla.org.
To view this discussion on the web visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/CAKKt98S-h--r4jGeghmjzeGKERiic_Jk%2B5c8eQxyUokJ2D8DtQ%40mail.gmail.com.
For more options, visit https://groups.google.com/a/openxla.org/d/optout.