Support IR conversion of proc aliases with proc-scoped channels flagged on.

richmckeever · copybara-github · commit 259d031e98cc · 2025-10-17T12:48:16.000-07:00
PiperOrigin-RevId: 820795684
diff --git a/xls/dslx/ir_convert/function_converter.cc b/xls/dslx/ir_convert/function_converter.cc
@@ -3586,10 +3586,6 @@ absl::Status FunctionConverter::HandleProcNextFunction(
     ScopedTypeInfoSwap stis(this, config_type_info);
 
     Function& config_fn = proc->config();
-    // This probably was checked already but might as well double-check it.
-    XLS_RET_CHECK(invocation != nullptr || !f->IsParametric())
-        << "Cannot lower a parametric proc without an invocation";
-
     proc_scoped_channel_scope = std::make_unique<ProcScopedChannelScope>(
         package_data_.conversion_info, import_data, options_, builder_ptr);
     proc_scoped_channel_scope->EnterFunctionContext(current_type_info_,
diff --git a/xls/dslx/ir_convert/get_conversion_records.cc b/xls/dslx/ir_convert/get_conversion_records.cc
@@ -55,12 +55,14 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
  public:
   ConversionRecordVisitor(Module* module, TypeInfo* type_info,
                           bool include_tests, ProcIdFactory proc_id_factory,
-                          AstNode* top)
+                          AstNode* top,
+                          std::optional<ResolvedProcAlias> resolved_proc_alias)
       : module_(module),
         type_info_(type_info),
         include_tests_(include_tests),
         proc_id_factory_(proc_id_factory),
-        top_(top) {}
+        top_(top),
+        resolved_proc_alias_(resolved_proc_alias) {}
 
   absl::StatusOr<ConversionRecord> InvocationToConversionRecord(
       const Function* f, const Invocation* invocation,
@@ -235,6 +237,30 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
     VLOG(5) << "HandleProc " << p->ToString();
     const Function* next_fn = &p->next();
 
+    if (top_ == next_fn && resolved_proc_alias_.has_value()) {
+      ProcId proc_id = proc_id_factory_.CreateProcId(
+          /*parent=*/std::nullopt, const_cast<Proc*>(p),
+          /*count_as_new_instance=*/false);
+      proc_id.alias_name = resolved_proc_alias_->name;
+      XLS_ASSIGN_OR_RETURN(
+          ConversionRecord config_record,
+          MakeConversionRecord(
+              const_cast<Function*>(&p->config()), top_->owner(),
+              resolved_proc_alias_->config_type_info, resolved_proc_alias_->env,
+              proc_id, /*invocation=*/nullptr,
+              /*is_top=*/false));
+      XLS_ASSIGN_OR_RETURN(
+          ConversionRecord next_record,
+          MakeConversionRecord(
+              const_cast<Function*>(&p->next()), top_->owner(),
+              resolved_proc_alias_->next_type_info, resolved_proc_alias_->env,
+              proc_id, /*invocation=*/nullptr,
+              /*is_top=*/true,
+              std::make_unique<ConversionRecord>(std::move(config_record))));
+      records_.push_back(std::move(next_record));
+      return absl::OkStatus();
+    }
+
     // This is required in order to process cross-module spawns; otherwise it
     // will never add procs from imported modules to the list of functions to
     // convert.
@@ -294,6 +320,9 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
   ProcIdFactory proc_id_factory_;
   AstNode* top_;
 
+  // The proc alias that was used to specify the top proc, if any.
+  std::optional<ResolvedProcAlias> resolved_proc_alias_;
+
   std::vector<ConversionRecord> records_;
 };
 
@@ -319,23 +348,27 @@ absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecords(
   // TODO: https://github.com/google/xls/issues/2078 - properly set
   // top instead of setting to nullptr.
   ConversionRecordVisitor visitor(module, type_info, include_tests,
-                                  proc_id_factory, /*top=*/nullptr);
+                                  proc_id_factory, /*top=*/nullptr,
+                                  /*resolved_proc_alias=*/std::nullopt);
   XLS_RETURN_IF_ERROR(module->Accept(&visitor));
 
   std::vector<ConversionRecord> records = visitor.records();
   return RemoveFunctionDuplicates(records);
 }
 
 absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecordsForEntry(
-    std::variant<Proc*, Function*> entry, TypeInfo* type_info) {
+    std::variant<Proc*, Function*> entry, TypeInfo* type_info,
+    std::optional<ResolvedProcAlias> resolved_proc_alias) {
   ProcIdFactory proc_id_factory;
   if (std::holds_alternative<Function*>(entry)) {
+    XLS_RET_CHECK(!resolved_proc_alias.has_value());
     Function* f = std::get<Function*>(entry);
     Module* m = f->owner();
     // We are only ever called for tests, so we set include_tests to
     // true, and make sure that this function is top.
     ConversionRecordVisitor visitor(m, type_info, /*include_tests=*/true,
-                                    proc_id_factory, f);
+                                    proc_id_factory, f,
+                                    /*resolved_proc_alias=*/std::nullopt);
     XLS_RETURN_IF_ERROR(m->Accept(&visitor));
 
     std::vector<ConversionRecord> records = visitor.records();
@@ -349,7 +382,8 @@ absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecordsForEntry(
   // We are only ever called for tests, so we set include_tests to true,
   // and make sure that this proc's next function is top.
   ConversionRecordVisitor visitor(m, new_ti, /*include_tests=*/true,
-                                  proc_id_factory, &p->next());
+                                  proc_id_factory, &p->next(),
+                                  resolved_proc_alias);
   XLS_RETURN_IF_ERROR(m->Accept(&visitor));
 
   std::vector<ConversionRecord> records = visitor.records();
diff --git a/xls/dslx/ir_convert/get_conversion_records.h b/xls/dslx/ir_convert/get_conversion_records.h
@@ -15,6 +15,7 @@
 #ifndef XLS_DSLX_IR_CONVERT_GET_CONVERSION_RECORDS_H_
 #define XLS_DSLX_IR_CONVERT_GET_CONVERSION_RECORDS_H_
 
+#include <optional>
 #include <variant>
 #include <vector>
 
@@ -39,7 +40,8 @@ absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecords(
 //  entry: Proc or Function to start from (the top)
 //  type_info: Mapping from node to type.
 absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecordsForEntry(
-    std::variant<Proc*, Function*> entry, TypeInfo* type_info);
+    std::variant<Proc*, Function*> entry, TypeInfo* type_info,
+    std::optional<ResolvedProcAlias> resolved_proc_alias);
 
 }  // namespace xls::dslx
 
