[Mlir-commits] [mlir] [mlir][tosa] Fix MulOp verifier handling for unranked operands (PR #141980)

[Luke Hutton]

================
@@ -1819,59 +1820,39 @@ LogicalResult tosa::MulOp::verify() {
   // Verify the op has same ranks for all main operands (excludes extra operands
   // such as shift of mul op, so this is the only difference with the built-in
   // `SameOperandsAndResultRank` trait) and results types, if known.
-
-  // delegate function that returns true if type is a shaped type with known
-  // rank
-  auto hasRank = [](const Type type) {
-    if (auto shaped_type = dyn_cast<ShapedType>(type))
-      return shaped_type.hasRank();
-
-    return false;
-  };
-
-  auto rankedOperandTypes =
-      llvm::to_vector(llvm::make_filter_range(getOperandTypes(), hasRank));
-
-  auto rankedResultTypes =
-      llvm::make_filter_range(getOperation()->getResultTypes(), hasRank);
-
-  // If all operands and results are unranked, then no further verification.
-  if (rankedOperandTypes.empty() && rankedResultTypes.empty())
+  TypeRange operandTypes = getOperandTypes();
+  ShapedType aType = cast<ShapedType>(operandTypes[0]);
+  ShapedType bType = cast<ShapedType>(operandTypes[1]);
+
+  const bool aHasRank = aType.hasRank();
+  const bool bHasRank = bType.hasRank();
+  if (aHasRank && bHasRank) {
+    const int64_t aRank = aType.getRank();
+    const int64_t bRank = bType.getRank();
+    if (aRank != bRank)
+      return emitOpError("a and b operands don't have matching ranks, got ")
+             << aRank << " and " << bRank;
+
+    // check for broadcast compatible shapes
+    SmallVector<int64_t> resultShape;
+    if (!mlir::OpTrait::util::getBroadcastedShape(
+            aType.getShape(), bType.getShape(), resultShape))
+      return emitOpError("a and b operands don't have broadcast-compatible "
+                         "shapes, got ")
+             << aType << " and " << bType;
+  }
+
+  ShapedType resultType = cast<ShapedType>(output.getType());
+  if (!resultType.hasRank())
----------------
lhutton1 wrote:

If the output rank is unknown, it's not yet known whether the op is invalid, so success is returned to allow for things like shape inference to happen. e.g. consider the case:
```
func.func @test_mul_different_operand_ranks(%arg0: tensor<3x1x3xf32>, %arg1: tensor<3x1x3xf32>) -> tensor<*xf32> {
  %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
  %0 = tosa.mul %arg0, %arg1, %shift : (tensor<3x1x3xf32>, tensor<3x1x3xf32>, tensor<1xi8>) -> tensor<*xf32>
  return %0 : tensor<*xf32>
}
```
construction like this should be valid, the output shape can then later be resolved by shape inference

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