Propagate reshapes through generics with reduction iterators

[IanWood1]

Removes the constraint in BubbleUpExpandShapes that prevents moving tensor reshape ops through reduction linalg.generic ops. This has the benefit of increasing the dimensionality of reduction ops (more fusion opportunities) as well as increasing the chance these ops will be moved to the edge of the program.

Closes #18854

[IanWood1]

Tracked regression to the following pattern (input to BubbleUpExpandShapes):

// Note %420 is used by %421 and %423
%420 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%collapsed_476 : tensor<32x4x1048576xf16>) outs(%358 : tensor<32x4x1048576xf32>) {
  ^bb0(%in: f16, %out: f32):
    %436 = arith.extf %in : f16 to f32
    linalg.yield %436 : f32
  } -> tensor<32x4x1048576xf32>
  %421 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0)>], iterator_types = ["parallel", "reduction", "reduction"]} ins(%420 : tensor<32x4x1048576xf32>) outs(%20 : tensor<32xf32>) {
  ^bb0(%in: f32, %out: f32):
    %436 = arith.addf %in, %out : f32
    linalg.yield %436 : f32
  } -> tensor<32xf32>
  %422 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%421 : tensor<32xf32>) outs(%19 : tensor<32xf32>) {
  ^bb0(%in: f32, %out: f32):
    %436 = arith.divf %in, %cst_3 : f32
    linalg.yield %436 : f32
  } -> tensor<32xf32>
  %423 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0)>, affine_map<(d0, d1, d2) -> (d0)>], iterator_types = ["parallel", "reduction", "reduction"]} ins(%420, %422 : tensor<32x4x1048576xf32>, tensor<32xf32>) outs(%20 : tensor<32xf32>) {
  ^bb0(%in: f32, %in_488: f32, %out: f32):
    %436 = arith.subf %in, %in_488 : f32
    %437 = arith.mulf %436, %436 : f32
    %438 = arith.addf %437, %out : f32
    linalg.yield %438 : f32
  } -> tensor<32xf32>
  %expanded_477 = tensor.expand_shape %422 [[0, 1]] output_shape [1, 32] : tensor<32xf32> into tensor<1x32xf32>
  %expanded_478 = tensor.expand_shape %423 [[0, 1]] output_shape [1, 32] : tensor<32xf32> into tensor<1x32xf32>

An tensor.expand_shape get caught in between %420 and %423 because %420 has multiple uses. I removed the check that prevents bubbling if the producer has multiple uses, and it seemed to clear up the problem. I'm currently testing locally to see what other regressions pop up

cc @MaheshRavishankar

[IanWood1]

Note: this is because the collapse_shape (from drop unit dims) now gets propagated down, and CollapseDimensionsPass won't collapse operations with multiple reduction iterators (waiting for codegen support)

[MaheshRavishankar]

I like having as little checks as possible here. THanks!