diff --git a/xls/dslx/ir_convert/ir_converter.cc b/xls/dslx/ir_convert/ir_converter.cc
@@ -403,7 +403,7 @@ absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecords(
   // lower_to_proc_scoped_channels is turned on everywhere, and call
   // GetConversionRecordsForEntry unconditionally.
   if (options.lower_to_proc_scoped_channels) {
-    return GetConversionRecordsForEntry(block, type_info);
+    return GetConversionRecordsForEntry(block, type_info, resolved_proc_alias);
   }
   return GetOrderForEntry(block, type_info, resolved_proc_alias);
 }
diff --git a/xls/dslx/ir_convert/ir_converter_test.cc b/xls/dslx/ir_convert/ir_converter_test.cc
@@ -5877,10 +5877,7 @@ pub proc main {
   XLS_ASSERT_OK_AND_ASSIGN(
       std::string converted,
       ConvertOneFunctionForTest(kProgram, "main", import_data,
-                                ConvertOptions{
-                                    .emit_positions = false,
-                                    .lower_to_proc_scoped_channels = true,
-                                }));
+                                kProcScopedChannelOptions));
   ExpectIr(converted);
 }
 
@@ -5911,11 +5908,88 @@ fn f() -> u32 {
 
   XLS_ASSERT_OK_AND_ASSIGN(
       std::string converted,
-      ConvertModuleForTest(
-          importer_program,
-          ConvertOptions{.emit_positions = false,
-                         .lower_to_proc_scoped_channels = true},
-          &import_data));
+      ConvertModuleForTest(importer_program, kProcScopedChannelOptions,
+                           &import_data));
+  ExpectIr(converted);
+}
+
+TEST_P(ProcScopedChannelsIrConverterTest, ProcScopedBasicProcAlias) {
+  constexpr std::string_view program = R"(
+proc Foo {
+  c: chan<u32> out;
+  init { u32:1 }
+  config(output_c: chan<u32> out) {
+    (output_c,)
+  }
+  next(i: u32) {
+    let tok = send(join(), c, i);
+    i + u32:2
+  }
+}
+
+pub proc FooAlias = Foo;
+)";
+
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertOneFunctionForTest(program, "FooAlias", import_data,
+                                kProcScopedChannelOptions));
+  ExpectIr(converted);
+}
+
+TEST_P(ProcScopedChannelsIrConverterTest, ProcScopedParametricProcAlias) {
+  constexpr std::string_view program = R"(
+proc Foo<N: u32> {
+  c: chan<uN[N]> out;
+  init { uN[N]:1 }
+  config(output_c: chan<uN[N]> out) {
+    (output_c,)
+  }
+  next(i: uN[N]) {
+    let tok = send(join(), c, i);
+    i + uN[N]:2
+  }
+}
+
+pub proc FooAlias = Foo<16>;
+)";
+
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertOneFunctionForTest(program, "FooAlias", import_data,
+                                kProcScopedChannelOptions));
+  ExpectIr(converted);
+}
+
+TEST_P(ProcScopedChannelsIrConverterTest, ProcScopedProcAliasToImportedProc) {
+  ImportData import_data = CreateImportDataForTest();
+
+  constexpr std::string_view imported = R"(
+pub proc Foo<N: u32> {
+  c: chan<uN[N]> out;
+  init { uN[N]:1 }
+  config(output_c: chan<uN[N]> out) {
+    (output_c,)
+  }
+  next(i: uN[N]) {
+    let tok = send(join(), c, i);
+    i + uN[N]:2
+  }
+}
+  )";
+  XLS_EXPECT_OK(
+      ParseAndTypecheck(imported, "imported.x", "imported", &import_data));
+
+  constexpr std::string_view program = R"(
+import imported;
+pub proc FooAlias = imported::Foo<16>;
+  )";
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertOneFunctionForTest(program, "FooAlias", import_data,
+                                kProcScopedChannelOptions));
   ExpectIr(converted);
 }
 
