# \[MLIR\]\[RFC\] Semantics of \`linalg.reduce\` with variadic operands

**URL:** <https://discourse.llvm.org/t/mlir-rfc-semantics-of-linalg-reduce-with-variadic-operands/83662>\
**Category:** MLIR\
**Created:** [December 13, 2024, 5:47pm UTC](https://discourse.llvm.org/t/mlir-rfc-semantics-of-linalg-reduce-with-variadic-operands/83662 "2024-12-13T17:47:15Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![oowekyala](https://sea1.discourse-cdn.com/flex021/user_avatar/discourse.llvm.org/oowekyala/32/23686_2.png) [@oowekyala](https://discourse.llvm.org/u/oowekyala)\
**Post date:** [December 13, 2024, 5:47pm UTC](https://discourse.llvm.org/t/mlir-rfc-semantics-of-linalg-reduce-with-variadic-operands/83662/1 "2024-12-13T17:47:15Z")

</div>

`linalg.reduce` currently has variadic input/output operands, like any DPS linalg op. The implementation of the verifier, as well as its documentation and tests, clearly show that the authors intended to support variadic operands in some form. However, there is confusion as to what a `reduce` with several operands should mean. In current linalg, this op also has the trait `SameVariadicOperandSize`, which adds to the confusion. This RFC is there to discuss

1. Whether `linalg.reduce` should even allow variadic operands, and if yes,
2. What the semantics of the op are when it has several operands.

@rengolin @MaheshRavishankar

# Current behavior

Current `linalg.reduce` allows having several inputs and outputs. In effect it allows you to write eg

```mlir
linalg.reduce
  ins(%input, %input2:memref<16x32x64xi32>, memref<16x32x64xi32>)
  outs(%init, %init2: memref<16x64xi32>, memref<16x64xi32>)
  dimensions = [1]
  (%in: i32, %in2: i32, %out: i32, %out2: i32) {
    %0 = arith.muli %in, %in2: i32
    %1 = arith.addi %out, %0: i32
    linalg.yield %1, %1: i32, i32
 }

```

The current semantics this code has can be seen more clearly when lowering to affine for instance:

```mlir
affine.for %i = 0 to 16 {
  affine.for %k = 0 to 32 {
    affine.for %j = 0 to 64 {
      %4 = affine.load %input[%i, %k, %j] : memref<16x32x64xi32>
      %5 = affine.load %input2[%i, %k, %j] : memref<16x32x64xi32>
      %6 = affine.load %init[%i, %j] : memref<16x64xi32>
      %7 = affine.load %init2[%i, %j] : memref<16x64xi32>
      %8 = arith.muli %4, %5 : i32
      %9 = arith.addi %6, %8 : i32
      affine.store %9, %init[%i, %j] : memref<16x64xi32>
      affine.store %9, %init2[%i, %j] : memref<16x64xi32>
    }
  }
}

```

So the current semantics that the op has in this case is, you can reduce with _one_ loop over same-size dimensions of two or more tensors. That seems to be well defined.

But there is a catch. The op requires to have the same numbers of inputs and outputs. That means if you want to contract two tensors like above, you _can_ use `linalg.reduce`, but you have to have a bogus second output that you will ignore. This is IMO an accident and not the intended behavior.

# Ways forward

In [#107005](https://github.com/llvm/llvm-project/pull/107005) I argue that that restriction is accidental and removed it. This change was pushed back as there is no consensus over whether `linalg.reduce` should be allowed to have several operands in the first place.

In any case there seems to be something wrong with the semantics of this op. Either

1. The accident is `SameVariadicOperandSize`, meaning the op should support variadic operands, but not have this restriction that num outputs == num inputs.
2. Or the accident is that this op should not be supporting variadic operands at all, and should be unary.

Option 1 seems to fit what the original authors intended, to me at least. Option 2 makes `reduce` more restricted, but maybe it fits better in the new tree-like design of [this RFC](https://discourse.llvm.org/t/rfc-mlir-linalg-operation-tree/83586).

As far as I can tell, with option 1,`reduce` with 2 inputs and one output is similar to the proposed [linalg.contract](https://discourse.llvm.org/t/mlir-rfc-introduce-linalg-contract/83589/2), but more general as it is not restricted to some combination of the _products_ of the input elements. It could be eg a combination of the sums or whatever. It seems that such a powerful op would fit between `generic` and `contract` in the abstraction tree. However it is probably useful to keep an operator for simple reduction over one dimension, as it seems to be how `linalg.reduce` is used in the wild (from a search on Github). So IMHO, a good compromise would be:

- restrict `linalg.reduce` to be unary;
- in the future, if needed, introduce a new operator between `generic` and `contract`, that supports an arbitrary combine/reduce region.

# Implementation

- PR [#107005](https://github.com/llvm/llvm-project/pull/107005) implemented option 1 above, but has been rejected for now.
- PR [#119871](https://github.com/llvm/llvm-project/pull/119871) sets the current behavior in stone by adding a test and fixing the verifier.
- If we want to make `linalg.reduce` unary, it might be a more involved change. Maybe we need to provide a new builder with just Value for the input and output for instance.

---

<div class="post-metadata">

**Author:** ![rengolin](https://sea1.discourse-cdn.com/flex021/user_avatar/discourse.llvm.org/rengolin/32/327_2.png) [@rengolin](https://discourse.llvm.org/u/rengolin)\
**Post date:** [December 13, 2024, 9:31pm UTC](https://discourse.llvm.org/t/mlir-rfc-semantics-of-linalg-reduce-with-variadic-operands/83662/2 "2024-12-13T21:31:13Z")

</div>

Thank you for the very detailed RFC!

> [@oowekyala](#):
>
> This is IMO an accident and not the intended behavior.

Agree. Even if someone is relying on this _“feature”_, this could be dangerous. I may be wrong, but this looks like an eager over-generalization without a lot of thought behind.

> [@oowekyala](#):
>
> As far as I can tell, with option 1,`reduce` with 2 inputs and one output is similar to the proposed [linalg.contract](https://discourse.llvm.org/t/mlir-rfc-introduce-linalg-contract/83589/2), but more general as it is not restricted to some combination of the _products_ of the input elements. It could be eg a combination of the sums or whatever

We discussed in the `contract` RFC that we do want to support that case, but not now, as there are no real users yet and it adds complexity. We’ll leave it as a TODO for now but definitely add as soon as the op is stable and the semantics is agreed upon.

> [@oowekyala](#):
>
> So IMHO, a good compromise would be:
> 
> - restrict `linalg.reduce` to be unary;
> - in the future, if needed, introduce a new operator between `generic` and `contract`, that supports an arbitrary combine/reduce region.

Agree! It can even be `contract` itself that supports it. But we can discuss that when time comes.

> [@oowekyala](#):
>
> If we want to make `linalg.reduce` unary, it might be a more involved change. Maybe we need to provide a new builder with just Value for the input and output for instance.

Unless there’s a real use case (that cannot be solved by other ways), it should be mostly a builder issue and update the tests.