[Proc-scoped channels] Support spawning the same proc with the same parametrics but different channels. This was done by getting all instantiations of the config function for a given parametric from TypeInfo. These are now keyed on both the instantiation and the parametric env; otherwise it was not picking the right ones.

xls/dslx/ir_convert/function_converter.cc

@@ -507,7 +507,8 @@ FunctionConverter::FunctionConverter(PackageData& package_data, Module* module,
                      : xls::Fileno(0)),
       proc_data_(proc_data),
       channel_scope_(channel_scope),
-      is_top_(is_top) {
+      is_top_(is_top),
+      name_uniquer_("_") {
   VLOG(5) << "Constructed IR converter: " << this;
 }
 
@@ -3371,14 +3372,21 @@ absl::Status FunctionConverter::HandleSpawn(const Spawn* node) {
     }
   }
 
-  XLS_RET_CHECK(package_data_.invocation_to_ir_proc.contains(invocation));
-  xls::Proc* ir_proc = package_data_.invocation_to_ir_proc[invocation];
+  XLS_RET_CHECK(package_data_.invocation_to_ir_proc.contains(
+      {invocation, GetParametricEnv()}))
+      << invocation->ToString() << " env " << GetParametricEnv();
+  xls::Proc* ir_proc =
+      package_data_.invocation_to_ir_proc[{invocation, GetParametricEnv()}];
   xls::Proc* current_proc = builder_ptr->proc();
-  XLS_RETURN_IF_ERROR(
-      current_proc
-          ->AddProcInstantiation(absl::StrFormat("%s_inst", ir_proc->name()),
-                                 absl::MakeSpan(channel_args), ir_proc)
-          .status());
+  // Give each proc_instantiation a unique name. Since this method may be called
+  // multiple times per config method with the same uniquifier, this guarantees
+  // uniqueness within a proc, but not across oprocs, which is acceptable.
+  XLS_RETURN_IF_ERROR(current_proc
+                          ->AddProcInstantiation(
+                              name_uniquer_.GetSanitizedUniqueName(
+                                  absl::StrFormat("%s_inst", ir_proc->name())),
+                              absl::MakeSpan(channel_args), ir_proc)
+                          .status());
   // Spawn is an Invocation, which is an Instantiation, which is
   // technically an Expr, so it needs to have an Ir value in the map, even
   // though that value is never read.
@@ -3734,10 +3742,22 @@ absl::Status FunctionConverter::HandleProcNextFunction(
   package_data_.ir_to_dslx[p] = f;
   if (record.config_record() != nullptr) {
     // The invocation is actually the "next", but we need the "config".
-    package_data_.invocation_to_ir_proc[record.config_record()->invocation()] =
-        p;
+    // Get all the invocations of the 'config' function given the current
+    // parametrics.
+    for (const auto& invocation_callee_data :
+         record.type_info()->GetInvocationCalleeData(
+             record.config_record()->f(), record.parametric_env())) {
+      package_data_
+          .invocation_to_ir_proc[{invocation_callee_data.invocation,
+                                  invocation_callee_data.caller_bindings}] = p;
+    }
+    // Always use the given invocation, because it may be outside this module,
+    // and the call to GetInvocationCalleeData won't return it.
+    package_data_.invocation_to_ir_proc[{record.config_record()->invocation(),
+                                         record.parametric_env()}] = p;
   }
-  package_data_.invocation_to_ir_proc[invocation] = p;
+  package_data_.invocation_to_ir_proc[{invocation, record.parametric_env()}] =
+      p;
   return absl::OkStatus();
 }
 

xls/dslx/ir_convert/function_converter.h

@@ -51,6 +51,7 @@
 #include "xls/ir/fileno.h"
 #include "xls/ir/function.h"
 #include "xls/ir/function_builder.h"
+#include "xls/ir/name_uniquer.h"
 #include "xls/ir/nodes.h"
 #include "xls/ir/package.h"
 #include "xls/ir/source_location.h"
@@ -95,7 +96,8 @@ absl::StatusOr<xls::Function*> EmitImplicitTokenEntryWrapper(
 struct PackageData {
   PackageConversionData* conversion_info;
   absl::flat_hash_map<xls::FunctionBase*, dslx::Function*> ir_to_dslx;
-  absl::flat_hash_map<const dslx::Invocation*, xls::Proc*>
+  absl::flat_hash_map<std::pair<const dslx::Invocation*, ParametricEnv>,
+                      xls::Proc*>
       invocation_to_ir_proc;
   absl::flat_hash_set<xls::Function*> wrappers;
 };
@@ -634,6 +636,10 @@ class FunctionConverter {
   // The last tuple converted. Used for mapping the return tuple of a proc
   // `config` method to actual proc members.
   std::vector<IrValue> last_tuple_;
+
+  // Used to uniquify proc instantiation node names when lowering to
+  // proc-scoped channels
+  NameUniquer name_uniquer_;
 };
 
 }  // namespace xls::dslx

