[Mlir-commits] [mlir] [MLIR][Linalg] Fix insert_slice fusion with rank reduction (PR #130961)

[Thomas Preud'homme]

RoboTux wrote:

> Thanks for the updates, Thomas!
> 
> I've left a few smallish comments, but before addressing those, it would be good to make sure that we do indeed need `tensor.collapse_shape` (as opposed to _a rank reducing_ `tensor.extract_slice`). I am sure that you thought it through and there's a good reason to insert `collapse_shape`, but it's not obvious to me just yet 😅

That's a very good question and I'll need to get back to you before I can answer fully. To be clear, the fusion does not insert a new extract_slice, this is done by another transform. I don't even think it was there in the earlier example and I don't know why there's a new linalg as well so I will look into understanding that. Normally there is an extract slice that might happen to collapse some unit dimension which then causes a type mismatch. This patch is to ensure we do the proper collapse_shape after the extract_slice is fused to reproduce that unit dimension collapsing and make the type match. By the way the if is because the pass requires the producing linalg and consuming linalg (the one that consume the extract_slice) to be in different block, the bug I was facing was a case of loop but I used a if here to make the code simpler.

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