Skip to content

[Explicit State Access] Codegen support for unread/multi-read Next nodes without optimization passes #4177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[Explicit State Access] Codegen support for unread/multi-read Next nodes without optimization passes #4177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions xls/codegen_v_1_5/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
55 changes: 55 additions & 0 deletions xls/codegen_v_1_5/convert_to_block_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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<StateRead>()->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
71 changes: 52 additions & 19 deletions xls/codegen_v_1_5/pipeline_register_insertion_pass.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -201,21 +201,54 @@ absl::StatusOr<ConcurrentStageGroups> CalculateConcurrentStages(
continue;
}
Next* next = node->As<Next>();
StateRead* state_read = node->As<Next>()->state_read()->As<StateRead>();
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<StateRead>();
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<StateRead* const> reads =
proc->GetStateReadsByStateElement(state_element);
std::optional<int64_t> 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) {
Expand Down Expand Up @@ -253,8 +286,7 @@ absl::StatusOr<bool> PipelineRegisterInsertionPass::InsertPipelineRegisters(
for (Node* node : block->nodes()) {
if (node->Is<Next>()) {
Next* next = node->As<Next>();
StateElement* state_element =
next->state_read()->As<StateRead>()->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);
Expand Down Expand Up @@ -351,8 +383,9 @@ absl::StatusOr<bool> 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<Next>() && user->As<Next>()->state_read() == node;
bool restore_state_read = user->Is<Next>() &&
user->As<Next>()->has_state_read() &&
user->As<Next>()->state_read() == node;
user->ReplaceOperand(node, live_node);
if (restore_state_read) {
XLS_RETURN_IF_ERROR(user->As<Next>()->ReplaceOperandNumber(
Expand Down
83 changes: 50 additions & 33 deletions xls/codegen_v_1_5/state_to_register_io_lowering_pass.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
#include <string>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/base/casts.h"
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
Expand Down Expand Up @@ -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<Next*> writes) {
absl::Span<StateRead* const> reads,
absl::Span<Next*> writes) {
// A token or empty tuple state element has no real functionality.
Type* type = state_element.type();
const bool has_data = type->GetFlatBitCount() > 0;
Expand All @@ -177,32 +179,40 @@ absl::Status LowerStateElement(ScheduledBlock* block,
state_element.name(), type->ToString()));
}

// Lower the read of the state element.
std::optional<Node*> 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<RegisterRead>(
read_stage_index, read->loc(), state_register,
/*name=*/state_register->name()));
} else {
XLS_ASSIGN_OR_RETURN(reg_read_or_zero, block->MakeNode<xls::Literal>(
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<Node*> reg_reads_or_zeros;
std::vector<int> 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<RegisterRead>(
read_stage_index, read->loc(), state_register,
/*name=*/state_register->name()));
} else {
XLS_ASSIGN_OR_RETURN(
reg_read_or_zero,
block->MakeNode<xls::Literal>(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;
Expand All @@ -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<Node*> all_write_conditions;
Expand All @@ -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;
}

Expand All @@ -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;
}

Expand Down Expand Up @@ -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];
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -382,11 +402,8 @@ absl::StatusOr<bool> LowerStateIoForBlock(ScheduledBlock* block) {

for (StateElement* state_element : source_proc->StateElements()) {
std::vector<StateRead*> 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])));
}

Expand Down
45 changes: 45 additions & 0 deletions xls/codegen_v_1_5/state_to_register_io_lowering_pass_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<a: bits[32] in, b: bits[32] in, result: bits[32] out>(__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
Loading
Loading