[MLIR][Shape] Canonicalize duplicate operands in `shape.cstr_broadcastable`

[frgossen]

Summary

Depends On D99156

[bkramer]

looks good

[frgossen]

Rebase

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

 }
 } // namespace
 
 void CstrBroadcastableOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *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.add<CstrBroadcastableEqOps>(context);
+  patterns.add<CstrBroadcastableEqOps,
+               RemoveDuplicateOperandsPattern<CstrBroadcastableOp>>(context);
 }
 
 // Return true if there is exactly one attribute not representing a scalar
 // broadcast.
 static bool hasAtMostSingleNonScalar(ArrayRef<Attribute> attributes) {
   bool nonScalarSeen = false;
   for (Attribute a : attributes) {
     if (!a || a.cast<DenseIntElementsAttr>().getNumElements() != 0) {

mlir/test/Dialect/Shape/canonicalize.mlir

 
 // -----
 // One scalar and one non-scalar and one unknown cannot be broadcasted at compile time
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %cs0 = shape.const_shape [8, 1] : !shape.shape
+  %cs0 = shape.const_shape [8, 1] : !shape.shape  
   %cs1 = shape.const_shape [1, 8] : !shape.shape
   %cs2 = shape.const_shape [1, -1] : !shape.shape
   %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // One scalar and two unknowns cannot be broadcasted at compile time
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [8, 1] : !shape.shape
   %cs1 = shape.const_shape [1, -1] : !shape.shape
-  %cs2 = shape.const_shape [1, -1] : !shape.shape
+  %cs2 = shape.const_shape [8, -1] : !shape.shape
   %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Broadcastable with scalars and a non-scalar can be constant folded
 // CHECK-LABEL: func @f
func @is_broadcastable_on_duplicate_shapes(%a : !shape.shape, %b : !shape.shape)
   // 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: @cstr_broadcastable_on_duplicate_shapes
+// CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape)
+func @cstr_broadcastable_on_duplicate_shapes(%a : !shape.shape,
+    %b : !shape.shape) -> !shape.witness {
+  // CHECK: %[[RES:.*]] = shape.cstr_broadcastable %[[A]], %[[B]] :
+  // CHECK: return %[[RES]]
+  %0 = shape.cstr_broadcastable %a, %b, %a, %a, %a, %b : !shape.shape,
+      !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape
+  return %0 : !shape.witness
+}
+
+// -----
+
 // 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