Question about reduce(reshape(x)) swap in algebraic_simplifier.cc HandleReduce

[apivovarov]

Hello
I have a question about AlgebraicSimplifierVisitor::HandleReduce.

Handle reduce(reshape(x)) cases 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
4195a90

Justin @jlebar, can you comment on this ? Thank you.

[jurahul]

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