[mlir][tosa] Add verifiers to ReduceOps, fix shape inference crash

[ubfx]

The tosa.reduce_* ops take an axis Attribute that determines along which dimension the reduction takes place. A crash can occur during shape inference when the input tensor rank is so low that the given axis doesn't exist.

Fix #68187

[llvmbot]

@llvm/pr-subscribers-mlir-tosa

@llvm/pr-subscribers-mlir

Author: Felix Schneider (ubfx)

Changes

The tosa.reduce_* ops take an axis Attribute that determines along which dimension the reduction takes place. A crash can occur during shape inference when the input tensor rank is so low that the given axis doesn't exist.

Fix #68187

---

Full diff: https://github.com/llvm/llvm-project/pull/69843.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Tosa/IR/TosaOps.cpp (+2-2)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index e03904a1611fc42..5292465477b1094 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -1109,14 +1109,14 @@ LogicalResult tosa::ScatterOp::inferReturnTypeComponents(
 static LogicalResult ReduceInferReturnTypes(
     ShapeAdaptor operandShape, Type inputType, IntegerAttr axis,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  if (!operandShape.hasRank() || operandShape.getRank() == 0) {
+  int64_t axisVal = axis.getValue().getSExtValue();
+  if (!operandShape.hasRank() || operandShape.getRank() <= axisVal) {
     inferredReturnShapes.push_back(ShapedTypeComponents(inputType));
     return success();
   }
 
   SmallVector<int64_t> outputShape;
   operandShape.getDims(outputShape);
-  int64_t axisVal = axis.getValue().getSExtValue();
   outputShape[axisVal] = 1;
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape, inputType));
   return success();

[joker-eph]

Can we get a test for this?

[ubfx]

When writing the test, I noticed that the ReduceOp folder has a similar issue where it crashes when the axis argument of the Op is too large for the shape.
What is the desired behaviour for a tosa.reduce_* Op when the axis is out of range? Is it supposed to be a No-Op, i.e. folded away or should we emit an error?

[GeorgeARM]

When writing the test, I noticed that the ReduceOp folder has a similar issue where it crashes when the axis argument of the Op is too large for the shape. What is the desired behaviour for a tosa.reduce_* Op when the axis is out of range? Is it supposed to be a No-Op, i.e. folded away or should we emit an error?

Behavior is debatable but my assumption would be that an error should be triggered if the axis parameter is >= than the input rank. Find it cleaner with well defined preconditions.

[ubfx]

my assumption would be that an error should be triggered if the axis parameter is >= than the input rank. Find it cleaner with well defined preconditions.

Makes sense, we could implement this as a verifier for the Op? Then we wouldn't emit the error in the shape inference pass itself, but in the verification afterwards.

[eric-k256]

The spec calls for an error if axis > the rank of the input tensor: https://www.mlplatform.org/tosa/tosa_spec.html#_reduce_all

ERROR_IF(axis < 0 || axis >= rank(shape1));

Adding it as a verifier for the op seems like a good option.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[eric-k256]

Thanks for the test updates, looks good to me. (assuming builds finish ok)