Fix shape mismatch caused by broadcast-to-TMS conversion when inserting on Cast edge

Author: ashokkumarkannan1

forge/csrc/passes/lower_to_mlir.cpp

@@ -170,6 +170,10 @@ class AttributeMapper
 
         add_op_mapping("upsample2d", "scale_factor", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
 
+        // Broadcast
+        add_op_mapping(
+            "broadcast", "broadcast_dimensions", AttributeRemap(std::nullopt, TargetType::DenseI64ArrayAttr));
+
         // Add more default mappings here
     }
 };
@@ -811,6 +815,7 @@ class MLIRGenerator
         lowering_handler_map["avg_pool2d"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::AvgPool2dOp>;
         lowering_handler_map["batchnorm"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::BatchNormOp>;
         lowering_handler_map["bitwise_and"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::BitwiseAndOp>;
+        lowering_handler_map["broadcast"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::BroadcastOp>;
         lowering_handler_map["cast"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::TypecastOp>;
         lowering_handler_map["clip"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::ClampScalarOp>;
         lowering_handler_map["concatenate"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::ConcatOp>;

forge/csrc/passes/pre_lowering_passes.cpp

@@ -18,17 +18,37 @@ using EdgeType = graphlib::EdgeType;
 
 void convert_broadcast_ops_to_tms(Graph *graph)
 {
-    std::vector<Node *> broadcast_ops = graph->nodes(
-        [](Node *node) -> bool
+    // Determines if the consumer node is unary to allow broadcast propagation.
+    // If not unary, the broadcast is lowered to TMs and inserted on the consumer edge.
+
+    std::vector<graphlib::Node *> broadcast_ops = graph->nodes(
+        [](graphlib::Node *node) -> bool
         {
             graphlib::OpNode *op = dynamic_cast<graphlib::OpNode *>(node);
             return op and op->op_type() == ops::OpType::Broadcast;
         });
 
     for (Node *node : broadcast_ops)
     {
-        graphlib::OpNode *op = node->as<graphlib::OpNode>();
-        ops::Op op_type = op->op();
+        graphlib::OpNode *op = dynamic_cast<graphlib::OpNode *>(node);
+        if (not op)
+            continue;
+
+        graphlib::OpType op_type = op->op_type();
+        bool has_unary_consumer = false;
+        for (graphlib::Edge user_edge : graph->user_data_edges(node))
+        {
+            graphlib::OpNode *consumer_op =
+                dynamic_cast<graphlib::OpNode *>(graph->node_by_id(user_edge.consumer_node_id));
+            if (consumer_op && consumer_op->is_eltwise_unary())
+            {
+                has_unary_consumer = true;
+                break;
+            }
+        }
+        if (has_unary_consumer)
+            continue;
+
         constexpr bool remove_node = true;
         graphlib::bypass_node(
             graph,

forge/forge/op/tm.py

@@ -387,8 +387,10 @@ def Broadcast(name: str, operandA: Tensor, dim: int, shape: int) -> Tensor:
     Tensor
         Forge tensor
     """
+    broadcast_dimensions = [1] * len(operandA.shape.dims)
+    broadcast_dimensions[dim] = shape
 
-    return op(OpType.Broadcast, name, operandA, dim=dim, size=shape).get_tensor()
+    return op("broadcast", name, operandA, broadcast_dimensions=broadcast_dimensions, dim=dim, size=shape).get_tensor()
 
 
 def Repeat(name: str, operandA: Tensor, repeats: List[int]) -> Tensor: