[MLIR][Shape] Add canonicalizations for `shape.broadcast`

[frgossen]

Summary

Eliminate empty shapes from the operands, partially fold all constant shape
operands, and fix normal folding.

[herhut]

Nice. Some comments to further generalize.

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

472

An alternative would be to filter all empty shapes, which simplifies the broadcast, as well.

491

This could be generalized to broadcasting all constant arguments. It would simplify the IR a little. Not sure it is worth it.

[tpopp]

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

472

I mostly agree with this. I think it would also make the above code easier to read by not requiring this stateful onlyPotentiallyNonEmptyShape.

My only concern with doing this in fold though is that I could imagine being confused if I create an op, and operands are missing due to folding upon Op creation.

[frgossen]

Address comments

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

472

I'm not sure if folding can/should create new operations anyways. If we were to replace a shape.broadcast op with a simplified version that sounds more like a canonicalization.

[frgossen]

Ping.

[tpopp]

Arc diff doesn't update the description, so FYI in case you didn't, update to mention the additional features you added.

[frgossen]

Thanks, updated it.

[tpopp]

A few places you checked that *.isa<ShapeType>() == *.isa<ShapeType>(). Out of curiosity, why not just compare the types?

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

453

Is this supposed to be ShapedType and ShapeType?

536–538

Is this change actually necessary for a reason that you know, or are you just being extra cautious?

540–543

nit: I think you can also replace this with something like nonEmptyShapeOperands = llvm::make_filter_range(op->getOperands(), isNonEmptyShape)

547

nit: maybe a different name like foldedConstantShape? I was confused at first thinking this should be a Vector of Values based on the name.

550–551

nit: if (auto constShape = shape.getDefiningOp<ConstShapeOp>()) {

565

Alternatively you could just check to see if foldedConstantOperands is empty.

[frgossen]

Address comments

[frgossen]

Thanks

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

536–538

Rather being cautious.

565

If there is exactly one constant operand foldedConstantOperands would not be empty but we don't want to rewrite because nothing changed.

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

void AssumingAllOp::build(OpBuilder &b, OperationState &state,
   build(b, state, b.getType<WitnessType>(), inputs);
 }
 
 //===----------------------------------------------------------------------===//
 // BroadcastOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
-  if (operands.size() == 1)
+  if (shapes().size() == 1)
     return shapes().front();
 
+  // If all but one of the operands are known to be empty shapes, return the
+  // remaining shape.
+  Value onlyPotentiallyNonEmptyShape;
+  Value value;
+  Attribute constOperand;
+  for (auto it : llvm::zip(shapes(), operands)) {
+    std::tie(value, constOperand) = it;
+    if (!constOperand || constOperand.cast<DenseIntElementsAttr>().size() > 0) {
+      if (!onlyPotentiallyNonEmptyShape) {
+        onlyPotentiallyNonEmptyShape = value;
+      } else {
+        onlyPotentiallyNonEmptyShape = nullptr;
+        break;
+      }
+    }
+  }
+  if (onlyPotentiallyNonEmptyShape &&
+      onlyPotentiallyNonEmptyShape.getType().isa<ShapeType>() ==
+          getType().isa<ShapeType>())
+    return onlyPotentiallyNonEmptyShape;
+
   // TODO: Support folding with more than 2 input shapes
   if (shapes().size() > 2)
     return nullptr;
 
-  if (!operands[1])
-    return nullptr;
-
-  auto rhsShape = llvm::to_vector<6>(
-      operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
-  if (rhsShape.empty())
-    return shapes()[0];
-
-  if (!operands[0])
+  if (!operands[0] || !operands[1])
     return nullptr;
-
   auto lhsShape = llvm::to_vector<6>(
       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
-  if (lhsShape.empty())
-    return shapes()[1];
-
+  auto rhsShape = llvm::to_vector<6>(
+      operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
   SmallVector<int64_t, 6> resultShape;
+
   // If the shapes are not compatible, we can't fold it.
   // TODO: Fold to an "error".
   if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
     return nullptr;
+
   Builder builder(getContext());
   return builder.getIndexTensorAttr(resultShape);
 }
 
 static LogicalResult verify(BroadcastOp op) {
   return verifyShapeOrExtentTensorOp(op);
 }
 
struct BroadcastForwardSingleOperandPattern
     : public OpRewritePattern<BroadcastOp> {
   using OpRewritePattern<BroadcastOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(BroadcastOp op,
                                 PatternRewriter &rewriter) const override {
     if (op.getNumOperands() == 1) {
       rewriter.replaceOp(op, op.shapes().front());
       return success();
     }
     return failure();
   }
 };
 } // namespace
 
 void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                               MLIRContext *context) {
   patterns.add<BroadcastForwardSingleOperandPattern,
                RemoveDuplicateOperandsPattern<BroadcastOp>>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ConcatOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ConcatOp::fold(ArrayRef<Attribute> operands) {
   if (!operands[0] || !operands[1])
     return nullptr;
   auto lhsShape = llvm::to_vector<6>(
       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
   auto rhsShape = llvm::to_vector<6>(
       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
   SmallVector<int64_t, 6> resultShape;
   resultShape.append(lhsShape.begin(), lhsShape.end());
   resultShape.append(rhsShape.begin(), rhsShape.end());
   Builder builder(getContext());
   return builder.getIndexTensorAttr(resultShape);
 }
 
 //===----------------------------------------------------------------------===//
 // ConstShapeOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, ConstShapeOp &op) {
   p << "shape.const_shape ";
   p.printOptionalAttrDict(op->getAttrs(), /*elidedAttrs=*/{"shape"});
   p << "[";

mlir/test/Dialect/Shape/canonicalize.mlir

func @f() -> !shape.shape {
   %0 = shape.const_shape [] : !shape.shape
   %1 = shape.const_shape [1, 2, 3] : !shape.shape
   %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
   return %2 : !shape.shape
 }
 
 // -----
 
+// All but one operands are known empty shapes.
+// CHECK-LABEL: @all_but_one_empty
+// CHECK-SAME:  (%[[ARG:.*]]: !shape.shape)
+func @all_but_one_empty(%arg0 : !shape.shape) -> !shape.shape {
+  // CHECK: return %[[ARG]]
+  %0 = shape.const_shape [] : !shape.shape
+  %1 = constant dense<[]> : tensor<0xindex>
+  %2 = shape.broadcast %0, %arg0, %1, %0 : !shape.shape, !shape.shape,
+      tensor<0xindex>, !shape.shape -> !shape.shape
+  return %2 : !shape.shape
+}
+
+// -----
+
 // Incompatible shapes. No folding.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
   // CHECK: shape.broadcast
   %0 = shape.const_shape [2] : !shape.shape
   %1 = shape.const_shape [7] : !shape.shape
   %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
   return %2 : !shape.shape