[MLIR] Canonicalize broadcast operations on single shapes

[frgossen]

Summary

This covers cases that are not folded away because the extent tensor type
becomes more concrete in the process.

[bkramer]

Thanks

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

if (unique.size() < op.getNumOperands()) {
       rewriter.replaceOpWithNewOp<OpTy>(op, op->getResultTypes(), unique,
                                         op->getAttrs());
       return success();
     }
 
     return failure();
   }
 };
+
+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(
     OwningRewritePatternList &patterns, MLIRContext *context) {
-  patterns.insert<RemoveDuplicateOperandsPattern<BroadcastOp>>(context);
+  patterns.insert<BroadcastForwardSingleOperandPattern,
+                  RemoveDuplicateOperandsPattern<BroadcastOp>>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ConcatOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ConcatOp::fold(ArrayRef<Attribute> operands) {
   if (!operands[0] || !operands[1])

mlir/test/Dialect/Shape/canonicalize.mlir

 func @broadcast_on_duplicate_shapes(%a : !shape.shape, %b : !shape.shape)
     -> !shape.shape {
   // CHECK: %[[RES:.*]] = shape.broadcast %[[A]], %[[B]] :
   // CHECK: return %[[RES]]
   %0 = shape.broadcast %a, %b, %a, %a, %a, %b : !shape.shape, !shape.shape,
       !shape.shape, !shape.shape, !shape.shape, !shape.shape -> !shape.shape
   return %0 : !shape.shape
 }
+
+// -----
+
+// CHECK-LABEL: @broadcast_on_single_operand
+// CHECK-SAME: (%[[A:.*]]: tensor<3xindex>)
+func @broadcast_on_single_operand(%a : tensor<3xindex>) {
+  // CHECK-NOT: broadcast
+  // CHECK: "use"(%[[A]])
+  %0 = shape.broadcast %a : tensor<3xindex> -> tensor<?xindex>
+  "use"(%0) : (tensor<?xindex>) -> ()
+  return
+}