From e8ed8ef212d046e8539dbec03cf5d8d17234816c Mon Sep 17 00:00:00 2001
From: Alex Light <allight@google.com>
Date: Thu, 16 Nov 2023 15:36:38 -0800
Subject: [PATCH] Ensure that we don't add narrow-nodes when narrowing is
 impossible.

For subtraction if the resulting value is negative we don't currently narrow. We were checking this after already placing some instrucitons however causing the pass to report no-changes despite changing the graph.

PiperOrigin-RevId: 583186368
---
 xls/passes/narrowing_pass.cc      | 10 +++++-----
 xls/passes/narrowing_pass_test.cc | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/xls/passes/narrowing_pass.cc b/xls/passes/narrowing_pass.cc
index ced3b9234f..181996eb73 100644
--- a/xls/passes/narrowing_pass.cc
+++ b/xls/passes/narrowing_pass.cc
@@ -863,11 +863,6 @@ class NarrowVisitor final : public DfsVisitorWithDefault {
     bool is_one = leading_ones != 0;
     int64_t known_leading = is_one ? leading_ones : leading_zeros;
     int64_t required_bits = bit_count - known_leading;
-    XLS_ASSIGN_OR_RETURN(Node * new_lhs, MaybeNarrow(lhs, required_bits));
-    XLS_ASSIGN_OR_RETURN(Node * new_rhs, MaybeNarrow(rhs, required_bits));
-    XLS_ASSIGN_OR_RETURN(Node * new_sub,
-                         sub->function_base()->MakeNode<BinOp>(
-                             sub->loc(), new_lhs, new_rhs, Op::kSub));
     if (is_one) {
       // XLS_ASSIGN_OR_RETURN(
       //     Node * all_ones,
@@ -880,6 +875,11 @@ class NarrowVisitor final : public DfsVisitorWithDefault {
       // known-negative results. Until we do though this is dead code.
       return NoChange();
     }
+    XLS_ASSIGN_OR_RETURN(Node * new_lhs, MaybeNarrow(lhs, required_bits));
+    XLS_ASSIGN_OR_RETURN(Node * new_rhs, MaybeNarrow(rhs, required_bits));
+    XLS_ASSIGN_OR_RETURN(Node * new_sub,
+                         sub->function_base()->MakeNode<BinOp>(
+                             sub->loc(), new_lhs, new_rhs, Op::kSub));
     XLS_RETURN_IF_ERROR(
         sub->ReplaceUsesWithNew<ExtendOp>(new_sub, bit_count, Op::kZeroExt)
             .status());
diff --git a/xls/passes/narrowing_pass_test.cc b/xls/passes/narrowing_pass_test.cc
index 74e5d75e08..3e0c5e2f8f 100644
--- a/xls/passes/narrowing_pass_test.cc
+++ b/xls/passes/narrowing_pass_test.cc
@@ -119,6 +119,26 @@ TEST_P(NarrowingPassSemanticsTest, NarrowSub) {
   TestNarrowedIsEquivalent(fb);
 }
 
+TEST_P(NarrowingPassTest, NoChangeIsNoChangeWithMaybeNarrowableSubNegative) {
+  if (analysis() == NarrowingPass::AnalysisType::kBdd) {
+    GTEST_SKIP() << "BDD Unable to determine sign of subtraction";
+  }
+  // Regression test for issue where we added nodes but the pass reported no
+  // changes were made.
+  auto p = CreatePackage();
+  FunctionBuilder fb(TestName(), p.get());
+  auto x = fb.Param("x", p->GetBitsType(4));
+  auto y = fb.Param("y", p->GetBitsType(4));
+  auto x_wide = fb.ZeroExtend(x, 32);
+  // y_wide is always larger than x
+  auto y_wide = fb.ZeroExtend(fb.Concat({fb.Literal(UBits(1, 1)), y}), 32);
+  fb.Subtract(x_wide, y_wide);
+  XLS_ASSERT_OK_AND_ASSIGN(Function * f, fb.Build());
+  int64_t num_nodes = f->node_count();
+  ASSERT_THAT(Run(p.get()), IsOkAndHolds(false));
+  EXPECT_EQ(f->node_count(), num_nodes);
+}
+
 TEST_P(NarrowingPassTest, NarrowableSubPositive) {
   if (analysis() == NarrowingPass::AnalysisType::kBdd) {
     GTEST_SKIP() << "BDD Unable to determine sign of subtraction";