[mlir] Canonicalization of shape.cstr_broadcastable

[tpopp]

Summary

This allows replacing of this op with a true_witness in the case of both
inputs being const_shapes and being found to be broadcastable.

[silvas]

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

265

As far as the constant folding logic, you can just use an i1 for a witness. You just need to add a case in ShapeDialect::materializeConstant know how to materialize it into a true_witness op or false_witness op (or just a single const_witness?).

Maybe you can rename this patch "introduce basic constant folding for witnesses" and use this fold as the first example?

[herhut]

I am fine with this to land and make the transition to folding later.

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

263

Would a simple op.lhs() == op.rhs() work?

272

getOperand(0)->lhs()?
getOperand(1)->rhs()?

[tpopp]

Switch to folding and also handle repeated inputs.

[tpopp]

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

265

It took me a long time to figure out that I could use a BoolAttr for this. In the spirit of small commits, I'm leaving the ConstWitnessOp creation separate, but I'm not folding this for constant false witnesses because that needs to be carefully thought through as that is removing an assert that can be triggered.

[silvas]

Looks great! Thanks!

[isuruf]

This commit breaks building MLIR for me. You need to create the directory IR in the build directory in mlir/lib/Dialect/Shape/CMakeLists.txt. Otherwise mlir-tblgen fails with a file not found error.

[isuruf]

https://reviews.llvm.org/D81283

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

def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result",
     [{ /* nothing to do */ }]>
   ];
 
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
-def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", []> {
+def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", [Commutative]> {
   let summary = "Determines if 2 shapes can be successfully broadcasted.";
   let description = [{
     Given 2 input shapes, return a witness specifying if they are broadcastable.
     This broadcastable follows the same logic as what shape.broadcast documents.
 
     "cstr" operations represent runtime assertions.
 
     Example:
     ```mlir
       %w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Success
       %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
     ```
   }];
 
   let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
   let results = (outs Shape_WitnessType:$result);
 
   let assemblyFormat = "$lhs `,` $rhs attr-dict";
+
+  let hasCanonicalizer = 1;
 }
 
 def Shape_CstrEqOp : Shape_Op<"cstr_eq", []> {
   let summary = "Determines if all input shapes are equal.";
   let description = [{
     Given 1 or more input shapes, determine if all shapes are the exact same.
 
     "cstr" operations represent runtime assertions.

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

ConstShapeOp::inferReturnTypes(MLIRContext *context,
                                Optional<Location> location, ValueRange operands,
                                DictionaryAttr attributes, RegionRange regions,
                                SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(ShapeType::get(context));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// CstrBroadcastableOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+// TODO: Add a case for unknown shapes that are still defined by the same
+// operation.
+// TODO: Once Witnesses are Attributes, replace this with folding.
+struct CstrBroadcastableTrue : public OpRewritePattern<CstrBroadcastableOp> {
+  using OpRewritePattern<CstrBroadcastableOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CstrBroadcastableOp op,
+                                PatternRewriter &rewriter) const override {
+    // Don't try to compare equality when the shapes are not constant.
+    auto lhs = op.getOperand(0).getDefiningOp<ConstShapeOp>();
+    auto rhs = op.getOperand(1).getDefiningOp<ConstShapeOp>();
+    if (!lhs || !rhs)
+      return failure();
+
+    SmallVector<int64_t, 6> resultShape;
+    // If the shapes are not compatible, we can't fold it.
+    if (!OpTrait::util::getBroadcastedShape(
+            llvm::to_vector<6>(lhs.shape().getValues<int64_t>()),
+            llvm::to_vector<6>(rhs.shape().getValues<int64_t>()), resultShape))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<TrueWitnessOp>(op);
+    return success();
+  }
+};
+}; // namespace
+
+void CstrBroadcastableOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  // If equal, return true op
+  patterns.insert<CstrBroadcastableTrue>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // ConstSizeOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ConstSizeOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                               ValueRange operands, DictionaryAttr attributes,
                               RegionRange regions,
                               SmallVectorImpl<Type> &inferredReturnTypes) {

mlir/test/Dialect/Shape/canonicalize.mlir

 // any is not yet replaced without a constant input.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> !shape.shape {
   // CHECK-NEXT: shape.any
   // CHECK-NEXT: return %[[CS]]
   %1 = shape.any %arg0, %arg1
   return %1 : !shape.shape
 }
+
+// -----
+// Broadcastable with broadcastable constant shapes can be removed.
+// CHECK-LABEL: func @f
+func @f() {
+  // CHECK-NEXT: true_witness
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %cs0 = shape.const_shape [3, 1]
+  %cs1 = shape.const_shape [1, 5]
+  %0 = shape.cstr_broadcastable %cs0, %cs1
+  "consume.witness"(%0) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+// Broadcastable with non-broadcastable constant shapes cannot be removed.
+// CHECK-LABEL: func @f
+func @f() {
+  // CHECK-NEXT: shape.const_shape
+  // CHECK-NEXT: shape.const_shape
+  // CHECK-NEXT: shape.cstr_broadcastable
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %cs0 = shape.const_shape [1, 3]
+  %cs1 = shape.const_shape [1, 5]
+  %0 = shape.cstr_broadcastable %cs0, %cs1
+  "consume.witness"(%0) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+// Broadcastable without constant shapes cannot be removed.
+// CHECK-LABEL: func @f
+func @f(%arg0 : !shape.shape) {
+  // CHECK-NEXT: shape.const_shape
+  // CHECK-NEXT: shape.cstr_broadcastable
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %cs0 = shape.const_shape [1,3]
+  %0 = shape.cstr_broadcastable %arg0, %cs0
+  "consume.witness"(%0) : (!shape.witness) -> ()
+  return
+}