[MLIR][Shape] Fix `shape.broadcast` to standard lowering

[tpopp]

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

598

typo. known

[frgossen]

Address comments

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

for (Value v : llvm::drop_begin(ranks, 1)) {
     maxRank = lb.create<SelectOp>(rankIsGreater, v, maxRank);
   }
 
   // Calculate the difference of ranks and the maximum rank for later offsets.
   llvm::append_range(rankDiffs, llvm::map_range(ranks, [&](Value v) {
                        return lb.create<SubIOp>(indexTy, maxRank, v);
                      }));
 
-  rewriter.replaceOp(
-      op, lb.create<tensor::GenerateOp>(
-                getExtentTensorType(lb.getContext()), ValueRange{maxRank},
-                [&](OpBuilder &b, Location loc, ValueRange args) {
-                  Value broadcastedDim = getBroadcastedDim(
-                      ImplicitLocOpBuilder(loc, b), transformed.shapes(),
-                      rankDiffs, args[0]);
+  Value replacement = lb.create<tensor::GenerateOp>(
+        Value broadcastedDim =
+            getBroadcastedDim(ImplicitLocOpBuilder(loc, b),
+                              transformed.shapes(), rankDiffs, args[0]);
 
-                  b.create<tensor::YieldOp>(loc, broadcastedDim);
-                })
-              ->getResults());
+      });
+  if (replacement.getType() != op.getType())
+    replacement = lb.create<tensor::CastOp>(op.getType(), replacement);
+  rewriter.replaceOp(op, replacement);
   return success();
 }
 
 namespace {
 class ConstShapeOpConverter : public OpConversionPattern<ConstShapeOp> {
 public:
   using OpConversionPattern<ConstShapeOp>::OpConversionPattern;
 

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 // CHECK:           } : tensor<?xindex>
 // CHECK:           return
 // CHECK:         }
   %0 = shape.broadcast %a, %b, %c
       : tensor<2xindex>, tensor<3xindex>, tensor<2xindex> -> tensor<?xindex>
   return
 }
 
+// ----
+
+// CHECK-LABEL: @broadcast_to_known_rank
+func @broadcast_to_known_rank(%a : tensor<1xindex>, %b : tensor<3xindex>)
+    -> tensor<3xindex> {
+  // CHECK: %[[RES:.*]] = tensor.cast %{{.*}} : tensor<?xindex> to tensor<3xindex>
+  // CHECK: return %[[RES]] : tensor<3xindex>
+  %0 = shape.broadcast %a, %b : tensor<1xindex>, tensor<3xindex> -> tensor<3xindex>
+  return %0 : tensor<3xindex>
+}
+
 // -----
 
 // Lower `split_at`
 // CHECK-LABEL: @split_at
 // CHECK-SAME: %[[SHAPE:.*]]: tensor<?xindex>, %[[INDEX:.*]]: index
 func @split_at(%shape: tensor<?xindex>, %index: index) -> (tensor<?xindex>, tensor<?xindex>) {
   // CHECK-NEXT: %[[C0:.*]] = constant 0 : index
   // CHECK-NEXT: %[[RANK:.*]] = memref.dim %[[SHAPE]], %[[C0]] : tensor<?xindex>