Add scheduled function/proc information to the extracted interface proto.

xls/dev_tools/BUILD

@@ -120,12 +120,14 @@ cc_library(
     hdrs = ["extract_interface.h"],
     deps = [
         "//xls/common:attribute_data",
+        "//xls/common/status:status_macros",
         "//xls/ir",
         "//xls/ir:channel",
         "//xls/ir:register",
         "//xls/ir:state_element",
         "//xls/ir:type",
         "//xls/ir:xls_ir_interface_cc_proto",
+        "//xls/scheduling:pipeline_schedule",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:statusor",
         "@com_google_protobuf//:protobuf",
@@ -164,6 +166,8 @@ cc_test(
         "//xls/ir:function_builder",
         "//xls/ir:ir_test_base",
         "//xls/ir:source_location",
+        "//xls/scheduling:run_pipeline_schedule",
+        "//xls/scheduling:scheduling_options",
         "@googletest//:gtest",
     ],
 )

xls/dev_tools/extract_interface.cc

@@ -14,13 +14,16 @@
 
 #include "xls/dev_tools/extract_interface.h"
 
+#include <optional>
 #include <string>
+#include <utility>
 #include <variant>
 
 #include "absl/log/check.h"
 #include "absl/status/statusor.h"
 #include "google/protobuf/text_format.h"
 #include "xls/common/attribute_data.h"
+#include "xls/common/status/status_macros.h"
 #include "xls/ir/block.h"
 #include "xls/ir/channel.h"
 #include "xls/ir/function.h"
@@ -130,13 +133,24 @@ PackageInterfaceProto::Block ExtractBlockInterface(Block* block) {
   return proto;
 }
 
-PackageInterfaceProto ExtractPackageInterface(Package* package) {
+PackageInterfaceProto ExtractPackageInterface(
+    Package* package, std::optional<PackageScheduleProto> schedule) {
+  auto res = TryExtractPackageInterface(package, schedule);
+  CHECK_OK(res);
+  return *std::move(res);
+}
+absl::StatusOr<PackageInterfaceProto> TryExtractPackageInterface(
+    Package* package, std::optional<PackageScheduleProto> schedule) {
   PackageInterfaceProto proto;
   // Basic information
   proto.set_name(package->name());
   for (const auto& [_, f] : package->fileno_to_name()) {
     *proto.add_files() = f;
   }
+  std::optional<PackageSchedule> sched;
+  if (schedule.has_value()) {
+    XLS_ASSIGN_OR_RETURN(sched, PackageSchedule::FromProto(package, *schedule));
+  }
   // Fill in channels
   for (const Channel* c : package->channels()) {
     auto* chan = proto.add_channels();
@@ -186,10 +200,23 @@ PackageInterfaceProto ExtractPackageInterface(Package* package) {
     }
   }
   for (const auto& f : package->functions()) {
-    *proto.add_functions() = ExtractFunctionInterface(f.get());
+    auto func_proto = ExtractFunctionInterface(f.get());
+    if (sched && sched->HasSchedule(f.get())) {
+      XLS_ASSIGN_OR_RETURN(auto scheduled_fb, sched->ToScheduledInterfaceProto(
+                                                  f.get(), func_proto));
+      *proto.add_scheduled_functions() =
+          std::move(scheduled_fb).scheduled_function();
+    }
+    *proto.add_functions() = std::move(func_proto);
   }
   for (const auto& p : package->procs()) {
-    *proto.add_procs() = ExtractProcInterface(p.get());
+    auto proc_proto = ExtractProcInterface(p.get());
+    if (sched && sched->HasSchedule(p.get())) {
+      XLS_ASSIGN_OR_RETURN(auto scheduled_fb, sched->ToScheduledInterfaceProto(
+                                                  p.get(), proc_proto));
+      *proto.add_scheduled_procs() = std::move(scheduled_fb).scheduled_proc();
+    }
+    *proto.add_procs() = std::move(proc_proto);
   }
   for (const auto& b : package->blocks()) {
     *proto.add_blocks() = ExtractBlockInterface(b.get());

xls/dev_tools/extract_interface.h

@@ -15,8 +15,11 @@
 #ifndef XLS_DEV_TOOLS_EXTRACT_INTERFACE_H_
 #define XLS_DEV_TOOLS_EXTRACT_INTERFACE_H_
 
+#include <optional>
+
 #include "xls/ir/package.h"
 #include "xls/ir/xls_ir_interface.pb.h"
+#include "xls/scheduling/pipeline_schedule.h"
 
 namespace xls {
 
@@ -25,8 +28,13 @@ PackageInterfaceProto::Proc ExtractProcInterface(Proc* proc);
 PackageInterfaceProto::Block ExtractBlockInterface(Block* block);
 
 // Create a PackageInterfaceProto based on the given Package.
-PackageInterfaceProto ExtractPackageInterface(Package* package);
-
+absl::StatusOr<PackageInterfaceProto> TryExtractPackageInterface(
+    Package* package,
+    std::optional<PackageScheduleProto> schedule = std::nullopt);
+// Create a PackageInterfaceProto based on the given Package.
+PackageInterfaceProto ExtractPackageInterface(
+    Package* package,
+    std::optional<PackageScheduleProto> schedule = std::nullopt);
 }  // namespace xls
 
 #endif  // XLS_DEV_TOOLS_EXTRACT_INTERFACE_H_

xls/dev_tools/extract_interface_main.cc

@@ -14,6 +14,7 @@
 
 #include <filesystem>
 #include <iostream>
+#include <optional>
 #include <string>
 #include <string_view>
 #include <vector>
@@ -38,14 +39,33 @@ const char kUsage[] = R"(
 
 ABSL_FLAG(bool, binary_proto, false, "Print as a binary proto");
 ABSL_FLAG(std::string, output_file, "/dev/stdout", "Output file");
+ABSL_FLAG(bool, binary_schedule_proto, false,
+          "schedule_file is a binary proto");
+ABSL_FLAG(std::optional<std::string>, schedule_file, std::nullopt,
+          "Optional schedule file to use.");
 
 namespace xls {
 namespace {
 
-absl::Status RealMain(bool binary_proto, const std::filesystem::path& path) {
+absl::Status RealMain(
+    bool binary_proto, const std::filesystem::path& path,
+    bool binary_schedule_proto,
+    const std::optional<std::filesystem::path>& schedule_path) {
   XLS_ASSIGN_OR_RETURN(auto ir_text, GetFileContents(path));
   XLS_ASSIGN_OR_RETURN(auto package, ParsePackage(ir_text, path.string()));
-  PackageInterfaceProto proto = ExtractPackageInterface(package.get());
+  std::optional<PackageScheduleProto> schedule_proto;
+  if (schedule_path.has_value()) {
+    XLS_ASSIGN_OR_RETURN(auto schedule_text, GetFileContents(*schedule_path));
+    schedule_proto.emplace();
+    if (binary_schedule_proto) {
+      schedule_proto->ParseFromString(schedule_text);
+    } else {
+      XLS_RET_CHECK(
+          google::protobuf::TextFormat::ParseFromString(schedule_text, &*schedule_proto));
+    }
+  }
+  PackageInterfaceProto proto =
+      ExtractPackageInterface(package.get(), schedule_proto);
 
   std::string output;
   if (binary_proto) {
@@ -69,6 +89,8 @@ int main(int argc, char** argv) {
     LOG(QFATAL) << "Expected invocation: " << argv[0] << " <ir_file>";
   }
 
-  return xls::ExitStatus(xls::RealMain(absl::GetFlag(FLAGS_binary_proto),
-                                       positional_arguments[0]));
+  return xls::ExitStatus(
+      xls::RealMain(absl::GetFlag(FLAGS_binary_proto), positional_arguments[0],
+                    absl::GetFlag(FLAGS_binary_schedule_proto),
+                    absl::GetFlag(FLAGS_schedule_file)));
 }

xls/dev_tools/extract_interface_test.cc

@@ -28,6 +28,8 @@
 #include "xls/ir/function_builder.h"
 #include "xls/ir/ir_test_base.h"
 #include "xls/ir/source_location.h"
+#include "xls/scheduling/run_pipeline_schedule.h"
+#include "xls/scheduling/scheduling_options.h"
 
 namespace xls {
 namespace {
@@ -137,6 +139,97 @@ top proc add(__token: token, init={token}) {
                   )pb"));
 }
 
+TEST_F(ExtractInterfaceTest, BasicProcWithSchedule) {
+  constexpr std::string_view kIr = R"(
+package sample
+
+file_number 0 "fake_file.x"
+
+chan sample__operand_0(bits[32], id=0, kind=streaming, ops=receive_only, flow_control=ready_valid)
+chan sample__operand_1(bits[32], id=1, kind=streaming, ops=receive_only, flow_control=ready_valid)
+chan sample__result(bits[32], id=2, kind=streaming, ops=send_only, flow_control=ready_valid)
+
+top proc add(__token: token, init={token}) {
+  receive.4: (token, bits[32]) = receive(__token, channel=sample__operand_0, id=4)
+  receive.7: (token, bits[32]) = receive(__token, channel=sample__operand_1, id=7)
+  tok_operand_0_val: token = tuple_index(receive.4, index=0, id=5, pos=[(0,14,9)])
+  tok_operand_1_val: token = tuple_index(receive.7, index=0, id=8, pos=[(0,15,9)])
+  operand_0_val: bits[32] = tuple_index(receive.4, index=1, id=6, pos=[(0,14,28)])
+  operand_1_val: bits[32] = tuple_index(receive.7, index=1, id=9, pos=[(0,15,28)])
+  tok_recv: token = after_all(tok_operand_0_val, tok_operand_1_val, id=10)
+  result_val: bits[32] = add(operand_0_val, operand_1_val, id=11, pos=[(0,18,35)])
+  tok_send: token = send(tok_recv, result_val, channel=sample__result, id=12)
+  after_all.14: token = after_all(__token, tok_operand_0_val, tok_operand_1_val, tok_recv, tok_send, id=14)
+  next (after_all.14)
+})";
+  XLS_ASSERT_OK_AND_ASSIGN(auto p, ParsePackage(kIr));
+
+  TestDelayEstimator delay_estimator;
+  XLS_ASSERT_OK_AND_ASSIGN(
+      auto schedule,
+      RunPipelineSchedule(
+          p->GetTop().value(), delay_estimator,
+          SchedulingOptions()
+              .pipeline_stages(2)
+              .worst_case_throughput(2)
+              .add_constraint(
+                  IOConstraint("sample__operand_0", IODirection::kReceive,
+                               "sample__result", IODirection::kSend,
+                               /*minimum_latency=*/1, /*maximum_latency=*/1))
+              .add_constraint(
+                  IOConstraint("sample__operand_1", IODirection::kReceive,
+                               "sample__result", IODirection::kSend,
+                               /*minimum_latency=*/1, /*maximum_latency=*/1))));
+  XLS_ASSERT_OK_AND_ASSIGN(auto proto, schedule.ToProto(delay_estimator));
+  PackageScheduleProto package_schedule_proto;
+  package_schedule_proto.mutable_schedules()->insert({"add", proto});
+  auto interface = ExtractPackageInterface(p.get(), package_schedule_proto);
+  RecordProperty("interface", interface.DebugString());
+  EXPECT_THAT(
+      interface,
+      ProtoEquivalent(
+          R"pb(
+            name: "sample"
+            files: "fake_file.x"
+            # Also includes channels and procs but those are tested by the
+            # BasicProc test.
+            scheduled_procs {
+              proc {
+                base { top: true name: "add" }
+                state {
+                  name: "__token"
+                  type { type_enum: TOKEN }
+                }
+                state_values {
+                  name {
+                    name: "__token"
+                    type { type_enum: TOKEN }
+                  }
+                  non_synthesizable: false
+                }
+              }
+              pipeline_info { pipeline_length: 2 initiation_interval: 2 }
+              sends {
+                stage: 1
+                channel_name: "sample__result"
+                is_blocking: false
+              }
+              recvs {
+                stage: 0
+                channel_name: "sample__operand_0"
+                is_blocking: true
+              }
+              recvs {
+                stage: 0
+                channel_name: "sample__operand_1"
+                is_blocking: true
+              }
+              state_reads { name: "__token" stage: 0 }
+              state_writes { name: "__token" stage: 1 }
+            }
+          )pb"));
+}
+
 TEST_F(ExtractInterfaceTest, BasicBlock) {
   auto p = CreatePackage();
   BlockBuilder bb(TestName(), p.get());

xls/ir/xls_ir_interface.proto

@@ -85,6 +85,52 @@ message PackageInterfaceProto {
     repeated string recv_channels = 4;
   }
 
+  // Common information for scheduled pipelines.
+  message PipelineInformation {
+    // How long the full pipeline is.
+    optional int32 pipeline_length = 1;
+    // What the initiation interval for this pipeline is. 1 is full throughput.
+    optional int32 initiation_interval = 2;
+  }
+
+  // Information about a function which has gone through scheduling.
+  message ScheduledFunction {
+    optional Function function = 1;
+    optional PipelineInformation pipeline_info = 2;
+  }
+
+  message ScheduledProc {
+    // Information about a send or receive.
+    message IO {
+      // The stage at which this IO is scheduled.
+      optional int32 stage = 1;
+      // Either a send/recv channel that is used in the given stage.
+      optional string channel_name = 2;
+      optional bool is_blocking = 4;
+    }
+    // Information about when state information is read or written which
+    // controls throughput and latency.
+    message StateAccess {
+      // The state element
+      optional string name = 1;
+      // The stage at which this state element is accessed (read or written).
+      optional int32 stage = 2;
+    }
+    optional Proc proc = 1;
+    optional PipelineInformation pipeline_info = 2;
+    repeated IO sends = 3;
+    repeated IO recvs = 4;
+    repeated StateAccess state_reads = 5;
+    repeated StateAccess state_writes = 6;
+  }
+
+  message ScheduledFunctionBase {
+    oneof type {
+      ScheduledFunction scheduled_function = 1;
+      ScheduledProc scheduled_proc = 2;
+    }
+  }
+
   message Block {
     optional FunctionBase base = 1;
     repeated NamedValue registers = 2;
@@ -126,4 +172,6 @@ message PackageInterfaceProto {
   repeated Function functions = 4;
   repeated Proc procs = 5;
   repeated Block blocks = 6;
+  repeated ScheduledFunction scheduled_functions = 7;
+  repeated ScheduledProc scheduled_procs = 8;
 }

xls/scheduling/BUILD

@@ -287,6 +287,7 @@ cc_library(
         "//xls/ir",
         "//xls/ir:channel",
         "//xls/ir:op",
+        "//xls/ir:xls_ir_interface_cc_proto",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",