xls/dslx/ir_convert/ir_converter_test.cc

@@ -4642,19 +4642,19 @@ pub proc main {
 TEST_P(ProcScopedChannelsIrConverterTest,
        ParametricMultipleSpawnDifferentParametrics) {
   constexpr std::string_view program = R"(
-proc spawnee1<M:u32> {
+proc second<M:u32> {
   init { }
   config(sender: chan<uN[M]> out, receiver: chan<uN[M]> in) { () }
   next(_: ()) { () }
 }
 
-pub proc main {
+pub proc first {
   init { }
   config() {
-    let (sender, receiver) = chan<u1>("data_0");
-    spawn spawnee1<u32:1>(sender, receiver);
-    let (sender, receiver) = chan<u2>("data_1");
-    spawn spawnee1<u32:2>(sender, receiver);
+    let (sender, receiver) = chan<u16>("data_0");
+    spawn second<u32:16>(sender, receiver);
+    let (sender, receiver) = chan<u64>("data_1");
+    spawn second<u32:64>(sender, receiver);
     ()
   }
   next(_: ()) { () }
@@ -4663,7 +4663,36 @@ pub proc main {
   auto import_data = CreateImportDataForTest();
   XLS_ASSERT_OK_AND_ASSIGN(
       std::string converted,
-      ConvertOneFunctionForTest(program, "main", import_data,
+      ConvertOneFunctionForTest(program, "first", import_data,
+                                kProcScopedChannelOptions));
+  ExpectIr(converted);
+}
+
+TEST_P(ProcScopedChannelsIrConverterTest,
+       ParametricMultipleSpawnSameParametrics) {
+  constexpr std::string_view program = R"(
+proc second<M:u32> {
+  init { }
+  config(sender: chan<uN[M]> out, receiver: chan<uN[M]> in) { () }
+  next(_: ()) { () }
+}
+
+pub proc first {
+  init { }
+  config() {
+    let (sender, receiver) = chan<u16>("data_0");
+    spawn second<u32:16>(sender, receiver);
+    let (sender, receiver) = chan<u16>("data_1");
+    spawn second<u32:16>(sender, receiver);
+    ()
+  }
+  next(_: ()) { () }
+})";
+
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertOneFunctionForTest(program, "first", import_data,
                                 kProcScopedChannelOptions));
   ExpectIr(converted);
 }
@@ -4672,26 +4701,26 @@ TEST(ProcScopedChannelsIrConverterTest, SpawnTree) {
   // This is a more complex case than the above test, since the spawn tree
   // needs to be "expanded" for each call of the intermediate proc.
   constexpr std::string_view program = R"(
-proc spawnee2<N:u32, M:u32> {
+proc third<N:u32, M:u32> {
   init { }
   config() { () }
   next(state: ()) { () }
 }
 
-proc spawnee1<N:u32> {
+proc second<N:u32> {
   init { }
   config() {
-    spawn spawnee2<N, u32:1>();
-    spawn spawnee2<N, u32:2>();
+    spawn third<N, u32:16>();
+    spawn third<N, u32:32>();
   }
   next(state: ()) { () }
 }
 
-pub proc main {
+pub proc first {
   init { }
   config() {
-    spawn spawnee1<u32:3>();
-    spawn spawnee1<u32:4>();
+    spawn second<u32:64>();
+    spawn second<u32:128>();
   }
   next(state: ()) { () }
 }
@@ -4700,7 +4729,41 @@ pub proc main {
   auto import_data = CreateImportDataForTest();
   XLS_ASSERT_OK_AND_ASSIGN(
       std::string converted,
-      ConvertOneFunctionForTest(program, "main", import_data,
+      ConvertOneFunctionForTest(program, "first", import_data,
+                                kProcScopedChannelOptions));
+  ExpectIr(converted);
+}
+
+TEST(ProcScopedChannelsIrConverterTest, ParametricNetwork) {
+  constexpr std::string_view program = R"(
+proc third<N:u32> {
+  init { }
+  config() { () }
+  next(state: ()) { () }
+}
+
+proc second<N:u32> {
+  init { }
+  config() {
+    spawn third<N>();
+  }
+  next(state: ()) { () }
+}
+
+pub proc first {
+  init { }
+  config() {
+    spawn second<u32:16>();
+    spawn second<u32:32>();
+  }
+  next(state: ()) { () }
+}
+)";
+
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertOneFunctionForTest(program, "first", import_data,
                                 kProcScopedChannelOptions));
   ExpectIr(converted);
 }
@@ -5735,15 +5798,73 @@ proc main {
   XLS_ASSERT_OK_AND_ASSIGN(
       std::string converted,
       ConvertOneFunctionForTest(program, "main", import_data,
-                                ConvertOptions{
-                                    .emit_positions = false,
-                                    .lower_to_proc_scoped_channels = true,
-                                }));
+                                kProcScopedChannelOptions));
+  ExpectIr(converted);
+}
+
+TEST_P(ProcScopedChannelsIrConverterTest,
+       ProcScopedChannelAttributesMultipleSpawns) {
+  constexpr std::string_view program = R"(#![feature(channel_attributes)]
+proc producer {
+  c: chan<u32> out;
+  init {
+    u32:0
+  }
+  config(input_c: chan<u32> out) {
+    (input_c,)
+  }
+  next(i: u32) {
+    let tok = send(join(), c, i);
+    i + u32:1
+  }
+}
+
+proc consumer {
+  c: chan<u32> in;
+  init {
+    u32:0
+  }
+  config(input_c: chan<u32> in) {
+    (input_c,)
+  }
+  next(i: u32) {
+    let (tok, i) = recv(join(), c);
+    i + i
+  }
+}
+
+proc main {
+  init { () }
+  config() {
+    let (p, c) =
+    #[channel(depth=0)]
+    chan<u32>("my_chan0");
+    spawn producer(p);
+    spawn consumer(c);
+
+    let (p, c) =
+    #[channel(depth=1, register_push_outputs=true, register_pop_outputs=true, bypass=false)]
+    chan<u32>("my_chan1");
+    spawn producer(p);
+    spawn consumer(c);
+
+    let (p, c) =
+    #[channel(depth=0, input_flop_kind=zero_latency, output_flop_kind=flop)]
+    chan<u32>("my_chan2");
+    spawn producer(p);
+    spawn consumer(c);
+    ()
+  }
+  next(state: ()) { () }
+}
+)";
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertModuleForTest(program, kProcScopedChannelOptions));
   ExpectIr(converted);
 }
 
-// This is not supported yet.
-TEST_P(ProcScopedChannelsIrConverterTest, DISABLED_MultipleSpawns) {
+TEST_P(ProcScopedChannelsIrConverterTest, MultipleSpawnsNoParametric) {
   constexpr std::string_view program = R"(
 proc producer {
   c: chan<u32> out;

xls/dslx/ir_convert/testdata/ir_converter_test_MultipleSpawnsNoParametric.ir

@@ -0,0 +1,36 @@
+package test_module
+
+file_number 0 "test_module.x"
+
+proc __test_module__producer_0_next<output_c: bits[32] out>(__state: bits[32], init={0}) {
+  chan_interface output_c(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  __state: bits[32] = state_read(state_element=__state, id=2)
+  literal.8: bits[32] = literal(value=1, id=8)
+  output_c: bits[32] = send_channel_end(id=4)
+  after_all.6: token = after_all(id=6)
+  literal.3: bits[1] = literal(value=1, id=3)
+  add.9: bits[32] = add(__state, literal.8, id=9)
+  __token: token = literal(value=token, id=1)
+  tuple.5: (bits[32]) = tuple(output_c, id=5)
+  tok: token = send(after_all.6, __state, predicate=literal.3, channel=output_c, id=7)
+  next_value.10: () = next_value(param=__state, value=add.9, id=10)
+}
+
+top proc __test_module__main_0_next<>(__state: (), init={()}) {
+  chan my_chan0(bits[32], id=0, kind=streaming, ops=send_receive, flow_control=ready_valid, strictness=proven_mutually_exclusive)
+  chan_interface my_chan0(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=none, flop_kind=none)
+  chan_interface my_chan0(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=none, flop_kind=none)
+  chan my_chan1(bits[32], id=1, kind=streaming, ops=send_receive, flow_control=ready_valid, strictness=proven_mutually_exclusive)
+  chan_interface my_chan1(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=none, flop_kind=none)
+  chan_interface my_chan1(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=none, flop_kind=none)
+  proc_instantiation __test_module__producer_0_next_inst(my_chan0, proc=__test_module__producer_0_next)
+  proc_instantiation __test_module__producer_0_next_inst_1(my_chan1, proc=__test_module__producer_0_next)
+  __state: () = state_read(state_element=__state, id=12)
+  tuple.17: () = tuple(id=17)
+  __token: token = literal(value=token, id=11)
+  literal.13: bits[1] = literal(value=1, id=13)
+  my_chan0: bits[32] = send_channel_end(id=14)
+  my_chan1: bits[32] = send_channel_end(id=15)
+  tuple.16: () = tuple(id=16)
+  next_value.18: () = next_value(param=__state, value=tuple.17, id=18)
+}