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

[Lukas Sommer]

sommerlukas wrote:


This PR is a breaking change because it changes the default behavior of the function for operations with operands that cannot be collapsed. 

In the default implementation, the following types are considered as not collapsible:
- `memref` with non-identify layout.
- `tensor` with an encoding.

Previously, in this case, the transformation would still drop the unit extent dimensions from the operation, but keep the shape of operands with a non-collapsible type and insert 0-entries into their indexing maps. 

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

func.func @drop_unit_dims_encoded_operand(%arg0: tensor<1x1x42xf32>, %arg1: tensor<1x1x42xf32, #encoding>) -> tensor<1x1x42xf32> {
  %0 = tensor.empty() : tensor<1x1x42xf32>
  %1 = linalg.generic {
    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
    iterator_types = ["parallel", "parallel", "parallel"]}
    ins(%arg0, %arg1 : tensor<1x1x42xf32>, tensor<1x1x42xf32, #encoding>) outs(%0 : tensor<1x1x42xf32>) {
      ^bb0(%in0: f32, %in1 : f32, %out : f32):
        %2 = arith.addf %in0, %in1 : f32
        linalg.yield %2 : f32
    } -> tensor<1x1x42xf32>
  return %1 : tensor<1x1x42xf32>
}
```

Previous to this PR, the transformation would produce the following output:
```mlir
#map = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (0, d0, d1)>
module {
  func.func @drop_unit_dims_encoded_operand(%arg0: tensor<1x1x42xf32>, %arg1: tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>) -> tensor<1x1x42xf32> {
    %0 = tensor.empty() : tensor<1x1x42xf32>
    %collapsed = tensor.collapse_shape %arg0 [[0, 1], [2]] : tensor<1x1x42xf32> into tensor<1x42xf32>
    %collapsed_0 = tensor.collapse_shape %0 [[0, 1], [2]] : tensor<1x1x42xf32> into tensor<1x42xf32>
    %1 = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]} ins(%collapsed, %arg1 : tensor<1x42xf32>, tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>) outs(%collapsed_0 : tensor<1x42xf32>) {
    ^bb0(%in: f32, %in_1: f32, %out: f32):
      %2 = arith.addf %in, %in_1 : f32
      linalg.yield %2 : f32
    } -> tensor<1x42xf32>
    %expanded = tensor.expand_shape %1 [[0, 1], [2]] output_shape [1, 1, 42] : tensor<1x42xf32> into tensor<1x1x42xf32>
    return %expanded : tensor<1x1x42xf32>
  }
}
```

With the changes from this PR, the transformation is now aborted in these cases, i.e., the operation is not changed.

The rationale for this change is that the purpose of this transformation is to produce a _normalized_ form of operations by dropping unit extent dimensions. After offline discussion with @hanhanW  and @MaheshRavishankar , the form of the operation with 0-entries in the indexing map and some operands being reshaped while others are not is not more normalized than the input, so the transformation is now aborted in these cases instead.

Due to the hooks exposed to customize the behavior it is also not possible to get back the old behavior after this PR.

Another important change resulting from the discussion with @hanhanW and @MaheshRavishankar  affects the compatibility with greedy rewrite drivers: In the case the transformation is aborted on a later operand, `collapse_shape` for earlier operands have already been inserted. These are trivially dead if the transformation is aborted, so can be can removed, but would affect a greedy driver from reaching a fix point. 

The rationale for this change is that a repeated application of the transformation is not necessary, as the first application already drops unit-extent dimensions if possible.

After this change, the `dropUnitDims` API should not be used with greedy rewriters and use `walkAndApply` instead. This was also documented in the API description.

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