[MLIR][Shape] Canonicalize `assuming_all` when all operands are `cstr_bcastable`

[bkramer]

We can do the same for cstr_eq.

[silvas]

This is not a correct xform because broadcastability is not transitive.

Consider:

A: tensor<1x4x1xf32>
B: tensor<1x1x4xf32>
C: tensor<1x3x4xf32>

cstr_broadcastable(A,B) && cstr_broadcastable(B,C) does not imply cstr_broadcastable(A,B,C).

(we can do this for cstr_eq though since that is transitive)

[frgossen]

Fix

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

for (Value v : yieldValues)
     assumingTypes.push_back(v.getType());
   result.addTypes(assumingTypes);
 }
 
 //===----------------------------------------------------------------------===//
 // AssumingAllOp
 //===----------------------------------------------------------------------===//
 
+namespace {
+struct AssumingAllToCstrBroadcastableCanonicalization
+    : public OpRewritePattern<AssumingAllOp> {
+  using OpRewritePattern<AssumingAllOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AssumingAllOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value, 8> shapes;
+    for (Value v : op.inputs()) {
+      auto cstrBcastableOp = v.getDefiningOp<CstrBroadcastableOp>();
+      if (!cstrBcastableOp)
+        return failure();
+      auto range = cstrBcastableOp.shapes();
+      shapes.append(range.begin(), range.end());
+    }
+    rewriter.replaceOpWithNewOp<CstrBroadcastableOp>(op, shapes);
+    return success();
+  }
+};
+} // namespace
+
 void AssumingAllOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                 MLIRContext *context) {
-  patterns.add<AssumingAllOneOp>(context);
+  patterns
+      .add<AssumingAllOneOp, AssumingAllToCstrBroadcastableCanonicalization>(
+          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;

mlir/test/Dialect/Shape/canonicalize.mlir

func @cstr_require_no_fold(%arg0: i1) {
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.cstr_require %arg0, "msg"
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
+// `assuming_all` with all `cstr_broadcastable` can be collapsed.
+// CHECK-LABEL: func @assuming_all_to_cstr_broadcastable
+// CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<3xindex>)
+func @assuming_all_to_cstr_broadcastable(%a : !shape.shape,
+    %b : tensor<?xindex>, %c : tensor<3xindex>) -> !shape.witness {
+  // CHECK: %[[RESULT:.*]] = shape.cstr_broadcastable %[[A]], %[[B]], %[[C]]
+  // CHECK: return %[[RESULT]]
+  %0 = shape.cstr_broadcastable %a, %b : !shape.shape, tensor<?xindex>
+  %1 = shape.cstr_broadcastable %b, %c : tensor<?xindex>, tensor<3xindex>
+  %2 = shape.assuming_all %0, %1
+  return %2 : !shape.witness
+}
+
+// -----
 // assuming_all with known passing witnesses can be folded
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.const_witness true
   %1 = shape.const_witness true