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

[Han-Chung Wang]

hanhanW wrote:

> In this context, I think this is the options struct should contain the following.
> 
> 1. The current control function. The control function already allows to control which dimensions are being dropped. That drives the calculation of the reassociation indices needed for expand and collapse. The only issue is w.r.t how to handle operands/results that have encodings. If there are no encodings, the transformation already knows what to do. If there is an encoding, we should

The sentence is not completed. I think if there is an encoding, we should remove the dims from the `allowedUnitDims` based on encoding's property. In this context, I think we can rename `controlFn` as `getDroppableDims`. By default, it can mark all the dimensions droppable, and remove relevant dims if an operand has an encoding. Reason: the upstream does not understand downstream encodings, so it is safer to drop them by default. In this context, the below code also needs an update:

https://github.com/llvm/llvm-project/blob/56beac9f0c2629b34d65dff2ea63605cc43dad39/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h#L527-L537

This is a precondition for user's other control functions, and it is user's responsibility to provide how to collapse/expand the operand in those functions.

> 2. The collapseValueFn callback should take as arguments the original operand and reassociation indices that tell you which dimensions of the original values to collapse. It should return the new value to use as operands to the newly created operation. In other words, without encoding the original value -> new value is obtained by using a `tensor.collapse_shape`.

I'm not sure if we want to move the op creation into the `collapseValueFn` or not. It may be a reasonable change because downstream project may have their own collapsing mechanism for downstream encodings. E.g., downstream_op.encode_and_maybe_cast.

It is no harm to expose the functionality, so I'm okay with this. In this context, you use the created value from the `collapseValueFn` callback for the new operation. Since we already know what dimensions are going to be dropped by (1), the transformation can construct the indexing map without any control function.

> 3. The expandValueFn callback should take as arguments the new result and the reassociation indices that tell you which dimensions of the new result value need to be expanded to get back the original rank. It should return the new value to use as replacements for the original result. In other words, without encoding the new result -> original value replacement is obtained using a `tensor.expand_shape`.

I think this step is similar to (2), so I'm not going to expand the context.

> 
> I think the change would be much more simple and targeted. It will be a breaking change, but thats OK. Its a one-time breaking change to use the options struct, after which the API wont break, and will be strictly additive.
> 
> Happy to chat on a video call if you have more questions.

I'm also happy to chat and see if I miss anything.

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