build_tools/bazel/iree_trace_runner_test.bzl

# Copyright 2021 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

"""Macros for defining tests that use a trace-runner."""

load("//build_tools/bazel:iree_bytecode_module.bzl", "iree_bytecode_module")
load("//build_tools/bazel:native_binary.bzl", "native_test")

def iree_trace_runner_test(
        name,
        src,
        module_name,
        target_backend,
        driver,
        trace_runner,
        trace,
        compiler_flags = [],
        runner_args = [],
        tags = [],
        target_cpu_features = None,
        timeout = None,
        **kwargs):
    """Creates a test running a custom trace-runner on a trace file (yaml).

    Args:
        name: Name of the target
        src: mlir source file to be compiled to an IREE module.
        target_backend: target backend to compile for.
        driver: driver to run the module with.
        compiler_flags: additional flags to pass to the compiler. Bytecode
            output format and backend flags are passed automatically.
        runner_args: additional args to pass to the trace-runner program. The
            driver and input file flags are passed automatically.
        tags: Additional labels to apply to the test. "driver=${DRIVER}" is
            added automatically.
        trace_runner: trace-runner program to run.
        trace: trace file input to the trace-runner program.
        module_name: specifies the  path to use for the enerated IREE module
            (.vmfb). Mandatory, unlike in iree_check_test, because trace files
            (.yaml) reference a specific module file path.
        timeout: timeout for the generated tests.
        target_cpu_features: target CPU features. Only for llvm-cpu backend.
        **kwargs: any additional attributes to pass to the underlying tests and
            test suite.
    """

    if target_cpu_features:
        fail("target_cpu_features must currently be empty")

    bytecode_module_name = name + "_bytecode_module"
    iree_bytecode_module(
        name = bytecode_module_name,
        module_name = module_name,
        src = src,
        flags = [
            "--iree-hal-target-backends=%s" % target_backend,
        ] + ([
            "--iree-llvmcpu-target-cpu-features=%s" % target_cpu_features,
        ] if target_cpu_features else []) + compiler_flags,
        visibility = ["//visibility:private"],
        testonly = True,
        **kwargs
    )

    native_test(
        name = name,
        args = [
            "--device=%s" % driver,
            "$(location :%s)" % trace,
        ] + runner_args,
        data = [
            ":%s" % bytecode_module_name,
            ":%s" % trace,
        ],
        src = trace_runner,
        tags = tags + ["driver=%s" % driver],
        timeout = timeout,
        **kwargs
    )

def iree_single_backend_generated_trace_runner_test(
        name,
        generator,
        trace_runner,
        target_backend,
        driver,
        generator_args = [],
        compiler_flags = [],
        runner_args = [],
        tags = [],
        target_cpu_features = None,
        timeout = None,
        **kwargs):
    """Generates an iree_trace_runner_test using a custom python generator script.

    The generator script produces a .mlir source and a .yaml trace, which are
    passed to iree_trace_runner_test.

    Args:
        name: Name of the target
        generator: Target to run to generate the source file and trace files.
            It will be invoked with the following standard flags, in addition
            to generator_args:
            --output_code=(current binary dir)/name.mlir
            --output_trace=(current binary dir)/name.yaml
            --module_path=(current binary dir)/name.vmfb
        generator_args: additional args to pass to the generator program.
        target_backend: target backend to compile for.
        driver: driver to run the module with.
        compiler_flags: additional flags to pass to the compiler. Bytecode
            output format and backend flags are passed automatically.
        runner_args: additional args to pass to the trace-runner program. The
            driver and input file flags are passed automatically.
        tags: Additional labels to apply to the test. "driver=${DRIVER}" is
            added automatically.
        trace_runner: trace-runner program to run.
        timeout: timeout for the generated tests.
        target_cpu_features: target CPU features. Only for llvm-cpu backend.
        **kwargs: any additional attributes to pass to the underlying tests and
            test suite.
    """

    src = "%s.mlir" % (name)
    trace = "%s.yaml" % (name)
    module_name = "%s.vmfb" % (name)
    native.genrule(
        name = "%s_generate" % (name),
        outs = [src, trace],
        cmd = " ".join([
            "$(location %s)" % (generator),
            " ".join([('"%s"' % arg) for arg in generator_args]),
            "--output_code=$(location %s)" % (src),
            "--output_trace=$(location %s)" % (trace),
            # Explanation for why "$(RULEDIR)/%s" instead of "$(location %s)" below:
            # module_path points to a file that does not yet exist as it will
            # be generated by iree_bytecode_module below iree_trace_runner_test.
            "--module_path=$(RULEDIR)/%s" % (module_name),
        ] + [('"%s"' % arg) for arg in generator_args]),
        tools = [generator],
        message = "Generating code and trace for test %s..." % (name),
        output_to_bindir = 1,
        testonly = True,
        **kwargs
    )
    iree_trace_runner_test(
        name = name,
        src = src,
        module_name = module_name,
        target_backend = target_backend,
        driver = driver,
        trace_runner = trace_runner,
        trace = trace,
        compiler_flags = compiler_flags,
        runner_args = runner_args,
        tags = tags,
        timeout = timeout,
        target_cpu_features = target_cpu_features,
        **kwargs
    )

