Add metadata mapping from ports to channels to the Block IR.

This metadata indicates which channel the ports where constructed for. Currently, this metadata is not load bearing. This is the first step in migrating this information from the channel to the block. Having the information on the block is more natural. It enables removal of channels after block creating and more naturally fits with the proc scoped channel paradigm.

Also, as part of this change add a node class PortNode which InputPort and OutputPort are derived from, and fix a minor bug having to do with port name tracking.

Serializing the metadata will be added in a followup change.

PiperOrigin-RevId: 721883772

xls/codegen/BUILD

@@ -1128,12 +1128,16 @@ cc_test(
     name = "port_legalization_pass_test",
     srcs = ["port_legalization_pass_test.cc"],
     deps = [
+        ":block_conversion",
+        ":codegen_options",
         ":codegen_pass",
         ":port_legalization_pass",
         "//xls/common:xls_gunit_main",
         "//xls/common/status:matchers",
         "//xls/ir",
         "//xls/ir:bits",
+        "//xls/ir:channel",
+        "//xls/ir:channel_ops",
         "//xls/ir:function_builder",
         "//xls/ir:ir_matcher",
         "//xls/ir:ir_test_base",

xls/codegen/block_conversion.cc

@@ -343,6 +343,15 @@ absl::Status UpdateChannelMetadata(const StreamingIOPipeline& io,
       down_cast<StreamingChannel*>(input.channel)
           ->AddBlockPortMapping(block->name(), data_name, valid_name,
                                 ready_name);
+
+      if (input.IsExternal()) {
+        XLS_RETURN_IF_ERROR(block->AddChannelPortMetadata(ChannelPortMetadata{
+            .channel_name = std::string{input.channel->name()},
+            .direction = PortDirection::kInput,
+            .data_port = input.port.value()->GetName(),
+            .valid_port = input.port_valid->GetName(),
+            .ready_port = input.port_ready->GetName()}));
+      }
     }
   }
 
@@ -364,6 +373,15 @@ absl::Status UpdateChannelMetadata(const StreamingIOPipeline& io,
       down_cast<StreamingChannel*>(output.channel)
           ->AddBlockPortMapping(block->name(), data_name, valid_name,
                                 ready_name);
+
+      if (output.IsExternal()) {
+        XLS_RETURN_IF_ERROR(block->AddChannelPortMetadata(ChannelPortMetadata{
+            .channel_name = std::string{output.channel->name()},
+            .direction = PortDirection::kOutput,
+            .data_port = output.port.value()->GetName(),
+            .valid_port = output.port_valid->GetName(),
+            .ready_port = output.port_ready->GetName()}));
+      }
     }
   }
 
@@ -373,6 +391,13 @@ absl::Status UpdateChannelMetadata(const StreamingIOPipeline& io,
 
     down_cast<SingleValueChannel*>(input.channel)
         ->AddBlockPortMapping(block->name(), input.port->name());
+
+    XLS_RETURN_IF_ERROR(block->AddChannelPortMetadata(
+        ChannelPortMetadata{.channel_name = std::string{input.channel->name()},
+                            .direction = PortDirection::kInput,
+                            .data_port = std::string{input.port->name()},
+                            .valid_port = std::nullopt,
+                            .ready_port = std::nullopt}));
   }
 
   for (const SingleValueOutput& output : io.single_value_outputs) {
@@ -381,6 +406,13 @@ absl::Status UpdateChannelMetadata(const StreamingIOPipeline& io,
 
     down_cast<SingleValueChannel*>(output.channel)
         ->AddBlockPortMapping(block->name(), output.port->name());
+
+    XLS_RETURN_IF_ERROR(block->AddChannelPortMetadata(
+        ChannelPortMetadata{.channel_name = std::string{output.channel->name()},
+                            .direction = PortDirection::kOutput,
+                            .data_port = std::string{output.port->name()},
+                            .valid_port = std::nullopt,
+                            .ready_port = std::nullopt}));
   }
 
   return absl::OkStatus();

xls/codegen/block_conversion_test.cc

@@ -91,6 +91,7 @@ using ::testing::_;
 using ::testing::AllOf;
 using ::testing::Each;
 using ::testing::ElementsAre;
+using ::testing::Eq;
 using ::testing::Ge;
 using ::testing::HasSubstr;
 using ::testing::IsEmpty;
@@ -725,6 +726,104 @@ proc my_proc(my_state: (), init={()}) {
               m::OutputPort("out", m::Neg(m::InputPort("in"))));
 }
 
