Enable visibility expression builder usage while filtering BFS on def-use DAG

xls/passes/BUILD

@@ -1080,6 +1080,7 @@ cc_library(
     hdrs = ["folding_graph.h"],
     deps = [
         ":visibility_analysis",
+        "//xls/common/status:ret_check",
         "//xls/ir",
         "//xls/ir:function_builder",
         "//xls/ir:op",

xls/passes/folding_graph.cc

@@ -30,6 +30,7 @@
 #include "absl/strings/str_format.h"
 #include "absl/types/span.h"
 #include "ortools/graph/cliques.h"
+#include "xls/common/status/ret_check.h"
 #include "xls/ir/function_base.h"
 #include "xls/ir/function_builder.h"
 #include "xls/ir/node.h"
@@ -399,10 +400,18 @@ absl::StatusOr<std::unique_ptr<NaryFoldingAction>> NaryFoldingAction::Clone(
     clone_to_edges.insert({clone_operand, clone_node});
   }
 
+  absl::flat_hash_set<Node*> clone_sinks;
+  clone_sinks.reserve(other.GetSinks().size());
+  for (Node* sink : other.GetSinks()) {
+    auto it = original_node_to_clone.find(sink);
+    XLS_RET_CHECK(it != original_node_to_clone.end());
+    clone_sinks.insert(it->second);
+  }
+
   // Create the cloned object
   auto clone_action = std::make_unique<NaryFoldingAction>(
       std::move(clone_froms), clone_to, std::move(clone_to_edges),
-      other.area_saved());
+      other.area_saved(), std::move(clone_sinks));
 
   return clone_action;
 }

xls/passes/folding_graph.h

@@ -85,14 +85,21 @@ class FoldingAction {
   // folding is performed
   double area_saved() const { return area_saved_; }
 
+  const absl::flat_hash_set<Node*>& GetSinks() const { return sinks_; }
+
  protected:
-  FoldingAction(Node* to, VisibilityEdges to_edges, double area_saved)
-      : to_(to), to_edges_(to_edges), area_saved_(area_saved) {}
+  FoldingAction(Node* to, VisibilityEdges to_edges, double area_saved,
+                absl::flat_hash_set<Node*> sinks = {})
+      : to_(to),
+        to_edges_(std::move(to_edges)),
+        area_saved_(area_saved),
+        sinks_(std::move(sinks)) {}
 
  private:
   Node* to_;
   VisibilityEdges to_edges_;
   double area_saved_;
+  absl::flat_hash_set<Node*> sinks_;
 };
 
 // This class represents a single folding action from an IR node into another IR
