MLIR crash in linalg::ReduceOp::verify

[oowekyala]

ReduceOp crashes during verification of the following op:

module {
  func.func @main(%arg0: tensor<32xi32>, %arg1: tensor<32xi32>) -> i32 {
      %c0_i32 = arith.constant 0 : i32
      %from_elements = tensor.from_elements %c0_i32 : tensor<i32>
      %reduced = linalg.reduce ins(%arg0, %arg1 : tensor<32xi32>, tensor<32xi32>) outs(%from_elements : tensor<i32>) dimensions = [0] 
        (%in0: i32, %in1: i32, %init: i32) {
          %44 = arith.muli %in0, %in1 : i32
          %5 = arith.addi %44, %init : i32
          linalg.yield %5 : i32
        }
      %ext = tensor.extract %reduced[] : tensor<i32>
      return %ext: i32
  }
}

This is at line 1656 here:

llvm-project/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Lines 1652 to 1658 in fb87e11

	LogicalResult ReduceOp::verify() {
	ArrayRef<int64_t> dimensionsRef = getDimensions();
	for (int64_t i = 1; i < getNumDpsInputs(); ++i) {
	if (llvm::cast<ShapedType>(getInputs()[i].getType()).getShape() !=
	llvm::cast<ShapedType>(getInputs()[0].getType()).getShape()) {
	return emitOpError() << "expects all inputs to have the same shapes. "

It appears the ODS-generated code for getInputs assumes that the variadic operands (inputs and inits here) are the same sizes. When computing the size of the variadic operand, it foordivs 3 by 2 and returns a range with 1 element. On the other hand the code in here is correct and returns an input count of 2, and an operand range with the two inputs. That's why here we enter the loop, but getInputs()[1] crashes.

One fix for the crash would be to use getDpsInputs instead of getInputs in the verifier. But this should ideally be patched in ODS I suppose. Thanks for reading.

---

Stack trace

cinm-opt: /home/clem/Documents/LLVM/llvm-18.1.6/llvm/include/llvm/ADT/STLExtras.h:1281: ReferenceT llvm::detail::indexed_accessor_range_base<mlir::OperandRange, mlir::OpOperand *, mlir::Value, mlir::Value, mlir::Value>::operator[](size_t) const [DerivedT = mlir::OperandRange, BaseT = mlir::OpOperand *, T = mlir::Value, PointerT = mlir::Value, ReferenceT = mlir::Value]: Assertion `Index < size() && "invalid index for value range"' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.	Program arguments: cinm-opt sandbox/matmul_tiled.mlir
 #0 0x000055eadabe7a4d llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) /home/clem/Documents/LLVM/llvm-18.1.6/llvm/lib/Support/Unix/Signals.inc:723:11
 #1 0x000055eadabe803b PrintStackTraceSignalHandler(void*) /home/clem/Documents/LLVM/llvm-18.1.6/llvm/lib/Support/Unix/Signals.inc:798:1
 #2 0x000055eadabe5ce3 llvm::sys::RunSignalHandlers() /home/clem/Documents/LLVM/llvm-18.1.6/llvm/lib/Support/Signals.cpp:105:5
 #3 0x000055eadabe8801 SignalHandler(int) /home/clem/Documents/LLVM/llvm-18.1.6/llvm/lib/Support/Unix/Signals.inc:413:1
 #4 0x00007f5e4b450ae0 (/usr/lib/libc.so.6+0x3cae0)
 #5 0x00007f5e4b4a8e44 (/usr/lib/libc.so.6+0x94e44)
 #6 0x00007f5e4b450a30 raise (/usr/lib/libc.so.6+0x3ca30)
 #7 0x00007f5e4b4384c3 abort (/usr/lib/libc.so.6+0x244c3)
 #8 0x00007f5e4b4383df (/usr/lib/libc.so.6+0x243df)
 #9 0x00007f5e4b448c67 (/usr/lib/libc.so.6+0x34c67)
#10 0x000055ead608a178 llvm::detail::indexed_accessor_range_base<mlir::OperandRange, mlir::OpOperand*, mlir::Value, mlir::Value, mlir::Value>::operator[](unsigned long) const /home/clem/Documents/LLVM/llvm-18.1.6/llvm/include/llvm/ADT/STLExtras.h:0:5
#11 0x000055ead7199ce5 mlir::linalg::ReduceOp::verify() /home/clem/Documents/LLVM/llvm-18.1.6/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:1584:32
#12 0x000055ead715bb94 mlir::Op<mlir::linalg::ReduceOp, mlir::OpTrait::OneRegion, mlir::OpTrait::VariadicResults, mlir::OpTrait::ZeroSuccessors, mlir::OpTrait::VariadicOperands, mlir::OpTrait::S

[llvmbot]

@llvm/issue-subscribers-mlir-linalg

Author: Clément Fournier (oowekyala)

`ReduceOp` crashes during verification of the following op: ```mlir module { func.func @main(%arg0: tensor<32xi32>, %arg1: tensor<32xi32>) -> i32 { %c0_i32 = arith.constant 0 : i32 %from_elements = tensor.from_elements %c0_i32 : tensor<i32> %reduced = linalg.reduce ins(%arg0, %arg1 : tensor<32xi32>, tensor<32xi32>) outs(%from_elements : tensor<i32>) dimensions = [0] (%in0: i32, %in1: i32, %init: i32) { %44 = arith.muli %in0, %in1 : i32 %5 = arith.addi %44, %init : i32 linalg.yield %5 : i32 } %ext = tensor.extract %reduced[] : tensor<i32> return %ext: i32 } } ``` This is at line 1656 here: https://github.com/llvm/llvm-project/blob/fb87e11e7253656ffe99726f45b679e08343bd5f/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp#L1652-L1658

It appears the ODS-generated code for getInputs assumes that the variadic operands (inputs and inits here) are the same sizes. When computing the size of the variadic operand, it foordivs 3 by 2 and returns a range with 1 element. On the other hand the code in here is correct and returns an init count of 2, and an operand range with the two inputs. That's why here we enter the loop, but getInputs()[1] crashes.

One fix for the crash would be to use getDpsInputs instead of getInputs in the verifier. But this should ideally be patched in ODS I suppose. Thanks for reading.

---

Stack trace

cinm-opt: /home/clem/Documents/LLVM/llvm-18.1.6/llvm/include/llvm/ADT/STLExtras.h:1281: ReferenceT llvm::detail::indexed_accessor_range_base<mlir::OperandRange, mlir::OpOperand *, mlir::Value, mlir::Value, mlir::Value>::operator[](size_t) const [DerivedT = mlir::OperandRange, BaseT = mlir::OpOperand *, T = mlir::Value, PointerT = mlir::Value, ReferenceT = mlir::Value]: Assertion `Index < size() && "invalid index for value range"' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.	Program arguments: cinm-opt sandbox/matmul_tiled.mlir
 #<!-- -->0 0x000055eadabe7a4d llvm::sys::PrintStackTrace(llvm::raw_ostream&amp;, int) /home/clem/Documents/LLVM/llvm-18.1.6/llvm/lib/Support/Unix/Signals.inc:723:11
 #<!-- -->1 0x000055eadabe803b PrintStackTraceSignalHandler(void*) /home/clem/Documents/LLVM/llvm-18.1.6/llvm/lib/Support/Unix/Signals.inc:798:1
 #<!-- -->2 0x000055eadabe5ce3 llvm::sys::RunSignalHandlers() /home/clem/Documents/LLVM/llvm-18.1.6/llvm/lib/Support/Signals.cpp:105:5
 #<!-- -->3 0x000055eadabe8801 SignalHandler(int) /home/clem/Documents/LLVM/llvm-18.1.6/llvm/lib/Support/Unix/Signals.inc:413:1
 #<!-- -->4 0x00007f5e4b450ae0 (/usr/lib/libc.so.6+0x3cae0)
 #<!-- -->5 0x00007f5e4b4a8e44 (/usr/lib/libc.so.6+0x94e44)
 #<!-- -->6 0x00007f5e4b450a30 raise (/usr/lib/libc.so.6+0x3ca30)
 #<!-- -->7 0x00007f5e4b4384c3 abort (/usr/lib/libc.so.6+0x244c3)
 #<!-- -->8 0x00007f5e4b4383df (/usr/lib/libc.so.6+0x243df)
 #<!-- -->9 0x00007f5e4b448c67 (/usr/lib/libc.so.6+0x34c67)
#<!-- -->10 0x000055ead608a178 llvm::detail::indexed_accessor_range_base&lt;mlir::OperandRange, mlir::OpOperand*, mlir::Value, mlir::Value, mlir::Value&gt;::operator[](unsigned long) const /home/clem/Documents/LLVM/llvm-18.1.6/llvm/include/llvm/ADT/STLExtras.h:0:5
#<!-- -->11 0x000055ead7199ce5 mlir::linalg::ReduceOp::verify() /home/clem/Documents/LLVM/llvm-18.1.6/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:1584:32
#<!-- -->12 0x000055ead715bb94 mlir::Op&lt;mlir::linalg::ReduceOp, mlir::OpTrait::OneRegion, mlir::OpTrait::VariadicResults, mlir::OpTrait::ZeroSuccessors, mlir::OpTrait::VariadicOperands, mlir::OpTrait::S

[oowekyala]

Upon closer inspection, this behavior is explicitly written in in the op definition:

llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

Line 314 in 2a1ea15

SameVariadicOperandSize,

I wrote oowekyala@6f89431 to lift this restriction.

[joker-eph]

Nothing can justify the compiler crashing, this is a bug regardless of one writes in ODS.

[oowekyala]

I agree. AFAICS there are two way forward:

1. make the verifier more defensive and reject programs where the inputs and outputs of linalg.reduce are not of the same size, or
2. remove this restriction, which would allow linalg.reduce to eg, reduce two tensors into a single one.

I am not sure what the original intentions of the authors were, but it might be option 1. I would prefer option 2 because it allows useful programs to be written, and the restriction seems a bit artificial (unexpected even). My patch implements that. Should I open a PR?

[CoTinker]

Hi, maybe you could open a PR. Thanks.

[rengolin]

Just got pointed to this issue via it's PR in the context of our restructuring of linalg.

The semantics you want in linalg.reduce is not a reduction (assumed unary), but an type of contraction. In the proposed tree, reduce is a data movement operation, as we already have other capable compute operations in flight.

I agree with @joker-eph there is a bug in the verifier, as the compiler should not crash, we need to fix it anyway. But increasing the semantics of an operation isn't the right fix.

Please, just make the verifier more strict on the original semantics and we can discuss the increase of semantics as an RFC.