From cf8373644768d0addb959400585efedfb2460268 Mon Sep 17 00:00:00 2001
From: Eric Astor <epastor@google.com>
Date: Fri, 26 Dec 2025 06:42:03 -0800
Subject: [PATCH] [codegen 1.5] Enable the simplest proc test for new block
 conversion

This required small patches to channel & state lowering, as well as to how we transfer Nodes into ScheduledBlocks:

1. Channel lowering now handles single-value channnels without crashing in some circumstances. (Whoops.)
2. State lowering now cleans up after itself.
3. ScheduledBlock conversion now takes ownership of the node name as well as the node.

We also add dataflow simplification to the pipeline (matching codegen 1.0).

PiperOrigin-RevId: 849163225
---
 xls/codegen_v_1_5/BUILD                       |  65 ++++++
 .../block_conversion_pass_pipeline.cc         |  21 +-
 .../block_conversion_pass_pipeline.h          |   5 +-
 .../block_conversion_pass_pipeline_test.cc    |  62 ++++--
 .../block_conversion_wrapper_pass.cc          |  32 +++
 .../block_conversion_wrapper_pass.h           |  74 +++++++
 .../block_conversion_wrapper_pass_test.cc     |  76 +++++++
 .../channel_to_port_io_lowering_pass.cc       |  61 +++++-
 .../channel_to_port_io_lowering_pass_test.cc  | 145 ++++++++++++++
 xls/codegen_v_1_5/convert_to_block.cc         |   4 +-
 .../flow_control_insertion_pass.cc            |  28 +++
 .../function_io_lowering_pass.cc              |  11 +-
 .../function_io_lowering_pass_test.cc         |  31 ++-
 xls/codegen_v_1_5/register_cleanup_pass.cc    | 188 ++++++++++++++++++
 xls/codegen_v_1_5/register_cleanup_pass.h     |  61 ++++++
 .../state_to_register_io_lowering_pass.cc     | 154 +++++++-------
 xls/ir/BUILD                                  |   1 +
 xls/ir/function_base.cc                       |   6 +
 xls/ir/function_base.h                        |   4 +-
 xls/ir/name_uniquer.cc                        |  10 +
 xls/ir/name_uniquer.h                         |  12 ++
 xls/ir/proc.cc                                |   4 +
 22 files changed, 936 insertions(+), 119 deletions(-)
 create mode 100644 xls/codegen_v_1_5/block_conversion_wrapper_pass.cc
 create mode 100644 xls/codegen_v_1_5/block_conversion_wrapper_pass.h
 create mode 100644 xls/codegen_v_1_5/block_conversion_wrapper_pass_test.cc
 create mode 100644 xls/codegen_v_1_5/register_cleanup_pass.cc
 create mode 100644 xls/codegen_v_1_5/register_cleanup_pass.h

diff --git a/xls/codegen_v_1_5/BUILD b/xls/codegen_v_1_5/BUILD
index 2cceef3c91..5eb74e6a0f 100644
--- a/xls/codegen_v_1_5/BUILD
+++ b/xls/codegen_v_1_5/BUILD
@@ -41,15 +41,20 @@ cc_library(
     hdrs = ["block_conversion_pass_pipeline.h"],
     deps = [
         ":block_conversion_pass",
+        ":block_conversion_wrapper_pass",
         ":block_finalization_pass",
         ":channel_to_port_io_lowering_pass",
         ":flow_control_insertion_pass",
         ":function_io_lowering_pass",
         ":idle_insertion_pass",
         ":pipeline_register_insertion_pass",
+        ":register_cleanup_pass",
         ":scheduled_block_conversion_pass",
         ":scheduling_pass",
         ":state_to_register_io_lowering_pass",
+        "//xls/passes:dataflow_simplification_pass",
+        "//xls/passes:dce_pass",
+        "//xls/passes:optimization_pass",
     ],
 )
 
@@ -93,6 +98,65 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "register_cleanup_pass",
+    srcs = ["register_cleanup_pass.cc"],
+    hdrs = ["register_cleanup_pass.h"],
+    deps = [
+        ":block_conversion_pass",
+        "//xls/common/status:status_macros",
+        "//xls/data_structures:transitive_closure",
+        "//xls/ir",
+        "//xls/ir:register",
+        "//xls/ir:value",
+        "//xls/ir:value_utils",
+        "//xls/passes:node_dependency_analysis",
+        "//xls/passes:partial_info_query_engine",
+        "//xls/passes:pass_base",
+        "//xls/passes:query_engine",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_library(
+    name = "block_conversion_wrapper_pass",
+    srcs = ["block_conversion_wrapper_pass.cc"],
+    hdrs = ["block_conversion_wrapper_pass.h"],
+    deps = [
+        ":block_conversion_pass",
+        "//xls/ir",
+        "//xls/passes:optimization_pass",
+        "//xls/passes:pass_base",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_test(
+    name = "block_conversion_wrapper_pass_test",
+    srcs = ["block_conversion_wrapper_pass_test.cc"],
+    deps = [
+        ":block_conversion_pass",
+        ":block_conversion_wrapper_pass",
+        "//xls/common:xls_gunit_main",
+        "//xls/common/status:matchers",
+        "//xls/ir",
+        "//xls/ir:function_builder",
+        "//xls/ir:ir_matcher",
+        "//xls/ir:ir_test_base",
+        "//xls/passes:dce_pass",
+        "//xls/passes:optimization_pass",
+        "//xls/passes:pass_base",
+        "@com_google_absl//absl/status:statusor",
+        "@googletest//:gtest",
+    ],
+)
+
 cc_library(
     name = "channel_to_port_io_lowering_pass",
     srcs = ["channel_to_port_io_lowering_pass.cc"],
@@ -170,6 +234,7 @@ cc_library(
         "//xls/common/status:status_macros",
         "//xls/estimators/delay_model:delay_estimator",
         "//xls/ir",
+        "//xls/passes:optimization_pass",
         "//xls/passes:pass_base",
         "//xls/scheduling:pipeline_schedule_cc_proto",
         "//xls/scheduling:scheduling_options",
diff --git a/xls/codegen_v_1_5/block_conversion_pass_pipeline.cc b/xls/codegen_v_1_5/block_conversion_pass_pipeline.cc
index ea3f42c842..52d6eaa968 100644
--- a/xls/codegen_v_1_5/block_conversion_pass_pipeline.cc
+++ b/xls/codegen_v_1_5/block_conversion_pass_pipeline.cc
@@ -20,20 +20,25 @@
 #include <memory>
 
 #include "xls/codegen_v_1_5/block_conversion_pass.h"
+#include "xls/codegen_v_1_5/block_conversion_wrapper_pass.h"
 #include "xls/codegen_v_1_5/block_finalization_pass.h"
 #include "xls/codegen_v_1_5/channel_to_port_io_lowering_pass.h"
 #include "xls/codegen_v_1_5/flow_control_insertion_pass.h"
 #include "xls/codegen_v_1_5/function_io_lowering_pass.h"
 #include "xls/codegen_v_1_5/idle_insertion_pass.h"
 #include "xls/codegen_v_1_5/pipeline_register_insertion_pass.h"
+#include "xls/codegen_v_1_5/register_cleanup_pass.h"
 #include "xls/codegen_v_1_5/scheduled_block_conversion_pass.h"
 #include "xls/codegen_v_1_5/scheduling_pass.h"
 #include "xls/codegen_v_1_5/state_to_register_io_lowering_pass.h"
+#include "xls/passes/dataflow_simplification_pass.h"
+#include "xls/passes/dce_pass.h"
+#include "xls/passes/optimization_pass.h"
 
 namespace xls::codegen {
 
-std::unique_ptr<BlockConversionCompoundPass>
-CreateBlockConversionPassPipeline() {
+std::unique_ptr<BlockConversionCompoundPass> CreateBlockConversionPassPipeline(
+    OptimizationContext& opt_context) {
   auto top = std::make_unique<BlockConversionCompoundPass>(
       "block_conversion", "Top level codegen v1.5 block conversion pipeline");
 
@@ -64,6 +69,18 @@ CreateBlockConversionPassPipeline() {
   // Lower scheduled block to standard block, inlining each stage.
   top->Add<BlockFinalizationPass>();
 
+  // Clean up unused registers & load-enable bits (including flow-control
+  // registers).
+  top->Add<RegisterCleanupPass>();
+
+  // Clean up unnecessary array/tuple manipulation.
+  top->Add<BlockConversionWrapperPass>(
+      std::make_unique<DataflowSimplificationPass>(), opt_context);
+
+  // Remove anything we created & then left dead.
+  top->Add<BlockConversionWrapperPass>(
+      std::make_unique<DeadCodeEliminationPass>(), opt_context);
+
   return top;
 }
 
diff --git a/xls/codegen_v_1_5/block_conversion_pass_pipeline.h b/xls/codegen_v_1_5/block_conversion_pass_pipeline.h
index 7541593a34..37d20dcbeb 100644
--- a/xls/codegen_v_1_5/block_conversion_pass_pipeline.h
+++ b/xls/codegen_v_1_5/block_conversion_pass_pipeline.h
@@ -18,13 +18,14 @@
 #include <memory>
 
 #include "xls/codegen_v_1_5/block_conversion_pass.h"
+#include "xls/passes/optimization_pass.h"
 
 namespace xls::codegen {
 
 // Returns a pipeline which converts an unscheduled IR package into a standard
 // block.
-std::unique_ptr<BlockConversionCompoundPass>
-CreateBlockConversionPassPipeline();
+std::unique_ptr<BlockConversionCompoundPass> CreateBlockConversionPassPipeline(
+    OptimizationContext& opt_context);
 
 }  // namespace xls::codegen
 
diff --git a/xls/codegen_v_1_5/block_conversion_pass_pipeline_test.cc b/xls/codegen_v_1_5/block_conversion_pass_pipeline_test.cc
index 5bff3c171c..12f7b3e567 100644
--- a/xls/codegen_v_1_5/block_conversion_pass_pipeline_test.cc
+++ b/xls/codegen_v_1_5/block_conversion_pass_pipeline_test.cc
@@ -54,7 +54,6 @@
 #include "xls/common/status/matchers.h"
 #include "xls/common/status/ret_check.h"
 #include "xls/common/status/status_macros.h"
-#include "xls/estimators/delay_model/delay_estimator.h"
 #include "xls/interpreter/block_evaluator.h"
 #include "xls/interpreter/block_interpreter.h"
 #include "xls/ir/bits.h"
@@ -97,6 +96,7 @@ using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::Ge;
 using ::testing::HasSubstr;
+using ::testing::IsEmpty;
 using ::testing::Optional;
 using ::testing::Pair;
 using ::testing::SizeIs;
@@ -282,7 +282,7 @@ TEST_F(BlockConversionTest, SimpleFunction) {
       m::OutputPort("out", m::Add(m::InputPort("x"), m::InputPort("y"))));
 }
 
-TEST_F(BlockConversionTest, DISABLED_SimpleFunctionWithNamedOutput) {
+TEST_F(BlockConversionTest, SimpleFunctionWithNamedOutput) {
   auto p = CreatePackage();
   FunctionBuilder fb(TestName(), p.get());
   BValue x = fb.Param("x", p->GetBitsType(32));
@@ -331,7 +331,7 @@ TEST_F(BlockConversionTest, ZeroWidthInputsAndOutput) {
   EXPECT_EQ(block->GetPorts().size(), 4);
 }
 
-TEST_F(BlockConversionTest, DISABLED_SimplePipelinedFunction) {
+TEST_F(BlockConversionTest, SimplePipelinedFunction) {
   auto p = CreatePackage();
   FunctionBuilder fb(TestName(), p.get());
   BValue x = fb.Param("x", p->GetBitsType(32));
@@ -348,11 +348,13 @@ TEST_F(BlockConversionTest, DISABLED_SimplePipelinedFunction) {
           SchedulingOptions().pipeline_stages(3)));
 
   EXPECT_THAT(GetOutputPort(block),
-              m::OutputPort(m::Neg(m::Register(m::Not(m::Register(
-                  m::Add(m::InputPort("x"), m::InputPort("y"))))))));
+              m::OutputPort(m::Neg(m::Register(m::Not(
+                  m::Register(m::Add(m::InputPort("x"), m::InputPort("y"))))))))
+      << "\n\nIR:\n"
+      << block->DumpIr();
 }
 
-TEST_F(BlockConversionTest, DISABLED_TrivialPipelinedFunctionNoFlopping) {
+TEST_F(BlockConversionTest, TrivialPipelinedFunctionNoFlopping) {
   auto p = CreatePackage();
   FunctionBuilder fb(TestName(), p.get());
   BValue x = fb.Param("x", p->GetBitsType(32));
@@ -373,7 +375,7 @@ TEST_F(BlockConversionTest, DISABLED_TrivialPipelinedFunctionNoFlopping) {
                   m::Add(m::InputPort("x"), m::InputPort("y"))))))));
 }
 
-TEST_F(BlockConversionTest, DISABLED_TrivialPipelinedFunctionFloppedInputs) {
+TEST_F(BlockConversionTest, TrivialPipelinedFunctionFloppedInputs) {
   auto p = CreatePackage();
   FunctionBuilder fb(TestName(), p.get());
   BValue x = fb.Param("x", p->GetBitsType(32));
@@ -395,7 +397,7 @@ TEST_F(BlockConversionTest, DISABLED_TrivialPipelinedFunctionFloppedInputs) {
           m::Register(m::InputPort("x")), m::Register(m::InputPort("y")))))))));
 }
 
