[opt] Hoist comparisons across selects when it is cheap to do so

xls/ir/node_util.cc

@@ -1587,4 +1587,28 @@ absl::StatusOr<Node*> GenericSelect::MakePredicateForDefault() const {
           absl::StrFormat("%s is not a select like operation.", n->ToString()));
   }
 }
+
+absl::StatusOr<Node*> GenericSelect::MakeSelectLikeWithNewArms(
+    absl::Span<Node* const> new_cases, std::optional<Node*> new_default_value,
+    const SourceInfo& loc) const {
+  XLS_RET_CHECK(valid());
+  XLS_RET_CHECK_EQ(new_cases.size(), cases().size());
+  FunctionBase* fb = AsNode()->function_base();
+  return std::visit(
+      Visitor{[&](Select* /*unused*/) -> absl::StatusOr<Node*> {
+                return fb->MakeNode<Select>(loc, selector(), new_cases,
+                                            new_default_value);
+              },
+              [&](PrioritySelect* /*unused*/) -> absl::StatusOr<Node*> {
+                XLS_RET_CHECK(new_default_value.has_value());
+                return fb->MakeNode<PrioritySelect>(loc, selector(), new_cases,
+                                                    *new_default_value);
+              },
+              [&](OneHotSelect* /*unused*/) -> absl::StatusOr<Node*> {
+                XLS_RET_CHECK(!new_default_value.has_value());
+                return fb->MakeNode<OneHotSelect>(loc, selector(), new_cases);
+              }},
+      sel_);
+}
+
 }  // namespace xls

xls/ir/node_util.h

