Polish glob_lit_tests() macro and related code.

Add a rule to provide the test file path and the data dependencies of a
lit test. Doing this in a rule allows other inputs than simple globs and
simplifies path handling.

Lit has a peculiar way to discover config and test files from the input
arguments.

In it's simplest form, lit finds the config files in a parent directory of
each input argument (relative to the current directory). It then uses the
config directory to execute the RUN commands in the test file at the
remainder of the path. If the config files are in a parent directory of the test
files, it is therefore sufficient to simply pass the test file paths as input
arguments.

The config files may override the execution directory by specifying a
config.test_exec_root, and the base path of the test files by specifying a
config.test_source_root. Each test is executed in the directory of the test
file (specified relative to config.test_source_root) relative to the
config.test_exec_root directory.

We use this to handle the case where the suite is not in a parent directory
of the test file. The input path is set to '<suite path>/<test file basename>',
and the config.test_source_root is set to the directory of the 'TEST_BINARY' environment variable that bazel provides. The config.test_exec_root is simply set to ".", i.e. use the standard bazel execroot.

PiperOrigin-RevId: 407648413

mlir_tests/lit.bzl

@@ -13,143 +13,107 @@
 # limitations under the License.
 
 # Test definitions for Lit, the LLVM test runner.
-# This file is derived from @org_tensorflow/compiler/mlir/glob_lit_test.bzl
 #
-# This is reusing the LLVM Lit test runner in the interim until the new build
-# rules are upstreamed.
-# TODO(b/136126535): remove this custom rule.
-"""Lit runner globbing test
-"""
-
-load("@bazel_skylib//lib:paths.bzl", "paths")
-
-# Default values used by the test runner.
-_default_test_file_exts = ["mlir", ".pbtxt", ".td"]
-_default_driver = "@llvm-project//mlir:run_lit.sh"
-_default_size = "small"
-_default_tags = []
-
-# These are patterns which we should never match, for tests, subdirectories, or
-# test input data files.
-_ALWAYS_EXCLUDE = [
-    "**/LICENSE.txt",
-    "**/README.txt",
-    "**/lit.local.cfg",
-    # Exclude input files that have spaces in their names, since bazel
-    # cannot cope with such "targets" in the srcs list.
-    "**/* *",
-    "**/* */**",
-]
-
-def _run_lit_test(name, data, size, tags, driver, features):
-    """Runs lit on all tests it can find in `data` under tensorflow/compiler/mlir.
-
-    Note that, due to Bazel's hermetic builds, lit only sees the tests that
-    are included in the `data` parameter, regardless of what other tests might
-    exist in the directory searched.
-
-    Args:
-      name: str, the name of the test, including extension.
-      data: [str], the data input to the test.
-      size: str, the size of the test.
-      tags: [str], tags to attach to the test.
-      driver: str, label of the driver shell script.
-              Note: use of a custom driver is not currently supported
-              and specifying a default driver will abort the tests.
-      features: [str], list of extra features to enable.
-    """
-    if driver != _default_driver:
-        fail("There is no present support for custom drivers. Please omit" +
-             " the driver parameter when running this test. If you require" +
-             " custom driver support, please file an issue to request it.")
-
-    # Disable tests on windows for now, to enable testing rest of all xla and mlir.
-    native.py_test(
-        name = name,
-        srcs = ["@llvm-project//llvm:lit"],
-        tags = tags + ["no_windows"],
-        args = [
-            "mlir_tests/" + paths.basename(data[-1]),
-            "--show-all",
-        ] + features,
-        data = data + [
-            "//mlir_tests:litcfgs",
-            "@llvm-project//llvm:FileCheck",
-            "@llvm-project//llvm:count",
-            "@llvm-project//llvm:not",
-        ],
-        size = size,
-        main = "lit.py",
-    )
+# TODO(b/136126535): consider upstreaming at least a subset of this file.
+"""Lit test macros"""
 
 def glob_lit_tests(
+        cfgs,
+        test_file_exts,
         exclude = [],
-        test_file_exts = _default_test_file_exts,
-        default_size = _default_size,
-        size_override = {},
         data = [],
         per_test_extra_data = {},
-        default_tags = _default_tags,
+        default_size = "small",
+        size_override = {},
+        default_tags = [],
         tags_override = {},
-        driver = _default_driver,
-        features = []):
-    """Creates all plausible Lit tests (and their inputs) under this directory.
+        features = [],
+        **kwargs):
+    """Creates all plausible Lit tests (and their inputs) in this package.
 
     Args:
-      exclude: [str], paths to exclude (for tests and inputs).
+      cfgs: label, lit config files.
       test_file_exts: [str], extensions for files that are tests.