-TEST_F(BlockConversionTest, DISABLED_TrivialPipelinedFunctionFloppedOutputs) {
+TEST_F(BlockConversionTest, TrivialPipelinedFunctionFloppedOutputs) {
   auto p = CreatePackage();
   FunctionBuilder fb(TestName(), p.get());
   BValue x = fb.Param("x", p->GetBitsType(32));
@@ -416,7 +418,7 @@ TEST_F(BlockConversionTest, DISABLED_TrivialPipelinedFunctionFloppedOutputs) {
                   m::Add(m::InputPort("x"), m::InputPort("y")))))))));
 }
 
-TEST_F(BlockConversionTest, DISABLED_TrivialPipelinedFunctionFloppedIo) {
+TEST_F(BlockConversionTest, TrivialPipelinedFunctionFloppedIo) {
   auto p = CreatePackage();
   FunctionBuilder fb(TestName(), p.get());
   BValue x = fb.Param("x", p->GetBitsType(32));
@@ -438,7 +440,7 @@ TEST_F(BlockConversionTest, DISABLED_TrivialPipelinedFunctionFloppedIo) {
                                      m::Register(m::InputPort("y"))))))))));
 }
 
-TEST_F(BlockConversionTest, DISABLED_ZeroWidthPipeline) {
+TEST_F(BlockConversionTest, ZeroWidthPipeline) {
   auto p = CreatePackage();
   FunctionBuilder fb(TestName(), p.get());
   BValue x = fb.Param("x", p->GetTupleType({}));
@@ -454,7 +456,28 @@ TEST_F(BlockConversionTest, DISABLED_ZeroWidthPipeline) {
               "clk"),
           SchedulingOptions().pipeline_stages(3)));
 
-  EXPECT_EQ(block->GetRegisters().size(), 4);
+  EXPECT_THAT(block->GetRegisters(), IsEmpty());
+}
+
+TEST_F(BlockConversionTest, ZeroWidthPipelineWithValidControl) {
+  auto p = CreatePackage();
+  FunctionBuilder fb(TestName(), p.get());
+  BValue x = fb.Param("x", p->GetTupleType({}));
+  BValue y = fb.Param("y", p->GetBitsType(0));
+  XLS_ASSERT_OK_AND_ASSIGN(Function * f,
+                           fb.BuildWithReturnValue(fb.Tuple({x, y})));
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Block * block,
+      ConvertToBlock(f,
+                     CodegenOptions()
+                         .flop_inputs(false)
+                         .flop_outputs(false)
+                         .clock_name("clk")
+                         .valid_control("inputs_valid", "output_valid"),
+                     SchedulingOptions().pipeline_stages(3)));
+
+  EXPECT_THAT(block->GetRegisters(), SizeIs(2));
 }
 
 // Verifies that an implicit token, as generated by the DSLX IR converter, is
@@ -473,11 +496,10 @@ fn __itok__implicit_token__main(__token: token, __activated: bits[1]) ->
 fn __implicit_token__main() -> () {
   after_all.9: token = after_all(id=9)
   literal.10: bits[1] = literal(value=1, id=10)
-  invoke.11: (token, ()) = invoke(after_all.9, literal.10,
-  to_apply=__itok__implicit_token__main, id=11) tuple_index.12: token =
-  tuple_index(invoke.11, index=0, id=12) invoke.13: (token, ()) =
-  invoke(tuple_index.12, literal.10, to_apply=__itok__implicit_token__main,
-  id=13) ret tuple_index.14: () = tuple_index(invoke.13, index=1, id=14)
+  invoke.11: (token, ()) = invoke(after_all.9, literal.10, to_apply=__itok__implicit_token__main, id=11)
+  tuple_index.12: token = tuple_index(invoke.11, index=0, id=12)
+  invoke.13: (token, ()) = invoke(tuple_index.12, literal.10, to_apply=__itok__implicit_token__main, id=13)
+  ret tuple_index.14: () = tuple_index(invoke.13, index=1, id=14)
 }
   )";
   XLS_ASSERT_OK_AND_ASSIGN(auto p, Parser::ParsePackage(kIrText));
@@ -489,7 +511,7 @@ fn __implicit_token__main() -> () {
   XLS_ASSERT_OK(VerifyBlock(block));
 }
 
-TEST_F(BlockConversionTest, DISABLED_SimpleProc) {
+TEST_F(BlockConversionTest, SimpleProc) {
   const std::string ir_text = R"(package test
 
 chan in(bits[32], id=0, kind=single_value, ops=receive_only)
@@ -513,7 +535,9 @@ proc my_proc(my_state: (), init={()}) {
       Block * block,
       ConvertToBlock(proc, codegen_options().generate_combinational(true)));
   EXPECT_THAT(FindNode("out", block),
-              m::OutputPort("out", m::Neg(m::InputPort("in"))));
+              m::OutputPort("out", m::Neg(m::InputPort("in"))))
+      << "\n\nIR:\n"
+      << block->DumpIr();
 }
 
 TEST_F(BlockConversionTest, DISABLED_StreamingChannelMetadataForSimpleProc) {
@@ -6394,7 +6418,7 @@ TEST_F(ProcConversionTestFixture,
   }
 }
 
