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