diff --git a/xls/dslx/ir_convert/testdata/ir_converter_test_ProcScopedBasicProcAlias.ir b/xls/dslx/ir_convert/testdata/ir_converter_test_ProcScopedBasicProcAlias.ir
@@ -0,0 +1,17 @@
+package test_module
+
+file_number 0 "test_module.x"
+
+top proc __test_module__FooAlias_next<output_c: bits[32] out>(__state: bits[32], init={1}) {
+  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=2, 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)
+}
diff --git a/xls/dslx/ir_convert/testdata/ir_converter_test_ProcScopedParametricProcAlias.ir b/xls/dslx/ir_convert/testdata/ir_converter_test_ProcScopedParametricProcAlias.ir
@@ -0,0 +1,18 @@
+package test_module
+
+file_number 0 "test_module.x"
+
+top proc __test_module__FooAlias_next<output_c: bits[16] out>(__state: bits[16], init={1}) {
+  chan_interface output_c(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  __state: bits[16] = state_read(state_element=__state, id=2)
+  literal.9: bits[16] = literal(value=2, id=9)
+  output_c: bits[16] = send_channel_end(id=5)
+  after_all.7: token = after_all(id=7)
+  literal.3: bits[1] = literal(value=1, id=3)
+  add.10: bits[16] = add(__state, literal.9, id=10)
+  __token: token = literal(value=token, id=1)
+  N: bits[32] = literal(value=16, id=4)
+  tuple.6: (bits[16]) = tuple(output_c, id=6)
+  tok: token = send(after_all.7, __state, predicate=literal.3, channel=output_c, id=8)
+  next_value.11: () = next_value(param=__state, value=add.10, id=11)
+}
diff --git a/xls/dslx/ir_convert/testdata/ir_converter_test_ProcScopedProcAliasToImportedProc.ir b/xls/dslx/ir_convert/testdata/ir_converter_test_ProcScopedProcAliasToImportedProc.ir
@@ -0,0 +1,18 @@
+package test_module
+
+file_number 0 "imported.x"
+
+top proc __imported__FooAlias_next<output_c: bits[16] out>(__state: bits[16], init={1}) {
+  chan_interface output_c(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
+  __state: bits[16] = state_read(state_element=__state, id=2)
+  literal.9: bits[16] = literal(value=2, id=9)
+  output_c: bits[16] = send_channel_end(id=5)
+  after_all.7: token = after_all(id=7)
+  literal.3: bits[1] = literal(value=1, id=3)
+  add.10: bits[16] = add(__state, literal.9, id=10)
+  __token: token = literal(value=token, id=1)
+  N: bits[32] = literal(value=16, id=4)
+  tuple.6: (bits[16]) = tuple(output_c, id=6)
+  tok: token = send(after_all.7, __state, predicate=literal.3, channel=output_c, id=8)
+  next_value.11: () = next_value(param=__state, value=add.10, id=11)
+}

Original file line number	Diff line number	Diff line change
`@@ -403,7 +403,7 @@ absl::StatusOr<std::vector<ConversionRecord>> GetConversionRecords(`
`403`	`403`	`// lower_to_proc_scoped_channels is turned on everywhere, and call`
`404`	`404`	`// GetConversionRecordsForEntry unconditionally.`
`405`	`405`	`if (options.lower_to_proc_scoped_channels) {`
`406`		`- return GetConversionRecordsForEntry(block, type_info);`
	`406`	`+ return GetConversionRecordsForEntry(block, type_info, resolved_proc_alias);`
`407`	`407`	`}`
`408`	`408`	`return GetOrderForEntry(block, type_info, resolved_proc_alias);`
`409`	`409`	`}`