Support IR conversion of a simple parametric ProcDef.

PiperOrigin-RevId: 919066739

Author: richmckeever

xls/dslx/ir_convert/conversion_record.cc

@@ -104,9 +104,10 @@ std::string ConversionRecord::ToString() const {
   }
   return absl::StrFormat(
       "ConversionRecord{m=%s, f=%s, top=%s, pid=%s, "
-      "parametric_env=%s, type_info=%p, config=%s impl=%s}",
+      "parametric_env=%s, type_info=%s, config=%s impl=%s}",
       module_->name(), f_->identifier(), is_top_ ? "true" : "false", proc_id,
-      parametric_env_.ToString(), type_info_, config,
+      parametric_env_.ToString(),
+      type_info_ == nullptr ? "<none>" : type_info_->name(), config,
       f_->impl().has_value() ? (*f_->impl())->struct_ref()->ToString()
                              : "none");
 }

xls/dslx/ir_convert/function_converter.cc

@@ -3766,16 +3766,21 @@ absl::Status FunctionConverter::HandleChannelDecl(const ChannelDecl* node) {
   return absl::OkStatus();
 }
 
-absl::Status FunctionConverter::InitProcDefBuilder(const ProcDef* proc_def) {
-  // TODO: https://github.com/google/xls/issues/4125 - Consider the parametrics
-  // and init values when invoking `MangleDslxName` here, using
-  // `HandleProcNextFunction` as a rough guide. ProcDef support is a WIP and we
-  // don't yet support parametrics.
-  ParametricEnv bindings;
-  XLS_ASSIGN_OR_RETURN(std::string mangled_name,
-                       MangleDslxName(module_->name(), proc_def->identifier(),
-                                      CallingConvention::kProcNext,
-                                      /*free_keys=*/{}, &bindings));
+absl::Status FunctionConverter::InitProcDefBuilder(const ProcDef* proc_def,
+                                                   const ParametricEnv& env) {
+  absl::btree_set<std::string> parametric_keys;
+
+  // Include parametric values in the mangled names of non-top procs only.
+  if (!is_top_) {
+    for (const ParametricBinding* binding : proc_def->parametric_bindings()) {
+      parametric_keys.insert(binding->identifier());
+    }
+  }
+
+  XLS_ASSIGN_OR_RETURN(
+      std::string mangled_name,
+      MangleDslxName(module_->name(), proc_def->identifier(),
+                     CallingConvention::kProcNext, parametric_keys, &env));
   auto unique_builder =
       std::make_unique<ProcBuilder>(NewStyleProc{}, mangled_name, package());
   ProcBuilder* builder = unique_builder.get();
@@ -4139,26 +4144,29 @@ absl::Status FunctionConverter::AddProcDefInstantiation(
 
 absl::Status FunctionConverter::ConvertProcDef(
     const ProcDef* proc_def, const InterpValue& canonical_initializer_value,
-    ProcId proc_id, TypeInfo* type_info) {
-  VLOG(5) << "Converting ProcDef: " << proc_def->identifier();
+    ProcId proc_id, TypeInfo* constructor_ti, TypeInfo* next_ti,
+    const ParametricEnv& env) {
+  VLOG(5) << "Converting ProcDef: " << proc_def->identifier()
+          << " using constructor TI " << constructor_ti->name()
+          << " and next TI " << next_ti->name();
 
   InterpValue::ProcInitializer canonical_initializer =
       canonical_initializer_value.GetProcInitializerOrDie();
-  ScopedTypeInfoSwap stis(this, type_info);
+  ScopedTypeInfoSwap stis(this, constructor_ti);
   proc_id_ = proc_id;
   proc_data_->id_to_members[proc_id] = {};
 
-  XLS_RETURN_IF_ERROR(InitProcDefBuilder(proc_def));
+  XLS_RETURN_IF_ERROR(InitProcDefBuilder(proc_def, env));
   XLS_RETURN_IF_ERROR(InitProcDefChannels(proc_def, canonical_initializer));
   XLS_RETURN_IF_ERROR(
       InitProcDefStateElements(proc_def, canonical_initializer));
 
   XLS_ASSIGN_OR_RETURN(std::vector<InterpValue> spawnees,
-                       type_info->GetProcDefSpawnsFrom(proc_def));
+                       constructor_ti->GetProcDefSpawnsFrom(proc_def));
   for (const InterpValue& spawnee : spawnees) {
-    const AstNode* definer = spawnee.GetProcInitializerOrDie().definer();
-    XLS_ASSIGN_OR_RETURN(InterpValue canonical_initializer,
-                         current_type_info_->GetConstExpr(definer));
+    XLS_ASSIGN_OR_RETURN(
+        InterpValue canonical_initializer,
+        current_type_info_->GetCanonicalProcInitializer(spawnee));
     XLS_RETURN_IF_ERROR(
         AddProcDefInstantiation(proc_def, spawnee, canonical_initializer));
   }
@@ -4168,7 +4176,9 @@ absl::Status FunctionConverter::ConvertProcDef(
 
   XLS_ASSIGN_OR_RETURN(std::unique_ptr<ProcDefInstance> instance,
                        CreateProcDefInstance(proc_def));
+
   SetNodeToIr((*next_fn)->params()[0]->name_def(), instance.get());
+  ScopedTypeInfoSwap stis_next(this, next_ti);
   XLS_RETURN_IF_ERROR(Visit((*next_fn)->body()));
 
   XLS_ASSIGN_OR_RETURN(
@@ -4179,7 +4189,7 @@ absl::Status FunctionConverter::ConvertProcDef(
   // structure to support an InterpValue::ProcInitializer (or its whole
   // contents + env) as a key.
   package_data_.ir_to_dslx[p] = *next_fn;
-  package_data_.callee_to_ir_proc[{*next_fn, ParametricEnv{}}] = p;
+  package_data_.callee_to_ir_proc[{*next_fn, env}] = p;
   package_data_
       .canonical_proc_def_initializer_to_ir_proc[canonical_initializer_value] =
       p;

xls/dslx/ir_convert/function_converter.h

@@ -161,7 +161,8 @@ class FunctionConverter {
 
   absl::Status ConvertProcDef(const ProcDef* proc_def,
                               const InterpValue& canonical_initializer_value,
-                              ProcId proc_id, TypeInfo* type_info);
+                              ProcId proc_id, TypeInfo* constructor_ti,
+                              TypeInfo* next_ti, const ParametricEnv& env);
 
   // Notes a constant-definition dependency for the function (so it can
   // participate in the IR conversion).
@@ -544,7 +545,8 @@ class FunctionConverter {
   absl::Status HandleProcDef(const ProcDef* proc_def,
                              const Function* constructor);
 
-  absl::Status InitProcDefBuilder(const ProcDef* proc_def);
+  absl::Status InitProcDefBuilder(const ProcDef* proc_def,
+                                  const ParametricEnv& env);
 
   template <typename NodeType>
   absl::Status DefineProcDefChannelOrArray(

xls/dslx/ir_convert/get_conversion_records.cc

@@ -63,14 +63,27 @@ absl::StatusOr<ConversionRecord> MakeConversionRecord(
 
 // Variant for a `ProcDef` as opposed to a legacy `Proc`.
 absl::StatusOr<ConversionRecord> MakeConversionRecord(
-    const ProcDef* p, ProcId proc_id, TypeInfo* ti, ParametricEnv env, bool top,
-    std::optional<InterpValue> initializer = std::nullopt) {
+    const ProcDef* p, ProcId proc_id, bool top,
+    const ProcInitializerWithTypeInfo& canonical_initializer) {
   std::optional<Function*> next_fn = GetProcNextFunction(p);
   XLS_RET_CHECK(next_fn.has_value());
-  return ConversionRecord::Make(*next_fn, p->owner(), ti, std::move(env),
-                                proc_id, top,
-                                /*config_record=*/nullptr,
-                                /*init_value=*/initializer, p);
+
+  // TODO: https://github.com/google/xls/issues/4125 - We are only using the
+  // "config record" here as a place to store the TypeInfo for the constructor.
+  // Legacy procs need more of the record. We should get rid of the separate
+  // config record concept once we get rid of legacy procs.
+  XLS_ASSIGN_OR_RETURN(
+      ConversionRecord config_record,
+      ConversionRecord::Make(*next_fn, p->owner(),
+                             canonical_initializer.constructor_type_info,
+                             canonical_initializer.constructor_env, proc_id,
+                             false, nullptr, std::nullopt, p));
+
+  return ConversionRecord::Make(
+      *next_fn, p->owner(), canonical_initializer.next_type_info,
+      canonical_initializer.constructor_env, proc_id, top,
+      std::make_unique<ConversionRecord>(std::move(config_record)),
+      /*init_value=*/canonical_initializer.initializer, p);
 }
 
 // An AstNodeVisitor that creates ConversionRecords from appropriate AstNodes
@@ -244,7 +257,8 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
       // We can still convert this function even though it's never been called.
       // Make sure we are using the right type info for imported functions.
       TypeInfo* invocation_ti = GetTypeInfo(f);
-      VLOG(5) << "Processing fn " << f->ToString();
+      VLOG(5) << "Processing fn " << f->ToString() << " with TI "
+              << invocation_ti->name();
       XLS_ASSIGN_OR_RETURN(ConversionRecord cr,
                            MakeConversionRecord(const_cast<Function*>(f),
                                                 f->owner(), invocation_ti,
@@ -373,24 +387,29 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
     XLS_RET_CHECK(next_fn.has_value());
 
     TypeInfo* proc_owner_ti = GetTypeInfo(p);
-    XLS_ASSIGN_OR_RETURN(std::vector<InterpValue> canonical_initializers,
-                         proc_owner_ti->GetCanonicalProcInitializers(p));
+    XLS_ASSIGN_OR_RETURN(
+        std::vector<ProcInitializerWithTypeInfo> canonical_initializers,
+        proc_owner_ti->GetCanonicalProcInitializers(p));
     if (p->IsParametric() && canonical_initializers.empty()) {
       VLOG(5) << "No calls to parametric proc " << p->name_def()->ToString();
       return absl::OkStatus();
     }
-    for (const InterpValue& canonical_initializer : canonical_initializers) {
+    for (const ProcInitializerWithTypeInfo& canonical_initializer :
+         canonical_initializers) {
       // TODO: https://github.com/google/xls/issues/4125 - Exclude test-only
       // procs, and those only used in test-only contexts, if desired.
 
       VLOG(5)
           << "Making conversion record for canonical initializer of ProcDef: "
-          << canonical_initializer.ToString();
+          << canonical_initializer.initializer.ToString()
+          << " with constructor TI "
+          << canonical_initializer.constructor_type_info->name()
+          << " and next() TI " << canonical_initializer.next_type_info->name();
 
       XLS_ASSIGN_OR_RETURN(
           ConversionRecord cr,
-          MakeConversionRecord(p, proc_id_factory_->CreateProcId(p), type_info_,
-                               ParametricEnv(), /*top=*/top_ == *next_fn,
+          MakeConversionRecord(p, proc_id_factory_->CreateProcId(p),
+                               /*top=*/top_ == *next_fn,
                                canonical_initializer));
       records_.push_back(std::move(cr));
     }

xls/dslx/ir_convert/ir_converter.cc

@@ -190,6 +190,8 @@ absl::Status ConvertOneFunctionInternal(PackageData& package_data,
                                         ProcConversionData* proc_data,
                                         ChannelScope* channel_scope,
                                         const ConvertOptions& options) {
+  VLOG(6) << "Convert record: " << record.ToString();
+
   // Validate the requested conversion looks sound in terms of provided
   // parametrics.
   XLS_RETURN_IF_ERROR(ConversionRecord::ValidateParametrics(
@@ -206,8 +208,10 @@ absl::Status ConvertOneFunctionInternal(PackageData& package_data,
   }
 
   if (record.proc_def().has_value()) {
-    return converter.ConvertProcDef(*record.proc_def(), *record.init_value(),
-                                    *record.proc_id(), record.type_info());
+    return converter.ConvertProcDef(
+        *record.proc_def(), *record.init_value(), *record.proc_id(),
+        record.config_record()->type_info(), record.type_info(),
+        record.parametric_env());
   }
 
   Function* f = record.f();

xls/dslx/ir_convert/ir_converter_test.cc

@@ -2444,6 +2444,70 @@ impl Main {
   ExpectIr(conv.DumpIr());
 }
 
+TEST_F(IrConverterTest, ParametricProcDef) {
+  constexpr std::string_view kModule = R"(
+#![feature(explicit_state_access)]
+
+proc Loopback<N: u32> {
+  c_in: chan<uN[N]> in,
+  c_out: chan<uN[N]> out,
+}
+
+impl Loopback<N> {
+  fn new(c_in: chan<uN[N]> in, c_out: chan<uN[N]> out) -> Self {
+    Loopback { c_in, c_out }
+  }
+
+  fn next(self) {
+    let (t, val) = recv(join(), self.c_in);
+    send(t, self.c_out, val);
+  }
+}
+
+proc Main {
+  c_in: chan<u32> in,
+  c_out: chan<u32> out,
+  c_in_from_loopback: chan<u32> in,
+  c_out_to_loopback: chan<u32> out,
+  i: u32,
+}
+
+impl Main {
+  fn new(c_in: chan<u32> in, c_out: chan<u32> out) -> Self {
+    let (out_to_loopback, loopback_in) = chan<u32>("main_to_loopback");
+    let (loopback_out, in_from_loopback) = chan<u32>("loopback_to_main");
+    Loopback<32>::new(loopback_in, loopback_out).spawn();
+
+    Main {
+      c_in: c_in,
+      c_out: c_out,
+      c_in_from_loopback: in_from_loopback,
+      c_out_to_loopback: out_to_loopback,
+      i: 1
+    }
+  }
+
+  fn next(self) {
+    let i_val = read(self.i);
+    let (_, j) = recv(join(), self.c_in);
+    let loopback_tok = send(join(), self.c_out_to_loopback, j);
+    let (_, loopback_val) = recv(loopback_tok, self.c_in_from_loopback);
+    send(join(), self.c_out, i_val + loopback_val);
+    write(self.i, i_val + loopback_val);
+  }
+}
+)";
+
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      TypecheckedModule tm,
+      ParseAndTypecheck(kModule, "test_module.x", "test_module", &import_data));
+  XLS_ASSERT_OK_AND_ASSIGN(PackageConversionData conv,
+                           ConvertModuleToPackage(tm.module, &import_data,
+                                                  kProcScopedChannelOptions));
+  ExpectIr(conv.DumpIr());
+}
+
 TEST_F(IrConverterTest, TopProcDefWithIndrectConstructorResult) {
   constexpr std::string_view program = R"(
 #![feature(explicit_state_access)]

xls/dslx/ir_convert/testdata/ir_converter_test_ParametricProcDef.ir

@@ -0,0 +1,68 @@
+package test_module
+
+file_number 0 "test_module.x"
+
+proc __test_module__Loopback__32_next<_c_in: bits[32] in, _c_out: bits[32] out>() {
+  chan_interface _c_in(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  chan_interface _c_out(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  after_all.3: token = after_all(id=3)
+  literal.2: bits[1] = literal(value=1, id=2)
+  receive.4: (token, bits[32]) = receive(after_all.3, predicate=literal.2, channel=_c_in, id=4)
+  t: token = tuple_index(receive.4, index=0, id=6)
+  val: bits[32] = tuple_index(receive.4, index=1, id=7)
+  __token: token = literal(value=token, id=1)
+  tuple_index.5: token = tuple_index(receive.4, index=0, id=5)
+  send.8: token = send(t, val, predicate=literal.2, channel=_c_out, id=8)
+  tuple.9: () = tuple(id=9)
+}
+
+proc __test_module__Main_next<_c_in: bits[32] in, _c_out: bits[32] out>(__i: bits[32], init={1}) {
+  chan_interface _c_in(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  chan_interface _c_out(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  chan _loopback_to_main(bits[32], id=0, kind=streaming, ops=send_receive, flow_control=ready_valid, strictness=proven_mutually_exclusive)
+  chan_interface _loopback_to_main(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=none, flop_kind=none)
+  chan_interface _loopback_to_main(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=none, flop_kind=none)
+  chan _main_to_loopback(bits[32], id=1, kind=streaming, ops=send_receive, flow_control=ready_valid, strictness=proven_mutually_exclusive)
+  chan_interface _main_to_loopback(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=none, flop_kind=none)
+  chan_interface _main_to_loopback(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=none, flop_kind=none)
+  proc_instantiation __test_module__Loopback__32_next_inst(_main_to_loopback, _loopback_to_main, proc=__test_module__Loopback__32_next)
+  after_all.23: token = after_all(id=23)
+  literal.11: bits[1] = literal(value=1, id=11)
+  receive.24: (token, bits[32]) = receive(after_all.23, predicate=literal.11, channel=_c_in, id=24)
+  after_all.28: token = after_all(id=28)
+  j: bits[32] = tuple_index(receive.24, index=1, id=27)
+  literal.12: bits[1] = literal(value=0, id=12)
+  loopback_tok: token = send(after_all.28, j, predicate=literal.11, channel=_main_to_loopback, id=29)
+  not.15: bits[1] = not(literal.11, id=15)
+  not.16: bits[1] = not(literal.12, id=16)
+  __i__1: bits[32] = state_read(state_element=__i, predicate=literal.11, id=20)
+  receive.30: (token, bits[32]) = receive(loopback_tok, predicate=literal.11, channel=_loopback_to_main, id=30)
+  __token: token = literal(value=token, id=10)
+  or.17: bits[1] = or(not.15, not.16, id=17)
+  not.38: bits[1] = not(literal.11, id=38)
+  or.19: bits[1] = or(literal.12, literal.11, id=19)
+  literal.13: bits[1] = literal(value=0, id=13)
+  i_val: bits[32] = identity(__i__1, id=21)
+  loopback_val: bits[32] = tuple_index(receive.30, index=1, id=33)
+  assert.18: token = assert(__token, or.17, message="State element read after read in same activation.", id=18)
+  or.39: bits[1] = or(not.38, or.19, id=39)
+  not.41: bits[1] = not(literal.11, id=41)
+  not.42: bits[1] = not(literal.13, id=42)
+  after_all.34: token = after_all(id=34)
+  add.35: bits[32] = add(i_val, loopback_val, id=35)
+  assert.40: token = assert(assert.18, or.39, message="State element written before read in same activation.", id=40)
+  or.43: bits[1] = or(not.41, not.42, id=43)
+  __i: bits[32] = state_read(state_element=__i, id=14)
+  add.37: bits[32] = add(i_val, loopback_val, id=37)
+  or.22: bits[1] = or(literal.12, literal.11, id=22)
+  tuple_index.25: token = tuple_index(receive.24, index=0, id=25)
+  tuple_index.26: token = tuple_index(receive.24, index=0, id=26)
+  tuple_index.31: token = tuple_index(receive.30, index=0, id=31)
+  tuple_index.32: token = tuple_index(receive.30, index=0, id=32)
+  send.36: token = send(after_all.34, add.35, predicate=literal.11, channel=_c_out, id=36)
+  assert.44: token = assert(assert.40, or.43, message="State element written after write in same activation.", id=44)
+  or.45: bits[1] = or(literal.13, literal.11, id=45)
+  next_value.46: () = next_value(param=__i, value=add.37, predicate=literal.11, id=46)
+  tuple.47: () = tuple(id=47)
+  tuple.48: () = tuple(id=48)
+}