+TEST_F(BlockConversionTest, StreamingChannelMetadataForSimpleProc) {
+  Package package(TestName());
+  Type* u32 = package.GetBitsType(32);
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Channel * ch_in,
+      package.CreateStreamingChannel("in", ChannelOps::kReceiveOnly, u32));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Channel * ch_out,
+      package.CreateStreamingChannel("out", ChannelOps::kSendOnly, u32));
+
+  TokenlessProcBuilder pb(TestName(), /*token_name=*/"tkn", &package);
+  BValue a = pb.Receive(ch_in);
+  pb.Send(ch_out, a);
+
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build({}));
+
+  XLS_ASSERT_OK_AND_ASSIGN(CodegenPassUnit unit,
+                           ProcToCombinationalBlock(proc, codegen_options()));
+  Block* block = FindBlock(TestName(), &package);
+
+  XLS_ASSERT_OK_AND_ASSIGN(ChannelPortMetadata in_metadata,
+                           block->GetChannelPortMetadata("in"));
+  EXPECT_EQ(in_metadata.channel_name, "in");
+  EXPECT_EQ(in_metadata.direction, PortDirection::kInput);
+  EXPECT_THAT(in_metadata.data_port, Optional(std::string{"in"}));
+  EXPECT_THAT(in_metadata.valid_port, Optional(std::string{"in_vld"}));
+  EXPECT_THAT(in_metadata.ready_port, Optional(std::string{"in_rdy"}));
+
+  EXPECT_THAT(block->GetDataPortForChannel("in"),
+              IsOkAndHolds(Optional(m::InputPort("in"))));
+  EXPECT_THAT(block->GetValidPortForChannel("in"),
+              IsOkAndHolds(Optional(m::InputPort("in_vld"))));
+  EXPECT_THAT(block->GetReadyPortForChannel("in"),
+              IsOkAndHolds(Optional(m::OutputPort("in_rdy"))));
+
+  XLS_ASSERT_OK_AND_ASSIGN(ChannelPortMetadata out_metadata,
+                           block->GetChannelPortMetadata("out"));
+  EXPECT_EQ(out_metadata.channel_name, "out");
+  EXPECT_EQ(out_metadata.direction, PortDirection::kOutput);
+  EXPECT_THAT(out_metadata.data_port, Optional(std::string{"out"}));
+  EXPECT_THAT(out_metadata.valid_port, Optional(std::string{"out_vld"}));
+  EXPECT_THAT(out_metadata.ready_port, Optional(std::string{"out_rdy"}));
+
+  EXPECT_THAT(block->GetDataPortForChannel("out"),
+              IsOkAndHolds(Optional(m::OutputPort("out"))));
+  EXPECT_THAT(block->GetValidPortForChannel("out"),
+              IsOkAndHolds(Optional(m::OutputPort("out_vld"))));
+  EXPECT_THAT(block->GetReadyPortForChannel("out"),
+              IsOkAndHolds(Optional(m::InputPort("out_rdy"))));
+}
+
+TEST_F(BlockConversionTest, SingleValueChannelMetadataForSimpleProc) {
+  Package package(TestName());
+  Type* u32 = package.GetBitsType(32);
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Channel * ch_in,
+      package.CreateSingleValueChannel("in", ChannelOps::kReceiveOnly, u32));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Channel * ch_out,
+      package.CreateSingleValueChannel("out", ChannelOps::kSendOnly, u32));
+
+  TokenlessProcBuilder pb(TestName(), /*token_name=*/"tkn", &package);
+  BValue a = pb.Receive(ch_in);
+  pb.Send(ch_out, a);
+
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build({}));
+
+  XLS_ASSERT_OK_AND_ASSIGN(CodegenPassUnit unit,
+                           ProcToCombinationalBlock(proc, codegen_options()));
+  Block* block = FindBlock(TestName(), &package);
+
+  XLS_ASSERT_OK_AND_ASSIGN(ChannelPortMetadata in_metadata,
+                           block->GetChannelPortMetadata("in"));
+  EXPECT_EQ(in_metadata.channel_name, "in");
+  EXPECT_EQ(in_metadata.direction, PortDirection::kInput);
+  EXPECT_THAT(in_metadata.data_port, Optional(std::string{"in"}));
+  EXPECT_THAT(in_metadata.valid_port, Eq(std::nullopt));
+  EXPECT_THAT(in_metadata.ready_port, Eq(std::nullopt));
+
+  EXPECT_THAT(block->GetDataPortForChannel("in"),
+              IsOkAndHolds(Optional(m::InputPort("in"))));
+  EXPECT_THAT(block->GetValidPortForChannel("in"), IsOkAndHolds(std::nullopt));
+  EXPECT_THAT(block->GetReadyPortForChannel("in"), IsOkAndHolds(std::nullopt));
+
+  XLS_ASSERT_OK_AND_ASSIGN(ChannelPortMetadata out_metadata,
+                           block->GetChannelPortMetadata("out"));
+  EXPECT_EQ(out_metadata.channel_name, "out");
+  EXPECT_EQ(out_metadata.direction, PortDirection::kOutput);
+  EXPECT_THAT(out_metadata.data_port, Optional(std::string{"out"}));
+  EXPECT_THAT(out_metadata.valid_port, Eq(std::nullopt));
+  EXPECT_THAT(out_metadata.ready_port, Eq(std::nullopt));
+
+  EXPECT_THAT(block->GetDataPortForChannel("out"),
+              IsOkAndHolds(Optional(m::OutputPort("out"))));
+  EXPECT_THAT(block->GetValidPortForChannel("out"), IsOkAndHolds(std::nullopt));
+  EXPECT_THAT(block->GetReadyPortForChannel("out"), IsOkAndHolds(std::nullopt));
+}
+
 TEST_F(BlockConversionTest, ProcWithVariousNextStateNodes) {
   // A block with corner-case next state nodes (e.g., not dependent on state
   // param, same as state param, and shared next state nodes).

xls/codegen/block_stitching_pass.cc

@@ -235,6 +235,14 @@ absl::Status ExposeStreamingOutput(
                               SourceInfo(), ext_ready, block_inst,
                               output->port_ready->GetName())
                           .status());
