Add peek channel operation

xls/dslx/bytecode/bytecode.cc

@@ -290,6 +290,10 @@ std::string OpToString(Bytecode::Op op) {
       return "negate";
     case Bytecode::Op::kOr:
       return "or";
+    case Bytecode::Op::kPeek:
+      return "peek";
+    case Bytecode::Op::kPeekNonBlocking:
+      return "peek_non_blocking";
     case Bytecode::Op::kPop:
       return "pop";
     case Bytecode::Op::kRange:
@@ -502,6 +506,15 @@ DEF_UNARY_BUILDER(Swap);
   return Bytecode(span, Op::kMatchArm, std::move(item));
 }
 
+/* static */ Bytecode Bytecode::MakePeek(Span span, ChannelData channel_data) {
+  return Bytecode(span, Op::kPeek, std::move(channel_data));
+}
+
+/* static */ Bytecode Bytecode::MakePeekNonBlocking(
+    Span span, ChannelData channel_data) {
+  return Bytecode(span, Op::kPeekNonBlocking, std::move(channel_data));
+}
+
 /* static */ Bytecode Bytecode::MakeRecv(Span span, ChannelData channel_data) {
   return Bytecode(span, Op::kRecv, std::move(channel_data));
 }

xls/dslx/bytecode/bytecode.h

