diff --git a/mlir/include/mlir/Dialect/Traits.h b/mlir/include/mlir/Dialect/Traits.h --- a/mlir/include/mlir/Dialect/Traits.h +++ b/mlir/include/mlir/Dialect/Traits.h @@ -47,6 +47,20 @@ bool getBroadcastedShape(ArrayRef shape1, ArrayRef shape2, SmallVectorImpl &resultShape); +/// Returns true if a broadcast between the 2 shapes is guaranteed to be +/// successful and not result in an error. False does not guarantee that the +/// shapes are not broadcastable; it just says they they are not or it is +/// unknown. +/// +/// Conceptually, this returns true if getBroadcastedShape would have returned +/// true and vice versa, with one exception. If a dimension is unknown in both +/// shapes, getBroadcastedShape would return true and have a result with unknown +/// dimension, while this function will return false because it's possible for +/// both shapes to have a dimension greater than 1 and different which would +/// fail to broadcast. +bool staticallyKnownBroadcastable(ArrayRef shape1, + ArrayRef shape2); + /// 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 diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp --- a/mlir/lib/Dialect/Shape/IR/Shape.cpp +++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp @@ -317,13 +317,88 @@ // CstrBroadcastableOp //===----------------------------------------------------------------------===// +namespace { +// Scalars can always be broadcasted +struct CstrBroadcastableScalars : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(CstrBroadcastableOp op, + PatternRewriter &rewriter) const override { + if (!hasScalarInput(op.lhs()) && !hasScalarInput(op.rhs())) + return failure(); + rewriter.replaceOpWithNewOp(op.getOperation(), true); + return success(); + } + + static bool hasScalarInput(const Value &input) { + // Only handle cases that aren't handled by folding. + if (auto inputOp = input.getDefiningOp()) { + auto type = inputOp.arg().getType().dyn_cast(); + if (!type.hasRank()) + return false; + if (type.getRank() == 0) + return true; + } + return false; + } +}; + +// For shapes that were created by some operations, we can obtain partial +// information on the shapes and sometimes determine if they will be +// broadcastable with that. +struct CstrBroadcastablePartialInfo + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(CstrBroadcastableOp op, + PatternRewriter &rewriter) const override { + SmallVector lhsShape, rhsShape; + if (auto lhsOp = op.lhs().getDefiningOp()) { + auto type = lhsOp.arg().getType().dyn_cast(); + if (!type.hasRank()) + return failure(); + lhsShape = llvm::to_vector<6>(type.getShape()); + } else if (auto lhsOp = op.lhs().getDefiningOp()) { + lhsShape = llvm::to_vector<6>(lhsOp.shape().getValues()); + } else { + return failure(); + } + + if (auto rhsOp = op.rhs().getDefiningOp()) { + auto type = rhsOp.arg().getType().dyn_cast(); + if (!type.hasRank()) + return failure(); + rhsShape = llvm::to_vector<6>(type.getShape()); + } else if (auto rhsOp = op.rhs().getDefiningOp()) { + rhsShape = llvm::to_vector<6>(rhsOp.shape().getValues()); + } else { + return failure(); + } + + if (!OpTrait::util::staticallyKnownBroadcastable(lhsShape, rhsShape)) + return failure(); + + rewriter.replaceOpWithNewOp(op.getOperation(), true); + return success(); + } +}; +} // namespace + void CstrBroadcastableOp::getCanonicalizationPatterns( OwningRewritePatternList &patterns, MLIRContext *context) { - // If inputs are equal, return passing witness - patterns.insert(context); + patterns.insert(context); } OpFoldResult CstrBroadcastableOp::fold(ArrayRef operands) { + // Both operands are not needed if one is a scalar. + if (operands[0] && + operands[0].cast().getNumElements() == 0) + return BoolAttr::get(true, getContext()); + if (operands[1] && + operands[1].cast().getNumElements() == 0) + return BoolAttr::get(true, getContext()); + if (!operands[0] || !operands[1]) return nullptr; auto lhsShape = llvm::to_vector<6>( @@ -331,7 +406,7 @@ auto rhsShape = llvm::to_vector<6>( operands[1].cast().getValues()); SmallVector resultShape; - if (OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape)) + if (OpTrait::util::staticallyKnownBroadcastable(lhsShape, rhsShape)) return BoolAttr::get(true, getContext()); // Because a failing witness result here represents an eventual assertion diff --git a/mlir/lib/Dialect/Traits.cpp b/mlir/lib/Dialect/Traits.cpp --- a/mlir/lib/Dialect/Traits.cpp +++ b/mlir/lib/Dialect/Traits.cpp @@ -13,6 +13,25 @@ using namespace mlir; +bool OpTrait::util::staticallyKnownBroadcastable(ArrayRef shape1, + ArrayRef shape2) { + // Two dimensions are compatible when + // 1. they are defined and equal, or + // 2. one of them is 1 + for (auto i1 = shape1.rbegin(), i2 = shape2.rbegin(), e1 = shape1.rend(), + e2 = shape2.rend(); + i1 != e1 && i2 != e2; ++i1, ++i2) { + if (*i1 == 1) + continue; + if (*i2 == 1) + continue; + if (*i1 == *i2 && *i1 != -1) + continue; + return false; + } + return true; +} + bool OpTrait::util::getBroadcastedShape(ArrayRef shape1, ArrayRef shape2, SmallVectorImpl &resultShape) { diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir --- a/mlir/test/Dialect/Shape/canonicalize.mlir +++ b/mlir/test/Dialect/Shape/canonicalize.mlir @@ -403,8 +403,8 @@ // ----- // Broadcastable with non-broadcastable constant shapes is always false -// CHECK-LABEL: func @f -func @f() { +// CHECK-LABEL: func @static_non_broadcastable +func @static_non_broadcastable() { // CHECK-NEXT: shape.const_shape // CHECK-NEXT: shape.const_shape // CHECK-NEXT: shape.cstr_broadcastable @@ -515,3 +515,49 @@ return %result : !shape.size } +// ----- + +// Canonicalize scalar cstr_broadcastable checks +// CHECK-LABEL: @cstr_broadcastable_scalar +func @cstr_broadcastable_scalar(%arg0 : tensor) { + // CHECK-NEXT: shape.const_witness true + // CHECK-NEXT: consume.witness + // CHECK-NEXT: return + %0 = shape.const_shape [] + %1 = shape.shape_of %arg0 : tensor + %2 = shape.cstr_broadcastable %0, %1 + "consume.witness"(%2) : (!shape.witness) -> () + return +} + +// ----- + +// Do not canonicalize cstr_broadcastable checks with 2 unknowns +// CHECK-LABEL: @cstr_broadcastable_unknown +func @cstr_broadcastable_unknown(%arg0 : tensor, %arg1 : tensor) { + // CHECK-NEXT: shape.shape_of %arg0 + // CHECK-NEXT: shape.shape_of %arg1 + // CHECK-NEXT: shape.cstr_broadcastable + // CHECK-NEXT: consume.witness + // CHECK-NEXT: return + %0 = shape.shape_of %arg0 : tensor + %1 = shape.shape_of %arg1 : tensor + %2 = shape.cstr_broadcastable %0, %1 + "consume.witness"(%2) : (!shape.witness) -> () + return +} + +// ----- + +// Scalars are safe to broadcast to unranked sizes. +// CHECK-LABEL: @cstr_broadcastable_scalar_unranked +func @cstr_broadcastable_scalar_unranked(%arg0 : tensor<*xf32>) { + // CHECK-NEXT: shape.const_witness true + // CHECK-NEXT: consume.witness + // CHECK-NEXT: return + %0 = shape.const_shape [] + %1 = shape.shape_of %arg0 : tensor<*xf32> + %2 = shape.cstr_broadcastable %0, %1 + "consume.witness"(%2) : (!shape.witness) -> () + return +}