@@ -117,8 +124,9 @@ class FoldingAction {
 class BinaryFoldingAction : public FoldingAction {
  public:
   BinaryFoldingAction(Node* from, Node* to, VisibilityEdges from_edges,
-                      VisibilityEdges to_edges, double area_saved)
-      : FoldingAction{to, to_edges, area_saved},
+                      VisibilityEdges to_edges, double area_saved,
+                      absl::flat_hash_set<Node*> sinks = {})
+      : FoldingAction{to, std::move(to_edges), area_saved, std::move(sinks)},
         from_{from},
         from_edges_{std::move(from_edges)} {}
 
@@ -159,8 +167,9 @@ class BinaryFoldingAction : public FoldingAction {
 class NaryFoldingAction : public FoldingAction {
  public:
   NaryFoldingAction(std::vector<std::pair<Node*, VisibilityEdges>>&& froms,
-                    Node* to, VisibilityEdges to_edges, double area_saved = 0.0)
-      : FoldingAction{to, std::move(to_edges), area_saved},
+                    Node* to, VisibilityEdges to_edges, double area_saved = 0.0,
+                    absl::flat_hash_set<Node*> sinks = {})
+      : FoldingAction{to, std::move(to_edges), area_saved, std::move(sinks)},
         from_(std::move(froms)) {}
 
   absl::Span<const std::pair<Node*, VisibilityEdges>> GetFrom() const {

xls/passes/resource_sharing_pass.cc

@@ -30,6 +30,7 @@
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/container/btree_set.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/functional/function_ref.h"
@@ -50,6 +51,7 @@
 #include "xls/ir/function_base.h"
 #include "xls/ir/function_builder.h"
 #include "xls/ir/node.h"
+#include "xls/ir/node_util.h"
 #include "xls/ir/nodes.h"
 #include "xls/ir/op.h"
 #include "xls/passes/bdd_query_engine.h"
@@ -294,13 +296,17 @@ ResourceSharingPass::ComputeFoldableActions(
 
     absl::flat_hash_set<OperandVisibilityAnalysis::OperandNode> one_edges;
     absl::flat_hash_set<OperandVisibilityAnalysis::OperandNode> other_edges;
+    absl::flat_hash_set<Node*> sinks;
     if (visibility.single_select.IsMutuallyExclusive(one_node, other_node)) {
       XLS_ASSIGN_OR_RETURN(
           one_edges,
           visibility.single_select.GetEdgesForVisibilityExpr(one_node));
       XLS_ASSIGN_OR_RETURN(
           other_edges,
           visibility.single_select.GetEdgesForVisibilityExpr(other_node));
+      // Specifying a sink indicates that visibility through ths node is
+      // equivalent to visibility through all terminal nodes.
+      sinks = {visibility.single_select.GetInfo(one_node)->select};
     } else {
       XLS_ASSIGN_OR_RETURN(
           one_edges,
@@ -326,15 +332,15 @@ ResourceSharingPass::ComputeFoldableActions(
       // Only n-ary foldings are evaluated with area saving thresholds, so we
       // defer computing area saved until then
       foldable_actions.push_back(std::make_unique<BinaryFoldingAction>(
-          one_node, other_node, one_edges, other_edges, 0.0));
+          one_node, other_node, one_edges, other_edges, 0.0, sinks));
     }
     if (can_other_to_one) {
       VLOG(4) << "Adding folding action: " << other_node << " into "
               << one_node;
       // Only n-ary foldings are evaluated with area saving thresholds, so we
       // defer computing area saved until then
       foldable_actions.push_back(std::make_unique<BinaryFoldingAction>(
-          other_node, one_node, other_edges, one_edges, 0.0));
+          other_node, one_node, other_edges, one_edges, 0.0, sinks));
     }
   }
 
@@ -423,11 +429,16 @@ ResourceSharingPass::MakeNaryFoldingAction(
           std::make_pair(binary_folding->GetFrom(), std::move(edges)));
     }
   }
+  absl::flat_hash_set<Node*> sinks;
+  for (const BinaryFoldingAction* action : subset_of_edges_to_n) {
+    sinks.insert(action->GetSinks().begin(), action->GetSinks().end());
+  }
   XLS_RET_CHECK_GT(froms.size(), 0);
   std::unique_ptr<NaryFoldingAction> new_action =
       std::make_unique<NaryFoldingAction>(
           std::move(froms), subset_of_edges_to_n[0]->GetTo(),
-          subset_of_edges_to_n[0]->GetToVisibilityEdges(), area_saved);
+          subset_of_edges_to_n[0]->GetToVisibilityEdges(), area_saved,
+          std::move(sinks));
   return new_action;
 }
 
@@ -1184,10 +1195,11 @@ ResourceSharingPass::LegalizeSequenceOfFolding(
 
     // The current n-ary folding is worth considering. Allocate a new n-ary
     // folding to capture it.
+    absl::flat_hash_set<Node*> sinks(folding->GetSinks());
     std::unique_ptr<NaryFoldingAction> new_folding =
         std::make_unique<NaryFoldingAction>(std::move(legal_froms), to_node,
                                             folding->GetToVisibilityEdges(),
-                                            area_saved);
+                                            area_saved, std::move(sinks));
 
     // Keep track of the current n-ary folding to legalize the next ones.
     prior_folding_of_destination[to_node] = new_folding.get();