@@ -136,6 +136,22 @@ class Bytecode {
     kPop,
     // Creates an array of values [TOS1, TOS0).
     kRange,
+    // Peeks a value from the channel at TOS1 if condition at TOS0 is fulfilled.
+    // If TOS0 is true, then
+    //   peeks a value from the channel or "blocks"
+    //   if empty: terminates execution at the opcode's PC. The interpreter can
+    //    be resumed/retried if/when a value becomes available.
+    // else
+    //   a tuple containing a tuple and zero value is pushed on the stack.
+    kPeek,
+    // Peeks a value off of the channel at TOS0, but does not block if empty.
+    // A tuple containing
+    //   0. A token.
+    //   1. Peeked value (or a zero value if the channel is empty).
+    //     and
+    //   2. A valid flag (false if the channel is empty).
+    // is pushed on the stack.
+    kPeekNonBlocking,
     // Pulls TOS0 (a condition) and TOS1 (a channel).
     // If TOS0 is true, then
     //   pulls a value off of the channel or "blocks"
@@ -390,6 +406,8 @@ class Bytecode {
   static Bytecode MakeLoad(Span span, SlotIndex slot_index);
   static Bytecode MakeLogicalOr(Span span);
   static Bytecode MakeMatchArm(Span span, MatchArmItem item);
+  static Bytecode MakePeek(Span span, ChannelData channel_data);
+  static Bytecode MakePeekNonBlocking(Span span, ChannelData channel_data);
   static Bytecode MakePop(Span span);
   static Bytecode MakeRecv(Span span, ChannelData channel_data);
   static Bytecode MakeRecvNonBlocking(Span span, ChannelData channel_data);

xls/dslx/bytecode/bytecode_emitter.cc

@@ -538,6 +538,79 @@ absl::Status BytecodeEmitter::HandleCast(const Cast* node) {
   return absl::OkStatus();
 }
 
+absl::Status BytecodeEmitter::HandleBuiltinPeek(const Invocation* node) {
+  Expr* token = node->args()[0];
+  Expr* channel = node->args()[1];
+
+  XLS_RETURN_IF_ERROR(token->AcceptExpr(this));
+  XLS_RETURN_IF_ERROR(channel->AcceptExpr(this));
+  // All receives need a predicate. Set to true for unconditional receive.
+  Add(Bytecode::MakeLiteral(node->span(), InterpValue::MakeUBits(1, 1)));
+  XLS_ASSIGN_OR_RETURN(
+      Bytecode::ChannelData channel_data,
+      CreateChannelData(channel, type_info_, options_.format_preference));
+  // Default value which is unused because the predicate is always
+  // true. Required because the `peek` bytecode has a predicate and
+  // a default value operand.
+  XLS_ASSIGN_OR_RETURN(InterpValue default_value,
+                       CreateZeroValueFromType(channel_data.payload_type()));
+  Add(Bytecode::MakeLiteral(node->span(), default_value));
+  Add(Bytecode::MakePeek(node->span(), std::move(channel_data)));
+  return absl::OkStatus();
+}
+
+absl::Status BytecodeEmitter::HandleBuiltinPeekNonBlocking(
+    const Invocation* node) {
+  Expr* token = node->args()[0];
+  Expr* channel = node->args()[1];
+  Expr* default_value = node->args()[2];
+
+  XLS_RETURN_IF_ERROR(token->AcceptExpr(this));
+  XLS_RETURN_IF_ERROR(channel->AcceptExpr(this));
+  Add(Bytecode::MakeLiteral(node->span(), InterpValue::MakeUBits(1, 1)));
+  XLS_RETURN_IF_ERROR(default_value->AcceptExpr(this));
+  XLS_ASSIGN_OR_RETURN(
+      Bytecode::ChannelData channel_data,
+      CreateChannelData(channel, type_info_, options_.format_preference));
+  Add(Bytecode::MakePeekNonBlocking(node->span(), std::move(channel_data)));
+  return absl::OkStatus();
+}
+
+absl::Status BytecodeEmitter::HandleBuiltinPeekIf(const Invocation* node) {
+  Expr* token = node->args()[0];
+  Expr* channel = node->args()[1];
+  Expr* condition = node->args()[2];
+  Expr* default_value = node->args()[3];
+
+  XLS_RETURN_IF_ERROR(token->AcceptExpr(this));
+  XLS_RETURN_IF_ERROR(channel->AcceptExpr(this));
+  XLS_RETURN_IF_ERROR(condition->AcceptExpr(this));
+  XLS_RETURN_IF_ERROR(default_value->AcceptExpr(this));
+  XLS_ASSIGN_OR_RETURN(
+      Bytecode::ChannelData channel_data,
+      CreateChannelData(channel, type_info_, options_.format_preference));
+  Add(Bytecode::MakePeek(node->span(), std::move(channel_data)));
+  return absl::OkStatus();
+}
+
+absl::Status BytecodeEmitter::HandleBuiltinPeekIfNonBlocking(
+    const Invocation* node) {
+  Expr* token = node->args()[0];
+  Expr* channel = node->args()[1];
+  Expr* condition = node->args()[2];
+  Expr* default_value = node->args()[3];
+
+  XLS_RETURN_IF_ERROR(token->AcceptExpr(this));
+  XLS_RETURN_IF_ERROR(channel->AcceptExpr(this));
+  XLS_RETURN_IF_ERROR(condition->AcceptExpr(this));
+  XLS_RETURN_IF_ERROR(default_value->AcceptExpr(this));
+  XLS_ASSIGN_OR_RETURN(
+      Bytecode::ChannelData channel_data,
+      CreateChannelData(channel, type_info_, options_.format_preference));
+  Add(Bytecode::MakePeekNonBlocking(node->span(), std::move(channel_data)));
+  return absl::OkStatus();
+}
+
 absl::Status BytecodeEmitter::HandleBuiltinRecv(const Invocation* node) {
   Expr* token = node->args()[0];
   Expr* channel = node->args()[1];
@@ -1142,6 +1215,18 @@ absl::Status BytecodeEmitter::HandleInvocation(const Invocation* node) {
     if (name_ref->identifier() == "join") {
       return HandleBuiltinJoin(node);
     }
+    if (name_ref->identifier() == "peek") {
+      return HandleBuiltinPeek(node);
+    }
+    if (name_ref->identifier() == "peek_non_blocking") {
+      return HandleBuiltinPeekNonBlocking(node);
+    }
+    if (name_ref->identifier() == "peek_if") {
+      return HandleBuiltinPeekIf(node);
+    }
+    if (name_ref->identifier() == "peek_if_non_blocking") {
+      return HandleBuiltinPeekIfNonBlocking(node);
+    }
     if (name_ref->identifier() == "recv") {
       return HandleBuiltinRecv(node);
     }

xls/dslx/bytecode/bytecode_emitter.h

@@ -167,6 +167,10 @@ class BytecodeEmitter : public ExprVisitor {
   absl::Status HandleBuiltinDecode(const Invocation* node);
   absl::Status HandleBuiltinElementCount(const Invocation* node);
   absl::Status HandleBuiltinJoin(const Invocation* node);
+  absl::Status HandleBuiltinPeek(const Invocation* node);
+  absl::Status HandleBuiltinPeekIf(const Invocation* node);
+  absl::Status HandleBuiltinPeekIfNonBlocking(const Invocation* node);
+  absl::Status HandleBuiltinPeekNonBlocking(const Invocation* node);
   absl::Status HandleBuiltinRecv(const Invocation* node);
   absl::Status HandleBuiltinRecvIf(const Invocation* node);
   absl::Status HandleBuiltinRecvIfNonBlocking(const Invocation* node);

xls/dslx/bytecode/bytecode_interpreter.cc

@@ -585,6 +585,14 @@ absl::Status BytecodeInterpreter::EvalNextInstruction() {
       XLS_RETURN_IF_ERROR(EvalOr(bytecode));
       break;
     }
+    case Bytecode::Op::kPeek: {
+      XLS_RETURN_IF_ERROR(EvalPeek(bytecode));
+      break;
+    }
+    case Bytecode::Op::kPeekNonBlocking: {
+      XLS_RETURN_IF_ERROR(EvalPeekNonBlocking(bytecode));
+      break;
+    }
     case Bytecode::Op::kPop: {
       XLS_RETURN_IF_ERROR(EvalPop(bytecode));
       break;
@@ -1326,6 +1334,83 @@ absl::Status BytecodeInterpreter::EvalOr(const Bytecode& bytecode) {
   });
 }
 
+absl::Status BytecodeInterpreter::EvalPeek(const Bytecode& bytecode) {
+  XLS_ASSIGN_OR_RETURN(InterpValue default_value, Pop());
+  XLS_ASSIGN_OR_RETURN(InterpValue condition, Pop());
+  XLS_ASSIGN_OR_RETURN(InterpValue channel_value, Pop());
+  XLS_ASSIGN_OR_RETURN(auto channel_reference,
+                       channel_value.GetChannelReference());
+  XLS_ASSIGN_OR_RETURN(const Bytecode::ChannelData* channel_data,
+                       bytecode.channel_data());
+
+  XLS_RET_CHECK(channel_reference.GetChannelId().has_value());
+  int64_t channel_id = channel_reference.GetChannelId().value();
+  XLS_RET_CHECK(channel_manager_.has_value());
+  InterpValueChannel& channel = (*channel_manager_)->GetChannel(channel_id);
+
+  if (condition.IsTrue()) {
+    if (channel.IsEmpty()) {
+      stack_.Push(channel_value);
+      stack_.Push(condition);
+      stack_.Push(default_value);
+      blocked_channel_info_ = BlockedChannelInfo{
+          .name = FormatChannelNameForTracing(*channel_data),
+          .span = bytecode.source_span(),
+      };
+      return absl::UnavailableError("Channel is empty.");
+    }
+
+    XLS_ASSIGN_OR_RETURN(InterpValue token, Pop());
+    InterpValue value = channel.Peek();
+    if (options_.trace_channels() && events_.has_value()) {
+      (*events_)->AddTraceChannelMessage(
+          import_data_->file_table(), bytecode.source_span(),
+          FormatChannelNameForTracing(*channel_data), value,
+          ChannelDirection::kIn, channel_data->value_fmt_desc());
+    }
+    stack_.Push(InterpValue::MakeTuple({token, std::move(value)}));
+  } else {
+    XLS_ASSIGN_OR_RETURN(InterpValue token, Pop());
+    stack_.Push(InterpValue::MakeTuple({token, default_value}));
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status BytecodeInterpreter::EvalPeekNonBlocking(
+    const Bytecode& bytecode) {
+  XLS_ASSIGN_OR_RETURN(InterpValue default_value, Pop());
+  XLS_ASSIGN_OR_RETURN(InterpValue condition, Pop());
+  XLS_ASSIGN_OR_RETURN(InterpValue channel_value, Pop());
+  XLS_ASSIGN_OR_RETURN(InterpValue::ChannelReference channel_reference,
+                       channel_value.GetChannelReference());
+  XLS_ASSIGN_OR_RETURN(InterpValue token, Pop());
+
+  XLS_RET_CHECK(channel_reference.GetChannelId().has_value());
+  int64_t channel_id = channel_reference.GetChannelId().value();
+  XLS_RET_CHECK(channel_manager_.has_value());
+  InterpValueChannel& channel = (*channel_manager_)->GetChannel(channel_id);
+
+  XLS_ASSIGN_OR_RETURN(const Bytecode::ChannelData* channel_data,
+                       bytecode.channel_data());
+  if (condition.IsTrue() && !channel.IsEmpty()) {
+    InterpValue value = channel.Peek();
+    if (options_.trace_channels() && events_.has_value()) {
+      (*events_)->AddTraceChannelMessage(
+          import_data_->file_table(), bytecode.source_span(),
+          FormatChannelNameForTracing(*channel_data), value,
+          ChannelDirection::kIn, channel_data->value_fmt_desc());
+    }
+    stack_.Push(InterpValue::MakeTuple(
+        {token, std::move(value), InterpValue::MakeBool(true)}));
+  } else {
+    stack_.Push(InterpValue::MakeTuple(
+        {token, default_value, InterpValue::MakeBool(false)}));
+  }
+
+  return absl::OkStatus();
+}
+
 absl::Status BytecodeInterpreter::EvalPop(const Bytecode& bytecode) {
   return Pop().status();
 }
@@ -1751,6 +1836,10 @@ absl::Status BytecodeInterpreter::RunBuiltinFn(const Bytecode& bytecode,
     case Builtin::kToken:
     case Builtin::kSend:
     case Builtin::kSendIf:
+    case Builtin::kPeek:
+    case Builtin::kPeekNonBlocking:
+    case Builtin::kPeekIf:
+    case Builtin::kPeekIfNonBlocking:
     case Builtin::kRecv:
     case Builtin::kRecvIf:
     case Builtin::kRecvNonBlocking:

xls/dslx/bytecode/bytecode_interpreter.h

@@ -146,6 +146,7 @@ class InterpValueChannel {
     queue_.pop_front();
     return result;
   }
+  InterpValue Peek() { return queue_.front(); }
   void Write(InterpValue v) { queue_.push_back(std::move(v)); }
 
  private:
@@ -271,6 +272,8 @@ class BytecodeInterpreter {
   absl::Status EvalNe(const Bytecode& bytecode);
   absl::Status EvalNegate(const Bytecode& bytecode);
   absl::Status EvalOr(const Bytecode& bytecode);
+  absl::Status EvalPeek(const Bytecode& bytecode);
+  absl::Status EvalPeekNonBlocking(const Bytecode& bytecode);
   absl::Status EvalPop(const Bytecode& bytecode);
   absl::Status EvalRange(const Bytecode& bytecode);
   absl::Status EvalRecv(const Bytecode& bytecode);

xls/dslx/dslx_builtins.h

@@ -69,6 +69,10 @@ namespace xls::dslx {
   X("token", kToken)                           \
   /* send/recv routines */                     \
   /* keep-sorted start */                   \
+  X("peek", kPeek)                             \
+  X("peek_if", kPeekIf)                        \
+  X("peek_if_nonblocking", kPeekIfNonBlocking) \
+  X("peek_nonblocking", kPeekNonBlocking)      \
   X("recv", kRecv)                             \
   X("recv_if", kRecvIf)                        \
   X("recv_if_nonblocking", kRecvIfNonBlocking) \

xls/dslx/frontend/ast.cc

@@ -248,6 +248,14 @@ std::string_view AstNodeKindToString(AstNodeKind kind) {
       return "index";
     case AstNodeKind::kRange:
       return "range";
+    case AstNodeKind::kPeek:
+      return "peek";
+    case AstNodeKind::kPeekNonBlocking:
+      return "peek-non-blocking";
+    case AstNodeKind::kPeekIf:
+      return "peek-if";
+    case AstNodeKind::kPeekIfNonBlocking:
+      return "peek-if-non-blocking";
     case AstNodeKind::kRecv:
       return "receive";
     case AstNodeKind::kRecvNonBlocking:

xls/dslx/frontend/ast_node.h

@@ -75,6 +75,10 @@ enum class AstNodeKind : uint8_t {
   kProcMember,
   kQuickCheck,
   kRange,
+  kPeek,
+  kPeekIf,
+  kPeekIfNonBlocking,
+  kPeekNonBlocking,
   kRecv,
   kRecvIf,
   kRecvIfNonBlocking,

xls/dslx/frontend/builtin_stubs.x

@@ -75,6 +75,14 @@ fn one_hot_sel<N: u32, M: u32, S: bool>(x: uN[N], y: xN[S][M][N]) -> xN[S][M];
 
 fn or_reduce<N: u32>(x: uN[N]) -> u1;
 
+fn peek_if_non_blocking<T: type>(tok: token, channel: chan<T> in, predicate: bool, value: T) -> (token, T, bool);
+
+fn peek_if<T: type>(tok: token, channel: chan<T> in, predicate: bool, value: T) -> (token, T);
+
+fn peek_non_blocking<T: type>(tok: token, channel: chan<T> in, value: T) -> (token, T, bool);
+
+fn peek<T: type>(tok: token, channel: chan<T> in) -> (token, T);
+
 fn priority_sel<N: u32, M: u32, S: bool>(x: uN[N], y: xN[S][M][N], z: xN[S][M]) -> xN[S][M];
 
 fn read<T: type>(source: State<T>) -> T;

xls/dslx/frontend/builtins_metadata.cc

@@ -111,6 +111,10 @@ const absl::flat_hash_map<std::string, BuiltinsData>& GetParametricBuiltins() {
           // proc scope.
           {"labeled_read", {}},
           {"labeled_write", {}},
+          {"peek", {}},
+          {"peek_if", {}},
+          {"peek_non_blocking", {}},
+          {"peek_if_non_blocking", {}},
           {"send", {}},
           {"send_if", {}},
           {"recv", {}},

xls/dslx/frontend/parser_test.cc

@@ -1460,6 +1460,36 @@ proc producer {
   RoundTrip(std::string(kModule));
 }
 
+TEST_F(ParserTest, ParsePeek) {
+  constexpr std::string_view kModule = R"(struct Packet {
+    id: u32,
+    data: u32,
+}
+proc PeekPacketFiller {
+    req_r: chan<Packet> in;
+    resp_s: chan<Packet> out;
+    config(req_r: chan<Packet> in, resp_s: chan<Packet> out) {
+        (req_r, resp_s)
+    }
+    init {
+        u32:0
+    }
+    next(current_id: u32) {
+        let (tok, packet) = peek(join(), req_r);
+        let (packet, next_state) = if current_id < packet.id {
+            (Packet { id: current_id, data: current_id }, current_id + u32:1)
+        } else {
+            let (tok, packet) = recv(tok, req_r);
+            (packet, packet.id + u32:1)
+        };
+        let tok = send(tok, resp_s, packet);
+        next_state
+    }
+})";
+
+  RoundTrip(std::string(kModule));
+}
+
 TEST_F(ParserTest, ParseSendIfAndRecvIf) {
   constexpr std::string_view kModule = R"(proc producer {
     c: chan<u32> in;