[mlir] Generalize broadcastable trait operands

[jpienaar]

Summary

Generalize broadcastable trait to variadic operands. Update the documentation that still talked about element type as part of broadcastable trait (that bug was already fixed). Also rename Broadcastable to ResultBroadcastableShape to be more explicit that the trait affects the result shape (it is possible for op to allow broadcastable operands but not have result shape that is broadcast compatible with operands).

Doing some intermediate work to have getBroadcastedType take an optional elementType as input and use that if specified, instead of the common element type of type1 and type2 in this function.

[mehdi_amini]

mlir/include/mlir/IR/OpBase.td

1333

To me Broadcastable implying element type seem fairly intuitive, have you looked into introducing a "BroadcastableDimensions" instead?

[merge_guards_bot]

Unit tests: pass. 61655 tests passed, 0 failed and 777 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[jpienaar]

mlir/include/mlir/IR/OpBase.td

1333

Broadcastable is only about shape* not element type: comparison operators have broadcasting behavior but boolean element type. Broadcastable should be the base and if we want to specialize it, it should get further constraints (e.g., BroadcastableWithSameElement is specialization of Broadcastable and clear from the name) rather than have specializations be more generic (e.g., BroadcastableDimensions is not Broadcastable).

• https://docs.scipy.org/doc/numpy/reference/ufuncs.html#broadcasting

[antiagainst]

mlir/include/mlir/Dialect/Traits.h

56

This is a bit counter-intuitive to me: the result type is naturally the same as type1 and type2 if not explicitly providing an elementType. What about: "If elementType is not provided, this requires type1 and type2 to have the same element type; otherwise, the given elementType will be used as the result's element type, ignoring element types from type1 and type2."

80

The current implementation of BroadcastableTwoOperandsOneResult only checks shape. I don't quite understand why don't we add a new trait to built on the existing trait to further constraint the element type if we want BroadcastableTwoOperandsOneResultWithSameElementType.

I do remember there are use cases in TensorFlow relying on BroadcastableTwoOperandsOneResult being not checking element types... This is gonna cause problem for those cases if we remove this.

[mehdi_amini]

mlir/include/mlir/IR/OpBase.td

1333

The Numpy link is using terminology like “broadcastable” to the same shape. I'd still like an explicit name to avoid confusion and possible error by misuse: what about ShapeBroadcastable?

[jpienaar]

Update summary as I realized the doc was out of date but we had already fixed the bug about element type. Instead generalized broadcastable to variadic operands, rename trait and addressed comments.

[jpienaar]

mlir/include/mlir/Dialect/Traits.h

80

Good point, remvoed. We don't want a new trait, the intention is to use two existing traits instead that define the orthogonal constraints.

Seems we had fixed the bug but not the documentation.

mlir/include/mlir/IR/OpBase.td

1333

Renamed to ResultsBroadcastableShape to emphasize that it affects the result (not just the operands as could have been the case) and is about shape. I might want to spell this differently still in future or add variant (e.g., certain operands can affect certain outputs, so while this is a convenient attribute for a common case, it is for special case and we might want to generalize it).

Open to alternatives :)

[merge_guards_bot]

Unit tests: pass. 61999 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

mlir/docs/Traits.md

 may be used directly by any dialect. The format of the header for each trait
 section goes as follows:
 
 *   `Header`
     -   (`C++ class` -- `ODS class`(if applicable))
 
 ### Broadcastable
 
-*   `OpTrait::BroadcastableTwoOperandsOneResult` -- `Broadcastable`
+*   `OpTrait::ResultsBroadcastableShape` -- `ResultsBroadcastableShape`
 
