[MLIR] Add canoncalization for `shape.is_broadcastable`

[frgossen]

Summary

Canonicalize is_broadcastable to constant true if fewer than 2 unique shape
operands. Eliminate redundant operands, otherwise.

[bkramer]

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

792

llvm::is_contained

[frgossen]

Address comment

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

let extraClassDeclaration = [{
     ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
     ::llvm::SmallVectorImpl<::mlir::Type>&inferredReturnTypes) {
       inferredReturnTypes.push_back(::mlir::IntegerType::get(context,
                                                              /*width=*/1));
       return success();
     };
   }];
 
+  let hasCanonicalizer = 1;
+
   let assemblyFormat = "$shapes attr-dict `:` type($shapes)";
   let verifier = [{ return ::verify(*this); }];
-
 }
 
 def Shape_RankOp : Shape_Op<"rank", [NoSideEffect]> {
   let summary = "Gets the rank of a shape";
   let description = [{
     Returns the rank of the shape or extent tensor, i.e. the number of extents.
   }];
 

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

 
 static LogicalResult verify(IsBroadcastableOp op) {
   // Ensure that AssumingAllOp contains at least one operand
   if (op.getNumOperands() < 2)
     return op.emitOpError("required at least 2 input shapes");
   return success();
 }
 
+namespace {
+struct IsBroadcastableCanonicalizationPattern
+    : public OpRewritePattern<IsBroadcastableOp> {
+  using OpRewritePattern<IsBroadcastableOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IsBroadcastableOp op,
+                                PatternRewriter &rewriter) const override {
+    // Find unique operands.
+    SmallVector<Value, 2> unique;
+    for (Value v : op.getOperands()) {
+      if (!llvm::is_contained(unique, v))
+        unique.push_back(v);
+    }
+
+    // Can always broadcast fewer than two shapes.
+    if (unique.size() < 2) {
+      rewriter.replaceOpWithNewOp<mlir::ConstantOp>(op,
+                                                    rewriter.getBoolAttr(true));
+      return success();
+    }
+
+    // Reduce op to equivalent with unique operands.
+    if (unique.size() < op.getNumOperands()) {
+      rewriter.replaceOpWithNewOp<IsBroadcastableOp>(op, rewriter.getI1Type(),
+                                                     unique);
+      return success();
+    }
+
+    return failure();
+  }
+};
+} // namespace
+
+void IsBroadcastableOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<IsBroadcastableCanonicalizationPattern>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // RankOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult shape::RankOp::fold(ArrayRef<Attribute> operands) {
   auto shape = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
   if (!shape)
     return {};

mlir/test/Dialect/Shape/canonicalize.mlir

 // CHECK-LABEL: @fold_tensor.cast_of_const_shape_returned_dynamic
 func @fold_tensor.cast_of_const_shape_returned_dynamic(%arg: i1) -> tensor<?xindex> {
   // CHECK: shape.const_shape [2] : tensor<?xindex>
   // CHECK-NOT: tensor.cast
   %0 = shape.const_shape [2] : tensor<1xindex>
   %1 = tensor.cast %0 : tensor<1xindex> to tensor<?xindex>
   return %1 : tensor<?xindex>
 }
+
+// -----
+
+// CHECK-LABEL: @is_broadcastable_on_same_shape
+func @is_broadcastable_on_same_shape(%shape : !shape.shape) -> i1 {
+  // CHECK-NOT: is_broadcastable
+  // CHECK: %[[RES:.*]] = constant true
+  // CHECK: return %[[RES]]
+  %0 = shape.is_broadcastable %shape, %shape, %shape
+      : !shape.shape, !shape.shape, !shape.shape
+  return %0 : i1
+}
+
+// -----
+
+// CHECK-LABEL: @is_broadcastable_on_duplicate_shapes
+// CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape)
+func @is_broadcastable_on_duplicate_shapes(%a : !shape.shape, %b : !shape.shape)
+    -> i1 {
+  // CHECK: %[[RES:.*]] = shape.is_broadcastable %[[A]], %[[B]]
+  // CHECK: return %[[RES]]
+  %0 = shape.is_broadcastable %a, %b, %a, %a, %a, %b : !shape.shape,
+      !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape
+  return %0 : i1
+}