<commit_before>import time

from collections.abc import Callable


# Log everything to stderr and flush so that we have a unified stream to match
# errors/info emitted by MLIR to stderr.
def log(*args):
  print(*args, file=sys.stderr)
  sys.stderr.flush()


def timed_invoke(execute: Callable, gflop_count: int, n_iters: int, args: list,
                 **kwargs):
  start = time.time()
  execute(*args, **kwargs, n_iters=n_iters)
  elapsed_s = time.time() - start
  elapsed_s_per_iter = elapsed_s / n_iters
  gflop_per_s_per_iter = gflop_count / (elapsed_s_per_iter)
  print(f'xxxxxxxxxx : {n_iters} iters time on {1} threads '
        f'in {elapsed_s_per_iter:.{4}}s per iter '
        f'sec ({gflop_per_s_per_iter:.{4}} GFlop/s) '
        f'total time {elapsed_s:.{4}}s ')


def setup_and_invoke(
    setup_fun: Callable,
    run_fun: Callable,
    n_iters: int,
    gflop_count: int,
    n_iters_dry_run=1,
    # Used in particular to pass a check_fun callable.
    **kwargs):
  tensors = setup_fun()

  module, execution_engine = None, None
  if 'compile_fun' in kwargs:
    module, execution_engine = kwargs['compile_fun'](*tensors)

  # Dry-run.
  if n_iters_dry_run > 0:
    timed_invoke(
        run_fun,
        gflop_count,
        n_iters_dry_run,
        tensors,
        module=module,
        execution_engine=execution_engine)

  if 'check_fun' in kwargs:
    kwargs['check_fun'](*tensors)

  # Run for ITERS and report timing.
  timed_invoke(<commit_msg>Call benchmark function once before timing it.<commit_after>import time

from collections.abc import Callable


# Log everything to stderr and flush so that we have a unified stream to match
# errors/info emitted by MLIR to stderr.
def log(*args):
  print(*args, file=sys.stderr)
  sys.stderr.flush()


def invoke(execute: Callable, n_iters: int, args: list, **kwargs):
  execute(*args, **kwargs, n_iters=n_iters)


def timed_invoke(execute: Callable, gflop_count: int, n_iters: int, args: list,
                 **kwargs):
  start = time.time()
  execute(*args, **kwargs, n_iters=n_iters)
  elapsed_s = time.time() - start
  elapsed_s_per_iter = elapsed_s / n_iters
  gflop_per_s_per_iter = gflop_count / (elapsed_s_per_iter)
  print(f'xxxxxxxxxx : {n_iters} iters time on {1} threads '
        f'in {elapsed_s_per_iter:.{4}}s per iter '
        f'sec ({gflop_per_s_per_iter:.{4}} GFlop/s) '
        f'total time {elapsed_s:.{4}}s ')


def setup_and_invoke(
    setup_fun: Callable,
    run_fun: Callable,
    n_iters: int,
    gflop_count: int,
    n_iters_dry_run=1,
    # Used in particular to pass a check_fun callable.
    **kwargs):
  tensors = setup_fun()

  module, execution_engine = None, None
  if 'compile_fun' in kwargs:
    module, execution_engine = kwargs['compile_fun'](*tensors)

  # Run the function with zero iterations to make subsequent timings accurate.
  invoke(run_fun, 0, tensors, module=module, execution_engine=execution_engine)

  # Dry-run.
  if n_iters_dry_run > 0:
    timed_invoke(
        run_fun,
        gflop_count,
        n_iters_dry_run,
        tensors,
        module=module,
        execution_engine=execution_engine)

  if 'check_fun' in kwargs:
    kwargs['check_fun'](*tensors)

  # Run for ITERS and report timing.
  timed_invoke(
{'fraction_non_alphanumeric': 0.10523076923076922, 'fraction_numerical': 0.004, 'mean_word_length': 3.4779614325068873, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6504069', 'n_tokens_mistral': 1215, 'n_tokens_neox': 1172, 'n_words': 290}