-This trait provides the API for operations that are known to have
+This trait adds the property that the operation is known to have
 [broadcast-compatible](https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
-operand and result types. Specifically, starting from the most varying
-dimension, each dimension pair of the two operands' types should either be the
-same or one of them is one. Also, the result type should have the corresponding
+operands and its result types' shape is the broadcast compatible with the shape
+of the broadcasted operands. Specifically, starting from the most varying
+dimension, each dimension pair of the two operands' shapes should either be the
+same or one of them is one. Also, the result shape should have the corresponding
 dimension equal to the larger one, if known. Shapes are checked partially if
 ranks or dimensions are not known. For example, an op with `tensor<?x2xf32>` and
 `tensor<2xf32>` as operand types and `tensor<3x2xf32>` as the result type is
 broadcast-compatible.
 
-Ths trait assumes the op has two operands and one result, and it asserts if the
-pre-condition is not satisfied.
+This trait requires that the operands are either vector or tensor types.
 
 ### Commutative
 
 *   `OpTrait::IsCommutative` -- `Commutative`
 
 This trait adds the property that the operation is commutative, i.e. `X op Y ==
 Y op X`
 

mlir/include/mlir/Dialect/Traits.h

 /// - Otherwise, incompatible shape.
 bool getBroadcastedShape(ArrayRef<int64_t> shape1, ArrayRef<int64_t> shape2,
                          SmallVectorImpl<int64_t> &resultShape);
 
 /// Returns the result broadcast composition type from the two given types by
 /// following NumPy broadcast semantics. Returned type may have dynamic shape if
 /// either of the input types has dynamic shape. Returns null type if the two
 /// given types are not broadcast-compatible.
-Type getBroadcastedType(Type type1, Type type2);
+///
+/// elementType, if specified, will be used as the element type of the
+/// broadcasted result type. Otherwise it is required that the element type of
+/// type1 and type2 is the same and this element type will be used as the
+/// resultant element type.
+Type getBroadcastedType(Type type1, Type type2, Type elementType = nullptr);
+
 } // namespace util
 
-/// This class provides the API for ops that are known to have broadcast-
-/// compatible operand and result types. Specifically,  starting from the
-/// most varying dimension, each dimension pair of the two operands' types
-/// should either be the same or one of them is one. Also, the result type
-/// should have the corresponding dimension equal to the larger one, if known.
-/// Shapes are checked partially if ranks or dimensions are not known. For
-/// example, an op with tensor<? x 2 x f32> and tensor <2 x f32> as operand
-/// types and tensor<3 x 2 x f32> as the result type is broadcast-compatible.
-///
-/// Ths trait assumes the op has two operands and one result, and it asserts
-/// if the pre-condition is not satisfied.
+/// Trait for ops that are known to have broadcast compatible operands and
+/// result types. Specifically, starting from the most varying dimension, each
+/// dimension pair of the operands' shapes should either be the same or one
+/// of them is one. Also, the results's shapes should have the corresponding
+/// dimension equal to the larger one, if known. Shapes are checked partially if
+/// ranks or dimensions are not known. For example, an op with tensor<?x2xf32>
+/// and tensor<2xf32> as operand types and tensor<5x3x2xi16> as the result
+/// type has broadcast compatible operands ns result types.
 template <typename ConcreteType>
-class BroadcastableTwoOperandsOneResult
-    : public TraitBase<ConcreteType, BroadcastableTwoOperandsOneResult> {
+class ResultsBroadcastableShape
+    : public TraitBase<ConcreteType, ResultsBroadcastableShape> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifyCompatibleOperandBroadcast(op);
   }
 };
 
 } // end namespace OpTrait
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_TRAITS

mlir/include/mlir/IR/OpBase.td

 
 // PredOpTrait is an op trait implemented by way of a predicate on the op.
 class PredOpTrait<string descr, Pred pred> : OpTrait {
   string description = descr;
   Pred predicate = pred;
 }
 
 // Op supports operand broadcast behavior.
-def Broadcastable  : NativeOpTrait<"BroadcastableTwoOperandsOneResult">;
+def ResultsBroadcastableShape :
+  NativeOpTrait<"ResultsBroadcastableShape">;
+// TODO: Alias of the above, remove post integrate.
+def Broadcastable : NativeOpTrait<"ResultsBroadcastableShape">;
 // X op Y == Y op X
 def Commutative  : NativeOpTrait<"IsCommutative">;
 // Op behaves like a function.
 def FunctionLike : NativeOpTrait<"FunctionLike">;
 // Op is isolated from above.
 def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;
 // Op results are float or vectors/tensors thereof.
 def ResultsAreFloatLike : NativeOpTrait<"ResultsAreFloatLike">;

mlir/lib/Dialect/Traits.cpp

 //===- Traits.cpp - Common op traits shared by dialects -------------------===//
 //
 // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/StandardTypes.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 
 bool OpTrait::util::getBroadcastedShape(ArrayRef<int64_t> shape1,
                                         ArrayRef<int64_t> shape2,
                                         SmallVectorImpl<int64_t> &resultShape) {
   // To compute the result broadcasted shape, we compare operand shapes
if (auto sType = type.dyn_cast<ShapedType>())
     return sType.getShape();
   return {};
 }
 
 /// Returns the result broadcast composition type from the two given types by
 /// following NumPy broadcast semantics. Returned type may have dynamic shape if
 /// either of the input types has dynamic shape. Returns null type if the two
 /// given types are not broadcast-compatible.