@@ -1781,9 +1793,17 @@ absl::StatusOr<bool> ResourceSharingPass::PerformFoldingActions(
     from_used_expressions.reserve(froms_to_use.size());
     for (const auto& [from_node, from_edges] : froms_to_use) {
       VLOG(4) << "        Source: " << from_node->ToString();
-      XLS_ASSIGN_OR_RETURN(
-          Node * from_used,
-          visibility_builder->BuildVisibilityIRExpr(f, from_node, from_edges));
+      VLOG(4) << "        From edges: " << from_edges.size();
+      for (const auto& edge : from_edges) {
+        VLOG(4) << "          " << edge.operand->ToString() << " -> "
+                << edge.node->ToString();
+      }
+      XLS_ASSIGN_OR_RETURN(Node * from_used,
+                           visibility_builder->BuildVisibilityIRExpr(
+                               f, from_node, from_edges, folding->GetSinks()));
+      // Check that visibility isn't always true (ideally we would double check
+      // that the expressions are mutually exclusive, but that is expensive)
+      XLS_RET_CHECK(!IsLiteralUnsignedOne(from_used));
       from_used_expressions.push_back(from_used);
       VLOG(4) << "        From used: "
               << ToMathNotation(from_used, [&](const Node* n) -> bool {

xls/passes/resource_sharing_pass_test.cc

@@ -1525,7 +1525,10 @@ TEST_F(ResourceSharingPassTest, CatchesCyclesBeforeTransforming) {
   BValue i = fb.Param("i", u8);
   BValue X = fb.UMul(i, i, 8, SourceInfo(), "X");
   BValue Y = fb.UMul(X, i, 8, SourceInfo(), "D");
-  XLS_ASSERT_OK_AND_ASSIGN(Function * f, fb.BuildWithReturnValue(Y));
+  BValue cond = fb.Param("cond", p->GetBitsType(2));
+  BValue sel =
+      fb.Select(cond, {X, Y, fb.Literal(UBits(0, 8)), fb.Literal(UBits(0, 8))});
+  XLS_ASSERT_OK_AND_ASSIGN(Function * f, fb.BuildWithReturnValue(sel));
   int64_t next_node_id = 10;
 
   BddQueryEngine bdd_engine;
@@ -1536,10 +1539,18 @@ TEST_F(ResourceSharingPassTest, CatchesCyclesBeforeTransforming) {
   XLS_ASSERT_OK(bpa.Populate(f));
   VisibilityBuilder visibility_builder(next_node_id, &bdd_engine, nda, bpa);
 
+  // Exercise producing a cycle by requesting to fold X into Y; yes, they are
+  // not mutually exclusive, but the transformation doesn't have a feasible way
+  // to check whether analyses were correct; however, it should detect the cycle
+  // that would result from this transformation.
+  VisibilityEdges edges = {
+      OperandVisibilityAnalysis::OperandNode(X.node(), sel.node()),
+      OperandVisibilityAnalysis::OperandNode(Y.node(), sel.node())};
   std::vector<std::pair<Node*, VisibilityEdges>> from_X = {
-      std::make_pair(X.node(), VisibilityEdges{})};
+      std::make_pair(X.node(), edges)};
   auto fold_X_into_Y = std::make_unique<NaryFoldingAction>(
-      std::move(from_X), Y.node(), VisibilityEdges{});
+      std::move(from_X), Y.node(), edges, /*area_saved=*/0.0,
+      /*sinks=*/absl::flat_hash_set<Node*>{sel.node()});
   std::vector<std::unique_ptr<NaryFoldingAction>> folding_actions_to_perform;
   folding_actions_to_perform.push_back(std::move(fold_X_into_Y));
 

xls/passes/visibility_expr_builder.cc

@@ -339,8 +339,8 @@ absl::StatusOr<Node*> VisibilityBuilder::BuildVisibilityIRExprFromEdges(
     const absl::flat_hash_set<OperandVisibilityAnalysis::OperandNode>&
         conditional_edges,
     absl::flat_hash_map<Node*, Node*>& node_to_visibility_ir_cache,
-    Literal* always_visible) {
-  if (node->users().empty()) {
+    Literal* always_visible, const absl::flat_hash_set<Node*>& sinks) {
+  if (node->users().empty() || sinks.contains(node)) {
     return always_visible;
   }
   if (auto it = node_to_visibility_ir_cache.find(node);
@@ -353,10 +353,16 @@ absl::StatusOr<Node*> VisibilityBuilder::BuildVisibilityIRExprFromEdges(
     if (user->id() > prior_existing_id_) {
       continue;
     }
-    XLS_ASSIGN_OR_RETURN(Node * user_is_used,
-                         BuildVisibilityIRExprFromEdges(
-                             func, user, source, conditional_edges,
-                             node_to_visibility_ir_cache, always_visible));
+    bool sink_depends_on_user = absl::c_any_of(
+        sinks, [&](Node* sink) { return nda_.IsDependent(user, sink); });
+    if (!sinks.empty() && !sink_depends_on_user) {
+      continue;
+    }
+    XLS_ASSIGN_OR_RETURN(
+        Node * user_is_used,
+        BuildVisibilityIRExprFromEdges(func, user, source, conditional_edges,
+                                       node_to_visibility_ir_cache,
+                                       always_visible, sinks));
     Node* user_uses_node = always_visible;
     if (conditional_edges.contains({node, user})) {
       XLS_ASSIGN_OR_RETURN(user_uses_node,
@@ -407,32 +413,26 @@ absl::StatusOr<Node*> VisibilityBuilder::BuildVisibilityIRExprFromEdges(
 absl::StatusOr<Node*> VisibilityBuilder::BuildVisibilityIRExpr(
     FunctionBase* func, Node* node,
     const absl::flat_hash_set<OperandVisibilityAnalysis::OperandNode>&
-        conditional_edges) {
+        conditional_edges,
+    const absl::flat_hash_set<Node*>& sinks) {
   XLS_ASSIGN_OR_RETURN(
       Literal * always_visible,
       func->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
-  if (conditional_edges.size() == 1) {
-    XLS_ASSIGN_OR_RETURN(
-        Node * user_uses_node,
-        BuildVisibilityExpr(conditional_edges.begin()->operand,
-                            conditional_edges.begin()->node, node, func));
-    return user_uses_node ? user_uses_node : always_visible;
-  }
   absl::flat_hash_map<Node*, Node*> node_to_visibility_ir_cache;
   absl::flat_hash_map<std::tuple<Op, Node*, Node*>, Node*> binary_op_cache;
   return BuildVisibilityIRExprFromEdges(func, node, node, conditional_edges,
                                         node_to_visibility_ir_cache,
-                                        always_visible);
+                                        always_visible, sinks);
 }
 
 absl::StatusOr<VisibilityEstimator::AreaDelay>
 VisibilityEstimator::GetAreaAndDelayOfVisibilityExpr(
     Node* node,
     const absl::flat_hash_set<OperandVisibilityAnalysis::OperandNode>&
         conditional_edges) {
-  XLS_ASSIGN_OR_RETURN(
-      Node * visibility_expr,
-      BuildVisibilityIRExpr(TmpFunc(), node, conditional_edges));
+  XLS_ASSIGN_OR_RETURN(Node * visibility_expr,
+                       BuildVisibilityIRExpr(TmpFunc(), node, conditional_edges,
+                                             /*sinks=*/{}));
   double area = area_analysis_.GetAreaThroughToNode(visibility_expr);
   int64_t delay = *delay_analysis_.GetInfo(visibility_expr);
   return VisibilityEstimator::AreaDelay{area, delay};

xls/passes/visibility_expr_builder.h

@@ -100,7 +100,8 @@ class VisibilityBuilder : public ExpressionBuilder {
   absl::StatusOr<Node*> BuildVisibilityIRExpr(
       FunctionBase* func, Node* node,
       const absl::flat_hash_set<OperandVisibilityAnalysis::OperandNode>&
-          conditional_edges);
+          conditional_edges,
+      const absl::flat_hash_set<Node*>& sinks);
 
  private:
   absl::StatusOr<Node*> MakeParamIfTmpFunc(Node* node, FunctionBase* func) {
@@ -142,7 +143,7 @@ class VisibilityBuilder : public ExpressionBuilder {
       const absl::flat_hash_set<OperandVisibilityAnalysis::OperandNode>&
           conditional_edges,
       absl::flat_hash_map<Node*, Node*>& node_to_visibility_ir_cache,
-      Literal* always_visible);
+      Literal* always_visible, const absl::flat_hash_set<Node*>& sinks);
 
   absl::StatusOr<Node*> GetNonRepeatedSourceOf(Node* operand,
                                                FunctionBase* func);

xls/passes/visibility_expr_builder_test.cc

@@ -78,8 +78,9 @@ class VisibilityExprBuilderTest : public IrTestBase {
     BitProvenanceAnalysis bpa;
     VisibilityEstimator estimator(last_node_id, bdd_engine.get(), nda, bpa, ae,
                                   de);
-    XLS_ASSIGN_OR_RETURN(Node * expr, estimator.BuildVisibilityIRExpr(
-                                          f, node, conditional_edges));
+    XLS_ASSIGN_OR_RETURN(
+        Node * expr, estimator.BuildVisibilityIRExpr(f, node, conditional_edges,
+                                                     /*sinks=*/{}));
     XLS_ASSIGN_OR_RETURN(
         auto area_delay,
         estimator.GetAreaAndDelayOfVisibilityExpr(node, conditional_edges));
@@ -260,5 +261,36 @@ TEST_F(VisibilityExprBuilderTest, NotAFunctionOfSelf) {
   EXPECT_THAT(is_x_used.first, m::Ne(m::Param("y"), m::Literal(7)));
 }
 
+TEST_F(VisibilityExprBuilderTest, VisibilityExpressionWithOneKeptEdge) {
+  auto p = CreatePackage();
+  FunctionBuilder fb(TestName(), p.get());
+  Type* u1 = p->GetBitsType(1);
+  BValue x = fb.Param("x", u1);
+  BValue y = fb.Param("y", u1);
+
+  // z == select(y, {false, x}) == and(y, x)
+  BValue literal_false = fb.Literal(UBits(0, 1));
+  BValue z = fb.Select(y, {literal_false, x});
+
+  BValue ret = fb.Xor(x, z);
+  XLS_ASSERT_OK_AND_ASSIGN(Function * f, fb.BuildWithReturnValue(ret));
+  std::pair<Node*, VisibilityEstimator::AreaDelay> is_x_used;
+  XLS_ASSERT_OK_AND_ASSIGN(is_x_used,
+                           BuildDefaultVisibilityExpr(f, x.node(), {}));
+
+  // The visibility of `x` should be `Literal(1)`, since:
+  //   XOR(x, z) == XOR(x, AND(y, x)),
+  // and our analysis can't determine the visibility of `x` through `XOR(x, z)`.
+  //
+  // In case we later improve the analysis, the correct visibility expression
+  // for `x` is `y == 0`. We can see this with a casewise analysis:
+  //   If y == 0, then z == AND(y, x) == 0, so XOR(x, z) == XOR(x, 0) == x.
+  //   If y == 1, then z == AND(y, x) == x, so XOR(x, z) == XOR(x, x) == 0.
+  // So the output is x if y == 0, and 0 if y == 1. Equivalently:
+  //   ret == x & !y
+  // As such, `x` is actually visible iff `y == 0`.
+  EXPECT_THAT(is_x_used.first, m::Literal(1));
+}
+
 }  // namespace
 }  // namespace xls