[MLIR] Add canonicalization for `shape.broadcast`

[frgossen]

Summary

Remove redundant operands and fold if only one left.

[tpopp]

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

395

Unless this is convention, I would remove the "Canonicalization" part of the name. nothing about this is limited to canonicalizations.

813–814

return BoolAttr::get(getContext(), true);

816

nullptr?

mlir/test/Dialect/Shape/canonicalize.mlir

1116

add a : or this isn't actually testing that anything was removed I believe.

[frgossen]

Address comments

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

def Shape_BroadcastOp : Shape_Op<"broadcast", [Commutative, NoSideEffect]> {
   let builders = [OpBuilder<(ins "::mlir::Type":$result,
                                 "::mlir::Value":$lhs, "::mlir::Value":$rhs,
                                 "/*optional*/ ::mlir::StringAttr":$error), [{
       build($_builder, $_state, result, ::llvm::makeArrayRef({lhs, rhs}), error);
     }]>
   ];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
   let verifier = [{ return ::verify(*this); }];
 }
 
 def Shape_ConstShapeOp : Shape_Op<"const_shape", [ConstantLike, NoSideEffect]> {
   let summary = "Creates a constant shape or extent tensor";
   let description = [{
     Creates a constant shape or extent tensor. The individual extents are given
     as the `shape` attribute. The number of these values equals the shape's
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 hasFolder = 1;
   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(AssumingAllOp op) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // BroadcastOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
-  if (!operands[1])
-    return nullptr;
+  if (operands.size() == 1)
+    return shapes().front();
 
   // TODO: Support folding with more than 2 input shapes
   if (shapes().size() > 2)
     return nullptr;
 
+  if (!operands[1])
+    return nullptr;
+
   auto rhsShape = llvm::to_vector<6>(
       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
   if (rhsShape.empty())
     return shapes()[0];
 
   if (!operands[0])
     return nullptr;
 
   auto lhsShape = llvm::to_vector<6>(
       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
   if (lhsShape.empty())
     return shapes()[1];
 
   SmallVector<int64_t, 6> resultShape;
   // If the shapes are not compatible, we can't fold it.
   // TODO: Fold to an "error".
   if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
     return nullptr;
   Builder builder(getContext());
   return builder.getIndexTensorAttr(resultShape);
 }
 
 static LogicalResult verify(BroadcastOp op) {
-  // Ensure that AssumingAllOp contains at least one operand
-  if (op.getNumOperands() < 2)
-    return op.emitOpError("required at least 2 input shapes");
-
   return verifyShapeOrExtentTensorOp(op);
 }
 
+namespace {
+template <typename OpTy>
+struct RemoveDuplicateOperandsPattern : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy 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);
+    }
+
+    // Reduce op to equivalent with unique operands.
+    if (unique.size() < op.getNumOperands()) {
+      rewriter.replaceOpWithNewOp<OpTy>(op, op->getResultTypes(), unique,
+                                        op.getAttrs());
+      return success();
+    }
+
+    return failure();
+  }
+};
+} // namespace
+
+void BroadcastOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<RemoveDuplicateOperandsPattern<BroadcastOp>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ConcatOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ConcatOp::fold(ArrayRef<Attribute> operands) {
   if (!operands[0] || !operands[1])
     return nullptr;
   auto lhsShape = llvm::to_vector<6>(
if (shape.getType().isa<ShapeType>()) {
     build(builder, result, builder.getIndexType(), shape, dim);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // IsBroadcastableOp
 //===----------------------------------------------------------------------===//
 
-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);
-    }
+void IsBroadcastableOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<RemoveDuplicateOperandsPattern<IsBroadcastableOp>>(context);
 
-    // 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();
+OpFoldResult IsBroadcastableOp::fold(ArrayRef<Attribute> operands) {
+  if (operands.size() < 2) {
+    return BoolAttr::get(getContext(), true);
   }
-};
-} // namespace
 
-void IsBroadcastableOp::getCanonicalizationPatterns(
-    OwningRewritePatternList &patterns, MLIRContext *context) {
-  patterns.insert<IsBroadcastableCanonicalizationPattern>(context);
+  return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // RankOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult shape::RankOp::fold(ArrayRef<Attribute> operands) {
   auto shape = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();

mlir/test/Dialect/Shape/canonicalize.mlir

 }
 
 // -----
 
 // 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: %[[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
 }
+
+// -----
+
+// CHECK-LABEL: @broadcast_on_same_shape
+// CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape)
+func @broadcast_on_same_shape(%shape : !shape.shape) -> !shape.shape {
+  // CHECK-NOT: broadcast
+  // CHECK: return %[[SHAPE]]
+  %0 = shape.broadcast %shape, %shape, %shape : !shape.shape, !shape.shape,
+      !shape.shape -> !shape.shape
+  return %0 : !shape.shape
+}
+
+// -----
+
+// CHECK-LABEL: @broadcast_on_duplicate_shapes
+// CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape)
+func @broadcast_on_duplicate_shapes(%a : !shape.shape, %b : !shape.shape)
+    -> !shape.shape {
+  // CHECK: %[[RES:.*]] = shape.broadcast %[[A]], %[[B]] :
+  // CHECK: return %[[RES]]
+  %0 = shape.broadcast %a, %b, %a, %a, %a, %b : !shape.shape, !shape.shape,
+      !shape.shape, !shape.shape, !shape.shape, !shape.shape -> !shape.shape
+  return %0 : !shape.shape
+}

mlir/test/Dialect/Shape/invalid.mlir

 
 // Test that shape function library is defined.
 
 // expected-error@+1 {{@fn not found}}
 module attributes {shape.lib = @fn} { }
 
 // -----
 
-func @fn(%arg: !shape.shape) -> i1 {
-  // expected-error@+1 {{required at least 2 input shapes}}
-  %0 = shape.is_broadcastable %arg : !shape.shape
-  return %0 : i1
-}
-
-// -----
-
 func @fn(%arg: !shape.shape) -> !shape.witness {
   // expected-error@+1 {{required at least 2 input shapes}}
   %0 = shape.cstr_broadcastable %arg : !shape.shape
   return %0 : !shape.witness
 }
-
-