[MLIR][Shape] Remove duplicate operands of `shape.assuming_all` op

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

for (Value w : op.inputs()) {
       if (!shapes.empty() && !cstrEqOp.shapes().empty() && disjointShapes)
         return failure();
       shapes.append(cstrEqOp.shapes().begin(), cstrEqOp.shapes().end());
     }
     rewriter.replaceOpWithNewOp<CstrEqOp>(op, shapes);
     return success();
   }
 };
+
+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 AssumingAllOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                 MLIRContext *context) {
-  patterns.add<AssumingAllOneOp, AssumingAllToCstrEqCanonicalization>(context);
+  patterns.add<AssumingAllOneOp, AssumingAllToCstrEqCanonicalization,
+               RemoveDuplicateOperandsPattern<AssumingAllOp>>(context);
 }
 
 OpFoldResult AssumingAllOp::fold(ArrayRef<Attribute> operands) {
   // Iterate in reverse to first handle all constant operands. They are
   // guaranteed to be the tail of the inputs because this is commutative.
   for (int idx = operands.size() - 1; idx >= 0; idx--) {
     Attribute a = operands[idx];
     // Cannot fold if any inputs are not constant;

 }
 
 static LogicalResult verify(BroadcastOp op) {
   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();
-  }
-};
-
-template <typename OpTy>
 struct RemoveEmptyShapeOperandsPattern : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy op,
                                 PatternRewriter &rewriter) const override {
     auto isPotentiallyNonEmptyShape = [](Value shape) {
       if (auto extentTensorTy = shape.getType().dyn_cast<RankedTensorType>()) {
         if (extentTensorTy.getDimSize(0) == 0)

mlir/test/Dialect/Shape/canonicalize.mlir

func @assuming_all_to_cstr_eq(%a : !shape.shape, %b : tensor<?xindex>,
   // CHECK: return %[[RESULT]]
   %0 = shape.cstr_eq %a, %b : !shape.shape, tensor<?xindex>
   %1 = shape.cstr_eq %b, %c : tensor<?xindex>, tensor<3xindex>
   %2 = shape.assuming_all %0, %1
   return %2 : !shape.witness
 }
 
 // -----
+// `assuming_all` with duplicate operands.
+// CHECK-LABEL: func @assuming_all_duplicate_operands
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<?xindex>)
+func @assuming_all_duplicate_operands(%arg0 : tensor<?xindex>,
+    %arg1 : tensor<?xindex>) -> !shape.witness {
+  // CHECK: %[[RES:.*]] = shape.cstr_broadcastable %[[ARG0]], %[[ARG1]]
+  // CHECK: return %[[RES]]
+  %0 = shape.cstr_broadcastable %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
+  %1 = shape.assuming_all %0, %0, %0
+  return %1 : !shape.witness
+}
+
+// -----
 // `assuming_all` with all `cstr_eq` but disjoint operands cannot be collapsed.
 // CHECK-LABEL: func @assuming_all_to_cstr_eq
 // CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<3xindex>, %[[D:.*]]: tensor<3xindex>)
 func @assuming_all_to_cstr_eq(%a : !shape.shape, %b : tensor<?xindex>,
     %c : tensor<3xindex>, %d : tensor<3xindex>) -> !shape.witness {
   // CHECK: %[[EQ0:.*]] = shape.cstr_eq %[[A]], %[[B]]
   // CHECK: %[[EQ1:.*]] = shape.cstr_eq %[[C]], %[[D]]
   // CHECK: %[[RESULT:.*]] = shape.assuming_all %[[EQ0]], %[[EQ1]]