+
+  XLS_RETURN_IF_ERROR(block->AddChannelPortMetadata(
+      ChannelPortMetadata{.channel_name = std::string{output->channel->name()},
+                          .direction = PortDirection::kOutput,
+                          .data_port = output->port.value()->GetName(),
+                          .valid_port = output->port_valid->GetName(),
+                          .ready_port = output->port_ready->GetName()}));
+
   return absl::OkStatus();
 }
 
@@ -279,6 +287,14 @@ absl::Status ExposeStreamingInput(
           .status());
   XLS_RETURN_IF_ERROR(
       block->AddOutputPort(input->port_ready->GetName(), block_ready).status());
+
+  XLS_RETURN_IF_ERROR(block->AddChannelPortMetadata(
+      ChannelPortMetadata{.channel_name = std::string{input->channel->name()},
+                          .direction = PortDirection::kInput,
+                          .data_port = input->port.value()->GetName(),
+                          .valid_port = input->port_valid->GetName(),
+                          .ready_port = input->port_ready->GetName()}));
+
   return absl::OkStatus();
 }
 
@@ -322,6 +338,12 @@ absl::Status StitchSingleValueChannel(
     XLS_ASSIGN_OR_RETURN(
         input_node, container->AddInputPort(subblock_input->port->GetName(),
                                             subblock_input->port->GetType()));
+    XLS_RETURN_IF_ERROR(container->AddChannelPortMetadata(
+        ChannelPortMetadata{.channel_name = std::string{channel->name()},
+                            .direction = PortDirection::kInput,
+                            .data_port = subblock_input->port->GetName(),
+                            .valid_port = std::nullopt,
+                            .ready_port = std::nullopt}));
   }
   if (has_input) {
     auto subblock_inst_iter = instantiations.find(
@@ -338,8 +360,15 @@ absl::Status StitchSingleValueChannel(
   }
   channel->AddBlockPortMapping(container->name(),
                                subblock_output->port->GetName());
-  return container->AddOutputPort(subblock_output->port->GetName(), input_node)
-      .status();
+  XLS_RETURN_IF_ERROR(
+      container->AddOutputPort(subblock_output->port->GetName(), input_node)
+          .status());
+  return container->AddChannelPortMetadata(
+      ChannelPortMetadata{.channel_name = std::string{channel->name()},
+                          .direction = PortDirection::kOutput,
+                          .data_port = subblock_output->port->GetName(),
+                          .valid_port = std::nullopt,
+                          .ready_port = std::nullopt});
 }
 
 // Stitch two ends of a streaming channel together.

xls/codegen/block_stitching_pass_test.cc

@@ -277,6 +277,55 @@ TEST_F(BlockStitchingPassTest, StitchNetworkWithFifos) {
           m::InstantiationOutput("ch2_valid", Eq(inst0)),
           m::InstantiationOutput("ch1_ready", Eq(inst0)),
       }));
