[Mlir-commits] [mlir] [mlir][Linalg] Allow more control in drop unit dims (PR #170104)

[Lukas Sommer]

sommerlukas wrote:

Thanks for your feedback!

> If preserving the encoding defaults to the previous behaviour, why isn't this the new default behaviour, instead of a test pass implementation?

Sorry, I didn't phrase that very well. 

The current default behavior also preserves encodings (i.e., it doesn't lose them), but it also doesn't try to collapse operands with an encoding.

One possible use of the new controls is to collapse operands if their encoding allows that. 

The difference is visible if we apply both variants of the test pass to the new test input.

This is the input: 
```mlir
#encoding = #test.tensor_encoding<"encoding">

func.func @drop_outermost_unit_dims_with_encoding(%arg0: tensor<1x1x42xf32, #encoding>) -> tensor<1x1x42xf32, #encoding> {
  %0 = tensor.empty() : tensor<1x1x42xf32, #encoding>
  %1 = 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(%arg0 : tensor<1x1x42xf32, #encoding>)
    outs(%0 : tensor<1x1x42xf32, #encoding>) {
      ^bb0(%b0: f32, %b1 : f32):
        %2 = arith.addf %b0, %b1 : f32
        linalg.yield %2 : f32
    } -> tensor<1x1x42xf32, #encoding>
  return %1 : tensor<1x1x42xf32, #encoding>
}
```

This is the output with the default behavior (before and after this PR):
```mlir
#map = affine_map<(d0) -> (0, 0, d0)>
module {
  func.func @drop_outermost_unit_dims_with_encoding(%arg0: tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>) -> tensor<1x1x42xf32, #test.tensor_encoding<"encoding">> {
    %0 = tensor.empty() : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
    %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>) outs(%0 : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>) {
    ^bb0(%in: f32, %out: f32):
      %2 = arith.addf %in, %out : f32
      linalg.yield %2 : f32
    } -> tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
    return %1 : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
  }
}
```

As you can see, the unit extent dimensions of the `linalg.generic` were folded, so we end up with only one iterator. However, the operands and output have not been folded and their type remains ` tensor<1x1x42xf32, #test.tensor_encoding<"encoding">` (so encoding is preserved), but their iteration map gets updated with `0`: `#map = affine_map<(d0) -> (0, 0, d0)>`.

With the additional controls, we can also fold the operands with encoding, as it's done by the new variant of the test pass:
```mlir
#map = affine_map<(d0) -> (d0)>
module {
  func.func @drop_outermost_unit_dims_with_encoding(%arg0: tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>) -> tensor<1x1x42xf32, #test.tensor_encoding<"encoding">> {
    %0 = tensor.empty() : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
    %collapsed = tensor.collapse_shape %arg0 [[0, 1, 2]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">> into tensor<42xf32, #test.tensor_encoding<"encoding">>
    %collapsed_0 = tensor.collapse_shape %0 [[0, 1, 2]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">> into tensor<42xf32, #test.tensor_encoding<"encoding">>
    %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%collapsed : tensor<42xf32, #test.tensor_encoding<"encoding">>) outs(%collapsed_0 : tensor<42xf32, #test.tensor_encoding<"encoding">>) {
    ^bb0(%in: f32, %out: f32):
      %2 = arith.addf %in, %out : f32
      linalg.yield %2 : f32
    } -> tensor<42xf32, #test.tensor_encoding<"encoding">>
    %expanded = tensor.expand_shape %1 [[0, 1, 2]] output_shape [1, 1, 42] : tensor<42xf32, #test.tensor_encoding<"encoding">> into tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
    return %expanded : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
  }
}
```
The operands and result are collapsed/expanded to match the new single iterator of the `linalg.generic`.

So, the default behavior has always been to preserve encoding and will continue to do so, but it doesn't collapse operands with encoding. The new controls enable to do that as one possible use case.

I'll rename the pass option for the test pass to better reflect that and also tweak the PR description a bit.

https://github.com/llvm/llvm-project/pull/170104