Skip to content

Commit

Permalink
[XLS] Remove support for old-style "next state elements"
Browse files Browse the repository at this point in the history 
We've finally consolidated on the new `next_value` nodes.

Fixes #1520

PiperOrigin-RevId: 726671610
  • Loading branch information
@ericastor @copybara-github
ericastor authored and copybara-github committed Feb 14, 2025
1 parent 533c748 commit 4e3a92b
Show file tree
Hide file tree
Showing 33 changed files with 151 additions and 754 deletions.
1 change: 0 additions & 1 deletion xls/codegen/clone_nodes_into_block_handler.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -393,7 +393,6 @@ class CloneNodesIntoBlockHandler {
XLS_ASSIGN_OR_RETURN(int64_t index,
proc->GetStateElementIndex(state_element));

CHECK_EQ(proc->GetNextStateElement(index), next->state_read());
StateRegister& state_register = *result_.state_registers.at(index);
state_register.next_values.push_back(
{.stage = stage,
Expand Down
127 changes: 47 additions & 80 deletions xls/contrib/mlir/tools/xls_translate/xls_translate_to_mlir.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1464,12 +1464,6 @@ absl::StatusOr<Operation*> translateFunction(::xls::Function& xls_func,
absl::StatusOr<Operation*> translateProc(::xls::Proc& xls_proc,
OpBuilder& builder, MLIRContext* ctx,
TranslationState& state) {
// Detect legacy procs with marked `next` node instead of `next_value` nodes.
// TODO: google/xls#1520 - remove this once fully transitioned over to
// `next_value` nodes.
bool is_legacy_proc =
(xls_proc.next_values().empty() && xls_proc.GetStateElementCount());

// New function base scope
state.newFunctionBaseContext(/*is_proc=*/true);

Expand Down Expand Up @@ -1513,24 +1507,11 @@ absl::StatusOr<Operation*> translateProc(::xls::Proc& xls_proc,
if (!err.ok()) {
return err;
}

if (is_legacy_proc) {
// Track all state elements for which this node defines their next
// value.
// TODO: google/xls#1520 - remove this once fully transitioned over to
// `next_value` nodes.
for (auto state_elem_idx : xls_proc.GetNextStateIndices(n)) {
state_elem_next_value[xls_proc.GetStateElement(state_elem_idx)
->name()] = *state_value;
}
}

continue; // StateRead get directly added to the ssa context and don't
// otherwise map to a MLIR operation.
}

if (n->Is<::xls::Next>()) {
assert(!is_legacy_proc && "Legacy procs don't use next_value nodes");
auto xls_next = n->As<::xls::Next>();
auto state_elem_name = xls_next->state_read()
->As<::xls::StateRead>()
Expand All @@ -1543,77 +1524,63 @@ absl::StatusOr<Operation*> translateProc(::xls::Proc& xls_proc,
if (!op.ok()) {
return op;
}

if (is_legacy_proc) {
// Track all state elements for which this node defines their next
// value:
// TODO: google/xls#1520 - remove this once fully transitioned over to
// `next_value` nodes.
for (auto state_elem_idx : xls_proc.GetNextStateIndices(n)) {
state_elem_next_value[xls_proc.GetStateElement(state_elem_idx)
->name()] = (*op)->getResult(0);
}
}
}

// Generate NextValueOps:
if (!is_legacy_proc) {
for (auto* state_elem : xls_proc.StateElements()) {
// All Next nodes that contribue to this state element:
std::vector<::xls::Next*> elem_next_value_ops =
next_value_ops[state_elem->name()];

if (elem_next_value_ops.empty()) {
return absl::InternalError(absl::StrCat(
"No `next_value` nodes for state element ", state_elem->name()));
for (auto* state_elem : xls_proc.StateElements()) {
// All Next nodes that contribute to this state element:
std::vector<::xls::Next*> elem_next_value_ops =
next_value_ops[state_elem->name()];

if (elem_next_value_ops.empty()) {
return absl::InternalError(absl::StrCat(
"No `next_value` nodes for state element ", state_elem->name()));
}

if (elem_next_value_ops.size() == 1 &&
!elem_next_value_ops[0]->predicate().has_value()) {
// State element defined by single Next node without predicate. No MLIR
// NextValueOp needed - directly feed YieldOp from value:
auto redundant_next_node = elem_next_value_ops[0];
auto next_value = state.getMlirValue(redundant_next_node->value()->id());
if (!next_value.ok()) {
return next_value.status();
}
state_elem_next_value[state_elem->name()] = *next_value;
} else {
// Generate new NextValueOp:
SmallVector<Value> values_vec{};
SmallVector<Value> predicates_vec{};

for (auto* next_value_op : elem_next_value_ops) {
if (!next_value_op->predicate().has_value()) {
return absl::InternalError(absl::StrCat(
"`next_value` nodes for state element ", state_elem->name(),
" lacks predicate but there are multiple `next_value` nodes."));
}

if (elem_next_value_ops.size() == 1 &&
!elem_next_value_ops[0]->predicate().has_value()) {
// State element defined by single Next node without predicate. No MLIR
// NextValueOp needed - directly feed YieldOp from value:
auto redundant_next_node = elem_next_value_ops[0];
auto next_value =
state.getMlirValue(redundant_next_node->value()->id());
if (!next_value.ok()) {
return next_value.status();
auto val = state.getMlirValue(next_value_op->value()->id());
if (!val.ok()) {
return val.status();
}
state_elem_next_value[state_elem->name()] = *next_value;
} else {
// Generate new NextValueOp:
SmallVector<Value> values_vec{};
SmallVector<Value> predicates_vec{};

for (auto* next_value_op : elem_next_value_ops) {
if (!next_value_op->predicate().has_value()) {
return absl::InternalError(absl::StrCat(
"`next_value` nodes for state element ", state_elem->name(),
" lacks predicate but there are multiple `next_value` nodes."));
}

auto val = state.getMlirValue(next_value_op->value()->id());
if (!val.ok()) {
return val.status();
}
values_vec.push_back(*val);

auto predicate =
state.getMlirValue(next_value_op->predicate().value()->id());
if (!predicate.ok()) {
return val.status();
}
predicates_vec.push_back(*predicate);
values_vec.push_back(*val);

auto predicate =
state.getMlirValue(next_value_op->predicate().value()->id());
if (!predicate.ok()) {
return val.status();
}
predicates_vec.push_back(*predicate);
}

auto elem_type = translateType(state_elem->type(), builder, ctx);
auto elem_type = translateType(state_elem->type(), builder, ctx);

builder.setInsertionPointToEnd(body);
auto next = builder.create<NextValueOp>(
builder.getUnknownLoc(), elem_type, ValueRange(predicates_vec),
ValueRange(values_vec));
builder.setInsertionPointToEnd(body);
auto next = builder.create<NextValueOp>(
builder.getUnknownLoc(), elem_type, ValueRange(predicates_vec),
ValueRange(values_vec));

state_elem_next_value[state_elem->name()] = next;
}
state_elem_next_value[state_elem->name()] = next;
}
}

Expand Down
34 changes: 4 additions & 30 deletions xls/dev_tools/ir_minimizer_main.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -425,19 +425,8 @@ absl::StatusOr<bool> RemoveDeadParameters(FunctionBase* f) {
if (f->IsProc()) {
Proc* p = f->AsProcOrDie();
absl::flat_hash_set<StateElement*> dead_state_elements;
absl::flat_hash_set<StateElement*> invariant_state_elements;
for (StateElement* state_element : p->StateElements()) {
StateRead* state_read = p->GetStateRead(state_element);
if (state_read->IsDead()) {
dead_state_elements.insert(state_element);
}
XLS_ASSIGN_OR_RETURN(int64_t index,
p->GetStateElementIndex(state_element));
if (state_read == p->GetNextStateElement(index)) {
invariant_state_elements.insert(state_element);
}
}

absl::flat_hash_set<StateElement*> invariant_state_elements(
p->StateElements().begin(), p->StateElements().end());
for (Next* next : p->next_values()) {
if (next->value() != next->state_read()) {
// This state param is not actually invariant.
Expand Down Expand Up @@ -581,20 +570,6 @@ absl::StatusOr<SimplificationResult> ReplaceImplicitUse(Node* node,
XLS_RETURN_IF_ERROR(f->set_return_value(replacement));
return SimplificationResult::kDidChange;
}
if (fb->IsProc()) {
Proc* p = fb->AsProcOrDie();
SimplificationResult changed = SimplificationResult::kDidNotChange;
if (!node->GetType()->IsEqualTo(replacement->GetType())) {
return SimplificationResult::kDidNotChange;
}
for (int64_t i = 0; i < p->GetStateElementCount(); ++i) {
if (p->GetNextStateElement(i) == node) {
XLS_RETURN_IF_ERROR(p->SetNextStateElement(i, replacement));
changed = SimplificationResult::kDidChange;
}
}
return changed;
}
return SimplificationResult::kDidNotChange;
}

Expand All @@ -604,9 +579,8 @@ std::vector<Node*> ImplicitlyUsed(FunctionBase* fb) {
return {f->return_value()};
}
if (fb->IsProc()) {
Proc* p = fb->AsProcOrDie();
absl::Span<Node* const> next_state = p->NextState();
return std::vector<Node*>(next_state.begin(), next_state.end());
// Procs have no implicit uses.
return {};
}
LOG(FATAL) << "ImplicitlyUsed only supports functions and procs";
}
Expand Down
14 changes: 0 additions & 14 deletions xls/fdo/iterative_sdc_scheduler.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -419,20 +419,6 @@ absl::StatusOr<ScheduleCycleMap> ScheduleByIterativeSDC(
}
}

if (f->IsProc()) {
Proc *proc = f->AsProcOrDie();
for (int64_t index = 0; index < proc->GetStateElementCount(); ++index) {
StateRead *const state_read = proc->GetStateRead(index);
Node *const next_state_element = proc->GetNextStateElement(index);

// The next-state element always has lifetime extended to the state-read
// node, since we can't store the new value in the state register until
// the old value's been used.
XLS_RETURN_IF_ERROR(
model.AddLifetimeConstraint(next_state_element, state_read));
}
}

XLS_RETURN_IF_ERROR(model.AddTimingConstraints(clock_period_ps));

int64_t min_pipeline_length = 1;
Expand Down
13 changes: 3 additions & 10 deletions xls/interpreter/proc_interpreter.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -314,16 +314,9 @@ absl::StatusOr<TickResult> ProcInterpreter::Tick(
}
}

// Proc completed execution of the Tick. Pass the next proc state to the
// continuation.
//
// TODO: Simplify this once fully transitioned over to `next_value` nodes.
std::vector<Value> next_state;
next_state.resize(proc()->GetStateElementCount());
for (int64_t index = 0; index < proc()->NextState().size(); ++index) {
next_state[index] =
ir_interpreter.ResolveAsValue(proc()->GetNextStateElement(index));
}
// Proc completed execution of the Tick. Let the active next_values update the
// proc state, then pass it to the continuation.
std::vector<Value> next_state = cont->GetState();
for (const auto& [state_element, next_values] : cont->GetActiveNextValues()) {
if (next_values.size() > 1) {
return absl::AlreadyExistsError(absl::StrFormat(
Expand Down
4 changes: 2 additions & 2 deletions xls/ir/function_builder_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -914,7 +914,7 @@ TEST(FunctionBuilderTest, ProcWithMultipleStateElements) {
ElementsAre(m::Next(m::StateRead("z"), m::StateRead("z"))));
}

TEST(FunctionBuilderTest, ProcWithNextStateElement) {
TEST(FunctionBuilderTest, ProcWithNextValue) {
Package p("p");
ProcBuilder pb("the_proc", &p);
BValue x = pb.StateElement("x", Value(UBits(1, 1)));
Expand All @@ -928,7 +928,7 @@ TEST(FunctionBuilderTest, ProcWithNextStateElement) {
/*predicate=*/m::StateRead("x")));
}

TEST(FunctionBuilderTest, ProcWithNextStateElementBadPredicate) {
TEST(FunctionBuilderTest, ProcWithNextValueBadPredicate) {
Package p("p");
ProcBuilder pb("the_proc", &p);
BValue x = pb.StateElement("x", Value(UBits(1, 32)));
Expand Down
10 changes: 1 addition & 9 deletions xls/ir/node.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -930,21 +930,13 @@ absl::Status Node::ReplaceUsesWith(Node* replacement,

absl::StatusOr<bool> Node::ReplaceImplicitUsesWith(Node* replacement) {
bool changed = false;
// Only functions have implicitly-used nodes, for their return value.
if (function_base()->IsFunction()) {
Function* function = function_base()->AsFunctionOrDie();
if (this == function->return_value()) {
XLS_RETURN_IF_ERROR(function->set_return_value(replacement));
changed = true;
}
} else if (function_base()->IsProc()) {
Proc* proc = function_base()->AsProcOrDie();
for (int64_t index : proc->GetNextStateIndices(this)) {
XLS_RETURN_IF_ERROR(proc->SetNextStateElement(index, replacement));
changed = true;
}
} else {
XLS_RET_CHECK(function_base()->IsBlock());
// Blocks have no implicit uses.
}
return changed;
}
Expand Down
Loading

0 comments on commit 4e3a92b

Please sign in to comment.