[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:

Yes, I think we should add an explicit ERROR_IF for the reduction ops `ERROR_IF(rank(shape1) != rank(shape))`. The text in the box was written before we added the ERROR_IF cases and we missed adding this check. Other ops such as CONV2D have ERROR_IF for their input/output shape checking. So I think you should be able to fail verification in this case.

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