Add new-style proc support to the interpreter and jit.

A previous change switched the interpreter and jit over to using an elaboration (instances) instead of the IR directly. However, only old-style procs were supported. This change adds supports for new-style procs. Requires passing  instance-specific information to the JITted function. This replaces the unused "user data" argument.

This also required a couple fixes to elaboration which wasn't properly distinguishing between channel names and channel reference names.

PiperOrigin-RevId: 597678974

Author: meheffernan

xls/interpreter/BUILD

@@ -386,7 +386,6 @@ cc_library(
         ":proc_evaluator",
         ":proc_runtime",
         "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
         "//xls/common/logging",
@@ -401,13 +400,20 @@ cc_test(
     name = "serial_proc_runtime_test",
     srcs = ["serial_proc_runtime_test.cc"],
     deps = [
+        ":channel_queue",
         ":interpreter_proc_runtime",
+        ":proc_evaluator",
         ":proc_interpreter",
+        ":proc_runtime",
         ":proc_runtime_test_base",
         ":serial_proc_runtime",
+        "@com_google_absl//absl/status:statusor",
         "//xls/common:xls_gunit",
         "//xls/common:xls_gunit_main",
+        "//xls/common/status:status_macros",
         "//xls/ir",
+        "//xls/ir:elaboration",
+        "//xls/ir:value",
         "//xls/jit:jit_channel_queue",
         "//xls/jit:jit_proc_runtime",
         "//xls/jit:proc_jit",
@@ -422,14 +428,16 @@ cc_library(
     deps = [
         ":channel_queue",
         ":proc_runtime",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
         "//xls/common/status:matchers",
+        "//xls/common/status:status_macros",
         "//xls/ir",
         "//xls/ir:bits",
         "//xls/ir:channel",
         "//xls/ir:channel_ops",
+        "//xls/ir:elaboration",
         "//xls/ir:function_builder",
         "//xls/ir:ir_parser",
         "//xls/ir:ir_test_base",

xls/interpreter/interpreter_proc_runtime.cc

@@ -29,34 +29,33 @@
 #include "xls/ir/value.h"
 
 namespace xls {
+namespace {
 
-absl::StatusOr<std::unique_ptr<SerialProcRuntime>>
-CreateInterpreterSerialProcRuntime(Package* package) {
-  // TODO(https://github.com/google/xls/issues/869): Support new-style procs.
-  XLS_ASSIGN_OR_RETURN(Elaboration elaboration,
-                       Elaboration::ElaborateOldStylePackage(package));
-
+absl::StatusOr<std::unique_ptr<SerialProcRuntime>> CreateRuntime(
+    Elaboration elaboration) {
   // Create a queue manager for the queues. This factory verifies that there an
   // receive only queue for every receive only channel.
   XLS_ASSIGN_OR_RETURN(std::unique_ptr<ChannelQueueManager> queue_manager,
                        ChannelQueueManager::Create(std::move(elaboration)));
 
   // Create a ProcInterpreter for each Proc.
   std::vector<std::unique_ptr<ProcEvaluator>> proc_interpreters;
-  for (auto& proc : package->procs()) {
+  for (Proc* proc : queue_manager->elaboration().procs()) {
     proc_interpreters.push_back(
-        std::make_unique<ProcInterpreter>(proc.get(), queue_manager.get()));
+        std::make_unique<ProcInterpreter>(proc, queue_manager.get()));
   }
 
   // Create a runtime.
-  XLS_ASSIGN_OR_RETURN(
-      std::unique_ptr<SerialProcRuntime> proc_runtime,
-      SerialProcRuntime::Create(package, std::move(proc_interpreters),
-                                std::move(queue_manager)));
+  XLS_ASSIGN_OR_RETURN(std::unique_ptr<SerialProcRuntime> proc_runtime,
+                       SerialProcRuntime::Create(std::move(proc_interpreters),
+                                                 std::move(queue_manager)));
 
-  // Inject initial values into channels.
-  for (Channel* channel : package->channels()) {
-    ChannelQueue& queue = proc_runtime->queue_manager().GetQueue(channel);
+  // Inject initial values into channel queues.
+  for (ChannelInstance* channel_instance :
+       proc_runtime->elaboration().channel_instances()) {
+    Channel* channel = channel_instance->channel;
+    ChannelQueue& queue =
+        proc_runtime->queue_manager().GetQueue(channel_instance);
     for (const Value& value : channel->initial_values()) {
       XLS_RETURN_IF_ERROR(queue.Write(value));
     }
@@ -65,4 +64,19 @@ CreateInterpreterSerialProcRuntime(Package* package) {
   return std::move(proc_runtime);
 }
 
+}  // namespace
+
+absl::StatusOr<std::unique_ptr<SerialProcRuntime>>
+CreateInterpreterSerialProcRuntime(Package* package) {
+  XLS_ASSIGN_OR_RETURN(Elaboration elaboration,
+                       Elaboration::ElaborateOldStylePackage(package));
+  return CreateRuntime(std::move(elaboration));
+}
+
+absl::StatusOr<std::unique_ptr<SerialProcRuntime>>
+CreateInterpreterSerialProcRuntime(Proc* top) {
+  XLS_ASSIGN_OR_RETURN(Elaboration elaboration, Elaboration::Elaborate(top));
+  return CreateRuntime(std::move(elaboration));
+}
+
 }  // namespace xls

xls/interpreter/interpreter_proc_runtime.h

@@ -23,10 +23,16 @@
 
 namespace xls {
 
-// Create a SerialProcRuntime composed of ProcInterpreters.
+// Create a SerialProcRuntime composed of ProcInterpreters. Supports old-style
+// procs.
 absl::StatusOr<std::unique_ptr<SerialProcRuntime>>
 CreateInterpreterSerialProcRuntime(Package* package);
 
+// Create a SerialProcRuntime composed of ProcInterpreters. Constructed from the
+// elaboration of the given proc. Supports new-style procs.
+absl::StatusOr<std::unique_ptr<SerialProcRuntime>>
+CreateInterpreterSerialProcRuntime(Proc* top);
+
 }  // namespace xls
 
 #endif  // XLS_INTERPRETER_INTERPRETER_PROC_RUNTIME_H_

xls/interpreter/proc_interpreter.cc

@@ -17,6 +17,7 @@
 #include <cstdint>
 #include <memory>
 #include <optional>
+#include <string_view>
 #include <utility>
 #include <vector>
 
@@ -47,23 +48,25 @@ namespace {
 class ProcIrInterpreter : public IrInterpreter {
  public:
   // Constructor args:
+  //   proc_instance: the instance of the proc which is being interpreted.
   //   state: is the value to use for the proc state in the tick being
   //     interpreted.
   //   node_values: map from Node to Value for already computed values in this
   //     tick of the proc. Used for continuations.
   //   events: events object to record events in (e.g, traces).
   //   queue_manager: manager for channel queues.
-  ProcIrInterpreter(absl::Span<const Value> state,
+  ProcIrInterpreter(ProcInstance* proc_instance, absl::Span<const Value> state,
                     absl::flat_hash_map<Node*, Value>* node_values,
                     InterpreterEvents* events,
                     ChannelQueueManager* queue_manager)
       : IrInterpreter(node_values, events),
+        proc_instance_(proc_instance),
         state_(state.begin(), state.end()),
         queue_manager_(queue_manager) {}
 
   absl::Status HandleReceive(Receive* receive) override {
-    XLS_ASSIGN_OR_RETURN(ChannelQueue * queue, queue_manager_->GetQueueByName(
-                                                   receive->channel_name()));
+    XLS_ASSIGN_OR_RETURN(ChannelQueue * queue,
+                         GetChannelQueue(receive->channel_name()));
 
     if (receive->predicate().has_value()) {
       const Bits& pred = ResolveAsBits(receive->predicate().value());
@@ -97,7 +100,7 @@ class ProcIrInterpreter : public IrInterpreter {
 
   absl::Status HandleSend(Send* send) override {
     XLS_ASSIGN_OR_RETURN(ChannelQueue * queue,
-                         queue_manager_->GetQueueByName(send->channel_name()));
+                         GetChannelQueue(send->channel_name()));
     if (send->predicate().has_value()) {
       const Bits& pred = ResolveAsBits(send->predicate().value());
       if (pred.IsZero()) {
@@ -139,6 +142,21 @@ class ProcIrInterpreter : public IrInterpreter {
   }
 
  private:
+  // Get the channel queue for the channel or channel reference of the given
+  // name.
+  absl::StatusOr<ChannelQueue*> GetChannelQueue(std::string_view name) {
+    if (proc_instance_->path().has_value()) {
+      // New-style proc-scoped channel.
+      XLS_ASSIGN_OR_RETURN(ChannelInstance * channel_instance,
+                           queue_manager_->elaboration().GetChannelInstance(
+                               name, *proc_instance_->path()));
+      return &queue_manager_->GetQueue(channel_instance);
+    }
+    // Old-style global channel.
+    return queue_manager_->GetQueueByName(name);
+  }
+
+  ProcInstance* proc_instance_;
   std::vector<Value> state_;
   ChannelQueueManager* queue_manager_;
 
@@ -167,8 +185,9 @@ absl::StatusOr<TickResult> ProcInterpreter::Tick(
   XLS_RET_CHECK_NE(cont, nullptr) << "ProcInterpreter requires a continuation "
                                      "of type ProcInterpreterContinuation";
 
-  ProcIrInterpreter ir_interpreter(cont->GetState(), &cont->GetNodeValues(),
-                                   &cont->GetEvents(), queue_manager_);
+  ProcIrInterpreter ir_interpreter(cont->proc_instance(), cont->GetState(),
+                                   &cont->GetNodeValues(), &cont->GetEvents(),
+                                   queue_manager_);
 
   // Resume execution at the node indicated in the continuation
   // (NodeExecutionIndex).

xls/interpreter/proc_runtime.cc

@@ -14,7 +14,6 @@
 
 #include "xls/interpreter/proc_runtime.h"
 
-#include <algorithm>
 #include <cstdint>
 #include <memory>
 #include <optional>
@@ -42,11 +41,9 @@
 namespace xls {
 
 ProcRuntime::ProcRuntime(
-    Package* package,
     absl::flat_hash_map<Proc*, std::unique_ptr<ProcEvaluator>>&& evaluators,
     std::unique_ptr<ChannelQueueManager>&& queue_manager)
-    : package_(package),
-      queue_manager_(std::move(queue_manager)),
+    : queue_manager_(std::move(queue_manager)),
       evaluators_(std::move(evaluators)) {
   for (ProcInstance* instance : elaboration().proc_instances()) {
     std::unique_ptr<ProcContinuation> continuation =
@@ -72,43 +69,57 @@ absl::Status ProcRuntime::Tick() {
 }
 
 absl::StatusOr<int64_t> ProcRuntime::TickUntilOutput(
-    absl::flat_hash_map<Channel*, int64_t> output_counts,
+    const absl::flat_hash_map<Channel*, int64_t>& output_counts,
+    std::optional<int64_t> max_ticks) {
+  absl::flat_hash_map<ChannelInstance*, int64_t> instance_output_counts;
+  for (const auto& [channel, count] : output_counts) {
+    XLS_ASSIGN_OR_RETURN(ChannelInstance * channel_instance,
+                         elaboration().GetUniqueInstance(channel));
+    instance_output_counts[channel_instance] = count;
+  }
+  return TickUntilOutput(instance_output_counts, max_ticks);
+}
+
+absl::StatusOr<int64_t> ProcRuntime::TickUntilOutput(
+    const absl::flat_hash_map<ChannelInstance*, int64_t>& output_counts,
     std::optional<int64_t> max_ticks) {
   XLS_VLOG(3) << absl::StreamFormat("TickUntilOutput on package %s",
-                                    package_->name());
+                                    package()->name());
   // Create a deterministically sorted vector of the output channels for
   // deterministic behavior and error messages.
-  std::vector<Channel*> output_channels;
-  for (auto [channel, _] : output_counts) {
-    output_channels.push_back(channel);
+  std::vector<ChannelInstance*> output_channels;
+  for (ChannelInstance* channel_instance : elaboration().channel_instances()) {
+    if (output_counts.contains(channel_instance)) {
+      output_channels.push_back(channel_instance);
+    }
   }
-  std::sort(output_channels.begin(), output_channels.end(),
-            [](Channel* a, Channel* b) { return a->name() < b->name(); });
 
   if (XLS_VLOG_IS_ON(3)) {
-    XLS_VLOG(3) << "Expected outputs produced for each channel:";
-    for (Channel* channel : output_channels) {
-      XLS_VLOG(3) << absl::StreamFormat("  %s : %d", channel->name(),
-                                        output_counts.at(channel));
+    XLS_VLOG(3) << "Expected outputs produced for each channel instance:";
+    for (ChannelInstance* channel_instance : output_channels) {
+      XLS_VLOG(3) << absl::StreamFormat("  %s : %d",
+                                        channel_instance->ToString(),
+                                        output_counts.at(channel_instance));
     }
   }
 
-  for (Channel* channel : output_channels) {
+  for (ChannelInstance* channel_instance : output_channels) {
+    Channel* channel = channel_instance->channel;
     if (channel->supported_ops() != ChannelOps::kSendOnly) {
       return absl::InvalidArgumentError(absl::StrFormat(
           "Channel `%s` is not a send-only channel", channel->name()));
     }
     if (channel->kind() == ChannelKind::kSingleValue &&
-        output_counts.at(channel) > 1) {
+        output_counts.at(channel_instance) > 1) {
       return absl::InvalidArgumentError(
           absl::StrFormat("Channel `%s` is single-value, expected number of "
                           "elements must be one or less, is %d",
-                          channel->name(), output_counts.at(channel)));
+                          channel->name(), output_counts.at(channel_instance)));
     }
   }
   int64_t ticks = 0;
   auto needs_more_output = [&]() {
-    for (Channel* ch : output_channels) {
+    for (ChannelInstance* ch : output_channels) {
       if (queue_manager().GetQueue(ch).GetSize() < output_counts.at(ch)) {
         return true;
       }
@@ -131,7 +142,7 @@ absl::StatusOr<int64_t> ProcRuntime::TickUntilOutput(
 absl::StatusOr<int64_t> ProcRuntime::TickUntilBlocked(
     std::optional<int64_t> max_ticks) {
   XLS_VLOG(3) << absl::StreamFormat("TickUntilBlocked on package %s",
-                                    package_->name());
+                                    package()->name());
   int64_t ticks = 0;
   while (!max_ticks.has_value() || ticks < max_ticks.value()) {
     XLS_ASSIGN_OR_RETURN(NetworkTickResult result, TickInternal());

xls/interpreter/proc_runtime.h

@@ -37,7 +37,6 @@ namespace xls {
 class ProcRuntime {
  public:
   ProcRuntime(
-      Package* package,
       absl::flat_hash_map<Proc*, std::unique_ptr<ProcEvaluator>>&& evaluators,
       std::unique_ptr<ChannelQueueManager>&& queue_manager);
 
@@ -53,14 +52,17 @@ class ProcRuntime {
   // error if no progress can be made due to a deadlock.
   absl::Status Tick();
 
-  // Tick the proc network until some output channels have produced at least a
-  // specified number of outputs as indicated by `output_counts`.
-  // `output_counts` must only contain output channels and need not contain all
-  // output channels. Returns the number of ticks executed before the conditions
-  // were met. `max_ticks` is the maximum number of ticks of the proc network
-  // before returning an error.
+  // Tick the proc network until some output channels (channel instances) have
+  // produced at least a specified number of outputs as indicated by
+  // `output_counts`. `output_counts` must only contain output channels and need
+  // not contain all output channels. Returns the number of ticks executed
+  // before the conditions were met. `max_ticks` is the maximum number of ticks
+  // of the proc network before returning an error.
+  absl::StatusOr<int64_t> TickUntilOutput(
+      const absl::flat_hash_map<Channel*, int64_t>& output_counts,
+      std::optional<int64_t> max_ticks = std::nullopt);
   absl::StatusOr<int64_t> TickUntilOutput(
-      absl::flat_hash_map<Channel*, int64_t> output_counts,
+      const absl::flat_hash_map<ChannelInstance*, int64_t>& output_counts,
       std::optional<int64_t> max_ticks = std::nullopt);
 
   // Tick until all procs with IO (send or receive nodes) are blocked on receive
@@ -111,6 +113,8 @@ class ProcRuntime {
     return queue_manager_->elaboration();
   }
 
+  Package* package() const { return elaboration().package(); }
+
  protected:
   // Execute (up to) a single iteration of every proc in the package.
   struct NetworkTickResult {
@@ -124,7 +128,6 @@ class ProcRuntime {
   };
   virtual absl::StatusOr<NetworkTickResult> TickInternal() = 0;
 
-  Package* package_;
   std::unique_ptr<ChannelQueueManager> queue_manager_;
   absl::flat_hash_map<Proc*, std::unique_ptr<ProcEvaluator>> evaluators_;
   absl::flat_hash_map<ProcInstance*, std::unique_ptr<ProcContinuation>>