[Proc-scoped channels] Update get_conversion_records to use all invocations of a proc instead of the proc functions themselves to determine if a proc should be added to the conversion records list.

Fix type info to properly record invocations.

PiperOrigin-RevId: 820224956

Author: dplassgit

xls/dslx/ir_convert/BUILD

@@ -507,6 +507,7 @@ cc_library(
         "//xls/dslx/type_system:parametric_env",
         "//xls/dslx/type_system:type_info",
         "//xls/public:status_macros",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

xls/dslx/ir_convert/conversion_record.cc

@@ -101,10 +101,11 @@ std::string ConversionRecord::ToString() const {
     config = config_record_->ToString();
   }
   return absl::StrFormat(
-      "ConversionRecord{m=%s, f=%s, top=%s, pid=%s, parametric_env=%s, "
-      "type_info=%p, config=%s}",
-      module_->name(), f_->identifier(), is_top_ ? "true" : "false", proc_id,
-      parametric_env_.ToString(), type_info_, config);
+      "ConversionRecord{m=%s, invocation=%p, f=%s, top=%s, pid=%s, "
+      "parametric_env=%s, type_info=%p, config=%s}",
+      module_->name(), invocation_, f_->identifier(),
+      is_top_ ? "true" : "false", proc_id, parametric_env_.ToString(),
+      type_info_, config);
 }
 
 }  // namespace xls::dslx

xls/dslx/ir_convert/get_conversion_records.cc

@@ -20,6 +20,7 @@
 #include <variant>
 #include <vector>
 
+#include "absl/container/flat_hash_set.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -70,12 +71,6 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
     } else {
       VLOG(5) << "Processing fn " << f->ToString();
     }
-    // TODO: davidplass - change this to gather invocations from *spawns*
-    // instead of *functions*. This will allow function_converter to emit
-    // multiple IR procs with the same parametrics but different config
-    // parameters. Then, HandleFunction would only have to handle procs
-    // that are not spawned explicitly, like test or top procs.
-
     // Note, it's possible there is no config invocation if it's a
     // top proc or some other reason.
     std::unique_ptr<ConversionRecord> config_record;