@@ -593,6 +593,8 @@ absl::StatusOr<Node*> RemoveFromTuple(Node* tuple,
 // select nodes.
 class GenericSelect {
  public:
+  enum class Kind { kSel, kPrioritySel, kOneHotSel };
+
   // noop constructor for lists and such
   GenericSelect() : sel_(static_cast<Select*>(nullptr)) {}
   explicit GenericSelect(Select* select) : sel_(select) {}
@@ -635,6 +637,13 @@ class GenericSelect {
   }
 
   bool valid() const { return AsNode() != nullptr; }
+  Kind kind() const {
+    return std::visit(
+        Visitor{[](Select*) { return Kind::kSel; },
+                [](PrioritySelect*) { return Kind::kPrioritySel; },
+                [](OneHotSelect*) { return Kind::kOneHotSel; }},
+        sel_);
+  }
   absl::Span<Node* const> cases() const {
     return std::visit(
         Visitor{
@@ -669,6 +678,16 @@ class GenericSelect {
   // Make and return a new node which is true if the default case is selected.
   absl::StatusOr<Node*> MakePredicateForDefault() const;
 
+  // Creates a new select-like node of the same kind as this GenericSelect.
+  //
+  // `new_cases` must have the same size as `cases()`.
+  //
+  // - For `priority_sel`, `new_default_value` must be present.
+  // - For `one_hot_sel`, `new_default_value` must be empty.
+  absl::StatusOr<Node*> MakeSelectLikeWithNewArms(
+      absl::Span<Node* const> new_cases, std::optional<Node*> new_default_value,
+      const SourceInfo& loc) const;
+
   friend bool operator==(const GenericSelect& lhs, const GenericSelect& rhs) {
     return lhs.AsNode() == rhs.AsNode();
   }

xls/passes/select_simplification_pass.cc

@@ -1202,9 +1202,115 @@ absl::StatusOr<bool> MaybeReorderSelect(Node* node,
   return true;
 }
 
+absl::StatusOr<bool> TryHoistCompareThroughSelectLike(
+    Node* node, const QueryEngine& query_engine) {
+  if (!node->Is<CompareOp>()) {
+    return false;
+  }
+  CompareOp* cmp = node->As<CompareOp>();
+
+  std::optional<GenericSelect> sel;
+  Node* other = nullptr;
+  bool select_like_is_lhs = false;
+  // Look for `cmp(sel_like, literal)` or `cmp(literal, sel_like)` where
+  // `sel_like` is any select-like node wrapped by `GenericSelect`. We keep
+  // track of which side the select-like node is on so we preserve operand
+  // ordering for non-commutative compares (e.g. ULT/UGT).
+  if (auto gs = GenericSelect::From(cmp->operand(0)); gs.ok()) {
+    sel = *gs;
+    other = cmp->operand(1);
+    select_like_is_lhs = true;
+  } else if (auto gs = GenericSelect::From(cmp->operand(1)); gs.ok()) {
+    sel = *gs;
+    other = cmp->operand(0);
+    select_like_is_lhs = false;
+  } else {
+    return false;
+  }
+  XLS_RET_CHECK(sel.has_value() && sel->valid());
+  if (sel->kind() == GenericSelect::Kind::kOneHotSel &&
+      !query_engine.ExactlyOneBitTrue(sel->selector())) {
+    // For `one_hot_sel`, hoisting the compare into the arms is only semantics-
+    // preserving when the selector is exactly-one-hot. Otherwise, `one_hot_sel`
+    // has "multiple active arms" behavior which is not equivalent to selecting
+    // exactly one arm and then comparing.
+    return false;
+  }
+
+  // Guardrails:
+  //
+  // - Single use: avoids duplicating the select tree under multiple
+  //   compares/users (this rewrite clones the compare into each arm). Use
+  //   `HasSingleUse()` so we account for implicit uses.
+  // - Literal other operand: maximizes folding opportunities after hoisting the
+  //   comparison through the select(-like) into the arms.
+  //   This keeps the rewrite roughly area-preserving in the intended cases:
+  //   all-but-(at most)-one arm becomes a constant after folding.
+  if (!HasSingleUse(sel->AsNode()) || !other->Is<Literal>()) {
+    return false;
+  }
+
+  auto count_non_literal_arms =
+      [&](absl::Span<Node* const> cases,
+          std::optional<Node*> default_value) -> int64_t {
+    int64_t non_literal_arm_count = 0;
+    for (Node* c : cases) {
+      if (!c->Is<Literal>()) {
+        ++non_literal_arm_count;
+      }
+    }
+    if (default_value.has_value() && !default_value.value()->Is<Literal>()) {
+      ++non_literal_arm_count;
+    }
+    return non_literal_arm_count;
+  };
+
+  auto make_cmp = [&](Node* arm) -> absl::StatusOr<Node*> {
+    // Preserve operand order for non-commutative compares by placing `arm` on
+    // the same side where the select-like result originally appeared.
+    Node* lhs = select_like_is_lhs ? arm : other;
+    Node* rhs = select_like_is_lhs ? other : arm;
+    return node->function_base()->MakeNode<CompareOp>(cmp->loc(), lhs, rhs,
+                                                      cmp->op());
+  };
+
+  // Profitability: we only rewrite when there is at most one non-literal arm
+  // (including the default). In these cases, hoisting the compare into the arms
+  // enables aggressive folding of constant arms and can reduce critical-path
+  // delay by moving the compare "closer" to the data.
+  if (count_non_literal_arms(sel->cases(), sel->default_value()) > 1) {
+    return false;
+  }
+
+  std::vector<Node*> new_cases;
+  new_cases.reserve(sel->cases().size());
+  for (Node* c : sel->cases()) {
+    XLS_ASSIGN_OR_RETURN(Node * new_case, make_cmp(c));
+    new_cases.push_back(new_case);
+  }
+  std::optional<Node*> new_default = std::nullopt;
+  if (sel->default_value().has_value()) {
+    XLS_ASSIGN_OR_RETURN(new_default, make_cmp(*sel->default_value()));
+  }
+
+  XLS_ASSIGN_OR_RETURN(
+      Node * replacement,
+      sel->MakeSelectLikeWithNewArms(new_cases, new_default, cmp->loc()));
+  VLOG(2) << "Hoisting compare through " << sel->AsNode()->op() << ": "
+          << node->ToString();
+  XLS_RETURN_IF_ERROR(node->ReplaceUsesWith(replacement));
+  return true;
+}
+
 absl::StatusOr<bool> SimplifyNode(Node* node, const QueryEngine& query_engine,
                                   BitProvenanceAnalysis& provenance,
                                   int64_t opt_level, bool range_analysis) {
+  XLS_ASSIGN_OR_RETURN(bool hoisted_cmp,
+                       TryHoistCompareThroughSelectLike(node, query_engine));
+  if (hoisted_cmp) {
+    return true;
+  }
+
   // Select with a constant selector can be replaced with the respective
   // case.
   if (node->Is<Select>() &&

xls/passes/select_simplification_pass_test.cc

@@ -163,6 +163,145 @@ TEST_P(SelectSimplificationPassTest, BinaryTuplePrioritySelect) {
                    /*default_value=*/m::TupleIndex(m::Tuple(), 1))));
 }
 
+TEST_P(SelectSimplificationPassTest,
+       PushCompareThroughSelectLikeWithMostlyLiteralArms) {
+  auto p = CreatePackage();
+  FunctionBuilder fb(TestName(), p.get());
+  BValue s = fb.Param("s", p->GetBitsType(2));
+  BValue x = fb.Param("x", p->GetBitsType(8));
+  BValue param_p = fb.Param("p", p->GetBitsType(1));
+
+  BValue lit0 = fb.Literal(Value(UBits(0, 8)));
+  BValue lit1 = fb.Literal(Value(UBits(1, 8)));
+  BValue lit2 = fb.Literal(Value(UBits(2, 8)));
+
+  BValue pr = fb.PrioritySelect(s, /*cases=*/{lit1, x}, /*default_value=*/lit2);
+  BValue eq_pr = fb.Eq(pr, lit0);
+
+  BValue se = fb.Select(s, /*cases=*/{lit1, x}, /*default_value=*/lit2);
+  BValue eq_se = fb.Eq(se, lit0);
+
+  // Selector is provably exactly-one-hot: concat(not(p), p).
+  BValue not_p = fb.Not(param_p);
+  BValue oh = fb.Concat({not_p, param_p});
+  BValue ohs = fb.OneHotSelect(oh, /*cases=*/{lit1, x});
+  BValue eq_ohs = fb.Eq(ohs, lit0);
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Function * f, fb.BuildWithReturnValue(fb.Tuple({eq_pr, eq_se, eq_ohs})));
+  EXPECT_TRUE(f->return_value()->Is<Tuple>());
+
+  solvers::z3::ScopedVerifyEquivalence stays_equivalent{f};
+  EXPECT_THAT(Run(f), IsOkAndHolds(true));
+  EXPECT_THAT(
+      f->return_value(),
+      m::Tuple(m::PrioritySelect(m::Param("s"),
+                                 {m::Eq(m::Literal(1), m::Literal(0)),
+                                  m::Eq(m::Param("x"), m::Literal(0))},
+                                 m::Eq(m::Literal(2), m::Literal(0))),
+               m::Select(m::Param("s"),
+                         {m::Eq(m::Literal(1), m::Literal(0)),
+                          m::Eq(m::Param("x"), m::Literal(0))},
+                         m::Eq(m::Literal(2), m::Literal(0))),
+               m::OneHotSelect(m::Concat(m::Not(m::Param("p")), m::Param("p")),
+                               {m::Eq(m::Literal(1), m::Literal(0)),
+                                m::Eq(m::Param("x"), m::Literal(0))})));
+}
+
+TEST_P(SelectSimplificationPassTest,
+       PushNonCommutativeCompareThroughSelectPreservesOrder) {
+  auto p = CreatePackage();
+  FunctionBuilder fb(TestName(), p.get());
+  BValue s = fb.Param("s", p->GetBitsType(2));
+  BValue x = fb.Param("x", p->GetBitsType(8));
+  BValue lit0 = fb.Literal(Value(UBits(0, 8)));
+  BValue lit1 = fb.Literal(Value(UBits(1, 8)));
+  BValue lit2 = fb.Literal(Value(UBits(2, 8)));
+
+  BValue se = fb.Select(s, /*cases=*/{lit1, x}, /*default_value=*/lit2);
+  XLS_ASSERT_OK_AND_ASSIGN(Function * f,
+                           fb.BuildWithReturnValue(fb.ULt(lit0, se)));
+  EXPECT_TRUE(f->return_value()->Is<CompareOp>());
+
+  solvers::z3::ScopedVerifyEquivalence stays_equivalent{f};
+  EXPECT_THAT(Run(f), IsOkAndHolds(true));
+  EXPECT_THAT(f->return_value(),
+              m::Select(m::Param("s"),
+                        {m::ULt(m::Literal(0), m::Literal(1)),
+                         m::ULt(m::Literal(0), m::Param("x"))},
+                        m::ULt(m::Literal(0), m::Literal(2))));
+}
+
+TEST_P(SelectSimplificationPassTest,
+       PushCompareThroughSelectLikeWithAllLiteralArms) {
+  auto p = CreatePackage();
+  FunctionBuilder fb(TestName(), p.get());
+  BValue s = fb.Param("s", p->GetBitsType(2));
+  BValue param_p = fb.Param("p", p->GetBitsType(1));
+
+  BValue lit0 = fb.Literal(Value(UBits(0, 8)));
+  BValue lit1 = fb.Literal(Value(UBits(1, 8)));
+  BValue lit2 = fb.Literal(Value(UBits(2, 8)));
+
+  BValue pr =
+      fb.PrioritySelect(s, /*cases=*/{lit0, lit1}, /*default_value=*/lit2);
+  BValue eq_pr = fb.Eq(pr, lit0);
+
+  BValue se = fb.Select(s, /*cases=*/{lit0, lit1}, /*default_value=*/lit2);
+  BValue eq_se = fb.Eq(se, lit0);
+
+  BValue not_p = fb.Not(param_p);
+  BValue oh = fb.Concat({not_p, param_p});
+  BValue ohs = fb.OneHotSelect(oh, /*cases=*/{lit0, lit1});
+  BValue eq_ohs = fb.Eq(ohs, lit0);
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Function * f, fb.BuildWithReturnValue(fb.Tuple({eq_pr, eq_se, eq_ohs})));
+  EXPECT_TRUE(f->return_value()->Is<Tuple>());
+
+  solvers::z3::ScopedVerifyEquivalence stays_equivalent{f};
+  EXPECT_THAT(Run(f), IsOkAndHolds(true));
+
+  // Key invariant: after pushdown, no compare should directly compare a
+  // select-like node (sel/priority_sel/one_hot_sel) against the literal.
+  for (Node* n : f->nodes()) {
+    if (!n->Is<CompareOp>()) {
+      continue;
+    }
+    if (n->operand(0)->OpIn({Op::kPrioritySel, Op::kSel, Op::kOneHotSel}) ||
+        n->operand(1)->OpIn({Op::kPrioritySel, Op::kSel, Op::kOneHotSel})) {
+      // The pushed-down compares can still exist (e.g. eq(lit0, lit0)), but the
+      // compare itself should no longer be comparing the select-like result to
+      // the literal.
+      FAIL() << "Unexpected compare of select-like value: " << n->ToString();
+    }
+  }
+}
+
+TEST_P(SelectSimplificationPassTest,
+       CompareThroughSelectLikeWithTwoNonLiteralArmsDoesNotHoist) {
+  auto p = CreatePackage();
+  FunctionBuilder fb(TestName(), p.get());
+  BValue s = fb.Param("s", p->GetBitsType(2));
+  BValue x = fb.Param("x", p->GetBitsType(8));
+  BValue y = fb.Param("y", p->GetBitsType(8));
+  BValue lit0 = fb.Literal(Value(UBits(0, 8)));
+  BValue lit2 = fb.Literal(Value(UBits(2, 8)));
+
+  // Two non-literal arms => profitability guard should block the rewrite.
+  BValue se = fb.Select(s, /*cases=*/{x, y}, /*default_value=*/lit2);
+  XLS_ASSERT_OK_AND_ASSIGN(Function * f,
+                           fb.BuildWithReturnValue(fb.Eq(se, lit0)));
+
+  solvers::z3::ScopedVerifyEquivalence stays_equivalent{f};
+  EXPECT_THAT(Run(f), IsOkAndHolds(false));
+  EXPECT_THAT(f->return_value(),
+              m::Eq(m::Select(m::Param("s"),
+                              /*cases=*/{m::Param("x"), m::Param("y")},
+                              /*default_value=*/m::Literal(2)),
+                    m::Literal(0)));
+}
+
 TEST_P(SelectSimplificationPassTest, FourWayTupleSelect) {
   auto p = CreatePackage();
   XLS_ASSERT_OK_AND_ASSIGN(Function * f, ParseFunction(R"(

xls/passes/stateless_query_engine.cc

@@ -262,6 +262,10 @@ bool StatelessQueryEngine::AtMostOneBitTrue(Node* node) const {
   if (node->Is<OneHot>()) {
     return true;
   }
+  // A <=1-bit value can never have more than one bit set.
+  if (node->GetType()->IsBits() && node->BitCountOrDie() <= 1) {
+    return true;
+  }
 
   return QueryEngine::AtMostOneBitTrue(node);
 }
@@ -277,6 +281,49 @@ bool StatelessQueryEngine::ExactlyOneBitTrue(Node* node) const {
     return true;
   }
 
+  // Fast pattern matches for selectors which are exactly-one-hot by
+  // construction.
+  //
+  // Note: This is intentionally local/structural reasoning; stateless query
+  // engine does not propagate facts through the graph.
+  if (node->op() == Op::kConcat && node->operand_count() == 2) {
+    Node* msb = node->operand(0);
+    Node* lsb = node->operand(1);
+    auto is_single_bit = [](Node* n) {
+      return n->GetType()->IsBits() && n->BitCountOrDie() == 1;
+    };
+    auto is_not_of = [](Node* n, Node* x) {
+      return n->op() == Op::kNot && n->operand_count() == 1 &&
+             n->operand(0) == x;
+    };
+
+    // concat(not(x), x) or concat(x, not(x)) is exactly-one-hot when x is a
+    // single-bit value.
+    if (is_single_bit(lsb) && is_not_of(msb, lsb)) {
+      return true;
+    }
+    if (is_single_bit(msb) && is_not_of(lsb, msb)) {
+      return true;
+    }
+
+    // OneHot(x) may be rewritten into concat(eq(x, 0), x) (e.g. by select
+    // simplification). This concat is exactly-one-hot when x is mutually
+    // exclusive (at most one bit set).
+    Node* maybe_eq = msb;
+    Node* x = lsb;
+    if (maybe_eq->op() == Op::kEq && maybe_eq->operand_count() == 2) {
+      Node* eq_lhs = maybe_eq->operand(0);
+      Node* eq_rhs = maybe_eq->operand(1);
+      // We only handle the literal-zero comparison in stateless mode.
+      if ((eq_lhs == x && IsAllZeros(eq_rhs)) ||
+          (eq_rhs == x && IsAllZeros(eq_lhs))) {
+        if (AtMostOneBitTrue(x)) {
+          return true;
+        }
+      }
+    }
+  }
+
   return QueryEngine::ExactlyOneBitTrue(node);
 }