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

Author: ashokkumarkannan1

forge/csrc/passes/lower_to_mlir.cpp

@@ -167,6 +167,10 @@ class AttributeMapper
         add_op_mapping("update_cache", "batch_offset", AttributeRemap(std::nullopt, TargetType::I32Attr));
         add_op_mapping("fill_cache", "batch_offset", AttributeRemap(std::nullopt, TargetType::I32Attr));
 
+        // Broadcast
+        add_op_mapping(
+            "broadcast", "broadcast_dimensions", AttributeRemap(std::nullopt, TargetType::DenseI64ArrayAttr));
+
         // Add more default mappings here
     }
 };
@@ -801,6 +805,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

@@ -16,19 +16,39 @@ using NodeType = graphlib::NodeType;
 using Edge = graphlib::Edge;
 using EdgeType = graphlib::EdgeType;
 
-void convert_broadcast_ops_to_tms(Graph *graph)
+void convert_broadcast_ops_to_tms(graphlib::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->new_op_type() == ops::OpType::Broadcast;
         });
 
     for (Node *node : broadcast_ops)
     {
-        graphlib::OpNode *op = node->as<graphlib::OpNode>();
+        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

@@ -381,8 +381,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("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: