carbon-lang/toolchain/lex/tokenized_buffer_benchmark.cpp

// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
// Exceptions. See /LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include <benchmark/benchmark.h>

#include <algorithm>
#include <utility>

#include "absl/random/random.h"
#include "common/check.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/StringExtras.h"
#include "toolchain/diagnostics/diagnostic_emitter.h"
#include "toolchain/diagnostics/null_diagnostics.h"
#include "toolchain/lex/token_kind.h"
#include "toolchain/lex/tokenized_buffer.h"

namespace Carbon::Lex {
namespace {

// A large value for measurement stability without making benchmarking too slow.
// Needs to be a multiple of 100 so we can easily divide it up into percentages,
// and 1% itself needs to not be too tiny. This makes 100,000 a great balance.
constexpr int NumTokens = 100'000;

auto IdentifierStartChars() -> llvm::ArrayRef<char> {
  static llvm::SmallVector<char> chars = [] {
    llvm::SmallVector<char> chars;
    chars.push_back('_');
    for (char c : llvm::seq_inclusive('A', 'Z')) {
      chars.push_back(c);
    }
    for (char c : llvm::seq_inclusive('a', 'z')) {
      chars.push_back(c);
    }
    return chars;
  }();
  return chars;
}

auto IdentifierChars() -> llvm::ArrayRef<char> {
  static llvm::SmallVector<char> chars = [] {
    llvm::ArrayRef<char> start_chars = IdentifierStartChars();
    llvm::SmallVector<char> chars(start_chars.begin(), start_chars.end());
    for (char c : llvm::seq_inclusive('0', '9')) {
      chars.push_back(c);
    }
    return chars;
  }();
  return chars;
}

// Generates a random identifier string of the specified length using the
// provided RNG BitGen.
auto GenerateRandomIdentifier(absl::BitGen& gen, int length) -> std::string {
  llvm::ArrayRef<char> start_chars = IdentifierStartChars();
  llvm::ArrayRef<char> chars = IdentifierChars();

  std::string id_result;
  llvm::raw_string_ostream os(id_result);
  llvm::StringRef id;
  do {
    // Erase any prior attempts to find an identifier.
    id_result.clear();
    os << start_chars[absl::Uniform<int>(gen, 0, start_chars.size())];
    for (int j : llvm::seq(0, length)) {
      static_cast<void>(j);
      os << chars[absl::Uniform<int>(gen, 0, chars.size())];
    }
    // Check if we ended up forming an integer type literal or a keyword, and
    // try again.
    id = llvm::StringRef(id_result);
  } while (
      llvm::any_of(TokenKind::KeywordTokens,
                   [id](auto token) { return id == token.fixed_spelling(); }) ||
      ((id.consume_front("i") || id.consume_front("u") ||
        id.consume_front("f")) &&
       llvm::all_of(id, [](const char c) { return llvm::isDigit(c); })));
  return id_result;
}

// Get a static pool of random identifiers with the desired distribution.
template <int MinLength = 1, int MaxLength = 64, bool Uniform = false>
auto GetRandomIdentifiers() -> const std::array<std::string, NumTokens>& {
  static_assert(MinLength <= MaxLength);
  static_assert(
      Uniform || MaxLength <= 64,
      "Cannot produce a meaningful non-uniform distribution of lengths longer "
      "than 64 as those are exceedingly rare in our observed data sets.");

  static const std::array<std::string, NumTokens> id_storage = [] {
    std::array<int, 64> id_length_counts;
    // For non-uniform distribution, we simulate a distribution roughly based on
    // the observed histogram of identifier lengths, but smoothed a bit and
    // reduced to small counts so that we cycle through all the lengths
    // reasonably quickly. We want sampling of even 10% of NumTokens from this
    // in a round-robin form to not be skewed overly much. This still inherently
    // compresses the long tail as we'd rather have coverage even though it
    // distorts the distribution a bit.
    //
    // The distribution here comes from a script that analyzes source code run
    // over a few directories of LLVM. The script renders a visual ascii-art
    // histogram along with the data for each bucket, and that output is
    // included in comments above each bucket size below to help visualize the
    // rough shape we're aiming for.
    //
    // 1 characters   [3976]  ███████████████████████████████▊
    id_length_counts[0] = 40;
    // 2 characters   [3724]  █████████████████████████████▊
    id_length_counts[1] = 40;
    // 3 characters   [4173]  █████████████████████████████████▍
    id_length_counts[2] = 40;
    // 4 characters   [5000]  ████████████████████████████████████████
    id_length_counts[3] = 50;
    // 5 characters   [1568]  ████████████▌
    id_length_counts[4] = 20;
    // 6 characters   [2226]  █████████████████▊
    id_length_counts[5] = 20;
    // 7 characters   [2380]  ███████████████████
    id_length_counts[6] = 20;
    // 8 characters   [1786]  ██████████████▎
    id_length_counts[7] = 18;
    // 9 characters   [1397]  ███████████▏
    id_length_counts[8] = 12;
    // 10 characters  [ 739]  █████▉
    id_length_counts[9] = 12;
    // 11 characters  [ 779]  ██████▎
    id_length_counts[10] = 12;
    // 12 characters  [1344]  ██████████▊
    id_length_counts[11] = 12;
    // 13 characters  [ 498]  ████
    id_length_counts[12] = 5;
    // 14 characters  [ 284]  ██▎
    id_length_counts[13] = 3;
    // 15 characters  [ 172]  █▍
    // 16 characters  [ 278]  ██▎
    // 17 characters  [ 191]  █▌
    // 18 characters  [ 207]  █▋
    for (int i : llvm::seq(14, 18)) {
      id_length_counts[i] = 2;
    }
    // 19 - 63 characters are all <100 but non-zero, and we map them to 1 for
    // coverage despite slightly over weighting the tail.
    for (int i : llvm::seq(18, 64)) {
      id_length_counts[i] = 1;
    }

    // Used to track the different count buckets when in a non-uniform
    // distribution.
    int length_bucket_index = 0;
    int length_count = 0;

    std::array<std::string, NumTokens> ids;
    absl::BitGen gen;
    for (auto [i, id] : llvm::enumerate(ids)) {
      if (Uniform) {
        // Rather than using randomness, for a uniform distribution rotate
        // lengths in round-robin to get a deterministic and exact size on every
        // run. We will then shuffle them at the end to produce a random
        // ordering.
        int length = MinLength + i % (1 + MaxLength - MinLength);
        id = GenerateRandomIdentifier(gen, length);
        continue;
      }

      // For non-uniform distribution, walk through each each length bucket
      // until our count matches the desired distribution, and then move to the
      // next.
      id = GenerateRandomIdentifier(gen, length_bucket_index + 1);

      if (length_count < id_length_counts[length_bucket_index]) {
        ++length_count;
      } else {
        length_bucket_index =
            (length_bucket_index + 1) % id_length_counts.size();
        length_count = 0;
      }
    }

    return ids;
  }();
  return id_storage;
}

// Compute a random sequence of just identifiers.
template <int MinLength = 1, int MaxLength = 64, bool Uniform = false>
auto RandomIdentifierSeq() -> std::string {
  // Get a static pool of identifiers with the desired distribution.
  const std::array<std::string, NumTokens>& ids =
      GetRandomIdentifiers<MinLength, MaxLength, Uniform>();

  // Shuffle tokens so we get exactly one of each identifier but in a random
  // order.
  std::array<llvm::StringRef, NumTokens> tokens;
  for (int i : llvm::seq(NumTokens)) {
    tokens[i] = ids[i];
  }
  std::shuffle(tokens.begin(), tokens.end(), absl::BitGen());
  return llvm::join(tokens, " ");
}

auto GetSymbolTokenTable() -> llvm::ArrayRef<TokenKind> {
  // Build our own table of symbols so we can use repetitions to skew the
  // distribution.
  static auto symbol_token_table_storage = [] {
    llvm::SmallVector<TokenKind> table;
#define CARBON_SYMBOL_TOKEN(TokenName, Spelling) \
  table.push_back(TokenKind::TokenName);
#define CARBON_OPENING_GROUP_SYMBOL_TOKEN(TokenName, Spelling, ClosingName)
#define CARBON_CLOSING_GROUP_SYMBOL_TOKEN(TokenName, Spelling, OpeningName)
#include "toolchain/lex/token_kind.def"
    table.insert(table.end(), 32, TokenKind::Semi);
    table.insert(table.end(), 16, TokenKind::Comma);
    table.insert(table.end(), 12, TokenKind::Period);
    table.insert(table.end(), 8, TokenKind::Colon);
    table.insert(table.end(), 8, TokenKind::Equal);
    table.insert(table.end(), 4, TokenKind::Amp);
    table.insert(table.end(), 4, TokenKind::ColonExclaim);
    table.insert(table.end(), 4, TokenKind::EqualEqual);
    table.insert(table.end(), 4, TokenKind::ExclaimEqual);
    table.insert(table.end(), 4, TokenKind::MinusGreater);
    table.insert(table.end(), 4, TokenKind::Star);
    return table;
  }();
  return symbol_token_table_storage;
}

// Compute a random sequence of mixed symbols, keywords, and identifiers, with
// percentages of each according to the parameters.
auto RandomMixedSeq(int symbol_percent, int keyword_percent) -> std::string {
  CARBON_CHECK(0 <= symbol_percent && symbol_percent <= 100)
      << "Must be a percent: [0, 100].";
  CARBON_CHECK(0 <= keyword_percent && keyword_percent <= 100)
      << "Must be a percent: [0, 100].";
  CARBON_CHECK((symbol_percent + keyword_percent) <= 100)
      << "Cannot have >100%.";
  static_assert((NumTokens % 100) == 0,
                "The number of tokens must be divisible by 100 so that we can "
                "easily scale integer percentages up to it.");

  // Get static pools of symbols, keywords, and identifiers.
  llvm::ArrayRef<TokenKind> symbols = GetSymbolTokenTable();
  llvm::ArrayRef<TokenKind> keywords = TokenKind::KeywordTokens;
  const std::array<std::string, NumTokens>& ids = GetRandomIdentifiers();

  // Build a list of StringRefs from the different types with the desired
  // distribution, then shuffle that list.
  std::array<llvm::StringRef, NumTokens> tokens;

  int num_symbols = (NumTokens / 100) * symbol_percent;
  int num_keywords = (NumTokens / 100) * keyword_percent;
  int num_identifiers = NumTokens - num_symbols - num_keywords;
  CARBON_CHECK(num_identifiers == 0 || num_identifiers > 500)
      << "We require at least 500 identifiers as we need to collect a "
         "reasonable number of samples to end up with a reasonable "
         "distribution of lengths.";

  for (int i : llvm::seq(num_symbols)) {
    tokens[i] = symbols[i % symbols.size()].fixed_spelling();
  }
  for (int i : llvm::seq(num_keywords)) {
    tokens[num_symbols + i] = keywords[i % keywords.size()].fixed_spelling();
  }
  for (int i : llvm::seq(num_identifiers)) {
    // We always have enough identifiers, so no need to mod here.
    tokens[num_symbols + num_keywords + i] = ids[i];
  }
  std::shuffle(tokens.begin(), tokens.end(), absl::BitGen());

  return llvm::join(tokens, " ");
}

class LexerBenchHelper {
 public:
  explicit LexerBenchHelper(llvm::StringRef text)
      : source_(MakeSourceBuffer(text)) {}

  auto Lex() -> TokenizedBuffer {
    DiagnosticConsumer& consumer = NullDiagnosticConsumer();
    return TokenizedBuffer::Lex(source_, consumer);
  }

  auto DiagnoseErrors() -> std::string {
    std::string result;
    llvm::raw_string_ostream out(result);
    StreamDiagnosticConsumer consumer(out);
    auto buffer = TokenizedBuffer::Lex(source_, consumer);
    consumer.Flush();
    CARBON_CHECK(buffer.has_errors())
        << "Asked to diagnose errors but none found!";
    return result;
  }

  auto source_text() -> llvm::StringRef { return source_.text(); }

 private:
  auto MakeSourceBuffer(llvm::StringRef text) -> SourceBuffer {
    CARBON_CHECK(fs_.addFile(filename_, /*ModificationTime=*/0,
                             llvm::MemoryBuffer::getMemBuffer(text)));
    return std::move(*SourceBuffer::CreateFromFile(
        fs_, filename_, ConsoleDiagnosticConsumer()));
  }

  llvm::vfs::InMemoryFileSystem fs_;
  std::string filename_ = "test.carbon";
  SourceBuffer source_;
};

void BM_ValidKeywords(benchmark::State& state) {
  absl::BitGen gen;
  std::array<llvm::StringRef, NumTokens> tokens;
  for (int i : llvm::seq(NumTokens)) {
    tokens[i] = TokenKind::KeywordTokens[i % TokenKind::KeywordTokens.size()]
                    .fixed_spelling();
  }
  std::shuffle(tokens.begin(), tokens.end(), gen);
  std::string source = llvm::join(tokens, " ");

  LexerBenchHelper helper(source);
  for (auto _ : state) {
    TokenizedBuffer buffer = helper.Lex();
    CARBON_CHECK(!buffer.has_errors());
  }

  state.SetBytesProcessed(state.iterations() * source.size());
  state.counters["tokens_per_second"] = benchmark::Counter(
      NumTokens, benchmark::Counter::kIsIterationInvariantRate);
}
BENCHMARK(BM_ValidKeywords);

template <int MinLength, int MaxLength, bool Uniform>
void BM_ValidIdentifiers(benchmark::State& state) {
  std::string source = RandomIdentifierSeq<MinLength, MaxLength, Uniform>();

  LexerBenchHelper helper(source);
  for (auto _ : state) {
    TokenizedBuffer buffer = helper.Lex();
    CARBON_CHECK(!buffer.has_errors()) << helper.DiagnoseErrors();
  }

  state.SetBytesProcessed(state.iterations() * source.size());
  state.counters["tokens_per_second"] = benchmark::Counter(
      NumTokens, benchmark::Counter::kIsIterationInvariantRate);
}
// Benchmark the non-uniform distribution we observe in C++ code.
BENCHMARK(BM_ValidIdentifiers<1, 64, /*Uniform=*/false>);

// Also benchmark a few uniform distribution ranges of identifier widths to
// cover different patterns that emerge with small, medium, and longer
// identifiers.
BENCHMARK(BM_ValidIdentifiers<1, 1, /*Uniform=*/true>);
BENCHMARK(BM_ValidIdentifiers<3, 5, /*Uniform=*/true>);
BENCHMARK(BM_ValidIdentifiers<3, 16, /*Uniform=*/true>);
BENCHMARK(BM_ValidIdentifiers<12, 64, /*Uniform=*/true>);

void BM_ValidMix(benchmark::State& state) {
  int symbol_percent = state.range(0);
  int keyword_percent = state.range(1);
  std::string source = RandomMixedSeq(symbol_percent, keyword_percent);

  LexerBenchHelper helper(source);
  for (auto _ : state) {
    TokenizedBuffer buffer = helper.Lex();

    // Ensure that lexing actually occurs for benchmarking and that it doesn't
    // hit errors that would skew the benchmark results.
    CARBON_CHECK(!buffer.has_errors()) << helper.DiagnoseErrors();
  }

  state.SetBytesProcessed(state.iterations() * source.size());
  state.counters["tokens_per_second"] = benchmark::Counter(
      NumTokens, benchmark::Counter::kIsIterationInvariantRate);
}
// The distributions between symbols, keywords, and identifiers here are
// guesses. Eventually, we should collect more data to help tune these, but
// hopefully the performance isn't too sensitive and we can just cover a wide
// range here.
BENCHMARK(BM_ValidMix)
    ->Args({10, 40})
    ->Args({25, 30})
    ->Args({50, 20})
    ->Args({75, 10});

// This is a speed-of-light benchmark that should reflect memory bandwidth
// (ideally) of simply reading all the source code. For speed-of-light we use
// `strcpy` -- this both examines ever byte of the input looking for a null to
// end the copy, and also writes to a data structure of roughly the same size as
// the input. This routine is one we expect to be *very* well optimized and give
// a good approximation of the fastest possible lexer given the physical
// constraints of the machine. Note that which particular source we use as input
// here isn't especially interesting, so we just pick one and should update it
// to reflect whatever distribution is most realistic long-term. The
// bytes/second throughput is the important output of this routine.
auto BM_SpeedOfLightStrCpy(benchmark::State& state) -> void {
  std::string source =
      RandomMixedSeq(/*symbol_percent=*/25, /*keyword_percent=*/30);

  // A buffer to write the null-terminated contents of `source` into.
  llvm::OwningArrayRef<char> buffer(source.size() + 1);

  for (auto _ : state) {
    const char* text = source.data();
    benchmark::DoNotOptimize(text);
    strcpy(buffer.data(), text);
    benchmark::DoNotOptimize(buffer.data());
  }

  state.SetBytesProcessed(state.iterations() * source.size());
  state.counters["tokens_per_second"] = benchmark::Counter(
      NumTokens, benchmark::Counter::kIsIterationInvariantRate);
}
BENCHMARK(BM_SpeedOfLightStrCpy);

// This is a speed-of-light benchmark that builds up a best-case byte-wise table
// dispatch using guaranteed tail recursion. The goal is both to ensure the
// general technique can reasonably hit the level of performance we need and to
// establish how far from this speed of light the actual lexer currently sits.
//
// A major impact on the observed performance of this technique is how many
// different functions are reached in this dispatch loop. This benchmark
// infrastructure tries to bracket the range of performance this technique
// affords with different numbers of dispatch target functions.
using DispatchPtrT = auto (*)(ssize_t& index, const char* text, char* buffer)
    -> void;
using DispatchTableT = std::array<DispatchPtrT, 256>;

template <const DispatchTableT& Table>
auto BasicDispatch(ssize_t& index, const char* text, char* buffer) -> void {
  *buffer = text[index];
  ++index;
  [[clang::musttail]] return Table[static_cast<unsigned char>(text[index])](
      index, text, buffer);
}

template <const DispatchTableT& Table, char C>
auto SpecializedDispatch(ssize_t& index, const char* text, char* buffer)
    -> void {
  CARBON_CHECK(C == text[index]);
  *buffer = C;
  ++index;
  [[clang::musttail]] return Table[static_cast<unsigned char>(text[index])](
      index, text, buffer);
}

// A sample of the symbol characters used in Carbon code. Doesn't need to be
// perfect, as we just need to have a reasonably large # of distinct dispatch
// functions.
constexpr char DispatchSpecializableSymbols[] = {
    '!', '%', '(', ')', '*', '+', ',', '-', '.', ':',
    ';', '<', '=', '>', '?', '[', ']', '{', '}', '~',
};

// Create an array of all the characters we can specialize dispatch over --
// [0-9A-Za-z] and the symbols above. Similar to the above symbols, doesn't need
// to be exhaustive.
constexpr std::array<char, 26 * 2 + 10 + sizeof(DispatchSpecializableSymbols)>
    DispatchSpecializableChars = []() constexpr {
      constexpr int Size = sizeof(DispatchSpecializableChars);
      std::array<char, Size> chars = {};
      int i = 0;
      for (char c = '0'; c <= '9'; ++c) {
        chars[i] = c;
        ++i;
      }
      for (char c = 'A'; c <= 'Z'; ++c) {
        chars[i] = c;
        ++i;
      }
      for (char c = 'a'; c <= 'z'; ++c) {
        chars[i] = c;
        ++i;
      }
      for (char c : DispatchSpecializableSymbols) {
        chars[i] = c;
        ++i;
      }
      CARBON_CHECK(i == Size);
      return chars;
    }();

// Instantiate a number of specialized dispatch functions for characters in the
// array above, and assign those function addresses to the character's entry in
// the provided table. The provided `tmp_table` is a temporary that will
// eventually initialize the provided `Table` constant, so the constant is what
// we propagate to the instantiated function and the temporary is the one we
// initialize.
template <const DispatchTableT& Table, size_t... Indices>
constexpr auto SpecializeDispatchTable(
    DispatchTableT& tmp_table, std::index_sequence<Indices...> /*indices*/)
    -> void {
  static_assert(sizeof...(Indices) <= sizeof(DispatchSpecializableChars));
  ((tmp_table[static_cast<unsigned char>(DispatchSpecializableChars[Indices])] =
        &SpecializedDispatch<Table, DispatchSpecializableChars[Indices]>),
   ...);
}

// The maximum number of dispatch targets is the size of the array + 1 (for the
// base case target).
constexpr int MaxDispatchTargets = sizeof(DispatchSpecializableChars) + 1;

// Dispatch tables with a provided number of distinct dispatch targets. There
// will always be one additional target for the null byte to end the loop.
template <int NumDispatchTargets>
constexpr DispatchTableT DispatchTable = []() constexpr {
  static_assert(NumDispatchTargets > 0, "Need at least one dispatch target.");
  static_assert(NumDispatchTargets <= MaxDispatchTargets,
                "Limited number of dispatch targets available.");

  DispatchTableT tmp_table = {};
  // Start with the basic dispatch target.
  for (int i = 0; i < 256; ++i) {
    tmp_table[i] = &BasicDispatch<DispatchTable<NumDispatchTargets>>;
  }
  if constexpr (NumDispatchTargets > 1) {
    // Add additional dispatch targets from our specializable array.
    SpecializeDispatchTable<DispatchTable<NumDispatchTargets>>(
        tmp_table, std::make_index_sequence<NumDispatchTargets - 1>());
  }
  // Special case the null byte index to end the tail-dispatch.
  tmp_table[0] =
      +[](ssize_t& index, const char* text, char* /*buffer*/) -> void {
    CARBON_CHECK(text[index] == '\0');
    return;
  };
  return tmp_table;
}();

template <int NumDispatchTargets>
auto BM_SpeedOfLightDispatch(benchmark::State& state) -> void {
  std::string source =
      RandomMixedSeq(/*symbol_percent=*/25, /*keyword_percent=*/30);

  // A buffer to write to, simulating some minimal write traffic.
  llvm::OwningArrayRef<char> buffer(source.size());

  for (auto _ : state) {
    const char* text = source.data();
    benchmark::DoNotOptimize(text);

    // Use `ssize_t` to minimize indexing overhead.
    ssize_t i = 0;
    // The dispatch table tail-recurses through the entire string.
    DispatchTable<NumDispatchTargets>[static_cast<unsigned char>(text[i])](
        i, text, buffer.data());
    CARBON_CHECK(i == static_cast<ssize_t>(source.size()));

    benchmark::DoNotOptimize(buffer.data());
  }

  state.SetBytesProcessed(state.iterations() * source.size());
  state.counters["tokens_per_second"] = benchmark::Counter(
      NumTokens, benchmark::Counter::kIsIterationInvariantRate);
}
BENCHMARK(BM_SpeedOfLightDispatch<1>);
BENCHMARK(BM_SpeedOfLightDispatch<2>);
BENCHMARK(BM_SpeedOfLightDispatch<4>);
BENCHMARK(BM_SpeedOfLightDispatch<8>);
BENCHMARK(BM_SpeedOfLightDispatch<16>);
BENCHMARK(BM_SpeedOfLightDispatch<32>);
BENCHMARK(BM_SpeedOfLightDispatch<MaxDispatchTargets>);

}  // namespace
}  // namespace Carbon::Lex