@@ -124,16 +119,13 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
     if (f->proc().has_value()) {
       proc_id = proc_id_factory_.CreateProcId(
           /*parent=*/std::nullopt, f->proc().value(),
-          // TODO: davidplass - For parametric procs we have to decide if this
-          // is a new instance if it has been called with the same parametrics
-          // before. Otherwise it needs a new procid.
           /*count_as_new_instance=*/false);
     }
     std::vector<InvocationCalleeData> calls =
         type_info_->GetUniqueInvocationCalleeData(f);
     if (f->IsParametric() && calls.empty()) {
       VLOG(5) << "No calls to parametric proc " << f->name_def()->ToString();
-      return absl::OkStatus();
+      return DefaultHandler(f);
     }
     for (auto& callee_data : calls) {
       XLS_ASSIGN_OR_RETURN(
@@ -184,7 +176,35 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
   }
 
   absl::Status HandleProc(const Proc* p) override {
-    return AddFunction(&p->next());
+    const Function* next_fn = &p->next();
+
+    if (p->IsParametric()) {
+      std::optional<ProcId> proc_id = proc_id_factory_.CreateProcId(
+          /*parent=*/std::nullopt, const_cast<Proc*>(p),
+          /*count_as_new_instance=*/false);
+
+      std::vector<InvocationCalleeData> next_calls =
+          type_info_->GetUniqueInvocationCalleeData(next_fn);
+      for (auto& callee_data : next_calls) {
+        XLS_ASSIGN_OR_RETURN(
+            ConversionRecord cr,
+            InvocationToConversionRecord(
+                next_fn, callee_data.invocation, callee_data.derived_type_info,
+                callee_data.callee_bindings, callee_data.caller_bindings,
+                // Since this proc is being spawned, it's certainly not top.
+                /* is_top= */ false, proc_id));
+        records_.push_back(std::move(cr));
+      }
+    }
+    if (top_ == next_fn || !p->IsParametric()) {
+      // "top" procs won't have spawns referencing them so they won't
+      // otherwise be added to the list, so we have to manually do it here.
+
+      // Similarly, if a proc is not parametric, while it might not have any
+      // spawns, we still want to convert it.
+      return AddFunction(next_fn);
+    }
+    return absl::OkStatus();
   }
 
   absl::Status HandleTestProc(const TestProc* tp) override {
@@ -215,6 +235,20 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
 
 }  // namespace
 
+// Filters duplicate conversion records from the given vector and returns a new
+// vector without duplicates.
+std::vector<ConversionRecord> RemoveFunctionDuplicates(
+    std::vector<ConversionRecord>& ready) {
+  absl::flat_hash_set<std::pair<Function*, ParametricEnv>> records;
+  std::vector<ConversionRecord> result;
+  for (auto& record : ready) {
+    if (records.emplace(record.f(), record.parametric_env()).second) {
+      result.push_back(std::move(record));
+    }
+  }
+  return result;
+}
+
 absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecords(
     Module* module, TypeInfo* type_info, bool include_tests) {
   ProcIdFactory proc_id_factory;
@@ -224,7 +258,8 @@ absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecords(
                                   proc_id_factory, /*top=*/nullptr);
   XLS_RETURN_IF_ERROR(module->Accept(&visitor));
 
-  return visitor.records();
+  std::vector<ConversionRecord> records = visitor.records();
+  return RemoveFunctionDuplicates(records);
 }
 
 absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecordsForEntry(
@@ -238,7 +273,9 @@ absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecordsForEntry(
     ConversionRecordVisitor visitor(m, type_info, /*include_tests=*/true,
                                     proc_id_factory, f);
     XLS_RETURN_IF_ERROR(m->Accept(&visitor));
-    return visitor.records();
+
+    std::vector<ConversionRecord> records = visitor.records();
+    return RemoveFunctionDuplicates(records);
   }
 
   Proc* p = std::get<Proc*>(entry);
@@ -250,6 +287,8 @@ absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecordsForEntry(
   ConversionRecordVisitor visitor(m, new_ti, /*include_tests=*/true,
                                   proc_id_factory, &p->next());
   XLS_RETURN_IF_ERROR(m->Accept(&visitor));
-  return visitor.records();
+
+  std::vector<ConversionRecord> records = visitor.records();
+  return RemoveFunctionDuplicates(records);
 }
 }  // namespace xls::dslx

xls/dslx/ir_convert/get_conversion_records_test.cc

@@ -164,17 +164,13 @@ proc top {
       GetConversionRecords(tm.module, tm.type_info, false));
   ASSERT_EQ(3, order.size());
   EXPECT_EQ(order[0].f()->identifier(), "P.next");
-  EXPECT_EQ(order[0].invocation()->ToString(),
-            "P.next<u32:2>(P.init<u32:2>())");
   const ConversionRecord* config_record = order[0].config_record();
   EXPECT_NE(config_record, nullptr);
   EXPECT_EQ(config_record->invocation()->ToString(), "P.config<u32:2>(u2:1)");
   EXPECT_EQ(order[0].parametric_env(),
             ParametricEnv(absl::flat_hash_map<std::string, InterpValue>{
                 {"N", InterpValue::MakeUBits(/*bit_count=*/32, /*value=*/2)}}));
   EXPECT_EQ(order[1].f()->identifier(), "P.next");
-  EXPECT_EQ(order[1].invocation()->ToString(),
-            "P.next<u32:4>(P.init<u32:4>())");
   config_record = order[1].config_record();
   EXPECT_NE(config_record, nullptr);
   EXPECT_EQ(config_record->invocation()->ToString(), "P.config<u32:4>(u4:2)");
@@ -388,8 +384,9 @@ fn my_test() -> bool { f<u32:8>(u8:1) == u32:8 }
   XLS_ASSERT_OK_AND_ASSIGN(
       TypecheckedModule tm,
       ParseAndTypecheck(kProgram, "test.x", "test", &import_data));
-  XLS_ASSERT_OK_AND_ASSIGN(std::vector<ConversionRecord> order,
-                           GetConversionRecords(tm.module, tm.type_info, true));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::vector<ConversionRecord> order,
+      GetConversionRecords(tm.module, tm.type_info, /*include_tests=*/true));
   ASSERT_EQ(2, order.size());
   EXPECT_EQ(order[0].f()->identifier(), "f");
   EXPECT_EQ(order[0].parametric_env(),
@@ -440,13 +437,15 @@ proc test {
   next(x: ()) { () }
 }
 )";
