[Mlir-commits] [mlir] [mlir][Linalg] Fix fusing of indexed linalg consumer with different axes (PR #140892)

[Simone Pellegrini]

simpel01 wrote:

Hi, thanks for the comments.

In order to trigger the compiler assert that this patch fixes, I need an N-D `linalg.generic` to be folded into some M-D `linalg.generic` with M<N. 

I haven't mentioned this, but without the fix the compiler would generate the following linalg (out of the example I added):
```
  %3 = "linalg.generic"(%2) <{indexing_maps = [affine_map<(d0) -> (d0)>], iterator_types = [#linalg.iterator_type<parallel>], operandSegmentSizes = array<i32: 0, 1>}> ({
  ^bb0(%arg1: i32):
    %4 = "linalg.index"() <{dim = 0 : i64}> : () -> index
    %5 = "linalg.index"() <{dim = 1 : i64}> : () -> index
    %6 = "affine.apply"(%4, %5) <{map = affine_map<(d0) -> (1)>}> : (index, index) -> index
    %7 = "arith.index_cast"(%6) : (index) -> i32
    "linalg.yield"(%7) : (i32) -> ()
  }) : (tensor<2xi32>) -> tensor<2xi32>
```
which is illegal as:
```
'linalg.index' op expected dim (1) to be lower than the number of loops (1) of the enclosing LinalgOp
```
and later the compiler hits the following assert:
```
mlir-opt: llvm-project/mlir/lib/IR/AffineMap.cpp:456: mlir::AffineMap mlir::AffineMap::partialConstantFold(llvm::ArrayRef<mlir::Attribute>, llvm::SmallVectorImpl<long int>*, bool*) const: Assertion `getNumInputs() == operandConstants.size()' failed.
```

The affine expression per se is not really important as long as is not trivial and requires remapping. For example using something trivial like this:
```
affine_map<(d0) -> (d0, d0)>
```
won't cause the assert to trigger.

The expression I used in my latest example is the simplest strictly affine map that causes the assert to trigger. The example that motivated this change was what I used in the original patch based on the `mod` and `floordiv` expressions.

I think it is going to be difficult to come up with something else that triggers the assert whilst remaining in the "strictly affine realm". But I am not an expert in this subject, so suggestions for suitable 1D into 2D strictly affine expressions are welcome! :)

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