chi_model.cpp

// SPDX-License-Identifier: MIT

#include "chi/chi_model.h"

#include <algorithm>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <utility>

namespace chi {

namespace {

constexpr std::array<Channel, kChannelCount> kChannels = {
    Channel::Req,
    Channel::Rsp,
    Channel::Dat,
    Channel::Snp,
};

constexpr std::size_t kMaxNodeOutputQueueDepth = 64;

void require_message(const bool condition, const std::string& reason, const ChiMessage& message) {
  if (!condition) {
    throw std::logic_error(reason + ": " + describe(message));
  }
}

} // namespace

ChiNode::ChiNode(std::string name, const NodeId id, NodePort* port, TraceLog* trace)
    : name_(std::move(name)), id_(id), port_(port), trace_(trace) {
  if (port_ == nullptr) {
    throw std::invalid_argument("a CHI node requires a port");
  }
}

void ChiNode::evaluate(const std::uint64_t cycle) {
  // Drive old queued work before handling new inputs. A response generated by
  // handle_message() therefore cannot leave the node until the next cycle.
  drive_outputs(cycle);
  accept_inputs(cycle);
}

bool ChiNode::has_pending_output() const {
  return std::any_of(output_queues_.begin(), output_queues_.end(),
                     [](const auto& queue) { return !queue.empty(); });
}

void ChiNode::enqueue(ChiMessage message) {
  require_message(message.src_id == id_, "node tried to send a message with the wrong SrcID",
                  message);
  auto& queue = output_queues_[channel_index(channel_for(message.opcode))];
  if (queue.size() >= kMaxNodeOutputQueueDepth) {
    throw std::overflow_error("transaction-layer output queue overflow at " + name());
  }
  queue.push_back(std::move(message));
}

void ChiNode::on_message_sent(const ChiMessage&, const std::uint64_t) {}

void ChiNode::drive_outputs(const std::uint64_t cycle) {
  for (const auto channel : kChannels) {
    auto& queue = output_queues_[channel_index(channel)];
    if (queue.empty()) {
      continue;
    }
    if (port_->node_to_fabric.try_send(queue.front(), cycle)) {
      const ChiMessage message = queue.front();
      if (trace_ != nullptr) {
        trace_->emit(cycle, name_, "TX", describe(message));
      }
      queue.pop_front();
      on_message_sent(message, cycle);
    }
  }
}

void ChiNode::accept_inputs(const std::uint64_t cycle) {
  for (const auto channel : kChannels) {
    const auto* message = port_->fabric_to_node.peek_received(channel);
    if (message == nullptr) {
      continue;
    }
    require_message(message->tgt_id == id_, "node received a message with the wrong TgtID",
                    *message);
    if (trace_ != nullptr) {
      trace_->emit(cycle, name_, "RX", describe(*message));
    }
    handle_message(*message, cycle);
    port_->fabric_to_node.consume_received(channel, cycle);
  }
}

RnF::RnF(const NodeId id, const NodeId home_id, NodePort* port, TraceLog* trace)
    : ChiNode(node_name(id), id, port, trace), home_id_(home_id) {}

TxnId RnF::issue_read_shared(const Address address) {
  return issue_read(address, Opcode::ReadShared);
}

TxnId RnF::issue_read_no_snp(const Address address) {
  return issue_read(address, Opcode::ReadNoSnp);
}

TxnId RnF::issue_read(const Address address, const Opcode opcode) {
  if (opcode != Opcode::ReadShared && opcode != Opcode::ReadNoSnp) {
    throw std::invalid_argument("unsupported read opcode");
  }

  const TxnId txn_id = next_txn_id_++;
  transactions_.emplace(txn_id, Transaction{TransactionKind::Read, address});

  ChiMessage request;
  request.opcode = opcode;
  request.src_id = id();
  request.tgt_id = home_id_;
  request.txn_id = txn_id;
  request.address = address;
  request.exp_comp_ack = true;
  enqueue(request);
  return txn_id;
}

TxnId RnF::issue_write_no_snp_full(const Address address, const Data data) {
  const TxnId txn_id = next_txn_id_++;
  Transaction transaction;
  transaction.kind = TransactionKind::Write;
  transaction.address = address;
  transaction.write_data = data;
  transactions_.emplace(txn_id, transaction);

  ChiMessage request;
  request.opcode = Opcode::WriteNoSnpFull;
  request.src_id = id();
  request.tgt_id = home_id_;
  request.txn_id = txn_id;
  request.address = address;
  enqueue(request);
  return txn_id;
}

bool RnF::transaction_complete(const TxnId txn_id) const {
  const auto iterator = transactions_.find(txn_id);
  if (iterator == transactions_.end()) {
    throw std::out_of_range("unknown RN-F transaction ID");
  }
  return iterator->second.complete;
}

Data RnF::read_data(const TxnId txn_id) const {
  const auto iterator = transactions_.find(txn_id);
  if (iterator == transactions_.end() || iterator->second.kind != TransactionKind::Read ||
      !iterator->second.complete) {
    throw std::logic_error("read data requested before a read transaction completed");
  }
  return iterator->second.read_data;
}

std::size_t RnF::outstanding_transactions() const {
  return static_cast<std::size_t>(
      std::count_if(transactions_.begin(), transactions_.end(),
                    [](const auto& entry) { return !entry.second.complete; }));
}

void RnF::handle_message(const ChiMessage& message, const std::uint64_t cycle) {
  switch (message.opcode) {
  case Opcode::CompData: {
    const auto iterator = transactions_.find(message.txn_id);
    require_message(iterator != transactions_.end(), "CompData has an unknown TxnID", message);
    auto& transaction = iterator->second;
    require_message(transaction.kind == TransactionKind::Read,
                    "CompData matched a non-read transaction", message);
    transaction.read_data = message.data;
    transaction.complete = true;

    if (message.exp_comp_ack) {
      require_message(message.has_home && message.has_dbid,
                      "CompData requiring CompAck lacks HomeNID or DBID", message);
      ChiMessage acknowledgement;
      acknowledgement.opcode = Opcode::CompAck;
      acknowledgement.src_id = id();
      acknowledgement.tgt_id = message.home_nid;
      acknowledgement.txn_id = message.dbid;
      enqueue(acknowledgement);
    }
    break;
  }
  case Opcode::CompDBIDResp: {
    const auto iterator = transactions_.find(message.txn_id);
    require_message(iterator != transactions_.end(), "CompDBIDResp has an unknown TxnID", message);
    auto& transaction = iterator->second;
    require_message(transaction.kind == TransactionKind::Write && message.has_dbid,
                    "invalid CompDBIDResp for RN-F write", message);
    transaction.has_dbid = true;
    transaction.dbid = message.dbid;
    write_dbid_to_txn_.emplace(message.dbid, message.txn_id);

    ChiMessage data;
    data.opcode = Opcode::NCBWrData;
    data.src_id = id();
    data.tgt_id = message.src_id;
    data.txn_id = message.dbid;
    data.address = transaction.address;
    data.data = transaction.write_data;
    enqueue(data);
    break;
  }
  default:
    require_message(false, "RN-F received an unsupported opcode", message);
  }

  (void)cycle;
}

void RnF::on_message_sent(const ChiMessage& message, const std::uint64_t cycle) {
  if (message.opcode != Opcode::NCBWrData) {
    return;
  }
  const auto dbid_iterator = write_dbid_to_txn_.find(message.txn_id);
  require_message(dbid_iterator != write_dbid_to_txn_.end(),
                  "RN-F sent write data with an unknown DBID", message);
  const TxnId original_txn_id = dbid_iterator->second;
  auto& transaction = transactions_.at(original_txn_id);
  transaction.complete = true;
  write_dbid_to_txn_.erase(dbid_iterator);
  if (trace() != nullptr) {
    trace()->emit(cycle, name(), "DONE", "write Txn=" + std::to_string(original_txn_id));
  }
}

HnF::HnF(const NodeId id, const NodeId subordinate_id, const bool enable_dmt, NodePort* port,
         TraceLog* trace)
    : ChiNode(node_name(id), id, port, trace), subordinate_id_(subordinate_id),
      enable_dmt_(enable_dmt) {}

void HnF::handle_message(const ChiMessage& message, const std::uint64_t cycle) {
  switch (message.opcode) {
  case Opcode::ReadShared:
  case Opcode::ReadNoSnp:
    accept_read_request(message, cycle);
    break;
  case Opcode::CompData:
    accept_read_data(message, cycle);
    break;
  case Opcode::CompAck:
    accept_comp_ack(message, cycle);
    break;
  case Opcode::WriteNoSnpFull:
    accept_write_request(message, cycle);
    break;
  case Opcode::NCBWrData:
    accept_write_data(message, cycle);
    break;
  case Opcode::CompDBIDResp:
    accept_subordinate_dbid(message, cycle);
    break;
  default:
    require_message(false, "HN-F received an unsupported opcode", message);
  }
}

void HnF::accept_read_request(const ChiMessage& message, const std::uint64_t cycle) {
  const TxnId home_txn_id = next_home_txn_id_++;
  const bool use_dmt = enable_dmt_;
  reads_.emplace(home_txn_id, ReadContext{message.src_id, message.txn_id, home_txn_id,
                                          message.address, use_dmt});

  // On a Home cache/directory miss, HN-F obtains memory data with ReadNoSnp.
  // DMT changes the return identifiers so SN-F sends CompData to the original
  // RN-F instead of returning it to HN-F first.
  ChiMessage downstream;
  downstream.opcode = Opcode::ReadNoSnp;
  downstream.src_id = id();
  downstream.tgt_id = subordinate_id_;
  downstream.txn_id = home_txn_id;
  downstream.address = message.address;
  // Without DMT, HN-F is the downstream requester and consumes SN-F's
  // completion locally. The original RN-F still acknowledges the separate
  // HN-F response, but SN-F must not wait for that acknowledgement.
  downstream.exp_comp_ack = use_dmt && message.exp_comp_ack;
  downstream.has_home = true;
  downstream.home_nid = id();
  downstream.has_dbid = true;
  downstream.dbid = home_txn_id;
  if (use_dmt) {
    downstream.has_return = true;
    downstream.return_nid = message.src_id;
    downstream.return_txn_id = message.txn_id;
  }
  enqueue(downstream);

  if (trace() != nullptr) {
    trace()->emit(cycle, name(), "ALLOC",
                  "read HomeTxn=" + std::to_string(home_txn_id) +
                      (use_dmt ? " path=DMT" : " path=via-HN"));
  }
}

void HnF::accept_read_data(const ChiMessage& message, const std::uint64_t) {
  const auto iterator = reads_.find(message.txn_id);
  require_message(iterator != reads_.end(), "HN-F received CompData for an unknown Home TxnID",
                  message);
  const auto& context = iterator->second;
  require_message(!context.dmt, "DMT CompData unexpectedly returned through HN-F", message);

  ChiMessage response;
  response.opcode = Opcode::CompData;
  response.src_id = id();
  response.tgt_id = context.requester_id;
  response.txn_id = context.requester_txn_id;
  response.address = context.address;
  response.data = message.data;
  response.exp_comp_ack = true;
  response.has_home = true;
  response.home_nid = id();
  response.has_dbid = true;
  response.dbid = context.home_txn_id;
  enqueue(response);
}

void HnF::accept_comp_ack(const ChiMessage& message, const std::uint64_t cycle) {
  const auto iterator = reads_.find(message.txn_id);
  require_message(iterator != reads_.end(), "HN-F received CompAck for an unknown DBID", message);
  if (trace() != nullptr) {
    trace()->emit(cycle, name(), "FREE",
                  "read HomeTxn=" + std::to_string(iterator->second.home_txn_id));
  }
  reads_.erase(iterator);
  ++completed_reads_;
}

void HnF::accept_write_request(const ChiMessage& message, const std::uint64_t cycle) {
  const TxnId subordinate_txn_id = next_home_txn_id_++;
  const TxnId requester_dbid = next_requester_dbid_++;
  writes_.emplace(subordinate_txn_id, WriteContext{message.src_id, message.txn_id, requester_dbid,
                                                   subordinate_txn_id, message.address});
  requester_dbid_to_write_.emplace(requester_dbid, subordinate_txn_id);

  // CompDBIDResp combines completion visibility with allocation of a Home data
  // buffer ID. RN-F must use this DBID, not its original request TxnID, on DAT.
  ChiMessage response;
  response.opcode = Opcode::CompDBIDResp;
  response.src_id = id();
  response.tgt_id = message.src_id;
  response.txn_id = message.txn_id;
  response.has_dbid = true;
  response.dbid = requester_dbid;
  enqueue(response);

  ChiMessage downstream;
  downstream.opcode = Opcode::WriteNoSnpFull;
  downstream.src_id = id();
  downstream.tgt_id = subordinate_id_;
  downstream.txn_id = subordinate_txn_id;
  downstream.address = message.address;
  enqueue(downstream);

  if (trace() != nullptr) {
    trace()->emit(cycle, name(), "ALLOC",
                  "write HomeTxn=" + std::to_string(subordinate_txn_id) +
                      " RN-DBID=" + std::to_string(requester_dbid));
  }
}

void HnF::accept_write_data(const ChiMessage& message, const std::uint64_t cycle) {
  const auto mapping = requester_dbid_to_write_.find(message.txn_id);
  require_message(mapping != requester_dbid_to_write_.end(),
                  "HN-F received write data for an unknown requester DBID", message);
  auto& context = writes_.at(mapping->second);
  context.has_data = true;
  context.data = message.data;
  forward_write_if_ready(mapping->second, cycle);
}

void HnF::accept_subordinate_dbid(const ChiMessage& message, const std::uint64_t cycle) {
  const auto iterator = writes_.find(message.txn_id);
  require_message(iterator != writes_.end() && message.has_dbid,
                  "HN-F received an invalid subordinate CompDBIDResp", message);
  iterator->second.subordinate_ready = true;
  iterator->second.subordinate_dbid = message.dbid;
  forward_write_if_ready(message.txn_id, cycle);
}

void HnF::forward_write_if_ready(const TxnId subordinate_txn_id, const std::uint64_t cycle) {
  const auto iterator = writes_.find(subordinate_txn_id);
  if (iterator == writes_.end()) {
    throw std::logic_error("unknown HN-F write context");
  }
  const auto context = iterator->second;
  if (!context.has_data || !context.subordinate_ready) {
    return;
  }

  ChiMessage data;
  data.opcode = Opcode::NCBWrData;
  data.src_id = id();
  data.tgt_id = subordinate_id_;
  data.txn_id = context.subordinate_dbid;
  data.address = context.address;
  data.data = context.data;
  enqueue(data);

  requester_dbid_to_write_.erase(context.requester_dbid);
  writes_.erase(iterator);
  ++completed_writes_;
  if (trace() != nullptr) {
    trace()->emit(cycle, name(), "FREE",
                  "write HomeTxn=" + std::to_string(subordinate_txn_id) + " after forwarding data");
  }
}

SnF::SnF(const NodeId id, NodePort* port, TraceLog* trace)
    : ChiNode(node_name(id), id, port, trace) {}

void SnF::store(const Address address, const Data data) { memory_[address] = data; }

Data SnF::load(const Address address) const {
  const auto iterator = memory_.find(address);
  if (iterator == memory_.end()) {
    throw std::out_of_range("SN-F read from an uninitialized address");
  }
  return iterator->second;
}

void SnF::handle_message(const ChiMessage& message, const std::uint64_t cycle) {
  switch (message.opcode) {
  case Opcode::ReadNoSnp: {
    ChiMessage response;
    response.opcode = Opcode::CompData;
    response.src_id = id();
    response.tgt_id = message.has_return ? message.return_nid : message.src_id;
    response.txn_id = message.has_return ? message.return_txn_id : message.txn_id;
    response.address = message.address;
    response.data = load(message.address);
    response.exp_comp_ack = message.exp_comp_ack;
    response.has_home = true;
    response.home_nid = message.has_home ? message.home_nid : message.src_id;
    response.has_dbid = true;
    response.dbid = message.has_dbid ? message.dbid : message.txn_id;
    enqueue(response);
    break;
  }
  case Opcode::WriteNoSnpFull: {
    const TxnId dbid = next_dbid_++;
    writes_by_dbid_.emplace(dbid, WriteContext{message.address, message.src_id, message.txn_id});

    ChiMessage response;
    response.opcode = Opcode::CompDBIDResp;
    response.src_id = id();
    response.tgt_id = message.src_id;
    response.txn_id = message.txn_id;
    response.has_dbid = true;
    response.dbid = dbid;
    enqueue(response);
    break;
  }
  case Opcode::NCBWrData: {
    const auto iterator = writes_by_dbid_.find(message.txn_id);
    require_message(iterator != writes_by_dbid_.end(),
                    "SN-F received write data for an unknown DBID", message);
    memory_[iterator->second.address] = message.data;
    if (trace() != nullptr) {
      std::ostringstream detail;
      detail << "memory[0x" << std::hex << iterator->second.address << "]=0x" << message.data;
      trace()->emit(cycle, name(), "STORE", detail.str());
    }
    writes_by_dbid_.erase(iterator);
    break;
  }
  default:
    require_message(false, "SN-F received an unsupported opcode", message);
  }
}

ChiFabric::ChiFabric(std::vector<PortBinding> ports, TraceLog* trace)
    : ports_(std::move(ports)), trace_(trace), round_robin_start_(ports_.size()) {
  if (ports_.empty()) {
    throw std::invalid_argument("CHI Fabric requires at least one port");
  }
  for (const auto& binding : ports_) {
    if (binding.port == nullptr) {
      throw std::invalid_argument("CHI Fabric cannot bind a null port");
    }
  }
}

void ChiFabric::evaluate(const std::uint64_t cycle) {
  // Each output channel independently selects one ingress. This is a small
  // round-robin crossbar: routing uses TgtID, while arbitration only resolves
  // simultaneous messages targeting the same output and channel.
  for (std::size_t output_index = 0; output_index < ports_.size(); ++output_index) {
    for (const auto channel : kChannels) {
      const auto channel_slot = channel_index(channel);
      const auto start = round_robin_start_[output_index][channel_slot] % ports_.size();

      for (std::size_t offset = 0; offset < ports_.size(); ++offset) {
        const auto input_index = (start + offset) % ports_.size();
        const auto* message = ports_[input_index].port->node_to_fabric.peek_received(channel);
        if (message == nullptr) {
          continue;
        }

        // Detect an invalid target even though arbitration is output-centric.
        (void)find_output_index(message->tgt_id);
        if (message->tgt_id != ports_[output_index].node_id) {
          continue;
        }
        if (!ports_[output_index].port->fabric_to_node.try_send(*message, cycle)) {
          break;
        }

        if (trace_ != nullptr) {
          trace_->emit(cycle, "FABRIC", "ROUTE", describe(*message));
        }
        ports_[input_index].port->node_to_fabric.consume_received(channel, cycle);
        round_robin_start_[output_index][channel_slot] = (input_index + 1) % ports_.size();
        break;
      }
    }
  }
}

std::size_t ChiFabric::find_output_index(const NodeId target) const {
  for (std::size_t index = 0; index < ports_.size(); ++index) {
    if (ports_[index].node_id == target) {
      return index;
    }
  }
  throw std::logic_error("Fabric has no route to " + node_name(target));
}

ChiSystem::ChiSystem(const ChiSystemConfig& config)
    : trace_(config.trace_enabled, config.trace_credits),
      rn_port_(node_name(kRnF0Id), config.credit_depth, config.forward_latency,
               config.return_latency, &trace_),
      hn_port_(node_name(kHnFId), config.credit_depth, config.forward_latency,
               config.return_latency, &trace_),
      sn_port_(node_name(kSnFId), config.credit_depth, config.forward_latency,
               config.return_latency, &trace_),
      rn_(kRnF0Id, kHnFId, &rn_port_, &trace_),
      hn_(kHnFId, kSnFId, config.enable_dmt, &hn_port_, &trace_), sn_(kSnFId, &sn_port_, &trace_),
      fabric_({{kRnF0Id, &rn_port_}, {kHnFId, &hn_port_}, {kSnFId, &sn_port_}}, &trace_) {}

void ChiSystem::step() {
  rn_.evaluate(cycle_);
  hn_.evaluate(cycle_);
  sn_.evaluate(cycle_);
  fabric_.evaluate(cycle_);
  commit_links();
  ++cycle_;
}

bool ChiSystem::run_until(const std::function<bool()>& predicate, const std::uint64_t max_cycles) {
  for (std::uint64_t iteration = 0; iteration < max_cycles; ++iteration) {
    if (predicate()) {
      return true;
    }
    step();
  }
  return predicate();
}

bool ChiSystem::links_idle() const {
  const std::array<const NodePort*, 3> ports = {&rn_port_, &hn_port_, &sn_port_};
  return std::all_of(ports.begin(), ports.end(), [](const auto* port) {
    return port->node_to_fabric.idle() && port->fabric_to_node.idle();
  });
}

void ChiSystem::print_link_stats(std::ostream& output) const {
  struct NamedLink {
    const char* name;
    const ChiLink* link;
  };
  const std::array<NamedLink, 6> links = {{{"RN-F0 -> Fabric", &rn_port_.node_to_fabric},
                                           {"Fabric -> RN-F0", &rn_port_.fabric_to_node},
                                           {"HN-F -> Fabric", &hn_port_.node_to_fabric},
                                           {"Fabric -> HN-F", &hn_port_.fabric_to_node},
                                           {"SN-F -> Fabric", &sn_port_.node_to_fabric},
                                           {"Fabric -> SN-F", &sn_port_.fabric_to_node}}};

  output << "\nLink/channel credit statistics\n";
  output << std::left << std::setw(20) << "Link" << std::setw(6) << "Ch" << std::setw(8) << "Depth"
         << std::setw(9) << "Credits" << std::setw(8) << "Flits" << "Stalls\n";
  for (const auto& named_link : links) {
    for (const auto channel : kChannels) {
      const auto stats = named_link.link->channel(channel).stats();
      if (stats.flits_launched == 0 && stats.credit_stalls == 0) {
        continue;
      }
      output << std::left << std::setw(20) << named_link.name << std::setw(6) << to_string(channel)
             << std::setw(8) << stats.depth << std::setw(9) << stats.sender_credits << std::setw(8)
             << stats.flits_launched << stats.credit_stalls << '\n';
    }
  }
}

void ChiSystem::commit_links() {
  const std::array<NodePort*, 3> ports = {&rn_port_, &hn_port_, &sn_port_};
  for (auto* port : ports) {
    port->node_to_fabric.commit(cycle_);
    port->fabric_to_node.commit(cycle_);
  }
}

CreditBenchmarkResult run_credit_benchmark(const std::size_t depth,
                                           const std::size_t forward_latency,
                                           const std::size_t return_latency,
                                           const std::size_t flit_count) {
  if (flit_count == 0) {
    throw std::invalid_argument("credit benchmark requires at least one flit");
  }

  TraceLog trace(false, false);
  CreditChannel channel({"BENCH.REQ", depth, forward_latency, return_latency}, &trace);
  std::size_t launched = 0;
  std::size_t received = 0;
  std::uint64_t cycle = 0;

  while (received < flit_count) {
    if (launched < flit_count) {
      ChiMessage message;
      message.opcode = Opcode::ReadNoSnp;
      message.src_id = kRnF0Id;
      message.tgt_id = kHnFId;
      message.txn_id = static_cast<TxnId>(launched);
      if (channel.try_send(message, cycle)) {
        ++launched;
      }
    }

    if (channel.peek_received() != nullptr) {
      channel.consume_received(cycle);
      ++received;
    }
    channel.commit(cycle);
    ++cycle;

    if (cycle > flit_count * (forward_latency + return_latency + depth + 8)) {
      throw std::runtime_error("credit benchmark did not make forward progress");
    }
  }

  const auto stats = channel.stats();
  return {
      depth,
      flit_count,
      cycle,
      stats.credit_stalls,
      static_cast<double>(flit_count) / static_cast<double>(cycle),
  };
}

} // namespace chi