build_tools/testing/generate_cmake_e2e_model_tests.py

#!/usr/bin/env python3
## Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Generates a CMake file to define e2e mdoel tests."""

import string
import sys
import pathlib
import argparse

# Add build_tools python dir to the search path.
sys.path.insert(0, str(pathlib.Path(__file__).parent / ".." / "python"))

import benchmark_suites.iree.benchmark_collections
import e2e_model_tests.cmake_generator

TEMPLATE_DIR = pathlib.Path(__file__).parent
GENERATED_E2E_MODEL_TESTS_CMAKE_TEMPLATE = string.Template(
    (TEMPLATE_DIR / "generated_e2e_model_tests_template.cmake").read_text())


def parse_arguments():
  """Parses command-line options."""

  parser = argparse.ArgumentParser()
  parser.add_argument("--output",
                      required=True,
                      help="Path to write the generated cmake file.")

  return parser.parse_args()


def main(args: argparse.Namespace):
  (gen_configs,
   _) = benchmark_suites.iree.benchmark_collections.generate_benchmarks()
  cmake_rules = e2e_model_tests.cmake_generator.generate_rules(
      module_generation_configs=gen_configs)
  output = GENERATED_E2E_MODEL_TESTS_CMAKE_TEMPLATE.substitute(
      __TEST_RULES="\n".join(cmake_rules))
  with open(args.output, "w") as output_file:
    output_file.write(output)


if __name__ == "__main__":
  main(parse_arguments())