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

[llvmlistbot at llvm.org]

https://github.com/MaheshRavishankar requested changes to this pull request.

Thanks @sommerlukas  for taking a stab at this. I should have probably provided a bit more context for you before sending you down this road :) . I think the options struct here is way to heavy-weight. The options struct should just be a container for the different ways in which the transformations can be controlled. I think here the options struct is now encompassing actual implementation details of the transformations as methods. The options struct should just have fields that represent the different options, and methods to set/retrieve those options. 

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
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`.
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 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.

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