[Mlir-commits] [mlir] [mlir][tosa] Fix crash in inferReturnTypes for ReduceOps (PR #69843)

[Eric Kunze]

================
@@ -1155,6 +1155,52 @@ REDUCE_SHAPE_INFER(tosa::ReduceSumOp)
 COMPATIBLE_RETURN_TYPES(tosa::ConcatOp)
 #undef COMPATIBLE_RETURN_TYPES
 
+template <typename T>
+static LogicalResult verifyReduceOp(T op) {
+  // All TOSA reduce Ops have input, output and axis.
+  TensorType inputType = op.getInput().getType();
+  TensorType outputType = op.getOutput().getType();
+  int32_t reduceAxis = op.getAxis();
+
+  if (reduceAxis < 0) {
+    op.emitOpError("reduce axis must not be negative");
+    return failure();
+  }
+  if (inputType.hasRank()) {
+    int64_t inputRank = inputType.getRank();
+    // We allow for a special case where the input shape has rank 0 and axis is
+    // also 0.
+    if (reduceAxis >= inputRank && !(reduceAxis == 0 && inputRank == 0)) {
+      op.emitOpError("expect input tensor rank (")
+          << inputType.getRank() << ") to be larger than reduce axis ("
+          << reduceAxis << ")";
+      return failure();
+    }
+  }
+  if (outputType.hasRank()) {
+    if (reduceAxis >= outputType.getRank()) {
----------------
eric-k256 wrote:

I think you will need the same special case here as for the input. It isn't caught in the existing test case because of the incorrect output type you pointed out.

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