Build trace-runner tests in Bazel. (#7499)
Build trace-runner tests in Bazel.
This enables the e2e matmul tests, and having them supported in Bazel means in particular that they run in Google internal CI, which will hopefully help us catch integration regressions earlier.
diff --git a/build_tools/bazel/iree_trace_runner_test.bzl b/build_tools/bazel/iree_trace_runner_test.bzl
index 4c2b50a..dbbb493 100644
--- a/build_tools/bazel/iree_trace_runner_test.bzl
+++ b/build_tools/bazel/iree_trace_runner_test.bzl
@@ -6,7 +6,224 @@
"""Macros for defining tests that use a trace-runner."""
-def iree_generated_trace_runner_test(**kwargs):
- # TODO: implement this. For now, it's only parsed by bazel_to_cmake.py, so
- # the iree_generated_check_test's are just omitted in bazel builds.
- pass
+load("//iree/tools:compilation.bzl", "iree_bytecode_module")
+load("//build_tools/bazel:run_binary_test.bzl", "run_binary_test")
+
+def iree_trace_runner_test(
+ name,
+ src,
+ module,
+ target_backend,
+ driver,
+ trace_runner,
+ trace,
+ compiler_flags = [],
+ runner_args = [],
+ opt_tool = "//iree/tools:iree-opt",
+ opt_flags = [],
+ tags = [],
+ 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
+ translation 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.
+ opt_tool: Defaulting to iree-opt. Tool used to preprocess the source files
+ if opt_flags is specified.
+ opt_flags: If specified, source files are preprocessed with opt_tool with
+ these flags.
+ trace_runner: trace-runner program to run.
+ trace: trace file input to the trace-runner program.
+ module: 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.
+ **kwargs: any additional attributes to pass to the underlying tests and test suite.
+ """
+
+ bytecode_module_name = name + "_bytecode_module"
+ iree_bytecode_module(
+ name = bytecode_module_name,
+ module = module,
+ src = src,
+ flags = [
+ "-iree-mlir-to-vm-bytecode-module",
+ "-mlir-print-op-on-diagnostic=false",
+ "-iree-hal-target-backends=%s" % target_backend,
+ ] + compiler_flags,
+ opt_tool = opt_tool,
+ opt_flags = opt_flags,
+ visibility = ["//visibility:private"],
+ **kwargs
+ )
+
+ run_binary_test(
+ name = name,
+ args = [
+ "--driver=%s" % driver,
+ "$(location :%s)" % trace,
+ ] + runner_args,
+ data = [
+ ":%s" % bytecode_module_name,
+ ":%s" % trace,
+ ],
+ test_binary = 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 = [],
+ opt_tool = "//iree/tools:iree-opt",
+ opt_flags = [],
+ tags = [],
+ 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
+ translation 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.
+ opt_tool: Defaulting to iree-opt. Tool used to preprocess the source files
+ if opt_flags is specified.
+ opt_flags: If specified, source files are preprocessed with opt_tool with
+ these flags.
+ trace_runner: trace-runner program to run.
+ timeout: timeout for the generated tests.
+ **kwargs: any additional attributes to pass to the underlying tests and test suite.
+ """
+
+ src = "%s.mlir" % (name)
+ trace = "%s.yaml" % (name)
+ module = "%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),
+ ] + [('"%s"' % arg) for arg in generator_args]),
+ tools = [generator],
+ message = "Generating code and trace for test %s..." % (name),
+ output_to_bindir = 1,
+ **kwargs
+ )
+ iree_trace_runner_test(
+ name = name,
+ src = src,
+ module = module,
+ target_backend = target_backend,
+ driver = driver,
+ trace_runner = trace_runner,
+ trace = trace,
+ compiler_flags = compiler_flags,
+ runner_args = runner_args,
+ opt_tool = opt_tool,
+ opt_flags = opt_flags,
+ tags = tags,
+ timeout = timeout,
+ **kwargs
+ )
+
+def iree_generated_trace_runner_test(
+ name,
+ generator,
+ trace_runner,
+ target_backends_and_drivers,
+ generator_args = [],
+ compiler_flags = [],
+ runner_args = [],
+ opt_tool = "//iree/tools:iree-opt",
+ opt_flags = [],
+ tags = [],
+ timeout = None,
+ **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
+ translation 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.
+ opt_tool: Defaulting to iree-opt. Tool used to preprocess the source files
+ if opt_flags is specified.
+ opt_flags: If specified, source files are preprocessed with opt_tool with
+ these flags.
+ trace_runner: trace-runner program to run.
+ timeout: timeout for the generated tests.
+ **kwargs: any additional attributes to pass to the underlying tests and test suite.
+ """
+
+ tests = []
+ for backend, driver in target_backends_and_drivers:
+ 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,
+ opt_tool = opt_tool,
+ opt_flags = opt_flags,
+ tags = tags,
+ timeout = timeout,
+ **kwargs
+ )
+ tests.append(suite_entry_name)
+ native.test_suite(
+ name = name,
+ tests = tests,
+ tags = tags,
+ **kwargs
+ )
diff --git a/build_tools/bazel_to_cmake/bazel_to_cmake_converter.py b/build_tools/bazel_to_cmake/bazel_to_cmake_converter.py
index 9abd641..c6fe264 100644
--- a/build_tools/bazel_to_cmake/bazel_to_cmake_converter.py
+++ b/build_tools/bazel_to_cmake/bazel_to_cmake_converter.py
@@ -15,7 +15,7 @@
# pylint: disable=exec-used
import itertools
-import textwrap
+import re
import bazel_to_cmake_targets
@@ -611,8 +611,11 @@
drivers = [it[1] for it in target_backends_and_drivers]
name_block = _convert_string_arg_block("NAME", name, quote=False)
+ # For now we assume that the generator target is a py_binary with a single
+ # source .py file named like it.
+ generator_py = f"{generator.split(':')[-1]}.py"
generator_block = _convert_string_arg_block("GENERATOR",
- generator,
+ generator_py,
quote=True)
generator_args_block = _convert_string_list_block("GENERATOR_ARGS",
generator_args)
@@ -639,6 +642,7 @@
f"{opt_flags_block}"
f")\n\n")
+
def iree_e2e_cartesian_product_test_suite(self,
name,
matrix,
diff --git a/iree/test/e2e/regression/BUILD b/iree/test/e2e/regression/BUILD
index 71f6df5..20e1f62 100644
--- a/iree/test/e2e/regression/BUILD
+++ b/iree/test/e2e/regression/BUILD
@@ -105,9 +105,14 @@
target_backend = "cuda",
)
+py_binary(
+ name = "generate_e2e_matmul_tests",
+ srcs = ["generate_e2e_matmul_tests.py"],
+)
+
[iree_generated_trace_runner_test(
name = "e2e_matmul_direct_%s_small" % lhs_rhs_type,
- generator = "generate_e2e_matmul_tests.py",
+ generator = ":generate_e2e_matmul_tests",
generator_args = [
"--lhs_rhs_type=%s" % lhs_rhs_type,
"--shapes=small",
@@ -124,7 +129,7 @@
[iree_generated_trace_runner_test(
name = "e2e_matmul_mmt4d_%s_small" % lhs_rhs_type,
- generator = "generate_e2e_matmul_tests.py",
+ generator = ":generate_e2e_matmul_tests",
generator_args = [
"--lhs_rhs_type=%s" % lhs_rhs_type,
"--shapes=small",
@@ -144,7 +149,7 @@
[iree_generated_trace_runner_test(
name = "e2e_matmul_mmt4d_%s_large" % lhs_rhs_type,
- generator = "generate_e2e_matmul_tests.py",
+ generator = ":generate_e2e_matmul_tests",
generator_args = [
"--lhs_rhs_type=%s" % lhs_rhs_type,
"--shapes=large",
diff --git a/iree/tools/BUILD b/iree/tools/BUILD
index 9a68ab2..152abc4 100644
--- a/iree/tools/BUILD
+++ b/iree/tools/BUILD
@@ -336,6 +336,25 @@
)
cc_binary(
+ name = "iree-e2e-matmul-test",
+ srcs = ["iree-e2e-matmul-test.c"],
+ deps = [
+ "//iree/base",
+ "//iree/base:core_headers",
+ "//iree/base:tracing",
+ "//iree/base/internal:file_path",
+ "//iree/base/internal:flags",
+ "//iree/hal",
+ "//iree/hal/drivers",
+ "//iree/modules/hal",
+ "//iree/tools/utils:trace_replay",
+ "//iree/tools/utils:yaml_util",
+ "//iree/vm",
+ "@com_github_yaml_libyaml//:yaml",
+ ],
+)
+
+cc_binary(
name = "iree-tblgen",
srcs = [
"//iree/compiler/Dialect/Util/Tools:GenSrcs",
diff --git a/iree/tools/compilation.bzl b/iree/tools/compilation.bzl
index ae4fd03..6722ba6 100644
--- a/iree/tools/compilation.bzl
+++ b/iree/tools/compilation.bzl
@@ -9,9 +9,11 @@
load("//build_tools/embed_data:build_defs.bzl", "c_embed_data")
# TODO(benvanik): port to a full starlark rule, document, etc.
+
def iree_bytecode_module(
name,
src,
+ module = None,
flags = ["-iree-mlir-to-vm-bytecode-module"],
translate_tool = "//iree/tools:iree-translate",
embedded_linker_tool = "@llvm-project//lld:lld",
@@ -19,7 +21,28 @@
opt_flags = [],
c_identifier = "",
**kwargs):
+ """Builds an IREE bytecode module.
+
+ Args:
+ name: Name of the target
+ src: mlir source file to be compiled to an IREE module.
+ flags: additional flags to pass to the compiler. Bytecode
+ translation and backend flags are passed automatically.
+ translate_tool: the compiler to use to generate the module.
+ Defaults to iree-translate.
+ embedded_linker_tool: the embedded linker to use.
+ Defaults to the lld from the llvm-project directory.
+ opt_tool: Defaulting to iree-opt. Tool used to preprocess the source file
+ if opt_flags is specified.
+ opt_flags: If specified, source files are preprocessed with opt_tool with
+ these flags.
+ module: Optional. Specifies the path to use for the enerated IREE module (.vmfb).
+ c_identifier: Optional. Enables embedding the module as C data.
+ **kwargs: any additional attributes to pass to the underlying rules.
+ """
+
translate_src = src
+
if opt_flags:
translate_src = "%s.opt.mlir" % (name)
native.genrule(
@@ -35,20 +58,24 @@
tools = [opt_tool],
message = "Transforming MLIR source for IREE module %s..." % (name),
output_to_bindir = 1,
+ **kwargs
)
+ if not module:
+ module = "%s.vmfb" % (name)
+
native.genrule(
name = name,
srcs = [translate_src],
outs = [
- "%s.vmfb" % (name),
+ module,
],
cmd = " && ".join([
" ".join([
"$(location %s)" % (translate_tool),
" ".join(flags),
"-iree-llvm-embedded-linker-path=$(location %s)" % (embedded_linker_tool),
- "-o $(location %s.vmfb)" % (name),
+ "-o $(location %s)" % (module),
"$(location %s)" % (translate_src),
]),
]),
@@ -63,7 +90,7 @@
c_embed_data(
name = "%s_c" % (name),
identifier = c_identifier,
- srcs = ["%s.vmfb" % (name)],
+ srcs = [module],
c_file_output = "%s_c.c" % (name),
h_file_output = "%s_c.h" % (name),
flatten = True,
diff --git a/iree/tools/iree-e2e-matmul-test.c b/iree/tools/iree-e2e-matmul-test.c
index 99222fb..bcfea0b 100644
--- a/iree/tools/iree-e2e-matmul-test.c
+++ b/iree/tools/iree-e2e-matmul-test.c
@@ -14,7 +14,6 @@
#include "iree/base/internal/flags.h"
#include "iree/base/target_platform.h"
#include "iree/hal/api.h"
-#include "iree/hal/buffer_view.h"
#include "iree/hal/drivers/init.h"
#include "iree/modules/hal/module.h"
#include "iree/tools/utils/trace_replay.h"