-      default_size: str, the test size for targets not in "size_override".
-      size_override: {str: str}, sizes to use for specific tests.
+      exclude: [str], paths to exclude (for tests and inputs).
       data: [str], additional input data to the test.
       per_test_extra_data: {str: [str]}, extra data to attach to a given file.
+      default_size: str, the test size for targets not in "size_override".
+      size_override: {str: str}, sizes to use for specific tests.
       default_tags: [str], additional tags to attach to the test.
       tags_override: {str: str}, tags to add to specific tests.
-      driver: str, label of the driver shell script.
-              Note: use of a custom driver is not currently supported
-              and specifying a default driver will abort the tests.
-      features: [str], list of extra features to enable.
+      features: [str], extra arguments to pass to lit.
+      **kwargs: arguments to pass on to lit_test()
     """
 
-    # Ignore some patterns by default for tests and input data.
-    exclude = _ALWAYS_EXCLUDE + exclude
-
-    tests = native.glob(
-        ["*." + ext for ext in test_file_exts],
-        exclude = exclude,
-    )
-
-    # Run tests individually such that errors can be attributed to a specific
-    # failure.
-    for i in range(len(tests)):
-        curr_test = tests[i]
-
-        # Instantiate this test with updated parameters.
-        lit_test(
-            name = curr_test,
-            data = data + per_test_extra_data.get(curr_test, []),
-            size = size_override.get(curr_test, default_size),
-            tags = default_tags + tags_override.get(curr_test, []),
-            driver = driver,
-            features = features,
+    # Add some default data/features.
+    data = data + [
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//llvm:count",
+        "@llvm-project//llvm:not",
+    ]
+    features = features + ["--verbose", "--show-all"]
+
+    include = ["**/*." + ext for ext in test_file_exts]
+    for test in native.glob(include, exclude):
+        input = test + ".input"
+
+        lit_input(
+            name = input,
+            srcs = [test],
+            cfgs = cfgs,
         )
 
-def lit_test(
-        name,
-        data = [],
-        size = _default_size,
-        tags = _default_tags,
-        driver = _default_driver,
-        features = []):
-    """Runs test files under lit.
+        native.py_test(
+            name = test + ".test",
+            srcs = ["@llvm-project//llvm:lit"],
+            args = ["@$(rootpath :%s)" % input] + features,
+            data = [input] + data + per_test_extra_data.get(test, []),
+            tags = default_tags + tags_override.get(test, []),
+            size = size_override.get(test, default_size),
+            main = "lit.py",
+            **kwargs
+        )
 
