main.cpp

// SPDX-License-Identifier: MIT

#include "chi/chi_model.h"

#include <cstdlib>
#include <exception>
#include <iomanip>
#include <iostream>
#include <set>
#include <stdexcept>
#include <string>
#include <vector>

namespace {

using chi::Address;
using chi::ChiSystem;
using chi::ChiSystemConfig;
using chi::Data;

struct Options {
  std::string scenario = "all";
  std::size_t credits = 4;
  std::size_t forward_latency = 1;
  std::size_t return_latency = 1;
  bool show_trace = true;
  bool trace_credits = false;
};

std::size_t parse_positive(const std::string& text, const std::string& option) {
  std::size_t consumed = 0;
  const auto value = std::stoull(text, &consumed, 0);
  if (consumed != text.size() || value == 0) {
    throw std::invalid_argument(option + " requires a positive integer");
  }
  return static_cast<std::size_t>(value);
}

void print_usage(const char* program) {
  std::cout << "Usage: " << program << " [options]\n\n"
            << "Options:\n"
            << "  --scenario all|read|dmt|write|credit\n"
            << "  --credits N             Per-hop, per-channel receive depth\n"
            << "  --forward-latency N     Flit latency in cycles\n"
            << "  --return-latency N      L-Credit return latency in cycles\n"
            << "  --trace-credits         Include every L-Credit pulse in the trace\n"
            << "  --no-trace              Print only results and link statistics\n"
            << "  --help                  Show this help\n";
}

Options parse_options(const int argc, char** argv) {
  Options options;
  for (int index = 1; index < argc; ++index) {
    const std::string argument = argv[index];
    const auto require_value = [&]() -> std::string {
      if (index + 1 >= argc) {
        throw std::invalid_argument(argument + " requires a value");
      }
      return argv[++index];
    };

    if (argument == "--scenario") {
      options.scenario = require_value();
    } else if (argument == "--credits") {
      options.credits = parse_positive(require_value(), argument);
    } else if (argument == "--forward-latency") {
      options.forward_latency = parse_positive(require_value(), argument);
    } else if (argument == "--return-latency") {
      options.return_latency = parse_positive(require_value(), argument);
    } else if (argument == "--trace-credits") {
      options.trace_credits = true;
    } else if (argument == "--no-trace") {
      options.show_trace = false;
    } else if (argument == "--help") {
      print_usage(argv[0]);
      std::exit(EXIT_SUCCESS);
    } else {
      throw std::invalid_argument("unknown option: " + argument);
    }
  }

  const std::vector<std::string> scenarios = {"all", "read", "dmt", "write", "credit"};
  bool valid_scenario = false;
  for (const auto& scenario : scenarios) {
    valid_scenario = valid_scenario || options.scenario == scenario;
  }
  if (!valid_scenario) {
    throw std::invalid_argument("unsupported scenario: " + options.scenario);
  }
  return options;
}

ChiSystemConfig system_config(const Options& options, const bool enable_dmt) {
  ChiSystemConfig config;
  config.credit_depth = options.credits;
  config.forward_latency = options.forward_latency;
  config.return_latency = options.return_latency;
  config.enable_dmt = enable_dmt;
  config.trace_enabled = options.show_trace;
  config.trace_credits = options.trace_credits;
  return config;
}

void print_result(ChiSystem& system, const Options& options) {
  if (options.show_trace) {
    system.trace().print(std::cout);
  }
  system.print_link_stats(std::cout);
}

void run_read(const Options& options, const bool dmt) {
  constexpr Address kAddress = 0x80000000;
  constexpr Data kExpectedData = 0x1122334455667788;

  std::cout << "\n=== ReadShared " << (dmt ? "with DMT" : "via HN-F") << " ===\n";
  ChiSystem system(system_config(options, dmt));
  system.sn().store(kAddress, kExpectedData);
  const auto txn_id = system.rn().issue_read_shared(kAddress);

  const bool completed = system.run_until(
      [&]() {
        return system.rn().transaction_complete(txn_id) && system.hn().outstanding_reads() == 0;
      },
      500);
  if (!completed) {
    throw std::runtime_error("read transaction timed out");
  }

  std::cout << "Completed in " << system.cycle() << " cycles, data=0x" << std::hex
            << system.rn().read_data(txn_id) << std::dec << '\n';
  print_result(system, options);
}

void run_write(const Options& options) {
  constexpr Address kAddress = 0x80000040;
  constexpr Data kInitialData = 0;
  constexpr Data kWriteData = 0xa5a55a5adeadbeef;

  std::cout << "\n=== WriteNoSnpFull with DBID flow ===\n";
  ChiSystem system(system_config(options, true));
  system.sn().store(kAddress, kInitialData);
  const auto txn_id = system.rn().issue_write_no_snp_full(kAddress, kWriteData);

  const bool completed = system.run_until(
      [&]() {
        return system.rn().transaction_complete(txn_id) && system.hn().outstanding_writes() == 0 &&
               !system.sn().has_pending_write() && system.sn().load(kAddress) == kWriteData;
      },
      500);
  if (!completed) {
    throw std::runtime_error("write transaction timed out");
  }

  std::cout << "Completed in " << system.cycle() << " cycles, memory[0x" << std::hex << kAddress
            << "]=0x" << system.sn().load(kAddress) << std::dec << '\n';
  print_result(system, options);
}

void run_credit_sweep(const Options& options) {
  constexpr std::size_t kFlitCount = 32;
  // With registered communicate/update phases, a returned credit is usable
  // one cycle after it arrives. Receiver consumption also takes one cycle.
  const auto full_rate_depth = options.forward_latency + options.return_latency + 2;
  const std::set<std::size_t> depths = {1, 2, 4, full_rate_depth, 8, options.credits};

  std::cout << "\n=== Credit depth throughput sweep ===\n";
  std::cout << "Forward latency=" << options.forward_latency
            << ", return latency=" << options.return_latency << ", flits=" << kFlitCount << "\n\n";
  std::cout << "Model full-rate estimate: credits >= forward + 1 + return + 1 = " << full_rate_depth
            << "\n\n";
  std::cout << std::left << std::setw(10) << "Credits" << std::setw(10) << "Cycles" << std::setw(12)
            << "Stalls" << "Flits/cycle\n";

  for (const auto depth : depths) {
    const auto result = chi::run_credit_benchmark(depth, options.forward_latency,
                                                  options.return_latency, kFlitCount);
    std::cout << std::left << std::setw(10) << result.depth << std::setw(10) << result.cycles
              << std::setw(12) << result.stall_attempts << std::fixed << std::setprecision(3)
              << result.flits_per_cycle << '\n';
  }
}

} // namespace

int main(const int argc, char** argv) {
  try {
    const auto options = parse_options(argc, argv);
    if (options.scenario == "all" || options.scenario == "read") {
      run_read(options, false);
    }
    if (options.scenario == "all" || options.scenario == "dmt") {
      run_read(options, true);
    }
    if (options.scenario == "all" || options.scenario == "write") {
      run_write(options);
    }
    if (options.scenario == "all" || options.scenario == "credit") {
      run_credit_sweep(options);
    }
    return EXIT_SUCCESS;
  } catch (const std::exception& error) {
    std::cerr << "error: " << error.what() << '\n';
    return EXIT_FAILURE;
  }
}