Register sharing pass

Add a pass to share registers between stages identified by register-chaining analysis.

PiperOrigin-RevId: 612635721

xls/codegen/BUILD

@@ -628,6 +628,28 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "register_combining_pass",
+    srcs = ["register_combining_pass.cc"],
+    hdrs = ["register_combining_pass.h"],
+    deps = [
+        ":codegen_pass",
+        ":register_chaining_analysis",
+        "//xls/common/logging",
+        "//xls/common/status:ret_check",
+        "//xls/common/status:status_macros",
+        "//xls/ir",
+        "//xls/ir:register",
+        "//xls/passes:pass_base",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
 cc_library(
     name = "register_legalization_pass",
     srcs = ["register_legalization_pass.cc"],
@@ -947,6 +969,35 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "register_combining_pass_test",
+    srcs = ["register_combining_pass_test.cc"],
+    deps = [
+        ":block_conversion",
+        ":codegen_options",
+        ":codegen_pass",
+        ":module_signature_cc_proto",
+        ":register_combining_pass",
+        "//xls/common:xls_gunit",
+        "//xls/common:xls_gunit_main",
+        "//xls/common/status:matchers",
+        "//xls/ir",
+        "//xls/ir:bits",
+        "//xls/ir:function_builder",
+        "//xls/ir:ir_matcher",
+        "//xls/ir:ir_test_base",
+        "//xls/ir:op",
+        "//xls/ir:register",
+        "//xls/ir:source_location",
+        "//xls/passes:pass_base",
+        "//xls/scheduling:pipeline_schedule",
+        "//xls/tools:codegen_flags_cc_proto",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+    ],
+)
+
 cc_test(
     name = "register_legalization_pass_test",
     srcs = ["register_legalization_pass_test.cc"],

xls/codegen/register_combining_pass.cc

@@ -0,0 +1,162 @@
+// Copyright 2024 The XLS Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "xls/codegen/register_combining_pass.h"
+
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/types/span.h"
+#include "xls/codegen/codegen_pass.h"
+#include "xls/codegen/register_chaining_analysis.h"
+#include "xls/common/logging/logging.h"
+#include "xls/common/status/ret_check.h"
+#include "xls/common/status/status_macros.h"
+#include "xls/ir/node.h"
+#include "xls/ir/nodes.h"
+#include "xls/ir/register.h"
+#include "xls/passes/pass_base.h"
+
+namespace xls::verilog {
+
+namespace {
+absl::Status CombineRegisters(absl::Span<const RegisterData> mutex_group,
+                              CodegenPassUnit* unit) {
+  XLS_RET_CHECK_GE(mutex_group.size(), 2)
+      << "Attempting to combine a single register is not meaningful. Single "
+         "element mutex groups should have been filtered out.";
+  // Registers are listed so that 'last' one is at the end.
+  // The register with a loop-back write (write from a later stage) is always
+  // at the front, if one exists.
+  // Merge from the front back.
+  const RegisterData& first = mutex_group.front();
+  std::vector<Node*> cleanup_nodes;
+  absl::flat_hash_set<Register*> cleanup_regs;
+
+  // No need to change load-enable bits, we're merging into the top which has
+  // the right bits already.
+  XLS_VLOG(2) << "Collapsing " << mutex_group.size() << " registers into "
+              << mutex_group.front().reg->ToString();
+  for (const RegisterData& merge : mutex_group.subspan(1)) {
+    XLS_RETURN_IF_ERROR(merge.read->ReplaceUsesWith(first.read));
+    cleanup_regs.insert(merge.reg);
+    cleanup_nodes.push_back(merge.read);
+    cleanup_nodes.push_back(merge.write);
+  }
+
+  // Do cleanup.
+  for (auto& stage : unit->streaming_io_and_pipeline.pipeline_registers) {
+    std::erase_if(stage, [&](const PipelineRegister& pr) {
+      return cleanup_regs.contains(pr.reg);
+    });
+  }
+  for (auto& state_reg : unit->streaming_io_and_pipeline.state_registers) {
+    CHECK(!state_reg || !cleanup_regs.contains(state_reg->reg))
+        << "Removed a state register: " << state_reg->reg->ToString();
+  }
+  for (Node* n : cleanup_nodes) {
+    XLS_RETURN_IF_ERROR(unit->block->RemoveNode(n)) << "can't remove " << n;
+  }
+  for (Register* r : cleanup_regs) {
+    XLS_RETURN_IF_ERROR(unit->block->RemoveRegister(r));
+  }
+  return absl::OkStatus();
+}
+}  // namespace
+
+absl::StatusOr<bool> RegisterCombiningPass::RunInternal(
+    CodegenPassUnit* unit, const CodegenPassOptions& options,
+    PassResults* results) const {
+  if (!unit->concurrent_stages) {
+    return false;
+  }
+  std::vector<RegisterData> candidate_registers;
+  candidate_registers.reserve(unit->block->GetRegisters().size());
+  // State registers (but not their valid/reset regs) are candidates for
+  // merging.
+  XLS_VLOG(2) << unit->block->DumpIr();
+  for (const auto& maybe_reg :
+       unit->streaming_io_and_pipeline.state_registers) {
+    if (maybe_reg) {
+      CHECK(!maybe_reg->next_values.empty());
+      auto write_stage =
+          absl::c_min_element(maybe_reg->next_values, [](const auto& l,
+                                                         const auto& r) {
+            return l.stage < r.stage;
+          })->stage;
+      if (maybe_reg->read_stage == write_stage) {
+        // Immediate back edge.
+        continue;
+      }
+      candidate_registers.push_back({.reg = maybe_reg->reg,
+                                     .read = maybe_reg->reg_read,
+                                     .read_stage = maybe_reg->read_stage,
+                                     .write = maybe_reg->reg_write,
+                                     .write_stage = write_stage});
+    }
+  }
+  // pipeline registers (but not their valid/reset regs) are candidates for
+  // merging.
+  for (const auto& stg_regs :
+       unit->streaming_io_and_pipeline.pipeline_registers) {
+    for (const auto& reg : stg_regs) {
+      CHECK(unit->streaming_io_and_pipeline.node_to_stage_map.contains(
+          reg.reg_read))
+          << reg.reg_read;
+      CHECK(unit->streaming_io_and_pipeline.node_to_stage_map.contains(
+          reg.reg_write))
+          << reg.reg_write;
+      Stage read_stage =
+          unit->streaming_io_and_pipeline.node_to_stage_map.at(reg.reg_read);
+      Stage write_stage =
+          unit->streaming_io_and_pipeline.node_to_stage_map.at(reg.reg_write);
+      CHECK_EQ(write_stage + 1, read_stage)
+          << "pipeline register skipping stage? " << reg.reg->ToString()
+          << "\nread: " << reg.reg_read << "\nwrite: " << reg.reg_write;
+      candidate_registers.push_back({
+          .reg = reg.reg,
+          .read = reg.reg_read,
+          .read_stage = read_stage,
+          .write = reg.reg_write,
+          .write_stage = write_stage,
+      });
+    }
+  }
+  // chains of registers which are possibly combinable.
+  RegisterChains reg_groups;
+
+  for (const RegisterData& rd : candidate_registers) {
+    reg_groups.InsertAndReduce(rd);
+  }
+  XLS_ASSIGN_OR_RETURN(
+      std::vector<std::vector<RegisterData>> mutex_chains,
+      reg_groups.SplitBetweenMutexRegions(*unit->concurrent_stages, options));
+  bool changed = !mutex_chains.empty();
+
+  for (const std::vector<RegisterData>& group : mutex_chains) {
+    XLS_RETURN_IF_ERROR(CombineRegisters(group, unit));
+  }
+
+  if (changed) {
+    unit->GcMetadata();
+  }
+
+  return changed;
+}
+
+}  // namespace xls::verilog

xls/codegen/register_combining_pass.h

@@ -0,0 +1,39 @@
+// Copyright 2024 The XLS Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef XLS_CODEGEN_REGISTER_COMBINING_PASS_H_
+#define XLS_CODEGEN_REGISTER_COMBINING_PASS_H_
+
+#include "absl/status/statusor.h"
+#include "xls/codegen/codegen_pass.h"
+#include "xls/passes/pass_base.h"
+
+namespace xls::verilog {
+// Eliminates (and removes) redundant registers by allowing registers to be
+// shared across many stages.
+class RegisterCombiningPass : public CodegenPass {
+ public:
+  RegisterCombiningPass()
+      : CodegenPass("register_combining",
+                    "Combine mutually exclusive registers") {}
+  ~RegisterCombiningPass() override = default;
+
+  absl::StatusOr<bool> RunInternal(CodegenPassUnit* unit,
+                                   const CodegenPassOptions& options,
+                                   PassResults* results) const override;
+};
+
+}  // namespace xls::verilog
+
+#endif  // XLS_CODEGEN_REGISTER_COMBINING_PASS_H_