[Mlir-commits] [mlir] [MLIR][Linalg] Add aggregate ops decomposition pass and softmax decom… (PR #97582)

[Petr Kurapov]

kurapov-peter wrote:

> > The implementation ends up generating the semantically the same code as the previous decomposition implementation (with minor deviation as you noted).
> 
> IMO this is key - please add that in the summary.
> 

Done.

> > > Also, what's the difference between "transform-op-decompose.mlir" and "decompose-named-ops.mlir"? Is it just "how" the decomposition is "driven"? (TD vs Pass) Couldn't that be one test instead?
> > 
> > 
> > The first one tests the new pass. The second one uses the transform interpreter and the decompose_interface op (which happen to partially rely on the same code now).
> 
> From what I can tell, both tests verify the decomposition of `linalg.softmax`:
> 
> ```mlir
> func.func @softmax(%arg0: tensor<2x16x32xf32>, %dst: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
>   %1 = linalg.softmax dimension(2) ins(%arg0 : tensor<2x16x32xf32>) outs(%dst: tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
>   return %1 : tensor<2x16x32xf32>
> }
> ```
> 
> Couldn't we re-use the input and the `CHECK` lines? To avoid duplication.
> 

Do I understand correctly that you suggest having a single lit with the body of `@softmax`, the transformation IR, and runs both the decomposition via the pass and transform interpreter?

> FWIW (without links to documentation), that IR is just an implementation detail of torch-mlir. In this PR we are discussing an implementation detail of Linalg. Are you saying that the implementation in Linalg should match torch-mlir? Why? What if the implementation in torch-mlir changes?
> 
> I'm trying to understand the motivation here and the overall design that we are converging towards.
> 

So the goal was to look into what frameworks actually do in the implementation. If it so happens that all of them lower softmax to the same sequence (and this is what we happen to have) - we can have it set as the default decomposition to avoid re-implementing the thing. The general idea and direction is to have an intermediate decomposition stage that deals with complex ops (such as softmax) to aid other transformations and analyses (this is also an easy route to adding more named ops, upstream and downstream, and not implement all the interfaces like tiling, but convert it to simpler ops instead and enjoy all the existing goodness). Note: I'm leaving the question of accumulation out for now, this should be addressed separately.

> > > Are you saying that this PR is changing the semantics of softmax in Linalg?
> > 
> > 
> > I'd say it sets it. The implementation follows the op description, so there's no real 'change'.
> 
> Key info, please highlight in the summary.

Done.

> I know what the options are, thanks. To me, your comment implies that `-linalg-decompose-named-ops` is meant to be followed by `-linalg-generalize-named-ops` ("aggregate ops are _decomposed_ first and then _generalized_")? Is that what you had in mind?

This just describes the transform op decomposition change (in other words it used to produce generics + fills, now it internally produces named ops sequence and then runs generalization), there's no strict requirement to run generalization after the decomposition of course.

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