[mlir] Add additional Canonicalization of shape.cstr_broadcastable.

[tpopp]

Summary

Added canonicalization and folding was:

• Folding when either input is an attribute indicating a scalar input

which can always be broadcasted.

• Canonicalization where it can be determined that either input shape is

a scalar.

• Canonicalization where the partially specified input shapes can be

proven to be broadcastable always.

[herhut]

mlir/include/mlir/Dialect/Traits.h

52

nit: they they. Also they may not ?

mlir/lib/Dialect/Shape/IR/Shape.cpp

335

Why is shape_of(%scalar) not folded to a constant?

[frgossen]

Nice :)

mlir/lib/Dialect/Shape/IR/Shape.cpp

329

nit: I think, you can pass op directly.

333

Why is this static?
Not saying it is wrong, I just want to learn more about the C++ way of doing things.

337

nit: return type.hasRank() && type.getRank() == 0 may be easier to read

361

This is constant folding, is it? Should this case be moved to fold?

367

Maybe you can extract this as a method and reuse it for lhsShape and rhsShape like you did above.

mlir/lib/Dialect/Traits.cpp

21

Does llvm have some utility for a zip iteration like this one?
That would be really cool but I could not find.

28

The -1 is used for dynamic sizes, right?
You could use ShapedType::kDynamicSize

[rriddle]

mlir/lib/Dialect/Shape/IR/Shape.cpp

333

nit: Value should be passed by value, i.e. non-const non-reference.

mlir/lib/Dialect/Traits.cpp

21

llvm::zip

[tpopp]

Refactor code. No new behavior was added.

[tpopp]

mlir/lib/Dialect/Shape/IR/Shape.cpp

333

I was just making it static because it didn't need any information specific to an instance of the object. I think more appropriate though is a function in an anonymous namespace outside of the class because this isn't really specific to that class' data.

335

I just wasn't able to test this properly. You're right in that it's folded.

mlir/include/mlir/Dialect/Traits.h

 ///   - If either dimension is 1, the other dimension is the result.
 ///   - Otherwise, the result dimension is unknown dimension.
 /// - If one of the dimension is 1, the other dimension is the result.
 /// - If two dimensions are the same, that's the result.
 /// - Otherwise, incompatible shape.
 bool getBroadcastedShape(ArrayRef<int64_t> shape1, ArrayRef<int64_t> shape2,
                          SmallVectorImpl<int64_t> &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 might guarantee that they are not
+/// broadcastable or it might mean that this function does not have enough
+/// information to know.
+///
+/// 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<int64_t> shape1,
+                                  ArrayRef<int64_t> 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
 /// given types are not broadcast-compatible.
 ///
 /// 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

mlir/lib/Dialect/Shape/IR/Shape.cpp

 }
 
 OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shapeAttr(); }
 
 //===----------------------------------------------------------------------===//
 // CstrBroadcastableOp
 //===----------------------------------------------------------------------===//
 
+namespace {
+// Given an input shape Value, try to obtain the shape's values.
+LogicalResult getShapeVec(Value input, SmallVectorImpl<int64_t> &shapeValues) {
+  if (auto inputOp = input.getDefiningOp<ShapeOfOp>()) {
+    auto type = inputOp.arg().getType().dyn_cast<ShapedType>();
+    if (!type.hasRank())
+      return failure();
+    shapeValues = llvm::to_vector<6>(type.getShape());
+    return success();
+  } else if (auto inputOp = input.getDefiningOp<ConstShapeOp>()) {
+    shapeValues = llvm::to_vector<6>(inputOp.shape().getValues<int64_t>());
+    return success();
+  } else {
+    return failure();
+  }
+}
+
+// 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<CstrBroadcastableOp> {
+  using OpRewritePattern<CstrBroadcastableOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CstrBroadcastableOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<int64_t, 6> lhsShape, rhsShape;
+    if (failed(getShapeVec(op.lhs(), lhsShape)))
+      return failure();
+    if (failed(getShapeVec(op.rhs(), rhsShape)))
+      return failure();
+    if (!OpTrait::util::staticallyKnownBroadcastable(lhsShape, rhsShape))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<ConstWitnessOp>(op.getOperation(), true);
+    return success();
+  }
+};
+
+// Scalars are always broadcastable.
+struct CstrBroadcastableScalar : public OpRewritePattern<CstrBroadcastableOp> {
+  using OpRewritePattern<CstrBroadcastableOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CstrBroadcastableOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<int64_t, 6> shape;
+    if (failed(getShapeVec(op.lhs(), shape)) || shape.size() > 0)
+      return failure();
+    if (failed(getShapeVec(op.rhs(), shape)) || shape.size() > 0)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<ConstWitnessOp>(op.getOperation(), true);
+    return success();
+  }
+};
+
+} // namespace
+
 void CstrBroadcastableOp::getCanonicalizationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