+
+  // Verify channel port metadata.
+  XLS_ASSERT_OK_AND_ASSIGN(Block * top_block, p->GetBlock("top_proc"));
+  EXPECT_THAT(top_block->GetChannelNames(), UnorderedElementsAre("ch0", "ch2"));
+  EXPECT_THAT(top_block->GetDataPortForChannel("ch0"),
+              IsOkAndHolds(Optional(m::InputPort("ch0_data"))));
+  EXPECT_THAT(top_block->GetValidPortForChannel("ch0"),
+              IsOkAndHolds(Optional(m::InputPort("ch0_valid"))));
+  EXPECT_THAT(top_block->GetReadyPortForChannel("ch0"),
+              IsOkAndHolds(Optional(m::OutputPort("ch0_ready"))));
+
+  EXPECT_THAT(top_block->GetDataPortForChannel("ch2"),
+              IsOkAndHolds(Optional(m::OutputPort("ch2_data"))));
+  EXPECT_THAT(top_block->GetValidPortForChannel("ch2"),
+              IsOkAndHolds(Optional(m::OutputPort("ch2_valid"))));
+  EXPECT_THAT(top_block->GetReadyPortForChannel("ch2"),
+              IsOkAndHolds(Optional(m::InputPort("ch2_ready"))));
+
+  XLS_ASSERT_OK_AND_ASSIGN(Block * block0, p->GetBlock("top_proc__1"));
+  EXPECT_THAT(block0->GetChannelNames(), UnorderedElementsAre("ch0", "ch1"));
+  EXPECT_THAT(block0->GetDataPortForChannel("ch0"),
+              IsOkAndHolds(Optional(m::InputPort("ch0_data"))));
+  EXPECT_THAT(block0->GetValidPortForChannel("ch0"),
+              IsOkAndHolds(Optional(m::InputPort("ch0_valid"))));
+  EXPECT_THAT(block0->GetReadyPortForChannel("ch0"),
+              IsOkAndHolds(Optional(m::OutputPort("ch0_ready"))));
+
+  EXPECT_THAT(block0->GetDataPortForChannel("ch1"),
+              IsOkAndHolds(Optional(m::OutputPort("ch1_data"))));
+  EXPECT_THAT(block0->GetValidPortForChannel("ch1"),
+              IsOkAndHolds(Optional(m::OutputPort("ch1_valid"))));
+  EXPECT_THAT(block0->GetReadyPortForChannel("ch1"),
+              IsOkAndHolds(Optional(m::InputPort("ch1_ready"))));
+
+  XLS_ASSERT_OK_AND_ASSIGN(Block * block1, p->GetBlock(proc1->name()));
+  EXPECT_THAT(block1->GetChannelNames(), UnorderedElementsAre("ch1", "ch2"));
+  EXPECT_THAT(block1->GetDataPortForChannel("ch1"),
+              IsOkAndHolds(Optional(m::InputPort("ch1_data"))));
+  EXPECT_THAT(block1->GetValidPortForChannel("ch1"),
+              IsOkAndHolds(Optional(m::InputPort("ch1_valid"))));
+  EXPECT_THAT(block1->GetReadyPortForChannel("ch1"),
+              IsOkAndHolds(Optional(m::OutputPort("ch1_ready"))));
+
+  EXPECT_THAT(block1->GetDataPortForChannel("ch2"),
+              IsOkAndHolds(Optional(m::OutputPort("ch2_data"))));
+  EXPECT_THAT(block1->GetValidPortForChannel("ch2"),
+              IsOkAndHolds(Optional(m::OutputPort("ch2_valid"))));
+  EXPECT_THAT(block1->GetReadyPortForChannel("ch2"),
+              IsOkAndHolds(Optional(m::InputPort("ch2_ready"))));
 }
 
 TEST_F(BlockStitchingPassTest, StitchNetworkWithDirectConnections) {
@@ -5082,6 +5131,7 @@ proc output_passthrough(state:bits[1], init={0}) {
       EvalBlock(top_block, unit.metadata, {{"in", {5, 10}}}),
       IsOkAndHolds(BlockOutputsEq({{"out0", {5, 10}}, {"out1", {5, 10}}})));
 }
+
 TEST_F(ProcInliningPassTest, NestedProcsUsingEmptyAfterAll) {
   // Uses tokens created using empty after_all.
   constexpr std::string_view ir_text = R"(package test

xls/codegen/port_legalization_pass.cc

@@ -15,6 +15,7 @@
 #include "xls/codegen/port_legalization_pass.h"
 
 #include <memory>
+#include <string>
 #include <variant>
 #include <vector>
 
@@ -72,11 +73,16 @@ absl::StatusOr<bool> PortLegalizationPass::RunInternal(
         InputPort* input_port = std::get<InputPort*>(port);
         if (input_port->GetType()->GetFlatBitCount() == 0) {
           VLOG(4) << "Removing zero-width input port " << input_port->name();
-          XLS_RETURN_IF_ERROR(input_port
-                                  ->ReplaceUsesWithNew<xls::Literal>(
-                                      ZeroOfType(input_port->GetType()))
-                                  .status());
+          XLS_ASSIGN_OR_RETURN(Node * replacement,
+                               input_port->ReplaceUsesWithNew<xls::Literal>(
+                                   ZeroOfType(input_port->GetType())));
+
+          // Ports have durable names which are not changed by the replacement
+          // process, so rename the replacement after removing the port.
+          std::string name = input_port->GetName();
           XLS_RETURN_IF_ERROR(block->RemoveNode(input_port));
+          replacement->SetName(name);
+
           changed = true;
         }
       } else if (std::holds_alternative<OutputPort*>(port)) {