[mlir] Fix broadcasting check with 1 values

[jpienaar]

Summary

The trait was inconsistent with the other broadcasting logic here. And
also fix printing here to use ? rather than -1 in the error.

mlir/lib/Dialect/Traits.cpp

 /// Returns a tuple corresponding to whether range has tensor or vector type.
 template <typename iterator_range>
 static std::tuple<bool, bool> hasTensorOrVectorType(iterator_range types) {
   return std::make_tuple(
       llvm::any_of(types, [](Type t) { return t.isa<TensorType>(); }),
       llvm::any_of(types, [](Type t) { return t.isa<VectorType>(); }));
 }
 
-static bool areCompatibleShapes(ArrayRef<int64_t> shape1,
-                                ArrayRef<int64_t> shape2) {
+static bool isCompatibleInferredReturnShape(ArrayRef<int64_t> inferred,
+                                            ArrayRef<int64_t> existing) {
   auto isCompatible = [](int64_t dim1, int64_t dim2) {
-    return dim1 == dim2 || dim1 == -1 || dim2 == -1;
+    // If the inferred and existing dim is the same, or one of them is unknown
+    // then it is compatible, else if the inferred dim is 1 then it is also
+    // compatible. But if the existing dim is 1 and the inferred is greater than
+    // 1 then flag.
+    return dim1 == dim2 || dim1 == -1 || dim2 == -1 || dim1 == 1;
   };
-  if (shape1.size() != shape2.size())
+  if (inferred.size() != existing.size())
     return false;
-  for (auto p : llvm::zip(shape1, shape2))
+  for (auto p : llvm::zip(inferred, existing))
     if (!isCompatible(std::get<0>(p), std::get<1>(p)))
       return false;
   return true;
 }
 
 static std::string getShapeString(ArrayRef<int64_t> shape) {
   // TODO: should replace with printing shape more uniformly across here and
   // when in type.
-  return std::string(
-      formatv("'{0:$[x]}'", llvm::make_range(shape.begin(), shape.end())));
+  std::string ret;
+  llvm::raw_string_ostream ss(ret);
+  ss << '\'';
+  llvm::interleave(
+      shape, ss,
+      [&](int64_t dim) {
+        if (ShapedType::isDynamic(dim))
+          ss << '?';
+        else
+          ss << dim;
+      },
+      "x");
+  ss << '\'';
+  return ss.str();
 }
 
 LogicalResult OpTrait::impl::verifyCompatibleOperandBroadcast(Operation *op) {
   // Ensure broadcasting only tensor or only vector types.
   auto operandsHasTensorVectorType =
       hasTensorOrVectorType(op->getOperandTypes());
   auto resultsHasTensorVectorType = hasTensorOrVectorType(op->getResultTypes());
   if ((std::get<0>(operandsHasTensorVectorType) ||
LogicalResult OpTrait::impl::verifyCompatibleOperandBroadcast(Operation *op) {
 
   // If all of the results are unranked then no further verification.
   if (rankedResults.empty())
     return success();
 
   for (auto type : rankedResults) {
     ArrayRef<int64_t> actualSuffix =
         getShape(type).take_back(resultShape.size());
-    if (!areCompatibleShapes(actualSuffix, resultShape))
+    if (!isCompatibleInferredReturnShape(resultShape, actualSuffix))
       return op->emitOpError()
              << "result type " << getShapeString(getShape(type))
              << " not broadcast compatible with broadcasted operands's shapes "
              << getShapeString(resultShape);
   }
   return success();
 }

mlir/test/Dialect/traits.mlir

 func @broadcast_tensor_tensor_tensor(tensor<4x3x2xi32>, tensor<?xi32>) -> tensor<4x3x2xi32> {
 ^bb0(%arg0: tensor<4x3x2xi32>, %arg1: tensor<?xi32>):
   %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4x3x2xi32>, tensor<?xi32>) -> tensor<4x3x2xi32>
   return %0 : tensor<4x3x2xi32>
 }
 
 // -----
 
+func @broadcast_tensor_tensor_tensor(%arg0: tensor<?x6x1xi32>, %arg1: tensor<*xi32>) -> tensor<?x6x6xi32> {
+  %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<?x6x1xi32>, tensor<*xi32>) -> tensor<?x6x6xi32>
+  return %0 : tensor<?x6x6xi32>
+}
+
+// -----
+
 // Unranked operands but ranked result
 func @broadcast_tensor_tensor_tensor(tensor<*xi32>, tensor<*xi32>) -> tensor<2xi32> {
 ^bb0(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>):
   %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<*xi32>, tensor<*xi32>) -> tensor<2xi32>
   return %0 : tensor<2xi32>
 }
 
 // -----