From 5346b4500a2f82ab733b6dc5122ebb6eba0e5e97 Mon Sep 17 00:00:00 2001 From: Angelo Matni Date: Mon, 6 Jul 2026 11:23:33 -0700 Subject: [PATCH] Enable visibility expression builder usage while filtering BFS on def-use DAG This change ensures `BuildVisibilityIRExpr` always returns a valid visibility expression regardless of `conditional_edges` provided. If none are provided, then the function would return always_visible. If every edge in the def-use DAG from `node` down to terminal nodes is provided, then the expression would be exact, assuming no BDD overflow, and conservative otherwise. Providing only a subset of the entire DAG would produce a smaller and conservative expression, perhaps less conservative if that subset is chosen wisely. BuildVisibilityIRExpr had a bug where the short-circuit path when you provided a single-element conditional_edges set assumed that that single edge was a sufficient visibility expression to account for every path through the DAG. This assumption was made because resource sharing was using BuildVisibilityIRExpr in this way if it could prove mutual exclusivity via SingleSelectVisibilityAnalysis. For more context: * BuildVisibilityIRExprFromEdges computes visibility as OR(vis_node_through_user_i AND vis_user_i) where if the edge (node->user) is NOT in conditional_edges we assume vis_node_through_user_i is always true. This function is conservative. Using conditional_edgescan help prevent BDD overflow. * The test helper BuildDefaultVisibilityExpr uses GetEdgesForMutuallyExclusiveVisibilityExpr to determine a subset of conditional_edges that helps BuildVisibilityIRExprFromEdgesnot overflow. PiperOrigin-RevId: 943398726 --- xls/passes/BUILD | 1 + xls/passes/folding_graph.cc | 11 ++++++- xls/passes/folding_graph.h | 21 +++++++++---- xls/passes/resource_sharing_pass.cc | 34 +++++++++++++++----- xls/passes/resource_sharing_pass_test.cc | 17 ++++++++-- xls/passes/visibility_expr_builder.cc | 36 +++++++++++----------- xls/passes/visibility_expr_builder.h | 5 +-- xls/passes/visibility_expr_builder_test.cc | 36 ++++++++++++++++++++-- 8 files changed, 122 insertions(+), 39 deletions(-) diff --git a/xls/passes/BUILD b/xls/passes/BUILD index 06f1308e37..021d4b8d31 100644 --- a/xls/passes/BUILD +++ b/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", diff --git a/xls/passes/folding_graph.cc b/xls/passes/folding_graph.cc index 3f421103a8..7ec902e36d 100644 --- a/xls/passes/folding_graph.cc +++ b/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> NaryFoldingAction::Clone( clone_to_edges.insert({clone_operand, clone_node}); } + absl::flat_hash_set 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( std::move(clone_froms), clone_to, std::move(clone_to_edges), - other.area_saved()); + other.area_saved(), std::move(clone_sinks)); return clone_action; } diff --git a/xls/passes/folding_graph.h b/xls/passes/folding_graph.h index da263d874c..a526ec4985 100644 --- a/xls/passes/folding_graph.h +++ b/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& 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 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 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 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>&& 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 sinks = {}) + : FoldingAction{to, std::move(to_edges), area_saved, std::move(sinks)}, from_(std::move(froms)) {} absl::Span> GetFrom() const { diff --git a/xls/passes/resource_sharing_pass.cc b/xls/passes/resource_sharing_pass.cc index 63bc4481f6..ddf5f4822c 100644 --- a/xls/passes/resource_sharing_pass.cc +++ b/xls/passes/resource_sharing_pass.cc @@ -30,6 +30,7 @@ #include #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,6 +296,7 @@ ResourceSharingPass::ComputeFoldableActions( absl::flat_hash_set one_edges; absl::flat_hash_set other_edges; + absl::flat_hash_set sinks; if (visibility.single_select.IsMutuallyExclusive(one_node, other_node)) { XLS_ASSIGN_OR_RETURN( one_edges, @@ -301,6 +304,9 @@ ResourceSharingPass::ComputeFoldableActions( 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,7 +332,7 @@ 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( - 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 " @@ -334,7 +340,7 @@ 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( - 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 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 new_action = std::make_unique( 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 sinks(folding->GetSinks()); std::unique_ptr new_folding = std::make_unique(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 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 { diff --git a/xls/passes/resource_sharing_pass_test.cc b/xls/passes/resource_sharing_pass_test.cc index 6ece9d209c..003d58b2bf 100644 --- a/xls/passes/resource_sharing_pass_test.cc +++ b/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> from_X = { - std::make_pair(X.node(), VisibilityEdges{})}; + std::make_pair(X.node(), edges)}; auto fold_X_into_Y = std::make_unique( - std::move(from_X), Y.node(), VisibilityEdges{}); + std::move(from_X), Y.node(), edges, /*area_saved=*/0.0, + /*sinks=*/absl::flat_hash_set{sel.node()}); std::vector> folding_actions_to_perform; folding_actions_to_perform.push_back(std::move(fold_X_into_Y)); diff --git a/xls/passes/visibility_expr_builder.cc b/xls/passes/visibility_expr_builder.cc index 807d542f16..886a774f85 100644 --- a/xls/passes/visibility_expr_builder.cc +++ b/xls/passes/visibility_expr_builder.cc @@ -339,8 +339,8 @@ absl::StatusOr VisibilityBuilder::BuildVisibilityIRExprFromEdges( const absl::flat_hash_set& conditional_edges, absl::flat_hash_map& node_to_visibility_ir_cache, - Literal* always_visible) { - if (node->users().empty()) { + Literal* always_visible, const absl::flat_hash_set& 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 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,22 +413,16 @@ absl::StatusOr VisibilityBuilder::BuildVisibilityIRExprFromEdges( absl::StatusOr VisibilityBuilder::BuildVisibilityIRExpr( FunctionBase* func, Node* node, const absl::flat_hash_set& - conditional_edges) { + conditional_edges, + const absl::flat_hash_set& sinks) { XLS_ASSIGN_OR_RETURN( Literal * always_visible, func->MakeNode(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_to_visibility_ir_cache; absl::flat_hash_map, Node*> binary_op_cache; return BuildVisibilityIRExprFromEdges(func, node, node, conditional_edges, node_to_visibility_ir_cache, - always_visible); + always_visible, sinks); } absl::StatusOr @@ -430,9 +430,9 @@ VisibilityEstimator::GetAreaAndDelayOfVisibilityExpr( Node* node, const absl::flat_hash_set& 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}; diff --git a/xls/passes/visibility_expr_builder.h b/xls/passes/visibility_expr_builder.h index 10a253c5e8..49dbd76886 100644 --- a/xls/passes/visibility_expr_builder.h +++ b/xls/passes/visibility_expr_builder.h @@ -100,7 +100,8 @@ class VisibilityBuilder : public ExpressionBuilder { absl::StatusOr BuildVisibilityIRExpr( FunctionBase* func, Node* node, const absl::flat_hash_set& - conditional_edges); + conditional_edges, + const absl::flat_hash_set& sinks); private: absl::StatusOr MakeParamIfTmpFunc(Node* node, FunctionBase* func) { @@ -142,7 +143,7 @@ class VisibilityBuilder : public ExpressionBuilder { const absl::flat_hash_set& conditional_edges, absl::flat_hash_map& node_to_visibility_ir_cache, - Literal* always_visible); + Literal* always_visible, const absl::flat_hash_set& sinks); absl::StatusOr GetNonRepeatedSourceOf(Node* operand, FunctionBase* func); diff --git a/xls/passes/visibility_expr_builder_test.cc b/xls/passes/visibility_expr_builder_test.cc index 81c0f87b32..32e53fe79b 100644 --- a/xls/passes/visibility_expr_builder_test.cc +++ b/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 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