From 0c3446e8afbc5fa5dc554088ed7468ef0f367188 Mon Sep 17 00:00:00 2001
From: Akash Levy <akash@silimate.com>
Date: Mon, 1 Jun 2026 19:03:52 -0700
Subject: [PATCH] Fixups for Greptile

---
 passes/opt/opt_balance_tree.cc | 36 +++++++++++++++++++++++++------
 tests/opt/opt_balance_tree.ys  | 39 ++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/passes/opt/opt_balance_tree.cc b/passes/opt/opt_balance_tree.cc
index 003dfc89c..123a6b6b1 100644
--- a/passes/opt/opt_balance_tree.cc
+++ b/passes/opt/opt_balance_tree.cc
@@ -272,17 +272,38 @@ struct OptBalanceTreeWorker {
 		return true;
 	}
 
-	bool has_downstream_add_sink(Cell *cell, SlicedAddContext &ctx)
+	bool operand_contains_full_child_output(const SigSpec &sig, Cell *child)
+	{
+		SigSpec y = sigmap(child->getPort(ID::Y));
+		int width = GetSize(y);
+		for (int pos = 0; pos + width <= GetSize(sig); pos++)
+		{
+			bool found = true;
+			for (int i = 0; i < width; i++)
+				if (sig[pos + i] != y[i]) {
+					found = false;
+					break;
+				}
+			if (found)
+				return true;
+		}
+		return false;
+	}
+
+	bool has_downstream_add_sink(Cell *cell, pool<Cell*> &consumed_cells, SlicedAddContext &ctx)
 	{
 		SigSpec y = sigmap(cell->getPort(ID::Y));
 		for (auto bit : y)
 			for (auto sink : ctx.bit_to_sink[bit])
-				if (sink != cell && is_unsigned_add(sink))
-					return true;
+				if (sink != cell && !consumed_cells.count(sink) && is_unsigned_add(sink))
+					for (IdString port : {ID::A, ID::B})
+						if (operand_contains_full_child_output(sigmap(sink->getPort(port)), cell))
+							return true;
 		return false;
 	}
 
-	bool sliced_cluster_has_external_fanout(Cell *head_cell, pool<Cell*> &cluster, SlicedAddContext &ctx)
+	bool sliced_cluster_has_external_fanout(Cell *head_cell, pool<Cell*> &cluster, pool<Cell*> &consumed_cells,
+			SlicedAddContext &ctx)
 	{
 		for (auto cell : cluster)
 		{
@@ -295,7 +316,7 @@ struct OptBalanceTreeWorker {
 				if (ctx.output_port_sigs.count(bit))
 					return true;
 				for (auto sink : ctx.bit_to_sink[bit])
-					if (!cluster.count(sink))
+					if (!cluster.count(sink) && !consumed_cells.count(sink))
 						return true;
 			}
 		}
@@ -305,7 +326,8 @@ struct OptBalanceTreeWorker {
 
 	bool try_sliced_add_tree(Cell *head_cell, pool<Cell*> &consumed_cells, SlicedAddContext &ctx)
 	{
-		if (!is_unsigned_add(head_cell) || consumed_cells.count(head_cell) || has_downstream_add_sink(head_cell, ctx))
+		if (!is_unsigned_add(head_cell) || consumed_cells.count(head_cell) ||
+				has_downstream_add_sink(head_cell, consumed_cells, ctx))
 			return false;
 
 		vector<SigSpec> summands;
@@ -315,7 +337,7 @@ struct OptBalanceTreeWorker {
 			return false;
 		if (!saw_sliced_edge || GetSize(cluster) <= 1 || GetSize(summands) <= 2)
 			return false;
-		if (sliced_cluster_has_external_fanout(head_cell, cluster, ctx))
+		if (sliced_cluster_has_external_fanout(head_cell, cluster, consumed_cells, ctx))
 			return false;
 
 		log_debug("  Creating sliced add tree for %s with %d summands and %d cells...\n",
diff --git a/tests/opt/opt_balance_tree.ys b/tests/opt/opt_balance_tree.ys
index e82923bb0..36f106a03 100644
--- a/tests/opt/opt_balance_tree.ys
+++ b/tests/opt/opt_balance_tree.ys
@@ -180,6 +180,45 @@ log -pop
 
 
 
+# Test 35
+log -header "Partial-slice downstream ADD sink does not block sliced root"
+log -push
+design -reset
+read_verilog -icells <<EOF
+module top (
+  input  wire [15:0] a,
+  input  wire [15:0] b,
+  input  wire [15:0] c,
+  input  wire [15:0] d,
+  output wire [31:0] y,
+  output wire [31:0] side
+);
+  wire [16:0] s0;
+  wire [31:0] y_int;
+
+  \$add #(.A_WIDTH(16), .B_WIDTH(16), .Y_WIDTH(17), .A_SIGNED(0), .B_SIGNED(0))
+    add0 (.A(a), .B({14'b0, b[15:14]}), .Y(s0));
+  \$add #(.A_WIDTH(31), .B_WIDTH(31), .Y_WIDTH(32), .A_SIGNED(0), .B_SIGNED(0))
+    add1 (.A({s0, b[13:0]}), .B({15'b0, c}), .Y(y_int));
+  \$add #(.A_WIDTH(31), .B_WIDTH(31), .Y_WIDTH(32), .A_SIGNED(0), .B_SIGNED(0))
+    add_side (.A(y_int[31:1]), .B({15'b0, d}), .Y(side));
+
+  assign y = y_int;
+endmodule
+EOF
+check -assert
+
+design -save preopt
+equiv_opt -assert opt_balance_tree
+design -load preopt
+opt_balance_tree
+select -assert-count 0 c:add0 c:add1
+select -assert-count 1 c:add_side
+design -reset
+log -pop
+
+
+
 # Test 2
 log -header "AND chain with intermediate outputs"
 log -push