[Proc-scoped channels] Spawn a proc with channel arrays. When figuring out the proc instantiation args, expand any channel arrays to the list of all the channel interfaces it contains.

PiperOrigin-RevId: 807775728

Author: dplassgit

xls/dslx/ir_convert/channel_scope.h

@@ -52,6 +52,18 @@ class ChannelArray {
   // For logging purposes.
   std::string ToString() const { return base_channel_name_; }
 
+  std::vector<ChannelRef> channels() const {
+    std::vector<ChannelRef> channels;
+    channels.reserve(flattened_names_in_order_.size());
+    for (auto& name : flattened_names_in_order_) {
+      std::optional<ChannelRef> channel = FindChannel(name);
+      if (channel.has_value()) {
+        channels.push_back(*channel);
+      }
+    }
+    return channels;
+  }
+
  private:
   explicit ChannelArray(std::string_view base_channel_name,
                         bool subarray = false)
@@ -69,7 +81,7 @@ class ChannelArray {
     flattened_name_to_channel_.emplace(flattened_name, channel);
   }
 
-  std::optional<ChannelRef> FindChannel(std::string_view flattened_name) {
+  std::optional<ChannelRef> FindChannel(std::string_view flattened_name) const {
     const auto it = flattened_name_to_channel_.find(flattened_name);
     if (it == flattened_name_to_channel_.end()) {
       return std::nullopt;

xls/dslx/ir_convert/function_converter.cc

@@ -480,6 +480,7 @@ absl::Status FunctionConverter::Visit(const AstNode* node) {
                                    ChannelDirectionToString(ci->direction()),
                                    ci);
           },
+          [](ChannelArray* ca) { return absl::StrFormat("%p", ca); },
       },
       value);
 }
