[mlir] Add canonicalizer to remove redundant shape.cstr_broadcastable ops

[ezhulenev]

Summary

Depends On D119025

[ezhulenev]

Rebased

[frgossen]

That's a nice way of doing this. Thanks for adding this.

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

497

nit: Eliminate cstr_broadcastable operands from assuming_all that are subsumed by others?

527

You can even apply this canonicalization if there are other witnesses in the list.

mlir/test/Dialect/Shape/canonicalize.mlir

576

Maybe, check that it consumes the right witness
Although somewhat enforced with the CHECK-NEXT this could be easier to read

596

Maybe, check that it consumes the right witnesses

[ezhulenev]

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

527

Added a TODO, will send a change in the next PR.

[ezhulenev]

Address comments

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

 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/InliningUtils.h"
+#include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::shape;
 
 #include "mlir/Dialect/Shape/IR/ShapeOpsDialect.cpp.inc"
if (operands.size() == op.getNumOperands())
       return failure();
 
     // Replace with a new `assuming_all` operation with merged constraints.
     rewriter.replaceOpWithNewOp<AssumingAllOp>(op, operands);
     return success();
   }
 };
 
+// Rewrite `assuming_all` operation with `cstr_broadcastable` operands, to an
+// `assuming_all` operation with a smaller number of broadcastable constraints.
+//
+//   %0 = shape.cstr_broadcastable %shape0, %shape1
+//   %1 = shape.cstr_broadcastable %shape0, %shape1, %shape2
+//
+//   %2 = shape.cstr_broadcastable %shape3, %shape4
+//   %3 = shape.cstr_broadcastable %shape3, %shape4, %shape5
+//
+//   %4 = shape.assuming_all %0, %1, %2, %3
+//
+// to:
+//
+//   %0 = shape.cstr_broadcastable %shape0, %shape1, %shape2
+//   %1 = shape.cstr_broadcastable %shape3, %shape4, %shape5
+//   %2 = shape.assuming_all %0, %1
+//
+// In this example if shapes [0, 1, 2] are broadcastable, then it means that
+// shapes [0, 1] are broadcastable too, and can be removed from the list of
+// constraints. If shapes [0, 1, 2] are not broadcastable, then it doesn't
+// matter if shapes [0, 1] are broadcastable (same for shapes [3, 4, 5]).
+struct AssumingAllOfCstrBroadcastable : public OpRewritePattern<AssumingAllOp> {
+  using OpRewritePattern<AssumingAllOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AssumingAllOp op,
+                                PatternRewriter &rewriter) const override {
+    // Collect all `CstrBroadcastableOp` operands first.
+    SetVector<CstrBroadcastableOp> operands;
+    for (Value operand : op.getInputs()) {
+      auto broadcastable = operand.getDefiningOp<CstrBroadcastableOp>();
+      if (!broadcastable)
+        return failure();
+
+      operands.insert(broadcastable);
+    }
+
+    // Skip trivial `assuming_all` operations.
+    if (operands.size() <= 1)
+      return failure();
+
+    // Collect shapes checked by `cstr_broadcastable` operands.
+    SmallVector<std::pair<CstrBroadcastableOp, DenseSet<Value>>> shapes;
+    for (auto cstr : operands) {
+      DenseSet<Value> shapes_set(cstr->operand_begin(), cstr->operand_end());
+      shapes.emplace_back(cstr, std::move(shapes_set));
+    }
+
+    // Sort by the number of shape operands (larger to smaller).
+    llvm::sort(shapes, [](auto a, auto b) {
+      return a.first.getNumOperands() > b.first.getNumOperands();
+    });
+
+    // We start from the `cst_broadcastable` operations with largest number of
+    // shape operands, and remove redundant `cst_broadcastable` operations. We
+    // do this until we find a set of `cst_broadcastable` operations with
+    // non-overlapping constraints.
+    SmallVector<CstrBroadcastableOp> marked_for_erase;
+
+    for (unsigned i = 0; i < shapes.size(); ++i) {
+      auto isSubset = [&](auto pair) {
+        return llvm::set_is_subset(pair.second, shapes[i].second);
+      };
+
+      // Keep redundant `cstr_broadcastable` operations to be erased.
+      auto *it = std::remove_if(shapes.begin() + i + 1, shapes.end(), isSubset);
+      for (auto *it0 = it; it0 < shapes.end(); ++it0)
+        marked_for_erase.push_back(it0->first);
+      shapes.erase(it, shapes.end());
+    }
+
+    // We didn't find any operands that could be removed.
+    if (marked_for_erase.empty())
+      return failure();
+
+    // Collect non-overlapping `cst_broadcastable` constraints.
+    SmallVector<Value> unique_constraints;
+    for (auto &shape : shapes)
+      unique_constraints.push_back(shape.first.getResult());
+
+    // Replace with a new `assuming_all` operation ...
+    rewriter.replaceOpWithNewOp<AssumingAllOp>(op, unique_constraints);
+
+    // ... and maybe erase `cstr_broadcastable` ops without uses.
+    for (auto &op : marked_for_erase)
+      if (op->use_empty())
+        rewriter.eraseOp(op);
+
+    return success();
+  }
+};
+
 struct AssumingAllToCstrEqCanonicalization
     : public OpRewritePattern<AssumingAllOp> {
   using OpRewritePattern<AssumingAllOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AssumingAllOp op,
                                 PatternRewriter &rewriter) const override {
     SmallVector<Value, 8> shapes;
     for (Value w : op.getInputs()) {
LogicalResult matchAndRewrite(OpTy op,
 
     return failure();
   }
 };
 } // namespace
 
 void AssumingAllOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                 MLIRContext *context) {
-  patterns.add<MergeAssumingAllOps, AssumingAllOneOp,
-               AssumingAllToCstrEqCanonicalization,
-               RemoveDuplicateOperandsPattern<AssumingAllOp>>(context);
+  patterns
+      .add<MergeAssumingAllOps, AssumingAllOneOp,
+           AssumingAllOfCstrBroadcastable, AssumingAllToCstrEqCanonicalization,
 }
 
 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 @f() {
   %2 = "test.source"() : () -> !shape.witness
   %3 = shape.assuming_all %0, %1, %2
   "consume.witness"(%3) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
+// merge cstr_broadcastable operations
+//
+// CHECK-LABEL: func @f
+// CHECK:         %[[ARG0:[a-z0-9]*]]: !shape.shape
+// CHECK-SAME:    %[[ARG1:[a-z0-9]*]]: !shape.shape
+// CHECK-SAME:    %[[ARG2:[a-z0-9]*]]: !shape.shape
+func @f(%arg0 : !shape.shape, %arg1 : !shape.shape, %arg2 : !shape.shape) {
+  // CHECK-NEXT: shape.cstr_broadcastable %[[ARG0]], %[[ARG1]], %[[ARG2]]
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
+  %1 = shape.cstr_broadcastable %arg0, %arg1, %arg2 : !shape.shape, !shape.shape, !shape.shape
+  %2 = shape.assuming_all %0, %1
+  "consume.witness"(%2) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+
+// do not merge cstr_broadcastable operations
+//
+// CHECK-LABEL: func @f
+// CHECK:         %[[ARG0:[a-z0-9]*]]: !shape.shape
+// CHECK-SAME:    %[[ARG1:[a-z0-9]*]]: !shape.shape
+// CHECK-SAME:    %[[ARG2:[a-z0-9]*]]: !shape.shape
+func @f(%arg0 : !shape.shape, %arg1 : !shape.shape, %arg2 : !shape.shape) {
+  // CHECK-NEXT: shape.cstr_broadcastable %[[ARG0]], %[[ARG1]]
+  // CHECK-NEXT: shape.cstr_broadcastable %[[ARG1]], %[[ARG2]]
+  // CHECK-NEXT: shape.assuming_all
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
+  %1 = shape.cstr_broadcastable %arg1, %arg2 : !shape.shape, !shape.shape
+  %2 = shape.assuming_all %0, %1
+  "consume.witness"(%2) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+
 // any can be replaced with a constant input if it has one.
 // CHECK-LABEL: func @f
 func @f(%arg : !shape.shape) -> !shape.shape {
   // CHECK-NEXT: %[[CS:.*]] = shape.const_shape
   // CHECK-NEXT: return %[[CS]]
   %0 = shape.const_shape [2, 3, 4] : !shape.shape
   %1 = shape.any %0, %arg : !shape.shape, !shape.shape -> !shape.shape
   return %1 : !shape.shape