-Type OpTrait::util::getBroadcastedType(Type type1, Type type2) {
-  // Returns the scalar type out of the given type.
-  auto getScalarType = [](Type type) -> Type {
-    if (auto shapedType = type.dyn_cast<ShapedType>())
-      return shapedType.getElementType();
-    return type;
-  };
-
-  // Make sure underlying scalar type is the same.
-  auto scalarType = getScalarType(type1);
-  if (scalarType != getScalarType(type2))
+///
+/// elementType, if specified, will be used as the element type of the
+/// broadcasted result type. Otherwise it is required that the element type of
+/// type1 and type2 is the same and this element type will be used as the
+/// resultant element type.
+Type OpTrait::util::getBroadcastedType(Type type1, Type type2,
+                                       Type elementType) {
+  // If the elementType is not specified, then the use the common element type
+  // of the inputs or fail if there is no common element type.
+  if (!elementType) {
+    elementType = getElementTypeOrSelf(type1);
+    if (elementType != getElementTypeOrSelf(type2))
       return {};
+  }
 
   // If one of the types is unranked tensor, then the other type shouldn't be
   // vector and the result should have unranked tensor type.
   if (type1.isa<UnrankedTensorType>() || type2.isa<UnrankedTensorType>()) {
     if (type1.isa<VectorType>() || type2.isa<VectorType>())
       return {};
-    return UnrankedTensorType::get(scalarType);
+    return UnrankedTensorType::get(elementType);
   }
 
   // Returns the type kind if the given type is a vector or ranked tensor type.
   // Returns llvm::None otherwise.
   auto getCompositeTypeKind = [](Type type) -> Optional<StandardTypes::Kind> {
     if (type.isa<VectorType>() || type.isa<RankedTensorType>())
       return static_cast<StandardTypes::Kind>(type.getKind());
     return llvm::None;
Type OpTrait::util::getBroadcastedType(Type type1, Type type2,
 
   // Get the shape of each type.
   SmallVector<int64_t, 4> resultShape;
   if (!getBroadcastedShape(getShape(type1), getShape(type2), resultShape))
     return {};
 
   // Compose the final broadcasted type
   if (resultCompositeKind == StandardTypes::Vector)
-    return VectorType::get(resultShape, scalarType);
+    return VectorType::get(resultShape, elementType);
   if (resultCompositeKind == StandardTypes::RankedTensor)
-    return RankedTensorType::get(resultShape, scalarType);
-  return scalarType;
+    return RankedTensorType::get(resultShape, elementType);
+  return elementType;
 }
 
-/// Returns true if the given types has both vector types and tensor types.
-static bool hasBothVectorAndTensorType(ArrayRef<Type> types) {
-  return llvm::any_of(types, [](Type t) { return t.isa<VectorType>(); }) &&
-         llvm::any_of(types, [](Type t) { return t.isa<TensorType>(); });
+/// 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) {
   auto isCompatible = [](int64_t dim1, int64_t dim2) {
     return dim1 == dim2 || dim1 == -1 || dim2 == -1;
   };
   if (shape1.size() != shape2.size())
     return false;
   for (auto p : llvm::zip(shape1, shape2))
     if (!isCompatible(std::get<0>(p), std::get<1>(p)))
       return false;
   return true;
 }
 
-LogicalResult OpTrait::impl::verifyCompatibleOperandBroadcast(Operation *op) {
-  assert(op->getNumOperands() == 2 &&
-         "only support broadcast check on two operands");
-  assert(op->getNumResults() == 1 &&
-         "only support broadcast check on one result");
-
-  auto type1 = op->getOperand(0).getType();
-  auto type2 = op->getOperand(1).getType();
-  auto retType = op->getResult(0).getType();
+static std::string getShapeString(ArrayRef<int64_t> shape) {
+  // TODO: should replace with printing shape more uniformly across here and
+  // when in type.
+  return formatv("'{0:$[x]}'", llvm::make_range(shape.begin(), shape.end()));
+}
 
-  // We forbid broadcasting vector and tensor.
-  if (hasBothVectorAndTensorType({type1, type2, retType}))
+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) ||
+       std::get<0>(resultsHasTensorVectorType)) &&
+      (std::get<1>(operandsHasTensorVectorType) ||
+       std::get<1>(resultsHasTensorVectorType)))
     return op->emitError("cannot broadcast vector with tensor");
 
-  if (retType.isa<UnrankedTensorType>())
-    return success();
-
-  bool isUnranked1 = type1.isa<UnrankedTensorType>();
-  bool isUnranked2 = type2.isa<UnrankedTensorType>();
+  auto rankedOperands = make_filter_range(
+      op->getOperandTypes(), [](Type t) { return t.isa<RankedTensorType>(); });
 
-  // If both operands are unranked, then all result shapes are possible.
-  if (isUnranked1 && isUnranked2)
+  // If all operands are unranked, then all result shapes are possible.
+  if (rankedOperands.empty())
     return success();
 
-  // If one of the operands is unranked, then the known dimensions in the result
-  // should be compatible with the other shaped operand.
-  if (isUnranked1 || isUnranked2) {
-    // Result should have higher rank than the shaped operand's rank and then
-    // the result's trailing dimensions should be compatible with the operand
-    // shape.
-    ArrayRef<int64_t> shape = getShape(!isUnranked1 ? type1 : type2);
-    ArrayRef<int64_t> actualSuffix = getShape(retType).take_back(shape.size());
-    if (!areCompatibleShapes(actualSuffix, shape))
-      return op->emitOpError()
-             << "result type " << retType
-             << " has shape incompatible with a ranked operand type";
-    return success();
-  }
-
-  // If both operands are shaped, then the computed broadcasted shape should be
-  // compatible with the result shape.
+  // Compute broadcasted shape of operands (which requires that operands are
+  // broadcast compatible). The results need to be broadcast compatible with
+  // this result shape.
   SmallVector<int64_t, 4> resultShape;
-  if (!util::getBroadcastedShape(getShape(type1), getShape(type2), resultShape))
+  (void)util::getBroadcastedShape(getShape(*rankedOperands.begin()), {},
+                                  resultShape);
+  for (auto other : make_early_inc_range(rankedOperands)) {
+    SmallVector<int64_t, 4> temp = resultShape;
+    if (!util::getBroadcastedShape(temp, getShape(other), resultShape))
       return op->emitOpError("operands don't have broadcast-compatible shapes");
+  }
 
-  if (!areCompatibleShapes(resultShape, getShape(retType)))
-    return op->emitOpError() << "result type " << retType
-                             << " does not have shape compatible with the one "
-                                "computed from the operand types";
+  auto rankedResults = make_filter_range(
+      op->getResultTypes(), [](Type t) { return t.isa<RankedTensorType>(); });
 
+  // If all of the results are unranked then no further verfication.
+  if (rankedResults.empty())
+    return success();
+
+  for (auto type : rankedResults) {
+    ArrayRef<int64_t> actualSuffix =
+        getShape(type).take_back(resultShape.size());
+    if (!areCompatibleShapes(actualSuffix, resultShape))
+      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

^bb0(%arg0: tensor<4x3x2xi32>, %arg1: tensor<3x3xi32>):
   return %0 : tensor<4x3x2xi32>
 }
 
 // -----
 
 // Check incompatible result type with known dimension
 func @broadcast_tensor_tensor_tensor(tensor<4x3x2xi32>, tensor<3x1xi32>) -> tensor<4x3x3xi32> {
 ^bb0(%arg0: tensor<4x3x2xi32>, %arg1: tensor<3x1xi32>):
-  // expected-error @+1 {{does not have shape compatible with the one computed}}
+  // expected-error @+1 {{op result type '4x3x3' not broadcast compatible with broadcasted operands's shapes '4x3x2'}}
   %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4x3x2xi32>, tensor<3x1xi32>) -> tensor<4x3x3xi32>
   return %0 : tensor<4x3x3xi32>
 }
 
 // -----
 
 // Check incompatible result type with known dimension
 func @broadcast_tensor_tensor_tensor(tensor<8x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x6x1xi32> {
 ^bb0(%arg0: tensor<8x1x6x1xi32>, %arg1: tensor<7x1x5xi32>):
-  // expected-error @+1 {{does not have shape compatible with the one computed}}
+  // expected-error @+1 {{op result type '8x7x6x1' not broadcast compatible with broadcasted operands's shapes '8x7x6x5'}}
   %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<8x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x6x1xi32>
   return %0 : tensor<8x7x6x1xi32>
 }
 
 // -----
 
 func @broadcast_tensor_tensor_tensor(tensor<2xi32>, tensor<2xi32>) -> tensor<*xi32> {
 ^bb0(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>):
^bb0(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>):
   return %0 : tensor<2xi32>
 }
 
 // -----
 
 // Unranked operand and compatible ranked result
 func @broadcast_tensor_tensor_tensor(tensor<3x2xi32>, tensor<*xi32>) -> tensor<4x3x2xi32> {
 ^bb0(%arg0: tensor<3x2xi32>, %arg1: tensor<*xi32>):
-  %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<3x2xi32>, tensor<*xi32>) -> tensor<4x3x2xi32>
+  %0 = "test.broadcastable"(%arg0, %arg0, %arg1) : (tensor<3x2xi32>, tensor<3x2xi32>, tensor<*xi32>) -> tensor<4x3x2xi32>
   return %0 : tensor<4x3x2xi32>
 }
 
 // -----
 
 func @broadcast_tensor_tensor_tensor(tensor<3x2xi32>, tensor<*xi32>) -> tensor<2xi32> {
 ^bb0(%arg0: tensor<3x2xi32>, %arg1: tensor<*xi32>):
-  // expected-error @+1 {{shape incompatible with a ranked operand type}}
+  // expected-error @+1 {{op result type '2' not broadcast compatible with broadcasted operands's shapes '3x2'}}
   %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<3x2xi32>, tensor<*xi32>) -> tensor<2xi32>
   return %0 : tensor<2xi32>
 }
 
 // -----
 
 func @broadcast_tensor_tensor_tensor(tensor<?x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x6x5xi32> {
 ^bb0(%arg0: tensor<?x1x6x1xi32>, %arg1: tensor<7x1x5xi32>):

mlir/test/lib/TestDialect/TestOps.td

def IfFirstOperandIsNoneThenSoIsSecond :
     TEST_Op<"if_first_operand_is_none_then_so_is_second", [PredOpTrait<
     "has either both none type operands or first is not none",
      Or<[
         And<[TypeIsPred<"x", NoneType>, TypeIsPred<"y", NoneType>]>,
         Neg<TypeIsPred<"x", NoneType>>]>>]> {
   let arguments = (ins AnyType:$x, AnyType:$y);
 }
 
-def BroadcastableOp : TEST_Op<"broadcastable", [Broadcastable]> {
-  let arguments = (ins AnyTensor, AnyTensor);
+def BroadcastableOp : TEST_Op<"broadcastable", [ResultsBroadcastableShape]> {
+  let arguments = (ins Variadic<AnyTensor>);
   let results = (outs AnyTensor);
 }
 
 // There the "HasParent" trait.
 def ParentOp : TEST_Op<"parent">;
 def ChildOp : TEST_Op<"child", [HasParent<"ParentOp">]>;
 
 

mlir/tools/mlir-tblgen/RewriterGen.cpp

if (isSameOperandsAndResultType || useFirstAttr) {
     // Then create the op.
     os.indent(6) << formatv(
         "{0} = rewriter.create<{1}>(loc, tblgen_values, tblgen_attrs);\n",
         valuePackName, resultOp.getQualCppClassName());
     os.indent(4) << "}\n";
     return resultValue;
   }
 
-  bool isBroadcastable =
-      resultOp.getTrait("OpTrait::BroadcastableTwoOperandsOneResult");
+  // TODO: Remove once broadcastable has been updated. This query here is not
+  // really about broadcastable or not, it is about which build method to invoke
+  // and that requires knowledge of whether ODS generated a builder that need
+  // not take return types. That knowledge should be captured in one place
+  // rather than duplicated.
+  bool isResultsBroadcastableShape =
+      resultOp.getTrait("OpTrait::ResultsBroadcastableShape");
   bool usePartialResults = valuePackName != resultValue;
 
-  if (isBroadcastable || usePartialResults || depth > 0 || resultIndex < 0) {
+  if (isResultsBroadcastableShape || usePartialResults || depth > 0 ||
+      resultIndex < 0) {
     // For these cases (broadcastable ops, op results used both as auxiliary
     // values and replacement values, ops in nested patterns, auxiliary ops), we
     // still need to supply the result types when building the op. But because
     // we don't generate a builder automatically with ODS for them, it's the
     // developer's responsibility to make sure such a builder (with result type
     // deduction ability) exists. We go through the separate-parameter builder
     // here given that it's easier for developers to write compared to
     // aggregate-parameter builders.