[Mlir-commits] [mlir] [mlir] Fix semantics of linalg::ReduceOp with several inputs (PR #107005)

[Clément Fournier]

oowekyala wrote:

Here is the affine version:
```mlir
    affine.for %i = 0 to 16 {
      affine.for %k = 0 to 32 {
        affine.for %j = 0 to 64 {
          %4 = affine.load %1[%i, %k, %j] : memref<16x32x64xi32>
          %5 = affine.load %0[%i, %k, %j] : memref<16x32x64xi32>
          %6 = affine.load %alloc[%i, %j] : memref<16x64xi32>
          %7 = arith.muli %4, %5 : i32
          %8 = arith.addi %6, %7 : i32
          affine.store %8, %alloc[%i, %j] : memref<16x64xi32>
        }
      }
    }
```

and the generic version:

```mlir
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>

    %0 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "reduction", "parallel"]} ins(%arg0, %arg1 : tensor<16x32x64xi32>, tensor<16x32x64xi32>) outs(%arg2 : tensor<16x64xi32>) {
    ^bb0(%in: i32, %in_0: i32, %out: i32):
      %1 = arith.muli %in, %in_0 : i32
      %2 = arith.addi %out, %1 : i32
      linalg.yield %2 : i32
    } -> tensor<16x64xi32>
 
```

> What are the constraints on the shapes of the inputs? Should they all be the same shape? Are they all accessed using the same indexing map and is this checked in the verifier?

The inputs should all have the same shape. This is because the shape of the output is the shape of the inputs with the reduction dimensions removed. So the reduction dimensions must match in extent, and the parallel dimensions must match because they determine the shape of the output, so in total the full shape must match. That is already checked by the verifier yes. The above code samples are generated with the `mlir-opt` built with this PR, so the way `linalg.reduce` maps to `linalg.generic` is not changed by this PR.


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