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

[Felix Schneider]

================
@@ -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()) {
----------------
ubfx wrote:

I see what you mean, we still need the special case check for the output side. I added both and changed the tests accordingly.

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