-    Args:
-      name: str, the name of the test.
-      data: [str], labels that should be provided as data inputs.
-      size: str, the size of the test.
-      tags: [str], tags to attach to the test.
-      driver: str, label of the driver shell script.
-              Note: use of a custom driver is not currently supported
-              and specifying a default driver will abort the tests.
-      features: [str], list of extra features to enable.
-    """
-    _run_lit_test(name + ".test", data + [name], size, tags, driver, features)
+def _lit_input_impl(ctx):
+    dirs = {
+        f.short_path[:-len(f.basename)]: True
+        for f in ctx.files.cfgs
+    }
+    if len(dirs) != 1:
+        fail("All 'cfgs' files must be in same directory")
+    cfgs_dir = dirs.keys()[0]
+
+    inputs = []
+    for src in ctx.attr.srcs:
+        # Use the cfgs directory so that the suite can be found even if it's not
+        # in a parent directory of the test file. Add just basename of the input
+        # so that tests are executed in the config.test_exec_root directory.
+        # This requires that config.test_source_root is set to the directory of
+        # the 'TEST_BINARY' environment variable that bazel provides.
+        inputs.extend([cfgs_dir + f.basename for f in src.files.to_list()])
+
+    output = ctx.actions.declare_file(ctx.label.name + ".out")
+    ctx.actions.write(output, "\n".join(inputs))
+
+    runfiles = [output] + ctx.files.cfgs + ctx.files.srcs
+    return [DefaultInfo(
+        files = depset([output]),
+        runfiles = ctx.runfiles(runfiles),
+    )]
+
+lit_input = rule(
+    implementation = _lit_input_impl,
+    attrs = {
+        "srcs": attr.label_list(
+            doc = "Files to test.",
+            allow_files = True,
+            mandatory = True,
+        ),
+        "cfgs": attr.label(
+            doc = "Suite cfg files.",
+            allow_files = True,
+            mandatory = True,
+        ),
+    },
+    doc = "Generates a file containing the lit tests to run.",
+)

mlir_tests/lit.cfg.py

@@ -14,13 +14,14 @@
 """Lit configuration."""
 
 from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
 
+import os
 import lit.formats
 from lit.llvm import llvm_config
 from lit.llvm.subst import ToolSubst
 
+cwd = os.getcwd()
+
 # pylint: disable=undefined-variable
 
 # name: The name of this test suite.
@@ -32,25 +33,39 @@
 # suffixes: A list of file extensions to treat as test files.
 config.suffixes = ['.mlir']
 
-# test_source_root: The root path where tests are located.
-config.test_source_root = config.tfrt_test_dir
+# test_source_root: Base path of the test files. The remainder (none in this
+# case) is the lit argument stripped by the path to the suite directory.
+config.test_source_root = os.path.dirname(os.environ['TEST_BINARY'])
+
+# test_exec_root: Base path to the execution directory. The remainder is the lit
+# argument directory stripped by the path to the suite directory.
+config.test_exec_root = '.'
 
-# test_exec_root: The root path where tests should be run.
-config.test_exec_root = config.runfile_srcdir
+config.llvm_tools_dir = os.path.join(cwd, '..', 'llvm-project', 'llvm')
+
+# pylint: enable=undefined-variable
 
 llvm_config.use_default_substitutions()
 
-llvm_config.config.substitutions.append(
-    ('%tfrt_bindir', 'tensorflow/compiler/aot'))
 
-tool_dirs = config.tfrt_tools_dirs + [config.llvm_tools_dir]
+def _AddToolSubstitutions(targets):
+  paths = [
+      t.lstrip('/').replace('@', '../').replace('//', '/').replace(':', '/')
+      for t in targets
+  ]
+  llvm_config.add_tool_substitutions([
+      ToolSubst(os.path.basename(p), os.path.join(cwd, p), unresolved='ignore')
+      for p in paths
+  ], [])
 
-tool_names = [
-    'bef_executor', 'bef_executor_lite', 'tfrt_translate', 'tfrt_opt',
-    'tfrt_gpu_translate', 'tfrt_gpu_opt', 'code_size_test_driver',
-    'bef_executor_debug_tracing'
-]
-tools = [ToolSubst(s, unresolved='ignore') for s in tool_names]
-llvm_config.add_tool_substitutions(tools, tool_dirs)
 
-# pylint: enable=undefined-variable
+_AddToolSubstitutions([
+    '//backends/gpu:tfrt_gpu_translate',
+    '//backends/gpu:tfrt_gpu_opt',
+    '//tools:bef_executor',
+    '//tools:bef_executor_debug_tracing',
+    '//tools:bef_executor_lite',
+    '//tools:code_size_test_driver',
+    '//tools:tfrt_opt',
+    '//tools:tfrt_translate',
+])

mlir_tests/lit.site.cfg.py

@@ -14,53 +14,16 @@
 """Lit site configuration."""
 
 from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
 
 import os
-import platform
-import sys
 import lit.llvm
 
-from pathlib import Path
-
-if platform.system() not in ['Linux', 'Windows']:
-  sys.exit('Currently TFRT only supports lit tests on Linux and Windows.')
-
-runfile_srcdir = os.environ['TEST_SRCDIR']
-tfrt_workspace = os.environ['TEST_WORKSPACE']
-
-# On Windows, bazel uses a different runfiles location
-# (https://bazel.build/designs/2016/09/05/build-python-on-windows.html)
-if platform.system() == 'Windows':
-  runfile_root = Path.cwd().parent
-  external_srcdir = os.path.join(
-      runfile_srcdir[:runfile_srcdir.find('mlir_tests')], 'external')
-  external_srcdir = Path(external_srcdir).absolute()
-else:
-  runfile_root = runfile_srcdir
-  external_srcdir = runfile_srcdir
+cfg_path = os.path.join('mlir_tests', 'lit.cfg.py')
 
 # pylint: disable=undefined-variable
 
-config.runfile_srcdir = os.path.join(runfile_srcdir, tfrt_workspace)
-
-config.llvm_tools_dir = os.path.join(external_srcdir, 'llvm-project', 'llvm')
-
-config.tfrt_tools_dirs = [
-    os.path.join(runfile_root, tfrt_workspace, 'tools'),
-    os.path.join(runfile_root, tfrt_workspace, 'backends', 'gpu'),
-]
-
-test_target_dir = os.environ['TEST_TARGET'].strip('/').rsplit(':')[0]
-config.tfrt_test_dir = os.path.join(runfile_root, tfrt_workspace,
-                                    test_target_dir)
-
-lit.llvm.initialize(lit_config, config)
-
 # Let the main config do the real work.
-lit_config.load_config(
-    config,
-    os.path.join(runfile_root, tfrt_workspace, 'mlir_tests', 'lit.cfg.py'))
+lit.llvm.initialize(lit_config, config)
+lit_config.load_config(config, cfg_path)
 
 # pylint: enable=undefined-variable