-TEST_F(ProcConversionTestFixture, DISABLED_SimpleFunctionWithProcsPresent) {
+TEST_F(ProcConversionTestFixture, SimpleFunctionWithProcsPresent) {
   XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr<Package> p,
                            CreateMultiProcPackage(/*with_functions=*/true));
   XLS_ASSERT_OK_AND_ASSIGN(Function * f0, p->GetFunction("f0"));
diff --git a/xls/codegen_v_1_5/block_conversion_wrapper_pass.cc b/xls/codegen_v_1_5/block_conversion_wrapper_pass.cc
new file mode 100644
index 0000000000..a3ce2b6054
--- /dev/null
+++ b/xls/codegen_v_1_5/block_conversion_wrapper_pass.cc
@@ -0,0 +1,32 @@
+// Copyright 2021 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_v_1_5/block_conversion_wrapper_pass.h"
+
+#include "absl/status/statusor.h"
+#include "xls/codegen_v_1_5/block_conversion_pass.h"
+#include "xls/ir/package.h"
+#include "xls/passes/optimization_pass.h"
+#include "xls/passes/pass_base.h"
+
+namespace xls::codegen {
+
+absl::StatusOr<bool> BlockConversionWrapperPass::RunInternal(
+    Package* package, const BlockConversionPassOptions& options,
+    PassResults* results) const {
+  return wrapped_pass_->Run(package, OptimizationPassOptions(options), results,
+                            opt_context_);
+}
+
+}  // namespace xls::codegen
diff --git a/xls/codegen_v_1_5/block_conversion_wrapper_pass.h b/xls/codegen_v_1_5/block_conversion_wrapper_pass.h
new file mode 100644
index 0000000000..caa9025638
--- /dev/null
+++ b/xls/codegen_v_1_5/block_conversion_wrapper_pass.h
@@ -0,0 +1,74 @@
+// Copyright 2021 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_V_1_5_BLOCK_CONVERSION_WRAPPER_PASS_H_
+#define XLS_CODEGEN_V_1_5_BLOCK_CONVERSION_WRAPPER_PASS_H_
+
+#include <memory>
+#include <utility>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/str_format.h"
+#include "absl/types/span.h"
+#include "xls/codegen_v_1_5/block_conversion_pass.h"
+#include "xls/ir/package.h"
+#include "xls/passes/optimization_pass.h"
+#include "xls/passes/pass_base.h"
+
+namespace xls::codegen {
+
+// A codegen pass wrapper which wraps a OptimizationFunctionBasePass. This is
+// useful for adding an optimization or transformation pass (most passes in
+// xls/passes are OptimizationFunctionBasePasses) to the codegen pipeline. The
+// wrapped pass is run on the block being lowered to Verilog.
+//
+// Takes the OptimizationContext object at construction, since it's specific to
+// optimization passes & cannot be passed via a codegen pass's Run function.
+class BlockConversionWrapperPass : public BlockConversionPass {
+ public:
+  explicit BlockConversionWrapperPass(
+      std::unique_ptr<OptimizationFunctionBasePass> wrapped_pass,
+      OptimizationContext& opt_context)
+      : BlockConversionPass(
+            absl::StrFormat("codegen_%s", wrapped_pass->short_name()),
+            absl::StrFormat("%s (codegen)", wrapped_pass->long_name())),
+        wrapped_pass_(std::move(wrapped_pass)),
+        opt_context_(opt_context) {}
+  ~BlockConversionWrapperPass() override = default;
+
+  bool IsCompound() const override { return wrapped_pass_->IsCompound(); }
+
+  absl::StatusOr<bool> RunInternal(Package* package,
+                                   const BlockConversionPassOptions& options,
+                                   PassResults* results) const final;
+
+  absl::StatusOr<bool> RunNested(Package* package,
+                                 const BlockConversionPassOptions& options,
+                                 PassResults* results,
+                                 PassInvocation& invocation,
+                                 absl::Span<const InvariantChecker* const>
+                                     invariant_checkers) const override {
+    return wrapped_pass_->RunNested(package, OptimizationPassOptions(options),
+                                    results, opt_context_, invocation,
+                                    /*invariant_checkers=*/{});
+  }
+
+ private:
+  std::unique_ptr<OptimizationFunctionBasePass> wrapped_pass_;
+  OptimizationContext& opt_context_;
+};
+
+}  // namespace xls::codegen
+
+#endif  // XLS_CODEGEN_V_1_5_BLOCK_CONVERSION_WRAPPER_PASS_H_
diff --git a/xls/codegen_v_1_5/block_conversion_wrapper_pass_test.cc b/xls/codegen_v_1_5/block_conversion_wrapper_pass_test.cc
new file mode 100644
index 0000000000..45600011a9
--- /dev/null
+++ b/xls/codegen_v_1_5/block_conversion_wrapper_pass_test.cc
@@ -0,0 +1,76 @@
+// Copyright 2021 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_v_1_5/block_conversion_wrapper_pass.h"
+
+#include <memory>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "absl/status/statusor.h"
+#include "xls/codegen_v_1_5/block_conversion_pass.h"
+#include "xls/common/status/matchers.h"
+#include "xls/ir/block.h"
+#include "xls/ir/function_builder.h"
+#include "xls/ir/ir_matcher.h"
+#include "xls/ir/ir_test_base.h"
+#include "xls/passes/dce_pass.h"
+#include "xls/passes/optimization_pass.h"
+#include "xls/passes/pass_base.h"
+
+namespace m = ::xls::op_matchers;
+
+namespace xls::codegen {
+namespace {
+
+using ::absl_testing::IsOkAndHolds;
+
+class BlockConversionWrapperPassTest : public IrTestBase {
+ protected:
+  absl::StatusOr<bool> Run(Block* block) {
+    PassResults results;
+    OptimizationContext optimization_context;
+    return BlockConversionWrapperPass(
+               std::make_unique<DeadCodeEliminationPass>(),
+               optimization_context)
+        .Run(block->package(), BlockConversionPassOptions(), &results);
+  }
+};
+
+TEST_F(BlockConversionWrapperPassTest, WrappedDce) {
+  auto p = CreatePackage();
+
+  BlockBuilder bb(TestName(), p.get());
+  XLS_ASSERT_OK(bb.block()->AddClockPort("clk"));
+  BValue a = bb.InputPort("a", p->GetBitsType(32));
+  BValue b = bb.InputPort("b", p->GetBitsType(32));
+  bb.Add(a, b);
+  bb.OutputPort("c", bb.Subtract(a, b));
+  XLS_ASSERT_OK_AND_ASSIGN(Block * block, bb.Build());
+
+  // The add should be eliminated.
+  EXPECT_EQ(block->node_count(), 5);
+  EXPECT_THAT(Run(block), IsOkAndHolds(true));
+  EXPECT_EQ(block->node_count(), 4);
+
+  EXPECT_THAT(FindNode("c", block),
+              m::OutputPort(m::Sub(m::InputPort("a"), m::InputPort("b"))));
+
+  // Pass should be idempotent.
+  EXPECT_THAT(Run(block), IsOkAndHolds(false));
+  EXPECT_EQ(block->node_count(), 4);
+}
+
+}  // namespace
+}  // namespace xls::codegen
diff --git a/xls/codegen_v_1_5/channel_to_port_io_lowering_pass.cc b/xls/codegen_v_1_5/channel_to_port_io_lowering_pass.cc
index 243894a632..bcc330c488 100644
--- a/xls/codegen_v_1_5/channel_to_port_io_lowering_pass.cc
+++ b/xls/codegen_v_1_5/channel_to_port_io_lowering_pass.cc
@@ -641,6 +641,7 @@ absl::Status ConnectReceivesToConnector(
     Stage& stage = block->stages()[stage_index];
     Node* token = receive->As<Receive>()->token();
     std::optional<Node*> predicate = receive->As<Receive>()->predicate();
+    bool is_blocking = receive->As<Receive>()->is_blocking();
 
     if (connector.ready.has_value()) {
       // The ready signal from this receive is:
@@ -679,11 +680,54 @@ absl::Status ConnectReceivesToConnector(
               .status());
     }
 
-    XLS_RETURN_IF_ERROR(
-        receive
-            ->ReplaceUsesWithNewInStage<Tuple>(
-                stage_index, absl::MakeConstSpan({token, connector.data}))
-            .status());
+    Node* data = connector.data;
+    if (options.codegen_options.gate_recvs()) {
+      std::vector<Node*> gate_conditions;
+      gate_conditions.reserve(2);
+      if (predicate.has_value()) {
+        gate_conditions.push_back(*predicate);
+      }
+      if (!is_blocking && connector.valid.has_value()) {
+        gate_conditions.push_back(*connector.valid);
+      }
+
+      Node* gate_condition = nullptr;
+      if (gate_conditions.size() == 1) {
+        gate_condition = gate_conditions.front();
+      } else if (gate_conditions.size() > 1) {
+        XLS_ASSIGN_OR_RETURN(gate_condition, block->MakeNodeInStage<NaryOp>(
+                                                 stage_index, receive->loc(),
+                                                 gate_conditions, Op::kAnd));
+      }
+
+      if (gate_condition != nullptr) {
+        XLS_ASSIGN_OR_RETURN(Node * zero,
+                             block->MakeNodeInStage<Literal>(
+                                 stage_index, receive->loc(),
+                                 ZeroOfType(connector.data->GetType())));
+        XLS_ASSIGN_OR_RETURN(data,
+                             block->MakeNodeInStage<Select>(
+                                 stage_index, receive->loc(), gate_condition,
+                                 absl::MakeConstSpan({zero, connector.data}),
+                                 /*default_value=*/std::nullopt));
+      }
+    }
+
+    std::vector<Node*> replacement_elements = {token, data};
+    if (!is_blocking) {
+      if (connector.valid.has_value()) {
+        replacement_elements.push_back(*connector.valid);
+      } else {
+        XLS_ASSIGN_OR_RETURN(
+            Node * valid, block->MakeNodeInStage<Literal>(
+                              stage_index, receive->loc(), Value(UBits(1, 1))));
+        replacement_elements.push_back(valid);
+      }
+    }
+    XLS_RETURN_IF_ERROR(receive
+                            ->ReplaceUsesWithNewInStage<Tuple>(
+                                stage_index, replacement_elements)
+                            .status());
     XLS_RETURN_IF_ERROR(block->RemoveNode(receive));
   }
   if (connector.ready.has_value()) {
@@ -1030,9 +1074,16 @@ absl::flat_hash_set<std::string> GetRamChannelNames(
 absl::StatusOr<FlopKind> GetFlopKind(ChannelRef channel,
                                      ChannelDirection direction,
                                      ScheduledBlock* block) {
+  if (ChannelRefKind(channel) == ChannelKind::kSingleValue) {
+    // NOTE: We control the flop insertion for single-value channels globally,
+    // ignoring this result. This matches the behavior in codegen v1.
+    return FlopKind::kNone;
+  }
+
   if (std::holds_alternative<Channel*>(channel)) {
     XLS_RET_CHECK(!block->package()->ChannelsAreProcScoped())
         << "For proc-scoped channels, the flop kind is set on the interface.";
+    XLS_RET_CHECK_EQ(ChannelRefKind(channel), ChannelKind::kStreaming);
     StreamingChannel* streaming_channel =
         down_cast<StreamingChannel*>(std::get<Channel*>(channel));
     switch (direction) {
diff --git a/xls/codegen_v_1_5/channel_to_port_io_lowering_pass_test.cc b/xls/codegen_v_1_5/channel_to_port_io_lowering_pass_test.cc
index c5a4b157e1..fc497b1b33 100644
--- a/xls/codegen_v_1_5/channel_to_port_io_lowering_pass_test.cc
+++ b/xls/codegen_v_1_5/channel_to_port_io_lowering_pass_test.cc
@@ -488,5 +488,150 @@ TEST_F(ChannelToPortIoLoweringPassTest, SingleValueOutputFlop) {
               Contains(m::RegisterRead(HasSubstr("__a_out_reg"))));
 }
 
+TEST_F(ChannelToPortIoLoweringPassTest, GateRecvsWithPredicate) {
+  auto p = std::make_unique<Package>("test");
+  ScheduledProcBuilder pb(NewStyleProc(), "test_main", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(auto a_in,
+                           pb.AddInputChannel("a_in", p->GetBitsType(32)));
+  BValue tkn = pb.Literal(Value::Token());
+  BValue pred = pb.Literal(Value(UBits(1, 1)));
+  pb.ReceiveIf(a_in, tkn, pred);
+  XLS_ASSERT_OK(pb.Build());
+
+  BlockConversionPassOptions options;
+  options.codegen_options.gate_recvs(true);
+
+  EXPECT_THAT(Run(p.get(), options), IsOkAndHolds(true));
+  XLS_ASSERT_OK_AND_ASSIGN(Block * block, p->GetBlock("test_main"));
+
+  EXPECT_THAT(
+      block->nodes(),
+      Contains(m::Select(m::Literal(1),
+                         /*cases=*/{m::Literal(0), m::InputPort("a_in")})));
+}
+
+TEST_F(ChannelToPortIoLoweringPassTest, GateRecvsNonBlocking) {
+  auto p = std::make_unique<Package>("test");
+  ScheduledProcBuilder pb(NewStyleProc(), "test_main", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(auto a_in,
+                           pb.AddInputChannel("a_in", p->GetBitsType(32)));
+  BValue tkn = pb.Literal(Value::Token());
+  pb.ReceiveNonBlocking(a_in, tkn);
+  XLS_ASSERT_OK(pb.Build());
+
+  BlockConversionPassOptions options;
+  options.codegen_options.gate_recvs(true);
+
+  EXPECT_THAT(Run(p.get(), options), IsOkAndHolds(true));
+  XLS_ASSERT_OK_AND_ASSIGN(Block * block, p->GetBlock("test_main"));
+
+  EXPECT_THAT(
+      block->nodes(),
+      Contains(m::Select(m::InputPort("a_in_vld"),
+                         /*cases=*/{m::Literal(0), m::InputPort("a_in")})));
+}
+
+TEST_F(ChannelToPortIoLoweringPassTest, GateRecvsNonBlockingWithPredicate) {
+  auto p = std::make_unique<Package>("test");
+  ScheduledProcBuilder pb(NewStyleProc(), "test_main", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(auto a_in,
+                           pb.AddInputChannel("a_in", p->GetBitsType(32)));
+  BValue tkn = pb.Literal(Value::Token());
+  BValue pred = pb.Literal(Value(UBits(1, 1)));
+  pb.ReceiveIfNonBlocking(a_in, tkn, pred);
+  XLS_ASSERT_OK(pb.Build());
+
+  BlockConversionPassOptions options;
+  options.codegen_options.gate_recvs(true);
+
+  EXPECT_THAT(Run(p.get(), options), IsOkAndHolds(true));
+  XLS_ASSERT_OK_AND_ASSIGN(Block * block, p->GetBlock("test_main"));
+
+  EXPECT_THAT(
+      block->nodes(),
+      Contains(m::Select(m::And(m::Literal(1), m::InputPort("a_in_vld")),
+                         /*cases=*/{m::Literal(0), m::InputPort("a_in")})));
+}
+
+TEST_F(ChannelToPortIoLoweringPassTest, GateRecvsBlockingNoPredicate) {
+  auto p = std::make_unique<Package>("test");
+  ScheduledProcBuilder pb(NewStyleProc(), "test_main", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(auto a_in,
+                           pb.AddInputChannel("a_in", p->GetBitsType(32)));
+  BValue tkn = pb.Literal(Value::Token());
+  pb.Receive(a_in, tkn);
+  XLS_ASSERT_OK(pb.Build());
+
+  BlockConversionPassOptions options;
+  options.codegen_options.gate_recvs(true);
+
+  EXPECT_THAT(Run(p.get(), options), IsOkAndHolds(true));
+  XLS_ASSERT_OK_AND_ASSIGN(Block * block, p->GetBlock("test_main"));
+
+  EXPECT_THAT(block->nodes(), Contains(m::Tuple(m::Literal(Value::Token()),
+                                                m::InputPort("a_in"))));
+}
+
+TEST_F(ChannelToPortIoLoweringPassTest, NoGateRecvsWithPredicate) {
+  auto p = std::make_unique<Package>("test");
+  ScheduledProcBuilder pb(NewStyleProc(), "test_main", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(auto a_in,
+                           pb.AddInputChannel("a_in", p->GetBitsType(32)));
+  BValue tkn = pb.Literal(Value::Token());
+  BValue pred = pb.Literal(Value(UBits(1, 1)));
+  pb.ReceiveIf(a_in, tkn, pred);
+  XLS_ASSERT_OK(pb.Build());
+
+  BlockConversionPassOptions options;
+  options.codegen_options.gate_recvs(false);
+
+  EXPECT_THAT(Run(p.get(), options), IsOkAndHolds(true));
+  XLS_ASSERT_OK_AND_ASSIGN(Block * block, p->GetBlock("test_main"));
+
+  EXPECT_THAT(block->nodes(), Contains(m::Tuple(m::Literal(Value::Token()),
+                                                m::InputPort("a_in"))));
+}
+
+TEST_F(ChannelToPortIoLoweringPassTest, NoGateRecvsNonBlocking) {
+  auto p = std::make_unique<Package>("test");
+  ScheduledProcBuilder pb(NewStyleProc(), "test_main", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(auto a_in,
+                           pb.AddInputChannel("a_in", p->GetBitsType(32)));
+  BValue tkn = pb.Literal(Value::Token());
+  pb.ReceiveNonBlocking(a_in, tkn);
+  XLS_ASSERT_OK(pb.Build());
+
+  BlockConversionPassOptions options;
+  options.codegen_options.gate_recvs(false);
+
+  EXPECT_THAT(Run(p.get(), options), IsOkAndHolds(true));
+  XLS_ASSERT_OK_AND_ASSIGN(Block * block, p->GetBlock("test_main"));
+
+  EXPECT_THAT(block->nodes(), Contains(m::Tuple(m::Literal(Value::Token()),
+                                                m::InputPort("a_in"),
+                                                m::InputPort("a_in_vld"))));
+}
+
+TEST_F(ChannelToPortIoLoweringPassTest, NoGateRecvsNonBlockingWithPredicate) {
+  auto p = std::make_unique<Package>("test");
+  ScheduledProcBuilder pb(NewStyleProc(), "test_main", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(auto a_in,
+                           pb.AddInputChannel("a_in", p->GetBitsType(32)));
+  BValue tkn = pb.Literal(Value::Token());
+  BValue pred = pb.Literal(Value(UBits(1, 1)));
+  pb.ReceiveIfNonBlocking(a_in, tkn, pred);
+  XLS_ASSERT_OK(pb.Build());
+
+  BlockConversionPassOptions options;
+  options.codegen_options.gate_recvs(false);
+
+  EXPECT_THAT(Run(p.get(), options), IsOkAndHolds(true));
+  XLS_ASSERT_OK_AND_ASSIGN(Block * block, p->GetBlock("test_main"));
+
+  EXPECT_THAT(block->nodes(), Contains(m::Tuple(m::Literal(Value::Token()),
+                                                m::InputPort("a_in"),
+                                                m::InputPort("a_in_vld"))));
+}
+
 }  // namespace
 }  // namespace xls::codegen
diff --git a/xls/codegen_v_1_5/convert_to_block.cc b/xls/codegen_v_1_5/convert_to_block.cc
index f59676ebcf..c5546b228c 100644
--- a/xls/codegen_v_1_5/convert_to_block.cc
+++ b/xls/codegen_v_1_5/convert_to_block.cc
@@ -26,6 +26,7 @@
 #include "xls/common/status/status_macros.h"
 #include "xls/estimators/delay_model/delay_estimator.h"
 #include "xls/ir/package.h"
+#include "xls/passes/optimization_pass.h"
 #include "xls/passes/pass_base.h"
 #include "xls/scheduling/scheduling_options.h"
 #include "xls/scheduling/scheduling_result.h"
@@ -56,8 +57,9 @@ absl::Status ConvertToBlock(
       .codegen_options = std::move(codegen_options),
       .package_schedule = schedule};
 
+  OptimizationContext opt_context;
   std::unique_ptr<BlockConversionCompoundPass> pipeline =
-      CreateBlockConversionPassPipeline();
+      CreateBlockConversionPassPipeline(opt_context);
   PassResults results;
   XLS_ASSIGN_OR_RETURN(bool result, pipeline->Run(p, options, &results));
   XLS_RET_CHECK(result);
diff --git a/xls/codegen_v_1_5/flow_control_insertion_pass.cc b/xls/codegen_v_1_5/flow_control_insertion_pass.cc
index 54712f11b7..48b5d80686 100644
--- a/xls/codegen_v_1_5/flow_control_insertion_pass.cc
+++ b/xls/codegen_v_1_5/flow_control_insertion_pass.cc
@@ -66,6 +66,34 @@ absl::StatusOr<bool> FlowControlInsertionPass::InsertFlowControl(
     return changed;
   }
 
+  const bool uses_valid =
+      options.codegen_options.valid_control().has_value() &&
+      (options.codegen_options.valid_control()->has_input_name() ||
+       options.codegen_options.valid_control()->has_output_name());
+  if (block->source() != nullptr && block->source()->IsFunction() &&
+      !uses_valid) {
+    // If this block represents a function and neither gates on valid nor
+    // signals valid, then the pipeline registers can just be considered
+    // always-valid and always-ready.
+    for (Stage& stage : block->stages()) {
+      if (!IsLiteralUnsignedOne(stage.inputs_valid())) {
+        XLS_RETURN_IF_ERROR(
+            stage.inputs_valid()
+                ->ReplaceUsesWithNew<Literal>(Value(UBits(1, 1)))
+                .status());
+        changed = true;
+      }
+      if (!IsLiteralUnsignedOne(stage.outputs_ready())) {
+        XLS_RETURN_IF_ERROR(
+            stage.outputs_ready()
+                ->ReplaceUsesWithNew<Literal>(Value(UBits(1, 1)))
+                .status());
+        changed = true;
+      }
+    }
+    return changed;
+  }
+
   absl::FixedArray<Node*> stage_done(block->stages().size());
   for (int64_t stage_index = 0; stage_index < block->stages().size();
        ++stage_index) {
diff --git a/xls/codegen_v_1_5/function_io_lowering_pass.cc b/xls/codegen_v_1_5/function_io_lowering_pass.cc
index 997622d6a3..3e963095c5 100644
--- a/xls/codegen_v_1_5/function_io_lowering_pass.cc
+++ b/xls/codegen_v_1_5/function_io_lowering_pass.cc
@@ -182,10 +182,13 @@ absl::StatusOr<bool> FunctionIOLoweringPass::LowerReturnValue(
     return absl::UnimplementedError("Splitting outputs not supported.");
   }
 
-  // The return value is always expected to be in the last stage.
-  XLS_ASSIGN_OR_RETURN(int64_t stage_index,
-                       block->GetStageIndex(block->source_return_value()));
-  XLS_RET_CHECK_EQ(stage_index, block->stages().size() - 1);
+  // The return value is always assumed to be in the last stage - or possibly
+  // unscheduled, if (e.g.) returning a literal.
+  if (block->IsStaged(block->source_return_value())) {
+    XLS_ASSIGN_OR_RETURN(int64_t stage_index,
+                         block->GetStageIndex(block->source_return_value()));
+    XLS_RET_CHECK_EQ(stage_index, block->stages().size() - 1);
+  }
 
   std::optional<Node*> output_valid = std::nullopt;
   if (options.codegen_options.valid_control().has_value() &&
diff --git a/xls/codegen_v_1_5/function_io_lowering_pass_test.cc b/xls/codegen_v_1_5/function_io_lowering_pass_test.cc
index 216d5e3f5f..8f44a4d6bd 100644
--- a/xls/codegen_v_1_5/function_io_lowering_pass_test.cc
+++ b/xls/codegen_v_1_5/function_io_lowering_pass_test.cc
@@ -108,12 +108,11 @@ class FunctionIOLoweringPassTest : public IrTestBase {
 
   // Runs PipelineRegisterInsertionPass and BlockFinalizationPass, then
   // verify & return the output of block interpretation.
-  absl::StatusOr<Value> RunFunctionalTest(Package* p, Block* block,
-                                          absl::Span<const Value> inputs,
-                                          int64_t expected_latency,
-                                          FunctionalTestOptions options = {}) {
+  absl::StatusOr<Value> RunFunctionalTest(
+      Package* p, Block* block, const BlockConversionPassOptions& pass_options,
+      absl::Span<const Value> inputs, int64_t expected_latency,
+      FunctionalTestOptions options = {}) {
     std::string block_name = block->name();
-    BlockConversionPassOptions pass_options;
     PassResults pass_results;
     XLS_RETURN_IF_ERROR(PipelineRegisterInsertionPass()
                             .Run(p, pass_options, &pass_results)
@@ -221,7 +220,7 @@ TEST_F(FunctionIOLoweringPassTest, SingleStage) {
               ElementsAre(Property(&PortNode::GetName, "out")));
   EXPECT_EQ(sb->stages().size(), 1);
 
-  EXPECT_THAT(RunFunctionalTest(p.get(), sb, inputs,
+  EXPECT_THAT(RunFunctionalTest(p.get(), sb, options, inputs,
                                 /*expected_latency=*/1),
               IsOkAndHolds(expected_output));
 }
@@ -264,7 +263,7 @@ TEST_F(FunctionIOLoweringPassTest, SingleStageWithInputValid) {
   EXPECT_EQ(sb->stages().size(), 1);
 
   EXPECT_THAT(
-      RunFunctionalTest(p.get(), sb, inputs,
+      RunFunctionalTest(p.get(), sb, options, inputs,
                         /*expected_latency=*/1, {.input_valid = "x_valid"}),
       IsOkAndHolds(expected_output));
 }
@@ -308,7 +307,7 @@ TEST_F(FunctionIOLoweringPassTest, SingleStageWithIOValid) {
   EXPECT_EQ(sb->stages().size(), 1);
 
   EXPECT_THAT(RunFunctionalTest(
-                  p.get(), sb, inputs,
+                  p.get(), sb, options, inputs,
                   /*expected_latency=*/1,
                   {.input_valid = "x_valid", .output_valid = "out_valid"}),
               IsOkAndHolds(expected_output));
@@ -350,7 +349,7 @@ TEST_F(FunctionIOLoweringPassTest, MultiStage) {
               ElementsAre(Property(&PortNode::GetName, "out")));
   EXPECT_EQ(sb->stages().size(), 3);
 
-  EXPECT_THAT(RunFunctionalTest(p.get(), sb, inputs,
+  EXPECT_THAT(RunFunctionalTest(p.get(), sb, options, inputs,
                                 /*expected_latency=*/3),
               IsOkAndHolds(expected_output));
 }
@@ -396,7 +395,7 @@ TEST_F(FunctionIOLoweringPassTest, MultiStageWithInputValid) {
   EXPECT_EQ(sb->stages().size(), 3);
 
   EXPECT_THAT(
-      RunFunctionalTest(p.get(), sb, inputs,
+      RunFunctionalTest(p.get(), sb, options, inputs,
                         /*expected_latency=*/3, {.input_valid = "x_valid"}),
       IsOkAndHolds(expected_output));
 }
@@ -442,7 +441,7 @@ TEST_F(FunctionIOLoweringPassTest, MultiStageWithIOValidDelayed) {
                           Property(&PortNode::GetName, "out")));
   EXPECT_EQ(sb->stages().size(), 3);
 
-  EXPECT_THAT(RunFunctionalTest(p.get(), sb, inputs,
+  EXPECT_THAT(RunFunctionalTest(p.get(), sb, options, inputs,
                                 /*expected_latency=*/3,
                                 {.input_valid = "x_valid",
                                  .output_valid = "out_valid",
@@ -492,7 +491,7 @@ TEST_F(FunctionIOLoweringPassTest, MultiStageWithIOValid) {
   EXPECT_EQ(sb->stages().size(), 3);
 
   EXPECT_THAT(RunFunctionalTest(
-                  p.get(), sb, inputs,
+                  p.get(), sb, options, inputs,
                   /*expected_latency=*/3,
                   {.input_valid = "x_valid", .output_valid = "out_valid"}),
               IsOkAndHolds(expected_output));
@@ -541,7 +540,7 @@ TEST_F(FunctionIOLoweringPassTest, MultiStageMultiInputWithIOValid) {
   EXPECT_EQ(sb->stages().size(), 3);
 
   EXPECT_THAT(RunFunctionalTest(
-                  p.get(), sb, inputs,
+                  p.get(), sb, options, inputs,
                   /*expected_latency=*/3,
                   {.input_valid = "x_valid", .output_valid = "out_valid"}),
               IsOkAndHolds(expected_output));
@@ -587,7 +586,7 @@ TEST_F(FunctionIOLoweringPassTest, MultiStageMultiInputWithFloppedInputs) {
               ElementsAre(Property(&PortNode::GetName, "out")));
   EXPECT_EQ(sb->stages().size(), 3);
 
-  EXPECT_THAT(RunFunctionalTest(p.get(), sb, inputs,
+  EXPECT_THAT(RunFunctionalTest(p.get(), sb, options, inputs,
                                 /*expected_latency=*/4),
               IsOkAndHolds(expected_output));
 }
@@ -632,7 +631,7 @@ TEST_F(FunctionIOLoweringPassTest, MultiStageMultiInputWithFloppedOutputs) {
               ElementsAre(Property(&PortNode::GetName, "out")));
   EXPECT_EQ(sb->stages().size(), 3);
 
-  EXPECT_THAT(RunFunctionalTest(p.get(), sb, inputs,
+  EXPECT_THAT(RunFunctionalTest(p.get(), sb, options, inputs,
                                 /*expected_latency=*/4),
               IsOkAndHolds(expected_output));
 }
@@ -678,7 +677,7 @@ TEST_F(FunctionIOLoweringPassTest, MultiStageMultiInputWithFloppedIO) {
               ElementsAre(Property(&PortNode::GetName, "out")));
   EXPECT_EQ(sb->stages().size(), 3);
 
-  EXPECT_THAT(RunFunctionalTest(p.get(), sb, inputs,
+  EXPECT_THAT(RunFunctionalTest(p.get(), sb, options, inputs,
                                 /*expected_latency=*/5),
               IsOkAndHolds(expected_output));
 }
diff --git a/xls/codegen_v_1_5/register_cleanup_pass.cc b/xls/codegen_v_1_5/register_cleanup_pass.cc
new file mode 100644
index 0000000000..af97b37df5
--- /dev/null
+++ b/xls/codegen_v_1_5/register_cleanup_pass.cc
@@ -0,0 +1,188 @@
+// Copyright 2025 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_v_1_5/register_cleanup_pass.h"
+
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/statusor.h"
+#include "absl/types/span.h"
+#include "xls/codegen_v_1_5/block_conversion_pass.h"
+#include "xls/common/status/status_macros.h"
+#include "xls/data_structures/transitive_closure.h"
+#include "xls/ir/block.h"
+#include "xls/ir/node.h"
+#include "xls/ir/nodes.h"
+#include "xls/ir/package.h"
+#include "xls/ir/register.h"
+#include "xls/ir/value.h"
+#include "xls/ir/value_utils.h"
+#include "xls/passes/node_dependency_analysis.h"
+#include "xls/passes/partial_info_query_engine.h"
+#include "xls/passes/pass_base.h"
+#include "xls/passes/query_engine.h"
+
+namespace xls::codegen {
+
+namespace {
+
+absl::StatusOr<std::vector<RegisterWrite*>> GetRegisterWrites(Block* block,
+                                                              Register* reg) {
+  XLS_ASSIGN_OR_RETURN(absl::Span<RegisterWrite* const> writes,
+                       block->GetRegisterWrites(reg));
+  return std::vector<RegisterWrite*>(writes.begin(), writes.end());
+}
+
+}  // namespace
+
+absl::StatusOr<bool> RegisterCleanupPass::RemoveTrivialLoadEnables(
+    Block* block, QueryEngine& query_engine) const {
+  bool changed = false;
+  for (Register* reg : block->GetRegisters()) {
+    XLS_ASSIGN_OR_RETURN(std::vector<RegisterWrite*> writes,
+                         GetRegisterWrites(block, reg));
+    for (RegisterWrite* write : writes) {
+      if (write->load_enable().has_value() &&
+          query_engine.IsAllOnes(*write->load_enable())) {
+        XLS_RETURN_IF_ERROR(
+            write
+                ->ReplaceUsesWithNew<RegisterWrite>(write->data(), std::nullopt,
+                                                    write->reset(), reg)
+                .status());
+        XLS_RETURN_IF_ERROR(block->RemoveNode(write));
+        changed = true;
+      }
+    }
+  }
+  return changed;
+}
+
+absl::StatusOr<bool> RegisterCleanupPass::RemoveImpossibleWrites(
+    Block* block, QueryEngine& query_engine) const {
+  for (Register* reg : block->GetRegisters()) {
+    XLS_ASSIGN_OR_RETURN(std::vector<RegisterWrite*> writes,
+                         GetRegisterWrites(block, reg));
+    for (RegisterWrite* write : writes) {
+      bool write_disabled = write->load_enable().has_value() &&
+                            query_engine.IsAllZeros(*write->load_enable());
+      if (write_disabled) {
+        XLS_RETURN_IF_ERROR(
+            write->ReplaceUsesWithNew<Literal>(Value::Tuple({})).status());
+        XLS_RETURN_IF_ERROR(block->RemoveNode(write));
+      }
+    }
+
+    XLS_ASSIGN_OR_RETURN(writes, GetRegisterWrites(block, reg));
+    if (writes.empty() || (reg->reset_value().has_value() &&
+                           absl::c_all_of(writes, [&](RegisterWrite* write) {
+                             return query_engine.KnownValue(write->data()) ==
+                                    *reg->reset_value();
+                           }))) {
+      // Nothing can ever write a new value to this register, so it's equivalent
+      // to its reset value.
+      XLS_ASSIGN_OR_RETURN(RegisterRead * read, block->GetRegisterRead(reg));
+      XLS_RETURN_IF_ERROR(
+          read->ReplaceUsesWithNew<Literal>(
+                  reg->reset_value().value_or(ZeroOfType(reg->type())))
+              .status());
+      XLS_RETURN_IF_ERROR(block->RemoveNode(read));
+      XLS_RETURN_IF_ERROR(block->RemoveRegister(reg));
+    }
+  }
+  return false;
+}
+
+absl::StatusOr<bool> RegisterCleanupPass::RemoveUnreadRegisters(
+    Block* block, QueryEngine& query_engine) const {
+  absl::flat_hash_map<Register*, absl::flat_hash_set<Register*>>
+      can_receive_value_from;
+  absl::flat_hash_set<Register*> directly_read_registers;
+  NodeBackwardDependencyAnalysis nda;
+  XLS_RETURN_IF_ERROR(nda.Attach(block).status());
+  for (Register* reg : block->GetRegisters()) {
+    XLS_ASSIGN_OR_RETURN(RegisterRead * read, block->GetRegisterRead(reg));
+    for (Node* user : nda.NodesDependingOn(read)) {
+      if (user->Is<RegisterWrite>()) {
+        can_receive_value_from[reg].insert(
+            user->As<RegisterWrite>()->GetRegister());
+      } else if (user->Is<OutputPort>()) {
+        directly_read_registers.insert(reg);
+      }
+    }
+  }
+
+  absl::flat_hash_map<Register*, absl::flat_hash_set<Register*>>
+      transitive_receivers = TransitiveClosure(can_receive_value_from);
+  std::vector<Register*> unread_registers;
+  for (Register* reg : block->GetRegisters()) {
+    if (!directly_read_registers.contains(reg) &&
+        absl::c_none_of(transitive_receivers[reg], [&](Register* user) {
+          return directly_read_registers.contains(user);
+        })) {
+      unread_registers.push_back(reg);
+    }
+  }
+
+  for (Register* reg : unread_registers) {
+    XLS_ASSIGN_OR_RETURN(RegisterRead * read, block->GetRegisterRead(reg));
+    XLS_ASSIGN_OR_RETURN(absl::Span<RegisterWrite* const> writes,
+                         block->GetRegisterWrites(reg));
+    if (!read->IsDead()) {
+      XLS_RETURN_IF_ERROR(
+          read->ReplaceUsesWithNew<Literal>(
+                  reg->reset_value().value_or(ZeroOfType(reg->type())))
+              .status());
+    }
+    XLS_RETURN_IF_ERROR(block->RemoveNode(read));
+    for (RegisterWrite* write : writes) {
+      if (!write->IsDead()) {
+        XLS_RETURN_IF_ERROR(
+            write->ReplaceUsesWithNew<Literal>(Value::Tuple({})).status());
+      }
+      XLS_RETURN_IF_ERROR(block->RemoveNode(write));
+    }
+    XLS_RETURN_IF_ERROR(block->RemoveRegister(reg));
+  }
+  return !unread_registers.empty();
+}
+
+absl::StatusOr<bool> RegisterCleanupPass::RunInternal(
+    Package* package, const BlockConversionPassOptions& options,
+    PassResults* results) const {
+  bool changed = false;
+  for (const std::unique_ptr<Block>& block : package->blocks()) {
+    PartialInfoQueryEngine query_engine;
+    XLS_RETURN_IF_ERROR(query_engine.Populate(block.get()).status());
+
+    XLS_ASSIGN_OR_RETURN(bool changed_load_enables,
+                         RemoveTrivialLoadEnables(block.get(), query_engine));
+    changed |= changed_load_enables;
+
+    XLS_ASSIGN_OR_RETURN(bool changed_writes,
+                         RemoveImpossibleWrites(block.get(), query_engine));
+    changed |= changed_writes;
+
+    XLS_ASSIGN_OR_RETURN(bool changed_registers,
+                         RemoveUnreadRegisters(block.get(), query_engine));
+    changed |= changed_registers;
+  }
+  return changed;
+}
+
+}  // namespace xls::codegen
diff --git a/xls/codegen_v_1_5/register_cleanup_pass.h b/xls/codegen_v_1_5/register_cleanup_pass.h
new file mode 100644
index 0000000000..a2a827d8ec
--- /dev/null
+++ b/xls/codegen_v_1_5/register_cleanup_pass.h
@@ -0,0 +1,61 @@
+// Copyright 2025 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_V_1_5_REGISTER_CLEANUP_PASS_H_
+#define XLS_CODEGEN_V_1_5_REGISTER_CLEANUP_PASS_H_
+
+#include "absl/status/statusor.h"
+#include "xls/codegen_v_1_5/block_conversion_pass.h"
+#include "xls/ir/package.h"
+#include "xls/passes/pass_base.h"
+#include "xls/passes/query_engine.h"
+
+namespace xls::codegen {
+
+// A pass which cleans up registers and register writes in the block.
+//
+// Specific cleanup actions performed:
+// 1. Removes load-enable signals from register writes if the load-enable is
+//    statically determined to be always one (true).
+// 2. Removes register writes if the load-enable is statically determined to
+//    be always zero (false).
+// 3. Replaces registers which are never written to, or only ever have values
+//    identical to their reset value written to them, with literals.
+// 4. Removes registers which do not transitively drive any output port.
+//
+// This pass is intended to run late in the codegen pipeline after flow control
+// and pipeline registers have been inserted.
+class RegisterCleanupPass : public BlockConversionPass {
+ public:
+  RegisterCleanupPass()
+      : BlockConversionPass("register_cleanup",
+                            "Remove dead registers & unused load-enable bits") {
+  }
+
+ protected:
+  absl::StatusOr<bool> RemoveTrivialLoadEnables(
+      Block* block, QueryEngine& query_engine) const;
+  absl::StatusOr<bool> RemoveImpossibleWrites(Block* block,
+                                              QueryEngine& query_engine) const;
+  absl::StatusOr<bool> RemoveUnreadRegisters(Block* block,
+                                             QueryEngine& query_engine) const;
+
+  absl::StatusOr<bool> RunInternal(Package* package,
+                                   const BlockConversionPassOptions& options,
+                                   PassResults* results) const override;
+};
+
+}  // namespace xls::codegen
+
+#endif  // XLS_CODEGEN_V_1_5_REGISTER_CLEANUP_PASS_H_
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 b23715703c..6fa15dd7c6 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
@@ -49,37 +49,42 @@
 namespace xls::codegen {
 namespace {
 
-absl::Status ReplaceNode(Block& block, Node* old_node, Node* new_node,
-                         bool replace_implicit_uses = true) {
+// Replaces `old_node` with `new_node` in the given block. Returns whether
+// `old_node` was removed after the replacement.
+absl::StatusOr<bool> ReplaceNode(Block* block, Node* old_node, Node* new_node,
+                                 bool replace_implicit_uses = true) {
   XLS_RETURN_IF_ERROR(old_node->ReplaceUsesWith(
       new_node, [&](Node* n) { return n != new_node; }, replace_implicit_uses));
 
   if (old_node->Is<StateRead>()) {
-    XLS_RETURN_IF_ERROR(block.RemoveNodeFromStage(old_node).status());
-  } else if (old_node->IsDead()) {
-    XLS_RETURN_IF_ERROR(block.RemoveNode(old_node));
+    XLS_RETURN_IF_ERROR(block->RemoveNodeFromStage(old_node).status());
+    return false;
   }
-  return absl::OkStatus();
+  if (old_node->IsDead()) {
+    XLS_RETURN_IF_ERROR(block->RemoveNode(old_node));
+    return true;
+  }
+  return false;
 }
 
-absl::StatusOr<Node*> NodeOrOne(FunctionBase& fb, std::optional<Node*> node) {
+absl::StatusOr<Node*> NodeOrOne(FunctionBase* fb, std::optional<Node*> node) {
   if (node.has_value()) {
     return *node;
   }
-  return fb.MakeNode<xls::Literal>(SourceInfo(), Value(UBits(1, 1)));
+  return fb->MakeNode<xls::Literal>(SourceInfo(), Value(UBits(1, 1)));
 }
 
 absl::StatusOr<Register*> CreateFullRegisterAndRead(
-    Block& block, const StateElement& state_element,
+    Block* block, const StateElement& state_element,
     std::optional<Node*> read_predicate, const Next& last_write,
     int read_stage_index, Stage& read_stage) {
   XLS_ASSIGN_OR_RETURN(
       Register * reg_full,
-      block.AddRegister(absl::StrCat("__", state_element.name(), "_full"),
-                        block.package()->GetBitsType(1), Value(UBits(1, 1))));
+      block->AddRegister(absl::StrCat("__", state_element.name(), "_full"),
+                         block->package()->GetBitsType(1), Value(UBits(1, 1))));
 
   XLS_ASSIGN_OR_RETURN(RegisterRead * reg_full_read,
-                       block.MakeNodeWithNameInStage<RegisterRead>(
+                       block->MakeNodeWithNameInStage<RegisterRead>(
                            read_stage_index, last_write.loc(), reg_full,
                            /*name=*/reg_full->name()));
 
@@ -87,32 +92,33 @@ absl::StatusOr<Register*> CreateFullRegisterAndRead(
   if (read_predicate.has_value()) {
     XLS_ASSIGN_OR_RETURN(
         Node * not_predicate,
-        block.MakeNodeInStage<UnOp>(read_stage_index, SourceInfo{},
-                                    *read_predicate, Op::kNot));
+        block->MakeNodeInStage<UnOp>(read_stage_index, SourceInfo{},
+                                     *read_predicate, Op::kNot));
 
     XLS_ASSIGN_OR_RETURN(
         state_ready_or_not_needed,
-        block.MakeNodeInStage<NaryOp>(
+        block->MakeNodeInStage<NaryOp>(
             read_stage_index, SourceInfo(),
             std::vector<Node*>{reg_full_read, not_predicate}, Op::kOr));
   }
 
   XLS_ASSIGN_OR_RETURN(
       Node * new_active_inputs_valid,
-      NaryAndIfNeeded(&block,
+      NaryAndIfNeeded(block,
                       absl::MakeConstSpan({read_stage.active_inputs_valid(),
                                            state_ready_or_not_needed}),
                       /*name=*/"", SourceInfo{},
                       /*drop_literal_one_operands=*/true));
   XLS_RETURN_IF_ERROR(ReplaceNode(block, read_stage.active_inputs_valid(),
-                                  new_active_inputs_valid));
+                                  new_active_inputs_valid)
+                          .status());
   // Replacement is prone to dropping the read from the stage.
   XLS_RETURN_IF_ERROR(
-      block.AddNodeToStage(read_stage_index, reg_full_read).status());
+      block->AddNodeToStage(read_stage_index, reg_full_read).status());
   return reg_full;
 }
 
-absl::Status AddFullRegisterWrite(Block& block, Register* reg_full,
+absl::Status AddFullRegisterWrite(Block* block, Register* reg_full,
                                   std::optional<Node*> read_predicate,
                                   const Stage& read_stage,
                                   const SourceInfo& last_write_loc,
@@ -124,31 +130,31 @@ absl::Status AddFullRegisterWrite(Block& block, Register* reg_full,
   }
 
   XLS_ASSIGN_OR_RETURN(Node * full_from_read,
-                       NaryAndIfNeeded(&block, full_from_read_operands));
+                       NaryAndIfNeeded(block, full_from_read_operands));
 
   XLS_ASSIGN_OR_RETURN(Node * write_load_enable_or_1,
                        NodeOrOne(block, write_load_enable));
   XLS_ASSIGN_OR_RETURN(
       Node * full_load_enable,
-      NaryOrIfNeeded(&block, absl::MakeConstSpan(
-                                 {full_from_read, write_load_enable_or_1})));
+      NaryOrIfNeeded(block, absl::MakeConstSpan(
+                                {full_from_read, write_load_enable_or_1})));
 
   XLS_RETURN_IF_ERROR(
       block
-          .MakeNode<RegisterWrite>(last_write_loc, write_load_enable_or_1,
-                                   full_load_enable,
-                                   /*reset=*/block.GetResetPort(), reg_full)
+          ->MakeNode<RegisterWrite>(last_write_loc, write_load_enable_or_1,
+                                    full_load_enable,
+                                    /*reset=*/block->GetResetPort(), reg_full)
           .status());
   return absl::OkStatus();
 }
 
-absl::Status LowerStateElementToZero(Block& block,
+absl::Status LowerStateElementToZero(ScheduledBlock* block,
                                      const StateElement& state_element,
-                                     StateRead& read,
+                                     StateRead* read,
                                      absl::Span<Next*> writes) {
   Type* type = state_element.type();
   if (!state_element.type()->IsToken() &&
-      state_element.type() != block.package()->GetTupleType({})) {
+      state_element.type() != block->package()->GetTupleType({})) {
     return absl::UnimplementedError(
         absl::StrFormat("Proc has zero-width state element `%s`, but type is "
                         "not token or empty tuple, instead got %s.",
@@ -157,15 +163,22 @@ absl::Status LowerStateElementToZero(Block& block,
 
   XLS_ASSIGN_OR_RETURN(
       Literal * replacement,
-      block.MakeNode<xls::Literal>(read.loc(), ZeroOfType(type)));
-  XLS_RETURN_IF_ERROR(ReplaceNode(block, &read, replacement,
-                                  /*replace_implicit_uses=*/false));
+      block->MakeNode<xls::Literal>(read->loc(), ZeroOfType(type)));
+  XLS_ASSIGN_OR_RETURN(bool read_removed,
+                       ReplaceNode(block, read, replacement,
+                                   /*replace_implicit_uses=*/false));
+  XLS_RET_CHECK(!read_removed);
+
+  for (Next* write : writes) {
+    XLS_RETURN_IF_ERROR(block->RemoveNode(write));
+  }
+
   return absl::OkStatus();
 }
 
-absl::Status LowerStateElement(ScheduledBlock& block,
+absl::Status LowerStateElement(ScheduledBlock* block,
                                const StateElement& state_element,
-                               StateRead& read, absl::Span<Next*> writes) {
+                               StateRead* read, absl::Span<Next*> writes) {
   // A token or empty tuple state element has no real functionality.
   Type* type = state_element.type();
   if (type->IsToken() || type->GetFlatBitCount() == 0) {
@@ -175,21 +188,23 @@ absl::Status LowerStateElement(ScheduledBlock& block,
   }
 
   // Lower the read of the state element.
-  std::optional<Node*> read_predicate = read.predicate();
+  std::optional<Node*> read_predicate = read->predicate();
   std::string name =
-      block.UniquifyNodeName(absl::StrCat("__", state_element.name()));
+      block->UniquifyNodeName(absl::StrCat("__", state_element.name()));
   XLS_ASSIGN_OR_RETURN(Register * state_register,
-                       block.AddRegister(name, state_element.type(),
-                                         state_element.initial_value()));
+                       block->AddRegister(name, state_element.type(),
+                                          state_element.initial_value()));
 
-  XLS_ASSIGN_OR_RETURN(int read_stage_index, block.GetStageIndex(&read));
-  Stage& read_stage = block.stages()[read_stage_index];
+  XLS_ASSIGN_OR_RETURN(int read_stage_index, block->GetStageIndex(read));
+  Stage& read_stage = block->stages()[read_stage_index];
   XLS_ASSIGN_OR_RETURN(RegisterRead * reg_read,
-                       block.MakeNodeWithNameInStage<RegisterRead>(
-                           read_stage_index, read.loc(), state_register,
+                       block->MakeNodeWithNameInStage<RegisterRead>(
+                           read_stage_index, read->loc(), state_register,
                            /*name=*/state_register->name()));
-  XLS_RETURN_IF_ERROR(ReplaceNode(block, &read, reg_read,
-                                  /*replace_implicit_uses=*/false));
+  XLS_ASSIGN_OR_RETURN(bool read_removed,
+                       ReplaceNode(block, read, reg_read,
+                                   /*replace_implicit_uses=*/false));
+  XLS_RET_CHECK(!read_removed);
 
   Next* last_write = nullptr;
   Node* last_value = nullptr;
@@ -202,7 +217,7 @@ absl::Status LowerStateElement(ScheduledBlock& block,
     last_value = last_write->value();
     last_value_loc = last_value->loc();
     XLS_ASSIGN_OR_RETURN(last_write_stage_index,
-                         block.GetStageIndex(last_write));
+                         block->GetStageIndex(last_write));
   }
 
   // If the next state can be determined in a later cycle than the state read,
@@ -223,8 +238,8 @@ absl::Status LowerStateElement(ScheduledBlock& block,
   write_conditions.reserve(writes.size());
   values.reserve(writes.size());
   for (Next* write : writes) {
-    XLS_ASSIGN_OR_RETURN(int stage_index, block.GetStageIndex(write));
-    Stage& stage = block.stages()[stage_index];
+    XLS_ASSIGN_OR_RETURN(int stage_index, block->GetStageIndex(write));
+    Stage& stage = block->stages()[stage_index];
 
     std::vector<Node*> gate_operands{stage.inputs_valid(),
                                      stage.active_inputs_valid()};
@@ -234,9 +249,9 @@ absl::Status LowerStateElement(ScheduledBlock& block,
       gate_operands.push_back(*write->predicate());
       condition_operands.push_back(*write->predicate());
     }
-    XLS_ASSIGN_OR_RETURN(Node * gate, NaryAndIfNeeded(&block, gate_operands));
+    XLS_ASSIGN_OR_RETURN(Node * gate, NaryAndIfNeeded(block, gate_operands));
     XLS_ASSIGN_OR_RETURN(Node * condition,
-                         NaryAndIfNeeded(&block, condition_operands));
+                         NaryAndIfNeeded(block, condition_operands));
     gates.push_back(gate);
     write_conditions.push_back(condition);
     values.push_back(write->value());
@@ -252,31 +267,32 @@ absl::Status LowerStateElement(ScheduledBlock& block,
     value = values[0];
     write_load_enable = write_conditions[0];
   } else {
-    XLS_ASSIGN_OR_RETURN(Node * selector,
-                         block.MakeNode<xls::Concat>(last_write->loc(), gates));
+    XLS_ASSIGN_OR_RETURN(Node * selector, block->MakeNode<xls::Concat>(
+                                              last_write->loc(), gates));
 
     // Reverse the order of the values, so they match up to the selector.
     std::reverse(values.begin(), values.end());
-    XLS_ASSIGN_OR_RETURN(value, block.MakeNode<OneHotSelect>(last_write->loc(),
-                                                             selector, values));
+    XLS_ASSIGN_OR_RETURN(value, block->MakeNode<OneHotSelect>(
+                                    last_write->loc(), selector, values));
 
     XLS_ASSIGN_OR_RETURN(
         write_load_enable,
-        block.MakeNode<NaryOp>(last_write->loc(), write_conditions, Op::kOr));
+        block->MakeNode<NaryOp>(last_write->loc(), write_conditions, Op::kOr));
   }
 
   // Create the register write, and replace the Next nodes with empty tuples.
   XLS_RETURN_IF_ERROR(block
-                          .MakeNode<RegisterWrite>(
+                          ->MakeNode<RegisterWrite>(
                               last_value_loc, value,
                               /*load_enable=*/write_load_enable,
-                              /*reset=*/block.GetResetPort(), state_register)
+                              /*reset=*/block->GetResetPort(), state_register)
                           .status());
   XLS_ASSIGN_OR_RETURN(
       Node * empty_tuple,
-      block.MakeNode<Tuple>(SourceInfo{}, absl::Span<Node*>{}));
+      block->MakeNode<Tuple>(SourceInfo{}, absl::Span<Node*>{}));
   for (Next* write : writes) {
-    XLS_RETURN_IF_ERROR(ReplaceNode(block, write, empty_tuple));
+    XLS_ASSIGN_OR_RETURN(bool removed, ReplaceNode(block, write, empty_tuple));
+    XLS_RET_CHECK(removed);
   }
 
   // Write the full bit if needed.
@@ -289,18 +305,18 @@ absl::Status LowerStateElement(ScheduledBlock& block,
   return absl::OkStatus();
 }
 
-absl::StatusOr<bool> LowerStateIoForBlock(ScheduledBlock& block) {
-  XLS_RET_CHECK(block.source() != nullptr);
-  Proc& source_proc = *block.source()->AsProcOrDie();
-  if (source_proc.StateElements().empty()) {
+absl::StatusOr<bool> LowerStateIoForBlock(ScheduledBlock* block) {
+  XLS_RET_CHECK(block->source() != nullptr);
+  Proc* source_proc = block->source()->AsProcOrDie();
+  if (source_proc->StateElements().empty()) {
     return false;
   }
 
   absl::flat_hash_map<StateElement*, std::vector<StateRead*>> state_reads;
   absl::flat_hash_map<StateElement*, std::vector<Next*>> next_values;
 
-  for (int stage = 0; stage < block.stages().length(); ++stage) {
-    for (Node* node : block.stages()[stage]) {
+  for (int stage = 0; stage < block->stages().length(); ++stage) {
+    for (Node* node : block->stages()[stage]) {
       if (node->Is<StateRead>()) {
         StateRead* state_read = node->As<StateRead>();
         state_reads[state_read->state_element()].push_back(state_read);
@@ -312,17 +328,17 @@ absl::StatusOr<bool> LowerStateIoForBlock(ScheduledBlock& block) {
     }
   }
 
-  for (StateElement* state_element : source_proc.StateElements()) {
+  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[0],
                           absl::MakeSpan(next_values[state_element])));
   }
 
-  XLS_RETURN_IF_ERROR(source_proc.RemoveAllStateElements());
+  XLS_RETURN_IF_ERROR(source_proc->RemoveAllStateElements());
   return true;
 }
 
@@ -334,11 +350,11 @@ absl::StatusOr<bool> StateToRegisterIoLoweringPass::RunInternal(
   bool changed = false;
   for (FunctionBase* fb : package->GetFunctionBases()) {
     if (fb->IsBlock() && fb->IsScheduled()) {
-      ScheduledBlock& sb = *down_cast<ScheduledBlock*>(fb);
-      if (sb.source() != nullptr && sb.source()->IsProc()) {
+      ScheduledBlock* sb = down_cast<ScheduledBlock*>(fb);
+      if (sb->source() != nullptr && sb->source()->IsProc()) {
         XLS_ASSIGN_OR_RETURN(
             bool block_changed,
-            LowerStateIoForBlock(*down_cast<ScheduledBlock*>(fb)));
+            LowerStateIoForBlock(down_cast<ScheduledBlock*>(fb)));
         changed |= block_changed;
       }
     }
diff --git a/xls/ir/BUILD b/xls/ir/BUILD
index f854573097..b1a3205769 100644
--- a/xls/ir/BUILD
+++ b/xls/ir/BUILD
@@ -1761,6 +1761,7 @@ cc_library(
         "//xls/common/status:ret_check",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
diff --git a/xls/ir/function_base.cc b/xls/ir/function_base.cc
index b16f018c2b..7b7c8437bb 100644
--- a/xls/ir/function_base.cc
+++ b/xls/ir/function_base.cc
@@ -257,6 +257,12 @@ void FunctionBase::TakeOwnershipOfNode(std::unique_ptr<Node>&& node) {
   }
 
   node->function_base_ = this;
+  if (node->HasAssignedName()) {
+    std::string unique_name = UniquifyNodeName(node->GetNameView());
+    if (unique_name != node->GetName()) {
+      node->SetNameDirectly(unique_name);
+    }
+  }
   AddNodeInternal(std::move(node));
 }
 
diff --git a/xls/ir/function_base.h b/xls/ir/function_base.h
index 81eb5180fe..786ffb0965 100644
--- a/xls/ir/function_base.h
+++ b/xls/ir/function_base.h
@@ -467,7 +467,9 @@ class FunctionBase {
   // channels, or proc instantiations. The move includes the staging of the
   // logic, if this is a scheduled entity.
   void MoveLogicFrom(FunctionBase& other) {
-    MoveFrom(other, [](const Node* n) { return !n->Is<Param>(); });
+    MoveFrom(other, [](const Node* n) {
+      return !n->Is<Param>() && !n->Is<StateRead>();
+    });
     other.next_values_by_state_read_.clear();
     other.next_values_.clear();
   }
diff --git a/xls/ir/name_uniquer.cc b/xls/ir/name_uniquer.cc
index e64d6c0677..a8c3f633e6 100644
--- a/xls/ir/name_uniquer.cc
+++ b/xls/ir/name_uniquer.cc
@@ -21,10 +21,12 @@
 #include <string_view>
 
 #include "absl/container/flat_hash_set.h"
+#include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/ascii.h"
 #include "absl/strings/numbers.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
 #include "xls/common/status/ret_check.h"
 
 namespace xls {
@@ -140,6 +142,14 @@ absl::Status NameUniquer::ReleaseIdentifier(std::string_view sv) {
   }
   std::string_view root;
   std::optional<int64_t> suffix = ParseNumericSuffix(sv, separator_, &root);
+  if (!suffix.has_value()) {
+    LOG(DO_NOT_SUBMIT) << "Registered names: "
+                       << absl::StrJoin(
+                              generated_names_, ", ",
+                              [](std::string* out, const auto& entry) {
+                                absl::StrAppend(out, entry.first);
+                              });
+  }
   XLS_RET_CHECK(suffix)
       << "Name '" << sv
       << "' was not a bare identifier and didn't have a numeric suffix?";
diff --git a/xls/ir/name_uniquer.h b/xls/ir/name_uniquer.h
index f86687e121..df96ac8e8f 100644
--- a/xls/ir/name_uniquer.h
+++ b/xls/ir/name_uniquer.h
@@ -18,6 +18,7 @@
 #include <cstdint>
 #include <string>
 #include <string_view>
+#include <utility>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
@@ -40,6 +41,17 @@ class NameUniquer {
   NameUniquer(const NameUniquer&) = delete;
   NameUniquer operator=(const NameUniquer&) = delete;
 
+  NameUniquer(NameUniquer&& other) : NameUniquer(other.separator_, {}) {
+    reserved_names_.swap(other.reserved_names_);
+    generated_names_.swap(other.generated_names_);
+  }
+  NameUniquer& operator=(NameUniquer&& other) {
+    separator_ = other.separator_;
+    reserved_names_.swap(other.reserved_names_);
+    generated_names_.swap(other.generated_names_);
+    return *this;
+  }
+
   // Return a sanitized unique name which starts with the given (sanitized)
   // prefix. Names are uniqued by adding a numeric suffix if necessary separated
   // from the given prefix by "separator_". For example,
diff --git a/xls/ir/proc.cc b/xls/ir/proc.cc
index e16f1a9270..c7e610ac68 100644
--- a/xls/ir/proc.cc
+++ b/xls/ir/proc.cc
@@ -996,6 +996,10 @@ void Proc::MoveNonLogicFrom(Proc& other) {
   is_new_style_proc_ = other.is_new_style_proc_;
   state_elements_ = std::move(other.state_elements_);
   state_vec_ = std::move(other.state_vec_);
+  state_name_uniquer_ = std::move(other.state_name_uniquer_);
+  FunctionBase::MoveFrom(other,
+                         [](const Node* n) { return n->Is<StateRead>(); });
+  state_reads_ = std::move(other.state_reads_);
   channel_interfaces_ = std::move(other.channel_interfaces_);
   interface_ = std::move(other.interface_);
   proc_instantiations_ = std::move(other.proc_instantiations_);