+
 TEST(GetConversionRecordsTest, TestProc) {
   auto import_data = CreateImportDataForTest();
   XLS_ASSERT_OK_AND_ASSIGN(
       TypecheckedModule tm,
       ParseAndTypecheck(kTestProc, "test.x", "test", &import_data));
-  XLS_ASSERT_OK_AND_ASSIGN(std::vector<ConversionRecord> order,
-                           GetConversionRecords(tm.module, tm.type_info, true));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::vector<ConversionRecord> order,
+      GetConversionRecords(tm.module, tm.type_info, /*include_tests=*/true));
   ASSERT_EQ(2, order.size());
   EXPECT_EQ(order[0].f()->identifier(), "P.next");
   EXPECT_EQ(order[0].parametric_env(),
@@ -462,7 +461,7 @@ TEST(GetConversionRecordsTest, TestProcSkipped) {
       ParseAndTypecheck(kTestProc, "test.x", "test", &import_data));
   XLS_ASSERT_OK_AND_ASSIGN(
       std::vector<ConversionRecord> order,
-      GetConversionRecords(tm.module, tm.type_info, false));
+      GetConversionRecords(tm.module, tm.type_info, /*include_tests=*/false));
   // It still converts the parametric proc because there is still a spawn,
   // in the test proc.
   ASSERT_EQ(1, order.size());
@@ -511,6 +510,61 @@ impl S {
   EXPECT_EQ(order[0].parametric_env(), ParametricEnv());
 }
 
+TEST(GetConversionRecordsTest, SpawnTree) {
+  constexpr std::string_view kProgram = R"(
+proc spawnee2<N:u32, M:u16> {
+  init { }
+  config() { () }
+  next(state: ()) { () }
+}
+
+proc spawnee1<N:u32> {
+  init { }
+  config() {
+    spawn spawnee2<N, u16:1>();
+    spawn spawnee2<N, u16:2>();
+  }
+  next(state: ()) { () }
+}
+
+pub proc main {
+  init { }
+  config() {
+    spawn spawnee1<u32:3>();
+    spawn spawnee1<u32:4>();
+  }
+  next(state: ()) { () }
+}
+)";
+
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      TypecheckedModule tm,
+      ParseAndTypecheck(kProgram, "test.x", "main", &import_data, {},
+                        TypeInferenceVersion::kVersion2));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::vector<ConversionRecord> order,
+      GetConversionRecords(tm.module, tm.type_info, false));
+  // 4 of spawnee2, 2 of spawnee1 and 1 of main.
+  ASSERT_EQ(7, order.size());
+  EXPECT_EQ(order[0].parametric_env(),
+            ParametricEnv(absl::flat_hash_map<std::string, InterpValue>{
+                {"N", InterpValue::MakeUBits(/*bit_count=*/32, /*value=*/3)},
+                {"M", InterpValue::MakeUBits(/*bit_count=*/16, /*value=*/1)}}));
+  EXPECT_EQ(order[1].parametric_env(),
+            ParametricEnv(absl::flat_hash_map<std::string, InterpValue>{
+                {"N", InterpValue::MakeUBits(/*bit_count=*/32, /*value=*/3)},
+                {"M", InterpValue::MakeUBits(/*bit_count=*/16, /*value=*/2)}}));
+  EXPECT_EQ(order[2].parametric_env(),
+            ParametricEnv(absl::flat_hash_map<std::string, InterpValue>{
+                {"N", InterpValue::MakeUBits(/*bit_count=*/32, /*value=*/4)},
+                {"M", InterpValue::MakeUBits(/*bit_count=*/16, /*value=*/1)}}));
+  EXPECT_EQ(order[3].parametric_env(),
+            ParametricEnv(absl::flat_hash_map<std::string, InterpValue>{
+                {"N", InterpValue::MakeUBits(/*bit_count=*/32, /*value=*/4)},
+                {"M", InterpValue::MakeUBits(/*bit_count=*/16, /*value=*/2)}}));
+}
+
 }  // namespace
 
 }  // namespace xls::dslx