def iree_generated_trace_runner_test(
        name,
        generator,
        trace_runner,
        target_backends_and_drivers,
        generator_args = [],
        compiler_flags = [],
        runner_args = [],
        tags = [],
        timeout = None,
        target_cpu_features_variants = [],
        **kwargs):
    """Generates a suite of iree_trace_runner_test on multiple backends/drivers.

    Args:
        name: Name of the target
        generator: Target to run to generate the source file and trace files.
            It will be invoked with the following standard flags, in addition
            to generator_args:
            --output_code=(current binary dir)/name.mlir
            --output_trace=(current binary dir)/name.yaml
            --module_path=(current binary dir)/name.vmfb
        generator_args: additional args to pass to the generator program.
        target_backends_and_drivers: backend/driver pairs to compile and run
            the module.
        compiler_flags: additional flags to pass to the compiler. Bytecode
            output format and backend flags are passed automatically.
        runner_args: additional args to pass to the trace-runner program. The
            driver and input file flags are passed automatically.
        tags: Additional labels to apply to the test. "driver=${DRIVER}" is
            added automatically.
        trace_runner: trace-runner program to run.
        timeout: timeout for the generated tests.
        target_cpu_features_variants: list of target cpu features variants.
            Currently unimplemented in Bazel due to difficulty of specializing
            to target architecture in Bazel. The following describes the
            semantics that this should have if implemented. Each
            entry is either "default" for the architecture defaults, or a colon-
            separated triple "arch:name:cpu_features" where "arch" filters
            for a target CPU architecture (in IREE_ARCH format), "name" is a
            short name for the CPU features set (used to generate target names)
            and cpu_features is a comma-separated list of LLVM target attributes
            to enable. Example:
              x86_64:avx2_fma:+avx,+avx2,+fma
        **kwargs: any additional attributes to pass to the underlying tests and test suite.
    """

    tests = []
    for backend, driver in target_backends_and_drivers:
        # CUDA backend/driver not supported by Bazel build.
        if backend == "cuda" or driver == "cuda":
            continue
        suite_entry_name = "_".join([name, backend, driver])
        iree_single_backend_generated_trace_runner_test(
            name = suite_entry_name,
            generator = generator,
            trace_runner = trace_runner,
            driver = driver,
            target_backend = backend,
            generator_args = generator_args,
            compiler_flags = compiler_flags,
            runner_args = runner_args,
            tags = tags,
            timeout = timeout,
            **kwargs
        )
        tests.append(suite_entry_name)
    native.test_suite(
        name = name,
        tests = tests,
        tags = tags,
        **kwargs
    )