@@ -745,6 +746,23 @@ absl::Status FunctionConverter::HandleTupleIndex(const TupleIndex* node) {
 
 absl::Status FunctionConverter::HandleXlsTuple(const XlsTuple* node) {
   VLOG(5) << "FunctionConverter::HandleXlsTuple: " << node->ToString();
+  if (current_fn_tag_ == FunctionTag::kProcConfig) {
+    std::vector<IrValue> ir_operands;
+    std::vector<BValue> b_operands;
+    for (Expr* operand : node->members()) {
+      std::optional<IrValue> v = GetNodeToIr(operand);
+      XLS_RET_CHECK(v.has_value());
+      if (std::holds_alternative<BValue>(*v)) {
+        b_operands.push_back(std::get<BValue>(*v));
+      }
+      ir_operands.push_back(*v);
+    }
+    last_tuple_ = ir_operands;
+    Def(node, [this, &b_operands](const SourceInfo& loc) {
+      return function_builder_->Tuple(b_operands, loc);
+    });
+    return absl::OkStatus();
+  }
   std::vector<BValue> operands;
   for (Expr* o : node->members()) {
     XLS_ASSIGN_OR_RETURN(BValue v, Use(o));
@@ -753,9 +771,6 @@ absl::Status FunctionConverter::HandleXlsTuple(const XlsTuple* node) {
   Def(node, [this, &operands](const SourceInfo& loc) {
     return function_builder_->Tuple(operands, loc);
   });
-  if (current_fn_tag_ == FunctionTag::kProcConfig) {
-    last_tuple_ = operands;
-  }
   return absl::OkStatus();
 }
 
@@ -850,9 +865,7 @@ absl::Status FunctionConverter::HandleNameRef(const NameRef* node) {
         } else if (std::holds_alternative<ChannelArray*>(v)) {
           VLOG(4) << "Reference to Proc member: " << k << " : Chan array : "
                   << std::get<ChannelArray*>(v)->ToString();
-          // There is no IR equivalent for a channel array. The Index nodes
-          // that refer to it have to be lowered to refer to specific channels.
-          return absl::OkStatus();
+          SetNodeToIr(from, std::get<ChannelArray*>(v));
         } else if (std::holds_alternative<Channel*>(v)) {
           VLOG(4) << "Reference to Proc member: " << k
                   << " : Chan  : " << std::get<Channel*>(v)->ToString();
@@ -2635,6 +2648,10 @@ absl::StatusOr<ChanRef> FunctionConverter::IrValueToChannelRef(
             return absl::InvalidArgumentError(
                 "Unexpected ChannelInterface in IrValue.");
           },
+          [](ChannelArray* chan) -> absl::StatusOr<ChanRef> {
+            return absl::InvalidArgumentError(
+                "Unexpected ChannelArray in IrValue.");
+          },
       },
       ir_value);
 }
@@ -3198,20 +3215,30 @@ absl::Status FunctionConverter::HandleSpawn(const Spawn* node) {
 
   const Invocation* invocation = node->config();
   XLS_RET_CHECK(package_data_.invocation_to_ir_proc.contains(invocation));
-  xls::Proc* ir_proc = package_data_.invocation_to_ir_proc[invocation];
 
-  // Lookup the channel interface for each actual arg and add it to the args
-  // vector.
   std::vector<ChannelInterface*> channel_args;
+  channel_args.reserve(invocation->args().size());
   for (Expr* arg : invocation->args()) {
     XLS_RETURN_IF_ERROR(Visit(arg));
-    std::optional<IrValue> arg_value = GetNodeToIr(arg);
-    XLS_RET_CHECK(arg_value.has_value());
-    XLS_ASSIGN_OR_RETURN(ChannelInterface * channel_interface,
-                         IrValueToChannelInterface(*arg_value));
-    channel_args.push_back(channel_interface);
+    std::optional<IrValue> arg_value_opt = GetNodeToIr(arg);
+    XLS_RET_CHECK(arg_value_opt.has_value());
+    IrValue arg_value = *arg_value_opt;
+    if (std::holds_alternative<ChannelArray*>(arg_value)) {
+      // Get all the channel interfaces of this array and
+      // add them to the channel_args
+      ChannelArray* channel_array = std::get<ChannelArray*>(arg_value);
+      for (ChannelRef channel : channel_array->channels()) {
+        XLS_RET_CHECK(std::holds_alternative<ChannelInterface*>(channel));
+        channel_args.push_back(std::get<ChannelInterface*>(channel));
+      }
+    } else {
+      XLS_ASSIGN_OR_RETURN(ChannelInterface * channel_interface,
+                           IrValueToChannelInterface(arg_value));
+      channel_args.push_back(channel_interface);
+    }
   }
 
+  xls::Proc* ir_proc = package_data_.invocation_to_ir_proc[invocation];
   xls::Proc* current_proc = builder_ptr->proc();
   XLS_RETURN_IF_ERROR(
       current_proc
@@ -3449,8 +3476,8 @@ absl::Status FunctionConverter::HandleProcNextFunction(
             proc_scoped_channel_scope_->DefineBoundaryChannelOrArray(
                 param, type_info));
         XLS_RET_CHECK(std::holds_alternative<ChannelArray*>(channel_or_array));
-        // TODO: davidplass - assign the param name to an appropriate BValue
-        // for this channel or array, for the Def call.
+        SetNodeToIr(param->name_def(),
+                    std::get<ChannelArray*>(channel_or_array));
       } else {
         ChannelType* channel_type = dynamic_cast<ChannelType*>(type);
         XLS_RET_CHECK_NE(channel_type, nullptr)
@@ -3509,28 +3536,39 @@ absl::Status FunctionConverter::HandleProcNextFunction(
     XLS_RET_CHECK_EQ(last_tuple_.size(), proc->members().size());
     int i = 0;
     for (const ProcMember* member : proc->members()) {
-      BValue tuple_entry = last_tuple_[i++];
-      Def(member, [tuple_entry](const SourceInfo& loc) { return tuple_entry; });
-
-      // Store the ChannelInterface for this entry in the id_to_members map.
-      Node* node = tuple_entry.node();
-      switch (node->op()) {
-        case Op::kSendChannelEnd: {
-          SendChannelEnd* sce = node->As<SendChannelEnd>();
-          proc_data_->id_to_members.at(proc_id)[member->identifier()] =
-              sce->channel_interface();
-          break;
-        }
-        case Op::kRecvChannelEnd: {
-          RecvChannelEnd* rce = node->As<RecvChannelEnd>();
-          proc_data_->id_to_members.at(proc_id)[member->identifier()] =
-              rce->channel_interface();
-          break;
+      IrValue tuple_entry = last_tuple_[i++];
+      if (std::holds_alternative<ChannelArray*>(tuple_entry)) {
+        ChannelArray* channel_array = std::get<ChannelArray*>(tuple_entry);
+        SetNodeToIr(member->name_def(), tuple_entry);
+        proc_data_->id_to_members.at(proc_id)[member->identifier()] =
+            channel_array;
+        XLS_RETURN_IF_ERROR(channel_scope_->AssociateWithExistingChannelOrArray(
+            *proc_id_, member->name_def(), channel_array));
+      } else {
+        BValue bvalue = std::get<BValue>(tuple_entry);
+        Def(member->name_def(),
+            [bvalue](const SourceInfo& loc) { return bvalue; });
+
+        // Store the ChannelInterface for this entry in the id_to_members map.
+        Node* node = bvalue.node();
+        switch (node->op()) {
+          case Op::kSendChannelEnd: {
+            SendChannelEnd* sce = node->As<SendChannelEnd>();
+            proc_data_->id_to_members.at(proc_id)[member->identifier()] =
+                sce->channel_interface();
+            break;
+          }
+          case Op::kRecvChannelEnd: {
+            RecvChannelEnd* rce = node->As<RecvChannelEnd>();
+            proc_data_->id_to_members.at(proc_id)[member->identifier()] =
+                rce->channel_interface();
+            break;
+          }
+          default:
+            return absl::InternalError(absl::StrFormat(
+                "Cannot process config return tuple element %d of type %s",
+                i - 1, OpToString(node->op())));
         }
-        default:
-          return absl::InternalError(absl::StrFormat(
-              "Cannot process config return tuple element %d of type %s", i - 1,
-              OpToString(node->op())));
       }
     }
   }

xls/dslx/ir_convert/function_converter.h

@@ -225,7 +225,8 @@ class FunctionConverter {
   // Every AST node has an "IR value" that is either a function builder value
   // (BValue) or its IR-conversion-time-constant-decorated cousin (CValue), or
   // an inter-proc Channel.
-  using IrValue = std::variant<BValue, CValue, Channel*, ChannelInterface*>;
+  using IrValue =
+      std::variant<BValue, CValue, Channel*, ChannelInterface*, ChannelArray*>;
 
   // Helper for converting an IR value to its BValue pointer for use in
   // debugging.
@@ -629,7 +630,7 @@ class FunctionConverter {
 
   // The last tuple converted. Used for mapping the return tuple of a proc
   // `config` method to actual proc members.
-  std::vector<BValue> last_tuple_;
+  std::vector<IrValue> last_tuple_;
 
   std::unique_ptr<ProcScopedChannelScope> proc_scoped_channel_scope_;
 };

xls/dslx/ir_convert/ir_converter_test.cc

@@ -5016,6 +5016,150 @@ pub proc main {
                HasSubstr("Cannot have non-channel parameters")));
 }
 
+TEST_P(ProcScopedChannelsIrConverterTest, ChannelArrayMembers) {
+  constexpr std::string_view kProgram = R"(
+pub proc main {
+  ins: chan<u32>[3] in;
+  outs: chan<u16>[2] out;
+
+  init { }
+  config(in_chans: chan<u32>[3] in, out_chans: chan<u16>[2] out) {
+    (in_chans, out_chans)
+  }
+  next(state: ()) {
+    send(token(), outs[0], u16:42);
+    recv(token(), ins[1]);
+    ()
+  }
+}
+)";
+
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertOneFunctionForTest(kProgram, "main", import_data,
+                                ConvertOptions{
+                                    .emit_positions = false,
+                                    .lower_to_proc_scoped_channels = true,
+                                }));
+  ExpectIr(converted);
+}
+
+TEST_P(ProcScopedChannelsIrConverterTest, SpawnFromChannelArrayParams) {
+  if (GetParam() == TypeInferenceVersion::kVersion1) {
+    // v1 fails figuring out the send argument type for some reason.
+    return;
+  }
+
+  constexpr std::string_view kProgram = R"(
+pub proc spawnee {
+  ins: chan<u32>[3] in;
+  outs: chan<u16>[2] out;
+
+  init { }
+  config(in_chans: chan<u32>[3] in, out_chans: chan<u16>[2] out) {
+    (in_chans, out_chans)
+  }
+  next(state: ()) {
+    send(token(), outs[0], u16:42);
+    recv(token(), ins[1]);
+    ()
+  }
+}
+
+pub proc main {
+  init { }
+  config(in_chans: chan<u32>[3] in, out_chans: chan<u16>[2] out) {
+    spawn spawnee(in_chans, out_chans);
+    ()
+  }
+  next(state: ()) { () }
+}
+)";
+
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertOneFunctionForTest(kProgram, "main", import_data,
+                                ConvertOptions{
+                                    .emit_positions = false,
+                                    .lower_to_proc_scoped_channels = true,
+                                }));
+  ExpectIr(converted);
+}
+
+TEST_P(ProcScopedChannelsIrConverterTest, ChannelArrayAndChannelMembers) {
+  constexpr std::string_view kProgram = R"(
+pub proc main {
+  ins: chan<u32>[3] in;
+  outs: chan<u16> out;
+
+  init { }
+  config(in_chans: chan<u32>[3] in, out_chan: chan<u16> out) {
+    (in_chans, out_chan)
+  }
+  next(state: ()) {
+    send(token(), outs, u16:42);
+    recv(token(), ins[1]);
+    ()
+  }
+}
+)";
+
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertOneFunctionForTest(kProgram, "main", import_data,
+                                ConvertOptions{
+                                    .emit_positions = false,
+                                    .lower_to_proc_scoped_channels = true,
+                                }));
+  ExpectIr(converted);
+}
+
+TEST_P(ProcScopedChannelsIrConverterTest,
+       SpawnFromChannelArrayAndChannelParams) {
+  if (GetParam() == TypeInferenceVersion::kVersion1) {
+    // v1 fails figuring out the send argument type for some reason.
+    return;
+  }
+
+  constexpr std::string_view kProgram = R"(
+pub proc spawnee {
+  ins: chan<u32>[3] in;
+  outch: chan<u16> out;
+
+  init { }
+  config(in_chans: chan<u32>[3] in, out_chan: chan<u16> out) {
+    (in_chans, out_chan)
+  }
+  next(state: ()) {
+    send(token(), outch, u16:42);
+    recv(token(), ins[1]);
+    ()
+  }
+}
+
+pub proc main {
+  init { }
+  config(in_chans: chan<u32>[3] in, out_chans: chan<u16>[2] out) {
+    spawn spawnee(in_chans, out_chans[1]);
+    ()
+  }
+  next(state: ()) { () }
+}
+)";
+
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertOneFunctionForTest(kProgram, "main", import_data,
+                                ConvertOptions{
+                                    .emit_positions = false,
+                                    .lower_to_proc_scoped_channels = true,
+                                }));
+  ExpectIr(converted);
+}
 TEST_P(ProcScopedChannelsIrConverterTest, LoopbackChannelMember) {
   constexpr std::string_view kProgram = R"(
 proc main {

xls/dslx/ir_convert/testdata/ir_converter_test_ChannelArrayAndChannelMembers.ir

@@ -0,0 +1,23 @@
+package test_module
+
+file_number 0 "test_module.x"
+
+top proc __test_module__main_0_next<in_chans__0: bits[32] in, in_chans__1: bits[32] in, in_chans__2: bits[32] in, out_chan: bits[16] out>(__state: (), init={()}) {
+  chan_interface in_chans__0(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  chan_interface in_chans__1(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  chan_interface in_chans__2(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  chan_interface out_chan(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  literal.9: token = literal(value=token, id=9)
+  literal.3: bits[1] = literal(value=1, id=3)
+  out_chan: bits[16] = send_channel_end(id=4)
+  literal.6: token = literal(value=token, id=6)
+  literal.7: bits[16] = literal(value=42, id=7)
+  receive.10: (token, bits[32]) = receive(literal.9, predicate=literal.3, channel=in_chans__1, id=10)
+  __state: () = state_read(state_element=__state, id=2)
+  tuple.12: () = tuple(id=12)
+  __token: token = literal(value=token, id=1)
+  tuple.5: (bits[16]) = tuple(out_chan, id=5)
+  send.8: token = send(literal.6, literal.7, predicate=literal.3, channel=out_chan, id=8)
+  tuple_index.11: token = tuple_index(receive.10, index=0, id=11)
+  next_value.13: () = next_value(param=__state, value=tuple.12, id=13)
+}