[MLIR][Shape] Ensure to preserve op type of `shape.broadcast`

[frgossen]

Summary

Ensure to preserve the correct type during when folding and canonicalization.
shape.broadcast of of a single operand can only be folded away if the argument
type is correct.

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 (shapes().size() == 1)
+  if (shapes().size() == 1) {
+    // Otherwise, we need a cast which would be a canonicalization, not folding.
+    if (shapes().front().getType() != getType())
+      return nullptr;
     return shapes().front();
+  }
 
   // TODO: Support folding with more than 2 input shapes
   if (shapes().size() > 2)
     return nullptr;
 
   if (!operands[0] || !operands[1])
     return nullptr;
   auto lhsShape = llvm::to_vector<6>(

 struct BroadcastForwardSingleOperandPattern
     : public OpRewritePattern<BroadcastOp> {
   using OpRewritePattern<BroadcastOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(BroadcastOp op,
                                 PatternRewriter &rewriter) const override {
     if (op.getNumOperands() == 1) {
       Value uniqueShapeOperand = op.shapes().front();
+      if (uniqueShapeOperand.getType() == op.getType()) {
         rewriter.replaceOp(op, uniqueShapeOperand);
         return success();
       }
+    }
     return failure();
   }
 };
 
 struct BroadcastFoldConstantOperandsPattern
     : public OpRewritePattern<BroadcastOp> {
   using OpRewritePattern<BroadcastOp>::OpRewritePattern;
 

mlir/test/Dialect/Shape/canonicalize.mlir

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