From e65cd729ab1f664d29bf60025c78914860342cae Mon Sep 17 00:00:00 2001 From: Nelson Liang Date: Mon, 27 Apr 2026 13:01:40 -0700 Subject: [PATCH] [Explicit State Access] Codegen support for unread/multi-read Next nodes without optimization passes This change enables the codegen pipeline to handle `Next` nodes that reference a `StateElement` directly, without an associated `StateRead`. - Update `CalculateConcurrentStages` and `InsertPipelineRegisters` in `pipeline_register_insertion_pass.cc` to query `Proc` for state reads when a `Next` node doesn't have one. - Add a new test case `ProcWithExplicitStateAccessNewApi` in `convert_to_block_test.cc` using the new API and verify it completes codegen successfully. PiperOrigin-RevId: 906504274 --- xls/codegen_v_1_5/BUILD | 4 + xls/codegen_v_1_5/convert_to_block_test.cc | 55 ++++++++++++ .../pipeline_register_insertion_pass.cc | 71 +++++++++++----- .../state_to_register_io_lowering_pass.cc | 83 +++++++++++-------- ...state_to_register_io_lowering_pass_test.cc | 45 ++++++++++ ...ing_test_ProcWithMultipleReadsSameStage.ir | 42 ++++++++++ xls/scheduling/pipeline_schedule.cc | 50 ++++++----- xls/scheduling/pipeline_schedule_test.cc | 62 ++++++++++++++ xls/scheduling/schedule_bounds.cc | 17 ++-- xls/scheduling/schedule_graph.cc | 13 +-- xls/scheduling/schedule_util.cc | 13 ++- 11 files changed, 365 insertions(+), 90 deletions(-) create mode 100644 xls/codegen_v_1_5/testdata/state_to_register_io_lowering_test_ProcWithMultipleReadsSameStage.ir diff --git a/xls/codegen_v_1_5/BUILD b/xls/codegen_v_1_5/BUILD index ee77d4f481..b0ebdc9c28 100644 --- a/xls/codegen_v_1_5/BUILD +++ b/xls/codegen_v_1_5/BUILD @@ -295,8 +295,11 @@ cc_test( "//xls/codegen:codegen_options", "//xls/common:xls_gunit_main", "//xls/common/status:matchers", + "//xls/ir:bits", + "//xls/ir:channel", "//xls/ir:function_builder", "//xls/ir:ir_test_base", + "//xls/ir:value", "//xls/passes:optimization_pass", "//xls/passes:pass_base", "//xls/scheduling:scheduling_options", @@ -885,6 +888,7 @@ cc_library( "//xls/ir:value", "//xls/ir:value_utils", "//xls/passes:pass_base", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status", diff --git a/xls/codegen_v_1_5/convert_to_block_test.cc b/xls/codegen_v_1_5/convert_to_block_test.cc index a49927c45b..1d3b455e9a 100644 --- a/xls/codegen_v_1_5/convert_to_block_test.cc +++ b/xls/codegen_v_1_5/convert_to_block_test.cc @@ -18,8 +18,11 @@ #include "gtest/gtest.h" #include "xls/codegen/codegen_options.h" #include "xls/common/status/matchers.h" +#include "xls/ir/bits.h" +#include "xls/ir/channel.h" #include "xls/ir/function_builder.h" #include "xls/ir/ir_test_base.h" +#include "xls/ir/value.h" #include "xls/passes/optimization_pass.h" #include "xls/passes/pass_base.h" #include "xls/scheduling/scheduling_options.h" @@ -66,5 +69,57 @@ TEST_F(ConvertToBlockTest, SimpleFunction) { // TODO: https://github.com/google/xls/issues/3356 - assert stuff. } +TEST_F(ConvertToBlockTest, ProcWithExplicitStateAccessSimpleAddOldApi) { + auto p = CreatePackage(); + TokenlessProcBuilder pb(NewStyleProc(), TestName(), "tkn", p.get()); + XLS_ASSERT_OK_AND_ASSIGN(SendChannelInterface * out_ch_interface, + pb.AddOutputChannel("out_ch", p->GetBitsType(32))); + + BValue state_read = pb.StateElement("state", Value(UBits(0, 32))); + BValue current = pb.Identity(state_read); + BValue add_val = pb.Add(current, pb.Literal(UBits(1, 32))); + + pb.Send(out_ch_interface, add_val); + // // Using the OLD API (passing the StateRead BValue) + pb.Next(state_read, add_val); + XLS_ASSERT_OK_AND_ASSIGN(Proc * top, pb.Build()); + XLS_ASSERT_OK(p->SetTop(top)); + + OptimizationContext opt_context; + PassResults pass_results; + XLS_ASSERT_OK(ConvertToBlock( + p.get(), + codegen_options().clock_name("clk").reset("rst", false, false, false), + scheduling_options(), &delay_estimator_, &opt_context, &pass_results)); +} + +TEST_F(ConvertToBlockTest, ProcWithExplicitStateAccessNewApi) { + auto p = CreatePackage(); + TokenlessProcBuilder pb(NewStyleProc(), TestName(), "tkn", p.get()); + XLS_ASSERT_OK_AND_ASSIGN(SendChannelInterface * out_ch_interface, + pb.AddOutputChannel("out_ch", p->GetBitsType(32))); + + BValue state_read = pb.StateElement("state", Value(UBits(0, 32))); + BValue current = pb.Identity(state_read); + BValue add_val = pb.Add(current, pb.Literal(UBits(1, 32))); + + pb.Send(out_ch_interface, add_val); + + // Using the NEW API (passing the StateElement) + StateElement* se = state_read.node()->As()->state_element(); + pb.Next(se, add_val); + + XLS_ASSERT_OK_AND_ASSIGN(Proc * top, pb.Build()); + XLS_ASSERT_OK(p->SetTop(top)); + + OptimizationContext opt_context; + PassResults pass_results; + XLS_ASSERT_OK(ConvertToBlock( + p.get(), + codegen_options().clock_name("clk").reset("rst", false, false, false), + scheduling_options().opt_level(0), &delay_estimator_, &opt_context, + &pass_results)); +} + } // namespace } // namespace xls::codegen diff --git a/xls/codegen_v_1_5/pipeline_register_insertion_pass.cc b/xls/codegen_v_1_5/pipeline_register_insertion_pass.cc index cce5aeff7d..63fa3245b6 100644 --- a/xls/codegen_v_1_5/pipeline_register_insertion_pass.cc +++ b/xls/codegen_v_1_5/pipeline_register_insertion_pass.cc @@ -201,21 +201,54 @@ absl::StatusOr CalculateConcurrentStages( continue; } Next* next = node->As(); - StateRead* state_read = node->As()->state_read()->As(); - StateElement* state_element = state_read->state_element(); - if (state_read->predicate().has_value()) { - // If the state read is predicated, then it doesn't start a mutual - // exclusion zone. - continue; - } - XLS_ASSIGN_OR_RETURN(read_by_stage[state_element], - block->GetStageIndex(state_read)); - - XLS_ASSIGN_OR_RETURN(int64_t write_stage, block->GetStageIndex(next)); - auto [it, inserted] = - first_write_stage.try_emplace(state_element, write_stage); - if (!inserted) { - it->second = std::min(it->second, write_stage); + if (next->has_state_read()) { + StateRead* state_read = next->state_read()->As(); + StateElement* state_element = state_read->state_element(); + if (state_read->predicate().has_value()) { + // If the state read is predicated, then it doesn't start a mutual + // exclusion zone. + continue; + } + XLS_ASSIGN_OR_RETURN(read_by_stage[state_element], + block->GetStageIndex(state_read)); + + XLS_ASSIGN_OR_RETURN(int64_t write_stage, block->GetStageIndex(next)); + auto [it, inserted] = + first_write_stage.try_emplace(state_element, write_stage); + if (!inserted) { + it->second = std::min(it->second, write_stage); + } + } else { + StateElement* state_element = next->state_element(); + if (block->source() == nullptr || !block->source()->IsProc()) { + continue; + } + Proc* proc = block->source()->AsProcOrDie(); + absl::Span reads = + proc->GetStateReadsByStateElement(state_element); + std::optional min_read_stage; + for (StateRead* read : reads) { + // Predicated interactions don't safely map exclusion zones + if (read->predicate().has_value()) { + continue; + } + XLS_ASSIGN_OR_RETURN(int64_t read_stage, block->GetStageIndex(read)); + if (!min_read_stage.has_value() || read_stage < *min_read_stage) { + min_read_stage = read_stage; + } + } + // If no unconditional paths trigger, ignore constraint + if (!min_read_stage.has_value()) { + continue; + } + read_by_stage[state_element] = *min_read_stage; + + XLS_ASSIGN_OR_RETURN(int64_t write_stage, block->GetStageIndex(next)); + auto [it, inserted] = + first_write_stage.try_emplace(state_element, write_stage); + if (!inserted) { + it->second = std::min(it->second, write_stage); + } } } for (const auto& [state_element, read_stage] : read_by_stage) { @@ -253,8 +286,7 @@ absl::StatusOr PipelineRegisterInsertionPass::InsertPipelineRegisters( for (Node* node : block->nodes()) { if (node->Is()) { Next* next = node->As(); - StateElement* state_element = - next->state_read()->As()->state_element(); + StateElement* state_element = next->state_element(); int64_t stage = *block->GetStageIndex(next); auto [it, inserted] = earliest_next_stage.try_emplace(state_element, stage); @@ -351,8 +383,9 @@ absl::StatusOr PipelineRegisterInsertionPass::InsertPipelineRegisters( // corresponding StateRead, we need to avoid updating that specific // operand, since no data is being passed. For simplicity, we put it // back afterwards rather than actually avoiding the update. - bool restore_state_read = - user->Is() && user->As()->state_read() == node; + bool restore_state_read = user->Is() && + user->As()->has_state_read() && + user->As()->state_read() == node; user->ReplaceOperand(node, live_node); if (restore_state_read) { XLS_RETURN_IF_ERROR(user->As()->ReplaceOperandNumber( diff --git a/xls/codegen_v_1_5/state_to_register_io_lowering_pass.cc b/xls/codegen_v_1_5/state_to_register_io_lowering_pass.cc index f3d43ab8b6..e1678f6a58 100644 --- a/xls/codegen_v_1_5/state_to_register_io_lowering_pass.cc +++ b/xls/codegen_v_1_5/state_to_register_io_lowering_pass.cc @@ -19,6 +19,7 @@ #include #include +#include "absl/algorithm/container.h" #include "absl/base/casts.h" #include "absl/container/flat_hash_map.h" #include "absl/status/status.h" @@ -165,7 +166,8 @@ absl::Status AddFullRegisterWrite(Block* block, Register* reg_full, absl::Status LowerStateElement(ScheduledBlock* block, const StateElement& state_element, - StateRead* read, absl::Span writes) { + absl::Span reads, + absl::Span writes) { // A token or empty tuple state element has no real functionality. Type* type = state_element.type(); const bool has_data = type->GetFlatBitCount() > 0; @@ -177,32 +179,40 @@ absl::Status LowerStateElement(ScheduledBlock* block, state_element.name(), type->ToString())); } - // Lower the read of the state element. - std::optional read_predicate = read->predicate(); + // Lower the reads of the state element. std::string name = block->UniquifyNodeName(absl::StrCat("__", state_element.name())); - - XLS_ASSIGN_OR_RETURN(int read_stage_index, block->GetStageIndex(read)); - Stage& read_stage = block->stages()[read_stage_index]; - Node* reg_read_or_zero = nullptr; Register* state_register = nullptr; if (has_data) { XLS_ASSIGN_OR_RETURN(state_register, block->AddRegister(name, state_element.type(), state_element.initial_value())); - XLS_ASSIGN_OR_RETURN(reg_read_or_zero, - block->MakeNodeWithNameInStage( - read_stage_index, read->loc(), state_register, - /*name=*/state_register->name())); - } else { - XLS_ASSIGN_OR_RETURN(reg_read_or_zero, block->MakeNode( - read->loc(), ZeroOfType(type))); } - XLS_ASSIGN_OR_RETURN(bool read_removed, - ReplaceNode(block, read, reg_read_or_zero, - /*replace_implicit_uses=*/false)); - XLS_RET_CHECK(!read_removed); + std::vector reg_reads_or_zeros; + std::vector read_stage_indices; + reg_reads_or_zeros.reserve(reads.size()); + read_stage_indices.reserve(reads.size()); + for (StateRead* read : reads) { + XLS_ASSIGN_OR_RETURN(int read_stage_index, block->GetStageIndex(read)); + read_stage_indices.push_back(read_stage_index); + Node* reg_read_or_zero = nullptr; + if (has_data) { + XLS_ASSIGN_OR_RETURN(reg_read_or_zero, + block->MakeNodeWithNameInStage( + read_stage_index, read->loc(), state_register, + /*name=*/state_register->name())); + } else { + XLS_ASSIGN_OR_RETURN( + reg_read_or_zero, + block->MakeNode(read->loc(), ZeroOfType(type))); + } + XLS_ASSIGN_OR_RETURN(bool read_removed, + ReplaceNode(block, read, reg_read_or_zero, + /*replace_implicit_uses=*/false)); + XLS_RET_CHECK(!read_removed); + reg_reads_or_zeros.push_back(reg_read_or_zero); + } Next* last_write = nullptr; Node* last_value = nullptr; @@ -222,11 +232,18 @@ absl::Status LowerStateElement(ScheduledBlock* block, // we have a non-trivial backedge between initiations (II>1); use a "full" bit // to track whether the state is currently valid. Register* reg_full = nullptr; - if (last_write_stage_index > read_stage_index) { - XLS_ASSIGN_OR_RETURN( - reg_full, - CreateFullRegisterAndRead(block, state_element, read_predicate, - *last_write, read_stage_index, read_stage)); + if (!reads.empty() && !writes.empty()) { + int earliest_read_stage_index = read_stage_indices[0]; + Stage& earliest_read_stage = block->stages()[earliest_read_stage_index]; + if (last_write_stage_index > earliest_read_stage_index) { + // TODO(nelsonliang): Support II > 1 with multiple reads in different + // stages. + XLS_ASSIGN_OR_RETURN( + reg_full, + CreateFullRegisterAndRead(block, state_element, reads[0]->predicate(), + *last_write, earliest_read_stage_index, + earliest_read_stage)); + } } std::vector all_write_conditions; @@ -247,7 +264,8 @@ absl::Status LowerStateElement(ScheduledBlock* block, // If we have no full bit to update and this is just a self-assignment, we // can skip this write entirely. - if (!reg_full && write->value() == reg_read_or_zero) { + if (!reg_full && + absl::c_linear_search(reg_reads_or_zeros, write->value())) { continue; } @@ -267,7 +285,7 @@ absl::Status LowerStateElement(ScheduledBlock* block, // For data updates, we skip all self-assignments, as they're just // unchanged values. - if (write->value() == reg_read_or_zero) { + if (absl::c_linear_search(reg_reads_or_zeros, write->value())) { continue; } @@ -303,7 +321,7 @@ absl::Status LowerStateElement(ScheduledBlock* block, if (values.empty()) { // If we never change the state element, write the current value // unconditionally. - value = reg_read_or_zero; + value = reg_reads_or_zeros.back(); } else if (values.size() == 1) { value = values[0]; write_load_enable = write_conditions[0]; @@ -350,9 +368,11 @@ absl::Status LowerStateElement(ScheduledBlock* block, NaryOrIfNeeded(block, all_write_conditions, /*name=*/"", last_write_loc)); } - XLS_RETURN_IF_ERROR(AddFullRegisterWrite(block, reg_full, read_predicate, - read_stage, last_write_loc, - full_write_enable)); + int earliest_read_stage_index = read_stage_indices[0]; + Stage& earliest_read_stage = block->stages()[earliest_read_stage_index]; + XLS_RETURN_IF_ERROR(AddFullRegisterWrite( + block, reg_full, reads[0]->predicate(), earliest_read_stage, + last_write_loc, full_write_enable)); } return absl::OkStatus(); @@ -382,11 +402,8 @@ absl::StatusOr LowerStateIoForBlock(ScheduledBlock* block) { for (StateElement* state_element : source_proc->StateElements()) { std::vector reads_for_element = state_reads.at(state_element); - // Multiple reads are not yet supported. - XLS_RET_CHECK_EQ(reads_for_element.size(), 1) - << "Found multiple reads for state element: " << state_element->name(); XLS_RETURN_IF_ERROR( - LowerStateElement(block, *state_element, reads_for_element[0], + LowerStateElement(block, *state_element, reads_for_element, absl::MakeSpan(next_values[state_element]))); } diff --git a/xls/codegen_v_1_5/state_to_register_io_lowering_pass_test.cc b/xls/codegen_v_1_5/state_to_register_io_lowering_pass_test.cc index 327ff9af6f..de09a78811 100644 --- a/xls/codegen_v_1_5/state_to_register_io_lowering_pass_test.cc +++ b/xls/codegen_v_1_5/state_to_register_io_lowering_pass_test.cc @@ -488,5 +488,50 @@ top scheduled_block __test__P_0_next(clk: clock, rst: bits[1]) { ExpectEqualToGoldenFile(output); } +TEST_F(StateToRegisterIoLoweringPassTest, ProcWithMultipleReadsSameStage) { + XLS_ASSERT_OK_AND_ASSIGN(std::string output, RunPassAndRoundTripIrText(R"( +package test + +top scheduled_block __test__P_0_next(clk: clock, rst: bits[1]) { + #![reset(port="rst", asynchronous=false, active_low=false)] + source proc __test__P_0_next__src(__state: bits[32], init={0}) { + chan_interface a(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none) + chan_interface b(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none) + chan_interface result(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none) + } + literal.3: bits[1] = literal(value=1, id=3) + stage_inputs_valid_0: bits[1] = literal(value=1, id=20) + stage_outputs_ready_0: bits[1] = literal(value=1, id=21) + controlled_stage(stage_inputs_valid_0, stage_outputs_ready_0) { + active_inputs_valid active_inputs_valid_0: bits[1] = literal(value=1, id=22) + after_all.5: token = after_all(id=5) + receive.6: (token, bits[32]) = receive(after_all.5, predicate=literal.3, channel=a, id=6) + tok: token = tuple_index(receive.6, index=0, id=8) + receive.10: (token, bits[32]) = receive(tok, predicate=literal.3, channel=b, id=10) + a_value: bits[32] = tuple_index(receive.6, index=1, id=9) + tok__1: token = tuple_index(receive.10, index=0, id=12) + b_value: bits[32] = tuple_index(receive.10, index=1, id=13) + umul.14: bits[32] = umul(a_value, b_value, id=14) + ret stage_outputs_valid_0: bits[1] = and(stage_inputs_valid_0, active_inputs_valid_0, id=23) + } + stage_inputs_valid_1: bits[1] = literal(value=1, id=24) + stage_outputs_ready_1: bits[1] = literal(value=1, id=25) + controlled_stage(stage_inputs_valid_1, stage_outputs_ready_1) { + active_inputs_valid active_inputs_valid_1: bits[1] = literal(value=1, id=26) + __state: bits[32] = state_read(state_element=__state, id=2) + __state_2: bits[32] = state_read(state_element=__state, id=28) + result_value: bits[32] = add(umul.14, __state, id=15) + result_value_2: bits[32] = add(result_value, __state_2, id=29) + send.16: token = send(tok__1, result_value_2, predicate=literal.3, channel=result, id=16) + next_value.17: () = next_value(param=__state, value=result_value_2, id=17) + ret stage_outputs_valid_1: bits[1] = and(stage_inputs_valid_1, active_inputs_valid_1, id=27) + } + rst: bits[1] = input_port(name=rst, id=19) +} +)")); + + ExpectEqualToGoldenFile(output); +} + } // namespace } // namespace xls::codegen diff --git a/xls/codegen_v_1_5/testdata/state_to_register_io_lowering_test_ProcWithMultipleReadsSameStage.ir b/xls/codegen_v_1_5/testdata/state_to_register_io_lowering_test_ProcWithMultipleReadsSameStage.ir new file mode 100644 index 0000000000..06c8356d83 --- /dev/null +++ b/xls/codegen_v_1_5/testdata/state_to_register_io_lowering_test_ProcWithMultipleReadsSameStage.ir @@ -0,0 +1,42 @@ +package test + +top scheduled_block __test__P_0_next(clk: clock, rst: bits[1]) { + #![reset(port="rst", asynchronous=false, active_low=false)] + reg ____state(bits[32], reset_value=0) + + source proc __test__P_0_next__src() { + chan_interface a(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none) + chan_interface b(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none) + chan_interface result(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none) + } + literal.3: bits[1] = literal(value=1, id=3) + stage_inputs_valid_0: bits[1] = literal(value=1, id=20) + stage_outputs_ready_0: bits[1] = literal(value=1, id=21) + controlled_stage(stage_inputs_valid_0, stage_outputs_ready_0) { + active_inputs_valid active_inputs_valid_0: bits[1] = literal(value=1, id=22) + after_all.5: token = after_all(id=5) + receive.6: (token, bits[32]) = receive(after_all.5, predicate=literal.3, channel=a, id=6) + tok: token = tuple_index(receive.6, index=0, id=8) + receive.10: (token, bits[32]) = receive(tok, predicate=literal.3, channel=b, id=10) + a_value: bits[32] = tuple_index(receive.6, index=1, id=9) + tok__1: token = tuple_index(receive.10, index=0, id=12) + b_value: bits[32] = tuple_index(receive.10, index=1, id=13) + umul.14: bits[32] = umul(a_value, b_value, id=14) + ret stage_outputs_valid_0: bits[1] = and(stage_inputs_valid_0, active_inputs_valid_0, id=23) + } + stage_inputs_valid_1: bits[1] = literal(value=1, id=24) + stage_outputs_ready_1: bits[1] = literal(value=1, id=25) + controlled_stage(stage_inputs_valid_1, stage_outputs_ready_1) { + active_inputs_valid active_inputs_valid_1: bits[1] = literal(value=1, id=26) + ____state__1: bits[32] = register_read(register=____state, id=41) + result_value: bits[32] = add(umul.14, ____state__1, id=15) + result_value_2: bits[32] = add(result_value, ____state__1, id=29) + send.16: token = send(tok__1, result_value_2, predicate=literal.3, channel=result, id=16) + ret stage_outputs_valid_1: bits[1] = and(stage_inputs_valid_1, active_inputs_valid_1, id=27) + } + rst: bits[1] = input_port(name=rst, id=19) + and.42: bits[1] = and(stage_outputs_valid_1, stage_outputs_ready_1, id=42) + and.43: bits[1] = and(stage_inputs_valid_1, active_inputs_valid_1, id=43) + register_write.44: () = register_write(result_value_2, register=____state, load_enable=and.42, reset=rst, id=44) + tuple.45: () = tuple(id=45) +} diff --git a/xls/scheduling/pipeline_schedule.cc b/xls/scheduling/pipeline_schedule.cc index 9d663b0add..778b54bfb5 100644 --- a/xls/scheduling/pipeline_schedule.cc +++ b/xls/scheduling/pipeline_schedule.cc @@ -283,17 +283,17 @@ absl::Status PipelineSchedule::Verify() const { int64_t worst_case_throughput = proc->GetInitiationInterval().value_or(1); if (worst_case_throughput >= 1) { for (Next* next : proc->next_values()) { - Node* state_read = next->state_read(); - // Verify that no write happens before the corresponding read. - XLS_RET_CHECK_LE(cycle(state_read), cycle(next)) - << "Next node " << next << " scheduled before state read " - << state_read; - // Verify that we determine the new state within II cycles of accessing - // the current param. - XLS_RET_CHECK_LT(cycle(next), cycle(state_read) + worst_case_throughput) - << "Next node " << next << " scheduled too late after " - << state_read << " (stage " << cycle(next) << " is not less than " - << cycle(state_read) << " + WCT " << worst_case_throughput << ")"; + absl::Span reads = + proc->GetStateReadsByStateElement(next->state_element()); + for (StateRead* read : reads) { + XLS_RET_CHECK_LE(cycle(read), cycle(next)) + << "Next node " << next << " scheduled before state read " + << read; + XLS_RET_CHECK_LT(cycle(next), cycle(read) + worst_case_throughput) + << "Next node " << next << " scheduled too late after " << read + << " (stage " << cycle(next) << " is not less than " + << cycle(read) << " + WCT " << worst_case_throughput << ")"; + } } } } @@ -551,18 +551,22 @@ absl::Status PipelineSchedule::VerifyConstraints( } Proc* proc = function_base_->AsProcOrDie(); for (Next* next : proc->next_values()) { - StateRead* state_read = next->state_read()->As(); - int64_t backedge_length = - cycle_map_.at(next) - cycle_map_.at(state_read); - if (backedge_length > max_backedge) { - return absl::ResourceExhaustedError(absl::StrFormat( - "Scheduling constraint violated: state read %s was scheduled %d " - "cycle%s before node %s, its next value, which violates the " - "constraint that we can achieve a worst-case throughput of one " - "iteration per %d cycle%s without external stalls.", - state_read->GetName(), backedge_length, plural_s(backedge_length), - next->ToString(), worst_case_throughput.value_or(1), - plural_s(worst_case_throughput.value_or(1)))); + absl::Span reads = + proc->GetStateReadsByStateElement(next->state_element()); + for (StateRead* read : reads) { + int64_t backedge_length = cycle_map_.at(next) - cycle_map_.at(read); + if (backedge_length > max_backedge) { + return absl::ResourceExhaustedError(absl::StrFormat( + "Scheduling constraint violated: state read %s was scheduled " + "%d " + "cycle%s before node %s, its next value, which violates the " + "constraint that we can achieve a worst-case throughput of " + "one " + "iteration per %d cycle%s without external stalls.", + read->GetName(), backedge_length, plural_s(backedge_length), + next->ToString(), worst_case_throughput.value_or(1), + plural_s(worst_case_throughput.value_or(1)))); + } } } } else if (std::holds_alternative(constraint)) { diff --git a/xls/scheduling/pipeline_schedule_test.cc b/xls/scheduling/pipeline_schedule_test.cc index 9e12f8d804..efd62b65c5 100644 --- a/xls/scheduling/pipeline_schedule_test.cc +++ b/xls/scheduling/pipeline_schedule_test.cc @@ -2118,5 +2118,67 @@ TEST_F(PipelineScheduleTest, SerializeAndDeserializeWithSynchronousSchedule) { EXPECT_EQ(clone.GetSynchronousCycle(add2.node()), 42); } +TEST_F(PipelineScheduleTest, ProcWithExplicitStateAccess) { + Package p(TestName()); + Type* u32 = p.GetBitsType(32); + XLS_ASSERT_OK_AND_ASSIGN( + Channel * out_ch, + p.CreateStreamingChannel("out_ch", ChannelOps::kSendOnly, u32)); + + TokenlessProcBuilder pb("the_proc", "tkn", &p); + XLS_ASSERT_OK_AND_ASSIGN(StateElement * se, + pb.UnreadStateElement("state", Value(UBits(0, 32)))); + BValue current = pb.StateRead(se); + BValue add_val = pb.Add(current, pb.Literal(UBits(1, 32))); + + pb.Send(out_ch, add_val); + + pb.Next(se, add_val); + + XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build()); + + SchedulingOptions options(SchedulingStrategy::ASAP); + options.clock_period_ps(2); + + XLS_ASSERT_OK_AND_ASSIGN( + PipelineSchedule schedule, + RunPipelineSchedule(proc, TestDelayEstimator(), options)); + + EXPECT_EQ(schedule.length(), 1); +} + +TEST_F(PipelineScheduleTest, ProcWithMultipleStateReads) { + Package p(TestName()); + Type* u32 = p.GetBitsType(32); + XLS_ASSERT_OK_AND_ASSIGN( + Channel * out_ch, + p.CreateStreamingChannel("out_ch", ChannelOps::kSendOnly, u32)); + + TokenlessProcBuilder pb("the_proc", "tkn", &p); + XLS_ASSERT_OK_AND_ASSIGN(StateElement * se, + pb.UnreadStateElement("state", Value(UBits(0, 32)))); + BValue read1 = pb.StateRead(se); + + // Create a second read with a predicate + BValue cond = pb.Literal(UBits(1, 1)); + BValue read2 = pb.StateRead(se, cond); + + BValue add_val = pb.Add(read1, read2); + pb.Send(out_ch, add_val); + + pb.Next(se, add_val); + + XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build()); + + SchedulingOptions options(SchedulingStrategy::ASAP); + options.clock_period_ps(2); + + XLS_ASSERT_OK_AND_ASSIGN( + PipelineSchedule schedule, + RunPipelineSchedule(proc, TestDelayEstimator(), options)); + + EXPECT_EQ(schedule.length(), 1); +} + } // namespace } // namespace xls diff --git a/xls/scheduling/schedule_bounds.cc b/xls/scheduling/schedule_bounds.cc index b8b3897b7a..98bd0ccfdb 100644 --- a/xls/scheduling/schedule_bounds.cc +++ b/xls/scheduling/schedule_bounds.cc @@ -199,12 +199,17 @@ class ConstraintConverter { return absl::OkStatus(); } for (Next* next : fb_->AsProcOrDie()->next_values()) { - node_constraints_.emplace_back(NodeDifferenceConstraint{ - .anchor = next->state_read(), - .subject = next, - .min_after = 0, - .max_after = *ii_ - 1, - }); + absl::Span reads = + fb_->AsProcOrDie()->GetStateReadsByStateElement( + next->state_element()); + for (StateRead* read : reads) { + node_constraints_.emplace_back(NodeDifferenceConstraint{ + .anchor = read, + .subject = next, + .min_after = 0, + .max_after = *ii_ - 1, + }); + } } return absl::OkStatus(); } diff --git a/xls/scheduling/schedule_graph.cc b/xls/scheduling/schedule_graph.cc index 8dc9738a9f..f212d13ec7 100644 --- a/xls/scheduling/schedule_graph.cc +++ b/xls/scheduling/schedule_graph.cc @@ -157,11 +157,14 @@ std::string ScheduleGraph::ToString() const { // Proc state is represented as backedges in the graph. if (f->IsProc()) { for (Next* next : f->AsProcOrDie()->next_values()) { - StateRead* state_read = next->state_read()->As(); - backedges.push_back( - ScheduleBackedge{.source = next, - .destination = state_read, - .distance = LessThanInitiationInterval()}); + absl::Span reads = + f->AsProcOrDie()->GetStateReadsByStateElement(next->state_element()); + for (StateRead* read : reads) { + backedges.push_back( + ScheduleBackedge{.source = next, + .destination = read, + .distance = LessThanInitiationInterval()}); + } } } diff --git a/xls/scheduling/schedule_util.cc b/xls/scheduling/schedule_util.cc index 21155d5ef4..6957432ac1 100644 --- a/xls/scheduling/schedule_util.cc +++ b/xls/scheduling/schedule_util.cc @@ -83,10 +83,15 @@ absl::StatusOr> GetDeadAfterSynthesisNodes( Proc* proc = f->AsProcOrDie(); for (StateElement* state_element : proc->StateElements()) { VLOG(2) << "Considering state element: " << state_element->name(); - if (live_after_synthesis.contains( - proc->GetStateReadByStateElement(state_element))) { - for (Next* next : - proc->GetStateReadByStateElement(state_element)->GetNextValues()) { + bool state_is_live = false; + for (StateRead* read : proc->GetStateReadsByStateElement(state_element)) { + if (live_after_synthesis.contains(read)) { + state_is_live = true; + break; + } + } + if (state_is_live) { + for (Next* next : proc->next_values(state_element)) { mark_live(next); } }