-  // If inputs are equal, return passing witness
-  patterns.insert<CstrBroadcastableEqOps>(context);
+  // Canonicalization patterns have overlap with the considerations during
+  // folding in case additional shape information is inferred at some point that
+  // does not result in folding.
+  patterns.insert<CstrBroadcastableEqOps, CstrBroadcastablePartialInfo,
+                  CstrBroadcastableScalar>(context);
 }
 
 OpFoldResult CstrBroadcastableOp::fold(ArrayRef<Attribute> operands) {
-  if (!operands[0] || !operands[1])
-    return nullptr;
+  // Both operands are not needed if one is a scalar.
+  if (operands[0] &&
+      operands[0].cast<DenseIntElementsAttr>().getNumElements() == 0)
+    return BoolAttr::get(true, getContext());
+  if (operands[1] &&
+      operands[1].cast<DenseIntElementsAttr>().getNumElements() == 0)
+    return BoolAttr::get(true, getContext());
+
+  if (operands[0] && operands[1]) {
     auto lhsShape = llvm::to_vector<6>(
         operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
     auto rhsShape = llvm::to_vector<6>(
         operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
     SmallVector<int64_t, 6> resultShape;
-  if (OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
+    if (OpTrait::util::staticallyKnownBroadcastable(lhsShape, rhsShape))
+      return BoolAttr::get(true, getContext());
+  }
+
+  // Lastly, see if folding can be completed based on what constraints are known
+  // on the input shapes.
+  SmallVector<int64_t, 6> lhsShape, rhsShape;
+  if (failed(getShapeVec(lhs(), lhsShape)))
+    return nullptr;
+  if (failed(getShapeVec(rhs(), rhsShape)))
+    return nullptr;
+
+  if (OpTrait::util::staticallyKnownBroadcastable(lhsShape, rhsShape))
     return BoolAttr::get(true, getContext());
 
   // Because a failing witness result here represents an eventual assertion
   // failure, we do not replace it with a constant witness.
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Traits.cpp

 //===- Traits.cpp - Common op traits shared by dialects -------------------===//
 //
 // Part of the LLVM 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::staticallyKnownBroadcastable(ArrayRef<int64_t> shape1,
+                                                 ArrayRef<int64_t> shape2) {
+  // Two dimensions are compatible when
+  //   1. they are defined and equal, or
+  //   2. one of them is 1
+  return llvm::all_of(llvm::zip(llvm::reverse(shape1), llvm::reverse(shape2)),
+                      [](auto dimensions) {
+                        auto dim1 = std::get<0>(dimensions);
+                        auto dim2 = std::get<1>(dimensions);
+                        if (dim1 == 1 || dim2 == 1)
+                          return true;
+                        if (dim1 == dim2 && !ShapedType::isDynamic(dim1))
+                          return true;
+                        return false;
+                      });
+}
+
 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
   // element-wise: starting with the trailing dimensions, and working the
   // way backward. Two dimensions are compatible when
   //   1. they are equal, or
   //   2. one of them is 1

mlir/test/Dialect/Shape/canonicalize.mlir

func @f() {
   %cs1 = shape.const_shape [1, 5]
   %0 = shape.cstr_broadcastable %cs0, %cs1
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // 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
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [1, 3]
   %cs1 = shape.const_shape [1, 5]
   %0 = shape.cstr_broadcastable %cs0, %cs1

 // CHECK-SAME: (%[[SIZE:.*]]: !shape.size) -> !shape.size
 func @size_to_index_to_size(%size : !shape.size) -> !shape.size {
   // CHECK: return %[[SIZE]] : !shape.size
   %idx = shape.size_to_index %size
   %result = shape.index_to_size %idx
   return %result : !shape.size
 }
 
+// -----
+
+// Canonicalize scalar cstr_broadcastable checks
+// CHECK-LABEL: @cstr_broadcastable_scalar
+func @cstr_broadcastable_scalar(%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
+}
+
+// -----
+
+// Do not canonicalize cstr_broadcastable checks with 2 unknowns
+// CHECK-LABEL: @cstr_broadcastable_unknown
+func @cstr_broadcastable_unknown(%arg0 : tensor<?xf32>, %arg1 : tensor<?xf32>) {
+  // 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<?xf32>
+  %1 = shape.shape_of %arg1 : tensor<?xf32>
+  %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>, %arg1 : tensor<index>) {
+  // CHECK-NEXT: shape.const_witness true
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %0 = shape.shape_of %arg1 : tensor<index>
+  %1 = shape.shape_of %arg0 : tensor<*xf32>
+  %2 = shape.cstr_broadcastable %0, %1
+  "consume.witness"(%2) : (!shape.witness) -> ()
+  return
+}