Use the real lit tool for lit testing (#8166)

For historical reasons, we've used a custom shell script to run our lit
tests. We've made it less bad over time, but it's time to just use the
real thing. This gets the real lit tool running under Bazel (and CMake,
but that's less hard).

I hope to upstream some of this Bazel stuff, but LLVM's lit setup is
way more complicated, so checking this in here for now.

Fixes https://github.com/google/iree/issues/818
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3978118..e1a705e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -216,7 +216,7 @@
 include(iree_add_all_subdirs)
 include(iree_check_test)
 include(iree_trace_runner_test)
-include(iree_run_binary_test)
+include(iree_native_test)
 include(iree_benchmark_suite)
 include(iree_hal_cts_test_suite)
 
diff --git a/build_tools/bazel/iree_check_test.bzl b/build_tools/bazel/iree_check_test.bzl
index 5359e42..b34ef02 100644
--- a/build_tools/bazel/iree_check_test.bzl
+++ b/build_tools/bazel/iree_check_test.bzl
@@ -7,7 +7,7 @@
 """Macros for defining tests that run a module using iree-check-module."""
 
 load("//build_tools/bazel:iree_bytecode_module.bzl", "iree_bytecode_module")
-load("//build_tools/bazel:run_binary_test.bzl", "run_binary_test")
+load("//build_tools/bazel:native_binary.bzl", "native_test")
 
 ALL_TARGET_BACKENDS_AND_DRIVERS = [
     ("vmvx", "vmvx"),
@@ -47,7 +47,7 @@
           automatically.
       target_cpu_features: currently unimplemented (must be empty), will eventually allow specifying target CPU features.
       timeout: timeout for the generated tests.
-      **kwargs: any additional attributes to pass to the underlying run_binary_test.
+      **kwargs: any additional attributes to pass to the underlying native_test.
     """
 
     if target_cpu_features:
@@ -67,14 +67,14 @@
         visibility = ["//visibility:private"],
     )
 
-    run_binary_test(
+    native_test(
         name = name,
         args = [
             "--driver=%s" % driver,
             "$(location :%s)" % bytecode_module_name,
         ] + runner_args,
         data = [":%s" % bytecode_module_name],
-        test_binary = "//iree/tools:iree-check-module",
+        src = "//iree/tools:iree-check-module",
         tags = tags + ["driver=%s" % driver],
         timeout = timeout,
         **kwargs
diff --git a/build_tools/bazel/iree_lit_test.bzl b/build_tools/bazel/iree_lit_test.bzl
new file mode 100644
index 0000000..f71c887
--- /dev/null
+++ b/build_tools/bazel/iree_lit_test.bzl
@@ -0,0 +1,97 @@
+# Copyright 2019 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
+
+"""Bazel macros for running lit tests."""
+
+load(":lit_test.bzl", "lit_test", "lit_test_suite")
+
+def iree_lit_test(
+        name,
+        cfg = "//iree:lit.cfg.py",
+        tools = None,
+        env = None,
+        **kwargs):
+    """A thin wrapper around lit_test with some opinionated settings.
+
+    See the base lit_test for more details on argument meanings.
+
+    Args:
+      name: name for the test.
+      cfg: string. lit config file.
+      tools: label_list. tools that should be included on the PATH.
+        llvm-symbolizer is added by default.
+      env: string_dict. Environment variables available to the test at runtime.
+        FILECHECK_OPTS=--enable-var-scope is added if FILECHECK_OPTS is not
+        already set.
+      **kwargs: additional keyword args to forward to the underyling lit_test.
+    """
+
+    tools = tools or []
+    env = env or {}
+
+    # Always include llvm-symbolizer so we get useful stack traces. Maybe it
+    # would be better to force everyone to do this explicitly, but since
+    # forgetting wouldn't cause the test to fail, only make debugging harder
+    # when it does, I think better to hardcode it here.
+    llvm_symbolizer = "@llvm-project//llvm:llvm-symbolizer"
+    if llvm_symbolizer not in tools:
+        tools.append(llvm_symbolizer)
+
+    filecheck_env_var = "FILECHECK_OPTS"
+    if filecheck_env_var not in env:
+        env[filecheck_env_var] = "--enable-var-scope"
+
+    lit_test(
+        name = name,
+        cfg = cfg,
+        tools = tools,
+        env = env,
+        **kwargs
+    )
+
+def iree_lit_test_suite(
+        name,
+        cfg = "//iree:lit.cfg.py",
+        tools = None,
+        env = None,
+        **kwargs):
+    """A thin wrapper around lit_test_suite with some opinionated settings.
+
+    See the base lit_test for more details on argument meanings.
+
+    Args:
+      name: name for the test suite.
+      cfg: string. lit config file.
+      tools: label_list. tools that should be included on the PATH.
+        llvm-symbolizer is added by default.
+      env: string_dict. Environment variables available to the test at runtime.
+        FILECHECK_OPTS=--enable-var-scope is added if FILECHECK_OPTS is not
+        already set.
+      **kwargs: additional keyword args to forward to the underyling
+        lit_test_suite.
+    """
+    tools = tools or []
+    env = env or {}
+
+    # Always include llvm-symbolizer so we get useful stack traces. Maybe it
+    # would be better to force everyone to do this explicitly, but since
+    # forgetting wouldn't cause the test to fail, only make debugging harder
+    # when it does, I think better to hardcode it here.
+    llvm_symbolizer = "@llvm-project//llvm:llvm-symbolizer"
+    if llvm_symbolizer not in tools:
+        tools.append(llvm_symbolizer)
+
+    filecheck_env_var = "FILECHECK_OPTS"
+    if filecheck_env_var not in env:
+        env[filecheck_env_var] = "--enable-var-scope"
+
+    lit_test_suite(
+        name = name,
+        cfg = cfg,
+        tools = tools,
+        env = env,
+        **kwargs
+    )
diff --git a/build_tools/bazel/iree_trace_runner_test.bzl b/build_tools/bazel/iree_trace_runner_test.bzl
index f5e3127..c0e6b79 100644
--- a/build_tools/bazel/iree_trace_runner_test.bzl
+++ b/build_tools/bazel/iree_trace_runner_test.bzl
@@ -7,7 +7,7 @@
 """Macros for defining tests that use a trace-runner."""
 
 load("//build_tools/bazel:iree_bytecode_module.bzl", "iree_bytecode_module")
-load("//build_tools/bazel:run_binary_test.bzl", "run_binary_test")
+load("//build_tools/bazel:native_binary.bzl", "native_test")
 
 def iree_trace_runner_test(
         name,
@@ -70,7 +70,7 @@
         **kwargs
     )
 
-    run_binary_test(
+    native_test(
         name = name,
         args = [
             "--driver=%s" % driver,
@@ -80,7 +80,7 @@
             ":%s" % bytecode_module_name,
             ":%s" % trace,
         ],
-        test_binary = trace_runner,
+        src = trace_runner,
         tags = tags + ["driver=%s" % driver],
         timeout = timeout,
         **kwargs
diff --git a/build_tools/bazel/lit_test.bzl b/build_tools/bazel/lit_test.bzl
new file mode 100644
index 0000000..e77d7a0
--- /dev/null
+++ b/build_tools/bazel/lit_test.bzl
@@ -0,0 +1,192 @@
+# 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
+
+"""Rules for running lit tests with the upstream lit binary."""
+
+# This exists as a separate file from iree_lit_test.bzl because we anticipate
+# upstreaming it soon.
+
+load("@bazel_skylib//lib:paths.bzl", "paths")
+load(":native_binary.bzl", "native_test")
+
+def _tools_on_path_impl(ctx):
+    runfiles = ctx.runfiles()
+
+    # For Bazel 4.x support. Drop when Bazel 4.x is no longer supported
+    to_merge = [d[DefaultInfo].default_runfiles for d in ctx.attr.srcs]
+    if hasattr(runfiles, "merge_all"):
+        runfiles = runfiles.merge_all(to_merge)
+    else:
+        for m in to_merge:
+            runfiles = runfiles.merge(m)
+
+    runfiles_symlinks = {}
+
+    for src in ctx.attr.srcs:
+        exe = src[DefaultInfo].files_to_run.executable
+        if not exe:
+            fail("All targets used as tools by lit tests must have exactly one" +
+                 " executable, but {} has none".format(src))
+        bin_path = paths.join(ctx.attr.bin_dir, exe.basename)
+        if bin_path in runfiles_symlinks:
+            fail("All tools used by lit tests must have unique basenames, as" +
+                 " they are added to the path." +
+                 " {} and {} conflict".format(runfiles_symlinks[bin_path], exe))
+        runfiles_symlinks[bin_path] = exe
+
+    return [
+        DefaultInfo(runfiles = ctx.runfiles(
+            symlinks = runfiles_symlinks,
+        ).merge(runfiles)),
+    ]
+
+_tools_on_path = rule(
+    _tools_on_path_impl,
+    attrs = {
+        "srcs": attr.label_list(allow_files = True, mandatory = True),
+        "bin_dir": attr.string(mandatory = True),
+    },
+    doc = "Symlinks srcs into a single lit_bin directory. All basenames must be unique.",
+)
+
+def lit_test(
+        name,
+        test_file,
+        cfg,
+        tools = None,
+        args = None,
+        data = None,
+        visibility = None,
+        env = None,
+        **kwargs):
+    """Runs a single test file with LLVM's lit tool.
+
+    Args:
+      name: string. the name of the generated test target.
+      test_file: label. The file on which to run lit.
+      cfg: label. The lit config file. It must list the file extension of
+        `test_file` in config.suffixes and must be in a parent directory of
+        `test_file`.
+      tools: label list. Tools invoked in the lit RUN lines. These binaries will
+        be symlinked into a directory which is on the path. They must therefore
+        have unique basenames.
+      args: string list. Additional arguments to pass to lit. Note that the test
+        file, `-v`, and a `--path` argument for the directory to which `tools`
+        are symlinked are added automatically.
+      data: label list. Additional data dependencies of the test. Note that
+        targets in `cfg` and `tools`, as well as their data dependencies, are
+        added automatically.
+      visibility: visibility of the generated test target.
+      env: string_dict. Environment variables available during test execution.
+        See the common Bazel test attribute.
+      **kwargs: additional keyword arguments to pass to all generated rules.
+
+    See https://llvm.org/docs/CommandGuide/lit.html for details on lit
+    """
+    args = args or []
+    data = data or []
+    tools = tools or []
+
+    tools_on_path_target_name = "_{}_tools_on_path".format(name)
+
+    bin_dir = paths.join(
+        native.package_name(),
+        tools_on_path_target_name,
+        "lit_bin",
+    )
+
+    _tools_on_path(
+        name = tools_on_path_target_name,
+        testonly = True,
+        srcs = tools,
+        bin_dir = bin_dir,
+        visibility = ["//visibility:private"],
+        **kwargs
+    )
+
+    native_test(
+        name = name,
+        src = "@llvm-project//llvm:lit",
+        # out = name,
+        args = [
+            "-v",
+            "--path",
+            bin_dir,
+            "$(location {})".format(test_file),
+        ] + args,
+        data = [test_file, cfg, tools_on_path_target_name] + data,
+        visibility = visibility,
+        env = env,
+        **kwargs
+    )
+
+def lit_test_suite(
+        name,
+        srcs,
+        cfg,
+        tools = None,
+        args = None,
+        data = None,
+        visibility = None,
+        size = "small",
+        env = None,
+        **kwargs):
+    """Creates one lit test per source file and a test suite that bundles them.
+
+    Args:
+      name: string. the name of the generated test suite.
+      srcs: label_list. The files which contain the lit tests.
+      cfg: label. The lit config file. It must list the file extension of
+        the files in `srcs` in config.suffixes and must be in a parent directory
+        of `srcs`.
+      tools: label list. Tools invoked in the lit RUN lines. These binaries will
+        be symlinked into a directory which is on the path. They must therefore
+        have unique basenames.
+      args: string list. Additional arguments to pass to lit. Note that the test
+        file, `-v`, and a `--path` argument for the directory to which `tools`
+        are symlinked are added automatically.
+      data: label list. Additional data dependencies of the test. Note that
+        targets in `cfg` and `tools`, as well as their data dependencies, are
+        added automatically.
+      visibility: visibility of the generated test targets and test suite.
+      size: string. size of the generated tests.
+      env: string_dict. Environment variables available during test execution.
+        See the common Bazel test attribute.
+      **kwargs: additional keyword arguments to pass to all generated rules.
+
+    See https://llvm.org/docs/CommandGuide/lit.html for details on lit
+    """
+    # If there are kwargs that need to be passed to only some of the generated
+    # rules, they should be extracted into separate named arguments.
+
+    args = args or []
+    data = data or []
+    tools = tools or []
+
+    tests = []
+    for test_file in srcs:
+        # It's generally good practice to prefix any generated names with the
+        # macro name, but it's also nice to have the test name just match the
+        # file name.
+        test_name = "%s.test" % (test_file)
+        tests.append(test_name)
+        lit_test(
+            name = test_name,
+            test_file = test_file,
+            cfg = cfg,
+            tools = tools,
+            args = args,
+            data = data,
+            visibility = visibility,
+            env = env,
+            **kwargs
+        )
+
+    native.test_suite(
+        name = name,
+        tests = tests,
+        **kwargs
+    )
diff --git a/build_tools/bazel/native_binary.bzl b/build_tools/bazel/native_binary.bzl
new file mode 100644
index 0000000..5b047e6
--- /dev/null
+++ b/build_tools/bazel/native_binary.bzl
@@ -0,0 +1,99 @@
+# 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
+
+"""native_binary() and native_test() rule implementations.
+
+Rewritten from the Bazel Skylib version pending several fixes and improvements
+to that rule:
+
+- https://github.com/bazelbuild/bazel-skylib/pull/338
+- https://github.com/bazelbuild/bazel-skylib/pull/339
+- https://github.com/bazelbuild/bazel-skylib/pull/340
+- https://github.com/bazelbuild/bazel-skylib/pull/341
+
+These rules let you wrap a pre-built binary or script in a conventional binary
+and test rule respectively. They fulfill the same goal as sh_binary and sh_test
+do, but they run the wrapped binary directly, instead of through Bash, so they
+don't depend on Bash and work with --shell_exectuable="".
+"""
+
+def _shared_impl(ctx):
+    out = ctx.attr.out
+    if not out:
+        out = ctx.attr.name
+    output = ctx.actions.declare_file(out)
+    ctx.actions.symlink(
+        target_file = ctx.executable.src,
+        output = output,
+        is_executable = True,
+    )
+
+    runfiles = ctx.runfiles(files = ctx.files.data)
+
+    # For Bazel 4.x support. Drop when Bazel 4.x is no longer supported
+    to_merge = ([d[DefaultInfo].default_runfiles for d in ctx.attr.data] +
+                [ctx.attr.src[DefaultInfo].default_runfiles])
+    if hasattr(runfiles, "merge_all"):
+        runfiles = runfiles.merge_all(to_merge)
+    else:
+        for m in to_merge:
+            runfiles = runfiles.merge(m)
+    return DefaultInfo(
+        executable = output,
+        files = depset([output]),
+        runfiles = runfiles,
+    )
+
+def _native_binary_impl(ctx):
+    default_info = _shared_impl(ctx)
+    return [default_info]
+
+def _native_test_impl(ctx):
+    default_info = _shared_impl(ctx)
+    return [default_info, testing.TestEnvironment(ctx.attr.env)]
+
+# We have to manually set "env" on the test rule because the builtin one is only
+# available in native rules. See
+# https://docs.bazel.build/versions/main/be/common-definitions.html#test.env
+# We don't have "env" on native_binary because there is no BinaryEnvironment
+# mirroring TestEnvironment. See https://github.com/bazelbuild/bazel/issues/7364
+_SHARED_ATTRS = {
+    "src": attr.label(
+        executable = True,
+        allow_files = True,
+        mandatory = True,
+        cfg = "target",
+    ),
+    "data": attr.label_list(allow_files = True),
+    # "out" is attr.string instead of attr.output, so that it is select()'able.
+    "out": attr.string(),
+}
+
+native_binary = rule(
+    implementation = _native_binary_impl,
+    attrs = _SHARED_ATTRS,
+    executable = True,
+)
+
+_TEST_ATTRS = {
+    k: v
+    for k, v in _SHARED_ATTRS.items() + [
+        (
+            "env",
+            attr.string_dict(
+                doc = "Mirrors the common env attribute that otherwise is" +
+                      " only available on native rules. See" +
+                      " https://docs.bazel.build/versions/main/be/common-definitions.html#test.env",
+            ),
+        ),
+    ]
+}
+
+native_test = rule(
+    implementation = _native_test_impl,
+    attrs = _TEST_ATTRS,
+    test = True,
+)
diff --git a/build_tools/bazel/run_binary_test.bzl b/build_tools/bazel/run_binary_test.bzl
deleted file mode 100644
index 0b8b9aa..0000000
--- a/build_tools/bazel/run_binary_test.bzl
+++ /dev/null
@@ -1,53 +0,0 @@
-# Copyright 2020 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
-
-"""Creates a test from the binary output of another rule.
-
-The rule instantiation can pass additional arguments to the binary and provide
-it with additional data files (as well as the standard bazel test classification
-attributes). This allows compiling the binary once and not recompiling or
-relinking it for each test rule. It also avoids a wrapper shell script, which
-adds unnecessary shell dependencies and confuses some tooling about the type of
-the binary.
-
-Example usage:
-
-run_binary_test(
-    name = "my_test",
-    args = ["--module_file=$(location :data_file)"],
-    data = [":data_file"],
-    test_binary = ":some_cc_binary",
-)
-"""
-
-def _run_binary_test_impl(ctx):
-    ctx.actions.symlink(
-        target_file = ctx.executable.test_binary,
-        output = ctx.outputs.executable,
-        is_executable = True,
-    )
-
-    data_runfiles = ctx.runfiles(files = ctx.files.data)
-
-    binary_runfiles = ctx.attr.test_binary[DefaultInfo].default_runfiles
-
-    return [DefaultInfo(
-        executable = ctx.outputs.executable,
-        runfiles = data_runfiles.merge(binary_runfiles),
-    )]
-
-run_binary_test = rule(
-    _run_binary_test_impl,
-    attrs = {
-        "test_binary": attr.label(
-            mandatory = True,
-            executable = True,
-            cfg = "target",
-        ),
-        "data": attr.label_list(allow_empty = True, allow_files = True),
-    },
-    test = True,
-)
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 7c66d0c..c875d87 100644
--- a/build_tools/bazel_to_cmake/bazel_to_cmake_converter.py
+++ b/build_tools/bazel_to_cmake/bazel_to_cmake_converter.py
@@ -508,15 +508,23 @@
                             f"{tblgen_block}"
                             f")\n\n")
 
-  def iree_lit_test_suite(self, name, srcs, data, tags=None, **kwargs):
+  def iree_lit_test_suite(self,
+                          name,
+                          srcs,
+                          tools=None,
+                          data=None,
+                          tags=None,
+                          **kwargs):
     name_block = _convert_string_arg_block("NAME", name, quote=False)
     srcs_block = _convert_srcs_block(srcs)
+    tools_block = _convert_target_list_block("TOOLS", tools)
     data_block = _convert_target_list_block("DATA", data)
     labels_block = _convert_string_list_block("LABELS", tags)
 
     self.converter.body += (f"iree_lit_test_suite(\n"
                             f"{name_block}"
                             f"{srcs_block}"
+                            f"{tools_block}"
                             f"{data_block}"
                             f"{labels_block}"
                             f")\n\n")
@@ -711,17 +719,16 @@
                             f"{labels_block}"
                             f")\n\n")
 
-  def run_binary_test(self, name, test_binary, args=None, data=None, tags=None):
+  def native_test(self, name, src, args=None, data=None, tags=None):
     if data is not None:
-      self._convert_unimplemented_function("iree_run_binary_test",
-                                           name + " has data")
+      self._convert_unimplemented_function("native_test", name + " has data")
 
     name_block = _convert_string_arg_block("NAME", name)
-    test_binary_block = _convert_single_target_block("TEST_BINARY", test_binary)
+    test_binary_block = _convert_single_target_block("SRC", src)
     args_block = _convert_string_list_block("ARGS", args)
     labels_block = _convert_string_list_block("LABELS", tags)
 
-    self.converter.body += (f"iree_run_binary_test(\n"
+    self.converter.body += (f"iree_native_test(\n"
                             f"{name_block}"
                             f"{args_block}"
                             f"{test_binary_block}"
diff --git a/build_tools/cmake/iree_check_test.cmake b/build_tools/cmake/iree_check_test.cmake
index 0340ac6..97cf9fd 100644
--- a/build_tools/cmake/iree_check_test.cmake
+++ b/build_tools/cmake/iree_check_test.cmake
@@ -178,12 +178,12 @@
     "${_RUNNER_TARGET}"
   )
 
-  iree_run_binary_test(
+  iree_native_test(
     NAME
       "${_RULE_NAME}"
     DRIVER
       "${_RULE_DRIVER}"
-    TEST_BINARY
+    SRC
       "${_RUNNER_TARGET}"
     TEST_INPUT_FILE_ARG
       "${_MODULE_FILE_NAME}"
diff --git a/build_tools/cmake/iree_lit_test.cmake b/build_tools/cmake/iree_lit_test.cmake
index d82e0fc..13ef3aa 100644
--- a/build_tools/cmake/iree_lit_test.cmake
+++ b/build_tools/cmake/iree_lit_test.cmake
@@ -15,13 +15,11 @@
 # Parameters:
 # NAME: Name of the target
 # TEST_FILE: Test file to run with the lit runner.
+# TOOLS: Tools that should be included on the PATH
 # DATA: Additional data dependencies invoked by the test (e.g. binaries
 #   called in the RUN line)
 # LABELS: Additional labels to apply to the test. The package path is added
 #     automatically.
-#
-# TODO(gcmn): allow using alternative driver
-# A driver other than the default iree/tools/run_lit.sh is not currently supported.
 function(iree_lit_test)
   if(NOT IREE_BUILD_TESTS)
     return()
@@ -37,7 +35,7 @@
     _RULE
     ""
     "NAME;TEST_FILE"
-    "DATA;LABELS"
+    "DATA;TOOLS;LABELS"
     ${ARGN}
   )
 
@@ -52,10 +50,16 @@
   get_filename_component(_TEST_FILE_PATH ${_RULE_TEST_FILE} ABSOLUTE)
 
   list(TRANSFORM _RULE_DATA REPLACE "^::" "${_PACKAGE_NS}::")
+  list(TRANSFORM _RULE_TOOLS REPLACE "^::" "${_PACKAGE_NS}::")
   set(_DATA_DEP_PATHS)
-  foreach(_DATA_DEP ${_RULE_DATA})
+  foreach(_DATA_DEP IN LISTS _RULE_DATA _RULE_TOOLS)
     list(APPEND _DATA_DEP_PATHS $<TARGET_FILE:${_DATA_DEP}>)
-  endforeach(_DATA_DEP)
+  endforeach()
+
+  set(_LIT_PATH_ARGS)
+  foreach(_TOOL IN LISTS _RULE_TOOLS)
+    list(APPEND _LIT_PATH_ARGS "--path" "$<TARGET_FILE_DIR:${_TOOL}>")
+  endforeach()
 
   iree_package_ns(_PACKAGE_NS)
   string(REPLACE "::" "/" _PACKAGE_PATH ${_PACKAGE_NS})
@@ -67,9 +71,11 @@
       # We run all our tests through a custom test runner to allow setup
       # and teardown.
       "${CMAKE_SOURCE_DIR}/build_tools/cmake/run_test.${IREE_HOST_SCRIPT_EXT}"
-      "${CMAKE_SOURCE_DIR}/iree/tools/run_lit.${IREE_HOST_SCRIPT_EXT}"
+      "${Python3_EXECUTABLE}"
+      "${LLVM_SOURCE_DIR}/utils/lit/lit.py"
+      "-v"
+      ${_LIT_PATH_ARGS}
       ${_TEST_FILE_PATH}
-      ${_DATA_DEP_PATHS}
   )
 
   list(APPEND _RULE_LABELS "${_PACKAGE_PATH}")
@@ -93,13 +99,10 @@
 # Parameters:
 # NAME: Name of the target
 # SRCS: List of test files to run with the lit runner. Creates one test per source.
-# DATA: Additional data dependencies invoked by the test (e.g. binaries
-#   called in the RUN line)
+# TOOLS: Tools that should be included on the PATH
+# DATA: Additional data dependencies used by the test
 # LABELS: Additional labels to apply to the generated tests. The package path is
 #     added automatically.
-#
-# TODO(gcmn): allow using alternative driver
-# A driver other than the default iree/tools/run_lit.sh is not currently supported.
 function(iree_lit_test_suite)
   if(NOT IREE_BUILD_TESTS)
     return()
@@ -113,7 +116,7 @@
     _RULE
     ""
     "NAME"
-    "SRCS;DATA;LABELS"
+    "SRCS;DATA;TOOLS;LABELS"
     ${ARGN}
   )
 
@@ -126,6 +129,8 @@
         "${_TEST_FILE}"
       DATA
         "${_RULE_DATA}"
+      TOOLS
+        "${_RULE_TOOLS}"
       LABELS
         "${_RULE_LABELS}"
     )
diff --git a/build_tools/cmake/iree_run_binary_test.cmake b/build_tools/cmake/iree_native_test.cmake
similarity index 83%
rename from build_tools/cmake/iree_run_binary_test.cmake
rename to build_tools/cmake/iree_native_test.cmake
index 5d68900..7a6ddcb 100644
--- a/build_tools/cmake/iree_run_binary_test.cmake
+++ b/build_tools/cmake/iree_native_test.cmake
@@ -6,7 +6,7 @@
 
 include(CMakeParseArguments)
 
-# iree_run_binary_test()
+# iree_native_test()
 #
 # Creates a test that runs the specified binary with the specified arguments.
 #
@@ -17,15 +17,15 @@
 # DRIVER: If specified, will pass --driver=DRIVER to the test binary and adds
 #     a driver label to the test.
 # TEST_INPUT_FILE_ARG: If specified, the input file will be added to DATA and
-#     its device path appended to ARGS. Note that the device path may be different
-#     from the host path, so this parameter should be used to portably pass file arguments
-#     to tests.
-# DATA: Additional input files needed by the test binary. When running tests on a
-#     separate device (e.g. Android), these files will be pushed to the device.
-#     TEST_INPUT_FILE_ARG is automatically added if specified.
+#     its device path appended to ARGS. Note that the device path may be
+#     different from the host path, so this parameter should be used to portably
+#     pass file arguments to tests.
+# DATA: Additional input files needed by the test binary. When running tests on
+#     a separate device (e.g. Android), these files will be pushed to the
+#     device. TEST_INPUT_FILE_ARG is automatically added if specified.
 # ARGS: additional arguments passed to the test binary. TEST_INPUT_FILE_ARG and
 #     --driver=DRIVER are automatically added if specified.
-# TEST_BINARY: binary target to run as the test.
+# SRC: binary target to run as the test.
 # LABELS: Additional labels to apply to the test. The package path is added
 #     automatically.
 #
@@ -38,16 +38,16 @@
 #     requires_args_to_run
 #   ...
 # )
-# iree_run_binary_test(
+# iree_native_test(
 #   NAME
 #     requires_args_to_run_test
 #   ARGS
 #    --do-the-right-thing
-#   TEST_BINARY
+#   SRC
 #     ::requires_args_to_run
 # )
 
-function(iree_run_binary_test)
+function(iree_native_test)
   if(NOT IREE_BUILD_TESTS)
     return()
   endif()
@@ -55,7 +55,7 @@
   cmake_parse_arguments(
     _RULE
     ""
-    "NAME;TEST_BINARY;DRIVER;TEST_INPUT_FILE_ARG"
+    "NAME;SRC;DRIVER;TEST_INPUT_FILE_ARG"
     "ARGS;LABELS;DATA"
     ${ARGN}
   )
@@ -88,7 +88,7 @@
   endif()
 
   # Replace binary passed by relative ::name with iree::package::name
-  string(REGEX REPLACE "^::" "${_PACKAGE_NS}::" _TEST_BINARY_TARGET ${_RULE_TEST_BINARY})
+  string(REGEX REPLACE "^::" "${_PACKAGE_NS}::" _SRC_TARGET ${_RULE_SRC})
 
   if(ANDROID)
     # Define a custom target for pushing and running the test on Android device.
@@ -98,7 +98,7 @@
         ${_TEST_NAME}
       COMMAND
         "${CMAKE_SOURCE_DIR}/build_tools/cmake/run_android_test.${IREE_HOST_SCRIPT_EXT}"
-        "${_ANDROID_ABS_DIR}/$<TARGET_FILE_NAME:${_TEST_BINARY_TARGET}>"
+        "${_ANDROID_ABS_DIR}/$<TARGET_FILE_NAME:${_SRC_TARGET}>"
         ${_RULE_ARGS}
     )
     # Use environment variables to instruct the script to push artifacts
@@ -108,7 +108,7 @@
     set(
       _ENVIRONMENT_VARS
         "TEST_ANDROID_ABS_DIR=${_ANDROID_ABS_DIR}"
-        "TEST_EXECUTABLE=$<TARGET_FILE:${_TEST_BINARY_TARGET}>"
+        "TEST_EXECUTABLE=$<TARGET_FILE:${_SRC_TARGET}>"
         "TEST_DATA=${_DATA_SPACE_SEPARATED}"
         "TEST_TMPDIR=${_ANDROID_ABS_DIR}/test_tmpdir"
     )
@@ -119,7 +119,7 @@
         ${_TEST_NAME}
       COMMAND
         "${CMAKE_SOURCE_DIR}/build_tools/cmake/run_test.${IREE_HOST_SCRIPT_EXT}"
-        "$<TARGET_FILE:${_TEST_BINARY_TARGET}>"
+        "$<TARGET_FILE:${_SRC_TARGET}>"
         ${_RULE_ARGS}
     )
     set_property(TEST ${_TEST_NAME} PROPERTY ENVIRONMENT "TEST_TMPDIR=${CMAKE_BINARY_DIR}/${_NAME}_test_tmpdir")
diff --git a/build_tools/cmake/iree_trace_runner_test.cmake b/build_tools/cmake/iree_trace_runner_test.cmake
index a2457a4..b32fcdf 100644
--- a/build_tools/cmake/iree_trace_runner_test.cmake
+++ b/build_tools/cmake/iree_trace_runner_test.cmake
@@ -96,12 +96,12 @@
     "${_RULE_TRACE_RUNNER}"
   )
 
-  iree_run_binary_test(
+  iree_native_test(
     NAME
       "${_RULE_NAME}"
     DRIVER
       "${_RULE_DRIVER}"
-    TEST_BINARY
+    SRC
       "${_RULE_TRACE_RUNNER}"
     TEST_INPUT_FILE_ARG
       ${_RULE_TRACE}
diff --git a/docs/developers/developing_iree/testing_guide.md b/docs/developers/developing_iree/testing_guide.md
index e003a7e..8537e46 100644
--- a/docs/developers/developing_iree/testing_guide.md
+++ b/docs/developers/developing_iree/testing_guide.md
@@ -145,12 +145,12 @@
 "lit".
 
 ```bzl
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//iree/build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 
 iree_lit_test_suite(
     name = "lit",
     srcs = glob(["*.mlir"]),
-    data = [
+    tools = [
         "@llvm-project//llvm:FileCheck",
         "//iree/tools:iree-opt",
     ],
diff --git a/integrations/tensorflow/BUILD b/integrations/tensorflow/BUILD
new file mode 100644
index 0000000..9da145a
--- /dev/null
+++ b/integrations/tensorflow/BUILD
@@ -0,0 +1,13 @@
+# 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
+
+package(
+    default_visibility = ["//visibility:public"],
+    features = ["layering_check"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+exports_files(["lit.cfg.py"])
diff --git a/integrations/tensorflow/build_tools/bazel/BUILD b/integrations/tensorflow/build_tools/bazel/BUILD
index a65ce03..f27d209 100644
--- a/integrations/tensorflow/build_tools/bazel/BUILD
+++ b/integrations/tensorflow/build_tools/bazel/BUILD
@@ -9,5 +9,3 @@
     features = ["layering_check"],
     licenses = ["notice"],  # Apache 2.0
 )
-
-exports_files(["run_lit.sh"])
diff --git a/integrations/tensorflow/build_tools/bazel/iree_lit_test.bzl b/integrations/tensorflow/build_tools/bazel/iree_lit_test.bzl
new file mode 100644
index 0000000..952fc78
--- /dev/null
+++ b/integrations/tensorflow/build_tools/bazel/iree_lit_test.bzl
@@ -0,0 +1,97 @@
+# Copyright 2019 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
+
+"""Bazel macros for running lit tests."""
+
+load(":lit_test.bzl", "lit_test", "lit_test_suite")
+
+def iree_lit_test(
+        name,
+        cfg = "//:lit.cfg.py",
+        tools = None,
+        env = None,
+        **kwargs):
+    """A thin wrapper around lit_test with some opinionated settings.
+
+    See the base lit_test for more details on argument meanings.
+
+    Args:
+      name: name for the test.
+      cfg: string. lit config file.
+      tools: label_list. tools that should be included on the PATH.
+        llvm-symbolizer is added by default.
+      env: string_dict. Environment variables available to the test at runtime.
+        FILECHECK_OPTS=--enable-var-scope is added if FILECHECK_OPTS is not
+        already set.
+      **kwargs: additional keyword args to forward to the underyling lit_test.
+    """
+
+    tools = tools or []
+    env = env or {}
+
+    # Always include llvm-symbolizer so we get useful stack traces. Maybe it
+    # would be better to force everyone to do this explicitly, but since
+    # forgetting wouldn't cause the test to fail, only make debugging harder
+    # when it does, I think better to hardcode it here.
+    llvm_symbolizer = "@llvm-project//llvm:llvm-symbolizer"
+    if llvm_symbolizer not in tools:
+        tools.append(llvm_symbolizer)
+
+    filecheck_env_var = "FILECHECK_OPTS"
+    if filecheck_env_var not in env:
+        env[filecheck_env_var] = "--enable-var-scope"
+
+    lit_test(
+        name = name,
+        cfg = cfg,
+        tools = tools,
+        env = env,
+        **kwargs
+    )
+
+def iree_lit_test_suite(
+        name,
+        cfg = "//:lit.cfg.py",
+        tools = None,
+        env = None,
+        **kwargs):
+    """A thin wrapper around lit_test_suite with some opinionated settings.
+
+    See the base lit_test for more details on argument meanings.
+
+    Args:
+      name: name for the test suite.
+      cfg: string. lit config file.
+      tools: label_list. tools that should be included on the PATH.
+        llvm-symbolizer is added by default.
+      env: string_dict. Environment variables available to the test at runtime.
+        FILECHECK_OPTS=--enable-var-scope is added if FILECHECK_OPTS is not
+        already set.
+      **kwargs: additional keyword args to forward to the underyling
+        lit_test_suite.
+    """
+    tools = tools or []
+    env = env or {}
+
+    # Always include llvm-symbolizer so we get useful stack traces. Maybe it
+    # would be better to force everyone to do this explicitly, but since
+    # forgetting wouldn't cause the test to fail, only make debugging harder
+    # when it does, I think better to hardcode it here.
+    llvm_symbolizer = "@llvm-project//llvm:llvm-symbolizer"
+    if llvm_symbolizer not in tools:
+        tools.append(llvm_symbolizer)
+
+    filecheck_env_var = "FILECHECK_OPTS"
+    if filecheck_env_var not in env:
+        env[filecheck_env_var] = "--enable-var-scope"
+
+    lit_test_suite(
+        name = name,
+        cfg = cfg,
+        tools = tools,
+        env = env,
+        **kwargs
+    )
diff --git a/integrations/tensorflow/build_tools/bazel/lit_test.bzl b/integrations/tensorflow/build_tools/bazel/lit_test.bzl
index ba16815..e77d7a0 100644
--- a/integrations/tensorflow/build_tools/bazel/lit_test.bzl
+++ b/integrations/tensorflow/build_tools/bazel/lit_test.bzl
@@ -1,95 +1,192 @@
-# Copyright 2019 The IREE Authors
+# 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
 
-"""Bazel macros for running lit tests."""
+"""Rules for running lit tests with the upstream lit binary."""
 
-def iree_lit_test(
+# This exists as a separate file from iree_lit_test.bzl because we anticipate
+# upstreaming it soon.
+
+load("@bazel_skylib//lib:paths.bzl", "paths")
+load(":native_binary.bzl", "native_test")
+
+def _tools_on_path_impl(ctx):
+    runfiles = ctx.runfiles()
+
+    # For Bazel 4.x support. Drop when Bazel 4.x is no longer supported
+    to_merge = [d[DefaultInfo].default_runfiles for d in ctx.attr.srcs]
+    if hasattr(runfiles, "merge_all"):
+        runfiles = runfiles.merge_all(to_merge)
+    else:
+        for m in to_merge:
+            runfiles = runfiles.merge(m)
+
+    runfiles_symlinks = {}
+
+    for src in ctx.attr.srcs:
+        exe = src[DefaultInfo].files_to_run.executable
+        if not exe:
+            fail("All targets used as tools by lit tests must have exactly one" +
+                 " executable, but {} has none".format(src))
+        bin_path = paths.join(ctx.attr.bin_dir, exe.basename)
+        if bin_path in runfiles_symlinks:
+            fail("All tools used by lit tests must have unique basenames, as" +
+                 " they are added to the path." +
+                 " {} and {} conflict".format(runfiles_symlinks[bin_path], exe))
+        runfiles_symlinks[bin_path] = exe
+
+    return [
+        DefaultInfo(runfiles = ctx.runfiles(
+            symlinks = runfiles_symlinks,
+        ).merge(runfiles)),
+    ]
+
+_tools_on_path = rule(
+    _tools_on_path_impl,
+    attrs = {
+        "srcs": attr.label_list(allow_files = True, mandatory = True),
+        "bin_dir": attr.string(mandatory = True),
+    },
+    doc = "Symlinks srcs into a single lit_bin directory. All basenames must be unique.",
+)
+
+def lit_test(
         name,
         test_file,
-        data,
-        size = "small",
-        driver = "//iree/tools:run_lit.sh",
+        cfg,
+        tools = None,
+        args = None,
+        data = None,
+        visibility = None,
+        env = None,
         **kwargs):
-    """Creates a lit test from the specified source file.
+    """Runs a single test file with LLVM's lit tool.
 
     Args:
-      name: name of the generated test suite.
-      test_file: the test file with the lit test
-      data: binaries used in the lit tests.
-      size: size of the tests.
-      driver: the shell runner for the lit tests.
-      **kwargs: Any additional arguments that will be passed to the underlying sh_test.
+      name: string. the name of the generated test target.
+      test_file: label. The file on which to run lit.
+      cfg: label. The lit config file. It must list the file extension of
+        `test_file` in config.suffixes and must be in a parent directory of
+        `test_file`.
+      tools: label list. Tools invoked in the lit RUN lines. These binaries will
+        be symlinked into a directory which is on the path. They must therefore
+        have unique basenames.
+      args: string list. Additional arguments to pass to lit. Note that the test
+        file, `-v`, and a `--path` argument for the directory to which `tools`
+        are symlinked are added automatically.
+      data: label list. Additional data dependencies of the test. Note that
+        targets in `cfg` and `tools`, as well as their data dependencies, are
+        added automatically.
+      visibility: visibility of the generated test target.
+      env: string_dict. Environment variables available during test execution.
+        See the common Bazel test attribute.
+      **kwargs: additional keyword arguments to pass to all generated rules.
+
+    See https://llvm.org/docs/CommandGuide/lit.html for details on lit
     """
-    data = data if data else []
+    args = args or []
+    data = data or []
+    tools = tools or []
 
-    # Always include llvm-sybmolizer so we get useful stack traces. Maybe it
-    # would be better to force everyone to do this explicitly, but since
-    # forgetting wouldn't cause the test to fail, only make debugging harder
-    # when it does, I think better to hardcode it here.
-    llvm_symbolizer = "@llvm-project//llvm:llvm-symbolizer"
-    if llvm_symbolizer not in data:
-        data.append(llvm_symbolizer)
+    tools_on_path_target_name = "_{}_tools_on_path".format(name)
 
-    # First argument is the test. The rest are tools to add to the path.
-    data = [test_file] + data
+    bin_dir = paths.join(
+        native.package_name(),
+        tools_on_path_target_name,
+        "lit_bin",
+    )
 
-    native.sh_test(
-        name = name,
-        srcs = [driver],
-        size = size,
-        data = data,
-        args = ["$(location {})".format(file) for file in data],
+    _tools_on_path(
+        name = tools_on_path_target_name,
+        testonly = True,
+        srcs = tools,
+        bin_dir = bin_dir,
+        visibility = ["//visibility:private"],
         **kwargs
     )
 
-def iree_lit_test_suite(
+    native_test(
+        name = name,
+        src = "@llvm-project//llvm:lit",
+        # out = name,
+        args = [
+            "-v",
+            "--path",
+            bin_dir,
+            "$(location {})".format(test_file),
+        ] + args,
+        data = [test_file, cfg, tools_on_path_target_name] + data,
+        visibility = visibility,
+        env = env,
+        **kwargs
+    )
+
+def lit_test_suite(
         name,
         srcs,
-        data,
+        cfg,
+        tools = None,
+        args = None,
+        data = None,
+        visibility = None,
         size = "small",
-        driver = "//iree/tools:run_lit.sh",
-        tags = [],
+        env = None,
         **kwargs):
     """Creates one lit test per source file and a test suite that bundles them.
 
     Args:
-      name: name of the generated test suite.
-      data: binaries used in the lit tests.
-      srcs: test file sources.
-      size: size of the tests.
-      driver: the shell runner for the lit tests.
-      tags: tags to apply to the test. Note that as in standard test suites, manual
-            is treated specially and will also apply to the test suite itself.
-      **kwargs: Any additional arguments that will be passed to the underlying tests and test_suite.
+      name: string. the name of the generated test suite.
+      srcs: label_list. The files which contain the lit tests.
+      cfg: label. The lit config file. It must list the file extension of
+        the files in `srcs` in config.suffixes and must be in a parent directory
+        of `srcs`.
+      tools: label list. Tools invoked in the lit RUN lines. These binaries will
+        be symlinked into a directory which is on the path. They must therefore
+        have unique basenames.
+      args: string list. Additional arguments to pass to lit. Note that the test
+        file, `-v`, and a `--path` argument for the directory to which `tools`
+        are symlinked are added automatically.
+      data: label list. Additional data dependencies of the test. Note that
+        targets in `cfg` and `tools`, as well as their data dependencies, are
+        added automatically.
+      visibility: visibility of the generated test targets and test suite.
+      size: string. size of the generated tests.
+      env: string_dict. Environment variables available during test execution.
+        See the common Bazel test attribute.
+      **kwargs: additional keyword arguments to pass to all generated rules.
+
+    See https://llvm.org/docs/CommandGuide/lit.html for details on lit
     """
+    # If there are kwargs that need to be passed to only some of the generated
+    # rules, they should be extracted into separate named arguments.
+
+    args = args or []
+    data = data or []
+    tools = tools or []
+
     tests = []
     for test_file in srcs:
         # It's generally good practice to prefix any generated names with the
-        # macro name, but we're trying to match the style of the names that are
-        # used for LLVM internally.
+        # macro name, but it's also nice to have the test name just match the
+        # file name.
         test_name = "%s.test" % (test_file)
-        iree_lit_test(
+        tests.append(test_name)
+        lit_test(
             name = test_name,
             test_file = test_file,
-            size = size,
+            cfg = cfg,
+            tools = tools,
+            args = args,
             data = data,
-            driver = driver,
-            tags = tags,
+            visibility = visibility,
+            env = env,
             **kwargs
         )
-        tests.append(test_name)
 
     native.test_suite(
         name = name,
         tests = tests,
-        # Note that only the manual tag really has any effect here. Others are
-        # used for test suite filtering, but all tests are passed the same tags.
-        tags = tags,
-        # If there are kwargs that need to be passed here which only apply to
-        # the generated tests and not to test_suite, they should be extracted
-        # into separate named arguments.
         **kwargs
     )
diff --git a/integrations/tensorflow/build_tools/bazel/native_binary.bzl b/integrations/tensorflow/build_tools/bazel/native_binary.bzl
new file mode 100644
index 0000000..5b047e6
--- /dev/null
+++ b/integrations/tensorflow/build_tools/bazel/native_binary.bzl
@@ -0,0 +1,99 @@
+# 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
+
+"""native_binary() and native_test() rule implementations.
+
+Rewritten from the Bazel Skylib version pending several fixes and improvements
+to that rule:
+
+- https://github.com/bazelbuild/bazel-skylib/pull/338
+- https://github.com/bazelbuild/bazel-skylib/pull/339
+- https://github.com/bazelbuild/bazel-skylib/pull/340
+- https://github.com/bazelbuild/bazel-skylib/pull/341
+
+These rules let you wrap a pre-built binary or script in a conventional binary
+and test rule respectively. They fulfill the same goal as sh_binary and sh_test
+do, but they run the wrapped binary directly, instead of through Bash, so they
+don't depend on Bash and work with --shell_exectuable="".
+"""
+
+def _shared_impl(ctx):
+    out = ctx.attr.out
+    if not out:
+        out = ctx.attr.name
+    output = ctx.actions.declare_file(out)
+    ctx.actions.symlink(
+        target_file = ctx.executable.src,
+        output = output,
+        is_executable = True,
+    )
+
+    runfiles = ctx.runfiles(files = ctx.files.data)
+
+    # For Bazel 4.x support. Drop when Bazel 4.x is no longer supported
+    to_merge = ([d[DefaultInfo].default_runfiles for d in ctx.attr.data] +
+                [ctx.attr.src[DefaultInfo].default_runfiles])
+    if hasattr(runfiles, "merge_all"):
+        runfiles = runfiles.merge_all(to_merge)
+    else:
+        for m in to_merge:
+            runfiles = runfiles.merge(m)
+    return DefaultInfo(
+        executable = output,
+        files = depset([output]),
+        runfiles = runfiles,
+    )
+
+def _native_binary_impl(ctx):
+    default_info = _shared_impl(ctx)
+    return [default_info]
+
+def _native_test_impl(ctx):
+    default_info = _shared_impl(ctx)
+    return [default_info, testing.TestEnvironment(ctx.attr.env)]
+
+# We have to manually set "env" on the test rule because the builtin one is only
+# available in native rules. See
+# https://docs.bazel.build/versions/main/be/common-definitions.html#test.env
+# We don't have "env" on native_binary because there is no BinaryEnvironment
+# mirroring TestEnvironment. See https://github.com/bazelbuild/bazel/issues/7364
+_SHARED_ATTRS = {
+    "src": attr.label(
+        executable = True,
+        allow_files = True,
+        mandatory = True,
+        cfg = "target",
+    ),
+    "data": attr.label_list(allow_files = True),
+    # "out" is attr.string instead of attr.output, so that it is select()'able.
+    "out": attr.string(),
+}
+
+native_binary = rule(
+    implementation = _native_binary_impl,
+    attrs = _SHARED_ATTRS,
+    executable = True,
+)
+
+_TEST_ATTRS = {
+    k: v
+    for k, v in _SHARED_ATTRS.items() + [
+        (
+            "env",
+            attr.string_dict(
+                doc = "Mirrors the common env attribute that otherwise is" +
+                      " only available on native rules. See" +
+                      " https://docs.bazel.build/versions/main/be/common-definitions.html#test.env",
+            ),
+        ),
+    ]
+}
+
+native_test = rule(
+    implementation = _native_test_impl,
+    attrs = _TEST_ATTRS,
+    test = True,
+)
diff --git a/integrations/tensorflow/build_tools/bazel/run_lit.sh b/integrations/tensorflow/build_tools/bazel/run_lit.sh
deleted file mode 100755
index f1258e9..0000000
--- a/integrations/tensorflow/build_tools/bazel/run_lit.sh
+++ /dev/null
@@ -1,60 +0,0 @@
-#!/bin/bash
-
-# Copyright 2019 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
-set -e
-set -o pipefail
-
-EXPLICIT_PATH=""
-
-# First argument is the src file. Remaining arguments are tools that should
-# be on the path.
-src_file="$1"
-shift
-for tool_exe in "$@"
-do
-  EXEDIR="$(dirname $tool_exe)"
-  if ! [ -z "$cygpath" ]; then
-    EXEDIR="$($cygpath -u "$EXEDIR")"
-  fi
-  EXPLICIT_PATH="${EXEDIR}:${EXPLICIT_PATH}"
-done
-
-echo "run_lit.sh: $src_file"
-echo "PWD=$(pwd)"
-echo "EXPLICIT_PATH=$EXPLICIT_PATH"
-
-# For each "// RUN:" line, run the command.
-runline_matches="$(egrep "^// RUN: " "$src_file")"
-if [ -z "$runline_matches" ]; then
-  echo "!!! No RUN lines found in test"
-  exit 1
-fi
-
-echo "$runline_matches" | while read -r runline
-do
-  echo "RUNLINE: $runline"
-  match="${runline%%// RUN: *}"
-  command="${runline##// RUN: }"
-  if [ -z "${command}" ]; then
-    echo "ERROR: Could not extract command from runline"
-    exit 1
-  fi
-
-  # Substitute any embedded '%s' with the file name.
-  full_command="${command//\%s/$src_file}"
-
-  # Run it.
-  export PATH="$EXPLICIT_PATH:$PATH"
-  echo "RUNNING TEST: $full_command"
-  echo "----------------"
-  if eval "$full_command"; then
-    echo "--- COMPLETE ---"
-  else
-    echo "!!! ERROR EVALUATING: $full_command"
-    exit 1
-  fi
-done
diff --git a/integrations/tensorflow/iree_tf_compiler/MHLO/test/BUILD b/integrations/tensorflow/iree_tf_compiler/MHLO/test/BUILD
index bbf6129..b34d04b 100644
--- a/integrations/tensorflow/iree_tf_compiler/MHLO/test/BUILD
+++ b/integrations/tensorflow/iree_tf_compiler/MHLO/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//build_tools/bazel:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -21,9 +21,8 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree_tf_compiler:iree-tf-opt",
         "@llvm-project//llvm:FileCheck",
     ],
-    driver = "//build_tools/bazel:run_lit.sh",
 )
diff --git a/integrations/tensorflow/iree_tf_compiler/TF/test/BUILD b/integrations/tensorflow/iree_tf_compiler/TF/test/BUILD
index a9bc5fa..bc7912a 100644
--- a/integrations/tensorflow/iree_tf_compiler/TF/test/BUILD
+++ b/integrations/tensorflow/iree_tf_compiler/TF/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//build_tools/bazel:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -30,9 +30,8 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree_tf_compiler:iree-tf-opt",
         "@llvm-project//llvm:FileCheck",
     ],
-    driver = "//build_tools/bazel:run_lit.sh",
 )
diff --git a/integrations/tensorflow/iree_tf_compiler/TFL/test/BUILD b/integrations/tensorflow/iree_tf_compiler/TFL/test/BUILD
index 556aa61..2bb735d 100644
--- a/integrations/tensorflow/iree_tf_compiler/TFL/test/BUILD
+++ b/integrations/tensorflow/iree_tf_compiler/TFL/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//build_tools/bazel:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -26,9 +26,8 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree_tf_compiler:iree-opt-tflite",
         "@llvm-project//llvm:FileCheck",
     ],
-    driver = "//build_tools/bazel:run_lit.sh",
 )
diff --git a/integrations/tensorflow/iree_tf_compiler/TFL/test/import/BUILD b/integrations/tensorflow/iree_tf_compiler/TFL/test/import/BUILD
index f63456f..65a378c 100644
--- a/integrations/tensorflow/iree_tf_compiler/TFL/test/import/BUILD
+++ b/integrations/tensorflow/iree_tf_compiler/TFL/test/import/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//build_tools/bazel:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -22,15 +22,15 @@
         ],
         include = ["*.mlir"],
     ),
-    data = glob(["*.tflite"]) + [
-        "//iree_tf_compiler:iree-import-tflite",
-        "@llvm-project//llvm:FileCheck",
-    ],
-    driver = "//build_tools/bazel:run_lit.sh",
+    data = glob(["*.tflite"]),
     # TODO: These tests have never passed
     tags = [
         "manual",
         "nokokoro",
         "notap",
     ],
+    tools = [
+        "//iree_tf_compiler:iree-import-tflite",
+        "@llvm-project//llvm:FileCheck",
+    ],
 )
diff --git a/integrations/tensorflow/lit.cfg.py b/integrations/tensorflow/lit.cfg.py
new file mode 100644
index 0000000..ae56105
--- /dev/null
+++ b/integrations/tensorflow/lit.cfg.py
@@ -0,0 +1,25 @@
+# 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
+"""Lit config for TensorFlow Integrations."""
+
+# Lint for undefined variables is disabled as config is not defined inside this
+# file, instead config is injected by way of evaluating runlit.cfg.py from
+# runlit.site.cfg.py which in turn is evaluated by lit.py.
+# pylint: disable=undefined-variable
+
+import os
+import tempfile
+
+import lit.formats
+
+config.name = "IREE TensorFlow Integrations"
+config.suffixes = [".mlir", ".txt"]
+config.test_format = lit.formats.ShTest(execute_external=True)
+
+# Use the most preferred temp directory.
+config.test_exec_root = (os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR") or
+                         os.environ.get("TEST_TMPDIR") or
+                         os.path.join(tempfile.gettempdir(), "lit"))
diff --git a/iree/BUILD.bazel b/iree/BUILD.bazel
index 54d8c62..1ce917f 100644
--- a/iree/BUILD.bazel
+++ b/iree/BUILD.bazel
@@ -12,6 +12,8 @@
     licenses = ["notice"],  # Apache 2.0
 )
 
+exports_files(["lit.cfg.py"])
+
 # Enables the debug service and other profiling features.
 # $ bazel build --define=IREE_DEBUG=1 :some_target
 config_setting(
diff --git a/iree/base/internal/BUILD b/iree/base/internal/BUILD
index 8504ff5..30fb9e3 100644
--- a/iree/base/internal/BUILD
+++ b/iree/base/internal/BUILD
@@ -8,9 +8,9 @@
 # These are not part of the IREE API. Though they may be used by external
 # projects their API may change at any time.
 
-load("//build_tools/bazel:run_binary_test.bzl", "run_binary_test")
+load("//build_tools/bazel:native_binary.bzl", "native_test")
 load("//iree:build_defs.oss.bzl", "iree_cmake_extra_content")
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -191,11 +191,11 @@
 iree_lit_test_suite(
     name = "flags_test",
     srcs = ["flags_test.txt"],
-    data = [
+    tags = ["hostonly"],
+    tools = [
         ":flags_demo",
         "@llvm-project//llvm:FileCheck",
     ],
-    tags = ["hostonly"],
 )
 
 cc_library(
@@ -294,10 +294,10 @@
     ],
 )
 
-run_binary_test(
+native_test(
     name = "synchronization_benchmark_test",
+    src = ":synchronization_benchmark",
     args = ["--benchmark_min_time=0"],
-    test_binary = ":synchronization_benchmark",
 )
 
 cc_test(
diff --git a/iree/base/internal/CMakeLists.txt b/iree/base/internal/CMakeLists.txt
index b7888af..65119e5 100644
--- a/iree/base/internal/CMakeLists.txt
+++ b/iree/base/internal/CMakeLists.txt
@@ -193,7 +193,7 @@
     flags_test
   SRCS
     "flags_test.txt"
-  DATA
+  TOOLS
     ::flags_demo
     FileCheck
   LABELS
@@ -308,12 +308,12 @@
   TESTONLY
 )
 
-iree_run_binary_test(
+iree_native_test(
   NAME
     "synchronization_benchmark_test"
   ARGS
     "--benchmark_min_time=0"
-  TEST_BINARY
+  SRC
     ::synchronization_benchmark
 )
 
diff --git a/iree/compiler/Bindings/Native/Transforms/test/BUILD b/iree/compiler/Bindings/Native/Transforms/test/BUILD
index d32c0a3..a043558 100644
--- a/iree/compiler/Bindings/Native/Transforms/test/BUILD
+++ b/iree/compiler/Bindings/Native/Transforms/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -21,7 +21,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Bindings/Native/Transforms/test/CMakeLists.txt b/iree/compiler/Bindings/Native/Transforms/test/CMakeLists.txt
index 19008de..9479fbc 100644
--- a/iree/compiler/Bindings/Native/Transforms/test/CMakeLists.txt
+++ b/iree/compiler/Bindings/Native/Transforms/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "wrap_entry_points.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Bindings/TFLite/Transforms/test/BUILD b/iree/compiler/Bindings/TFLite/Transforms/test/BUILD
index d5e27f8..ea1a218 100644
--- a/iree/compiler/Bindings/TFLite/Transforms/test/BUILD
+++ b/iree/compiler/Bindings/TFLite/Transforms/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -21,7 +21,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Bindings/TFLite/Transforms/test/CMakeLists.txt b/iree/compiler/Bindings/TFLite/Transforms/test/CMakeLists.txt
index a787975..7de2644 100644
--- a/iree/compiler/Bindings/TFLite/Transforms/test/CMakeLists.txt
+++ b/iree/compiler/Bindings/TFLite/Transforms/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "wrap_entry_points.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Codegen/Common/test/BUILD b/iree/compiler/Codegen/Common/test/BUILD
index b54dc92..3c47ed9 100644
--- a/iree/compiler/Codegen/Common/test/BUILD
+++ b/iree/compiler/Codegen/Common/test/BUILD
@@ -6,7 +6,7 @@
 
 # Tests for common transforms.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -38,7 +38,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Codegen/Common/test/CMakeLists.txt b/iree/compiler/Codegen/Common/test/CMakeLists.txt
index a2e7310..dd6ad5b 100644
--- a/iree/compiler/Codegen/Common/test/CMakeLists.txt
+++ b/iree/compiler/Codegen/Common/test/CMakeLists.txt
@@ -30,7 +30,7 @@
     "transpose_canonicalization.mlir"
     "vectorize_linalg_conv.mlir"
     "vectorize_linalg_mmt4d.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Codegen/Dialect/test/BUILD b/iree/compiler/Codegen/Dialect/test/BUILD
index 15797ef..22b576f 100644
--- a/iree/compiler/Codegen/Dialect/test/BUILD
+++ b/iree/compiler/Codegen/Dialect/test/BUILD
@@ -6,7 +6,7 @@
 
 # Tests for common transforms.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -23,7 +23,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Codegen/Dialect/test/CMakeLists.txt b/iree/compiler/Codegen/Dialect/test/CMakeLists.txt
index c260132..976a43b 100644
--- a/iree/compiler/Codegen/Dialect/test/CMakeLists.txt
+++ b/iree/compiler/Codegen/Dialect/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "lowering_config_attr.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Codegen/LLVMCPU/test/BUILD b/iree/compiler/Codegen/LLVMCPU/test/BUILD
index 511c649..d390ba2 100644
--- a/iree/compiler/Codegen/LLVMCPU/test/BUILD
+++ b/iree/compiler/Codegen/LLVMCPU/test/BUILD
@@ -6,7 +6,7 @@
 
 # Tests for common transforms.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -35,7 +35,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt b/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
index 742023d..cfd1235 100644
--- a/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
+++ b/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
@@ -27,7 +27,7 @@
     "tile_fuse_and_vectorize.mlir"
     "unfused_fma.mlir"
     "vector_contract_custom_kernels.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Codegen/LLVMGPU/test/BUILD b/iree/compiler/Codegen/LLVMGPU/test/BUILD
index cbe8cad..47b34ee 100644
--- a/iree/compiler/Codegen/LLVMGPU/test/BUILD
+++ b/iree/compiler/Codegen/LLVMGPU/test/BUILD
@@ -6,7 +6,7 @@
 
 # Tests for common transforms.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -32,7 +32,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt b/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt
index 421f3ae..02074dd 100644
--- a/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt
+++ b/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt
@@ -24,7 +24,7 @@
     "rocdl_pipeline_test.mlir"
     "tensorcore_vectorization.mlir"
     "vectorization.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Codegen/SPIRV/test/BUILD b/iree/compiler/Codegen/SPIRV/test/BUILD
index f9cbecc..1eda09a 100644
--- a/iree/compiler/Codegen/SPIRV/test/BUILD
+++ b/iree/compiler/Codegen/SPIRV/test/BUILD
@@ -6,7 +6,7 @@
 
 # Tests for common transforms.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -51,7 +51,7 @@
             "promote_workgroup_memory.mlir",
         ],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt b/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt
index d139690..4dcce06 100644
--- a/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt
+++ b/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt
@@ -39,7 +39,7 @@
     "vectorize_elementwise_ops.mlir"
     "vectorize_load_store.mlir"
     "vectorize_matmul.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Codegen/Sandbox/test/BUILD b/iree/compiler/Codegen/Sandbox/test/BUILD
index 728dfd3..3340d08 100644
--- a/iree/compiler/Codegen/Sandbox/test/BUILD
+++ b/iree/compiler/Codegen/Sandbox/test/BUILD
@@ -6,7 +6,7 @@
 
 # Tests for common transforms.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -24,7 +24,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Codegen/Sandbox/test/CMakeLists.txt b/iree/compiler/Codegen/Sandbox/test/CMakeLists.txt
index 35efc1e..ae4b2eb 100644
--- a/iree/compiler/Codegen/Sandbox/test/CMakeLists.txt
+++ b/iree/compiler/Codegen/Sandbox/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "fusion_expert.mlir"
     "single_tiling_expert.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/ConstEval/test/BUILD b/iree/compiler/ConstEval/test/BUILD
index d82d8cf..4b7046b 100644
--- a/iree/compiler/ConstEval/test/BUILD
+++ b/iree/compiler/ConstEval/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -21,7 +21,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/ConstEval/test/CMakeLists.txt b/iree/compiler/ConstEval/test/CMakeLists.txt
index 1ec22a2..3fe0d66 100644
--- a/iree/compiler/ConstEval/test/CMakeLists.txt
+++ b/iree/compiler/ConstEval/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "jit_globals.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/BUILD b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/BUILD
index 0f862cf..6e2c4bf 100644
--- a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/BUILD
+++ b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -25,7 +25,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/CMakeLists.txt b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/CMakeLists.txt
index 3a875b4..c4cd83e 100644
--- a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/CMakeLists.txt
@@ -19,7 +19,7 @@
     "extract_slice.mlir"
     "from_elements.mlir"
     "insert_slice.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Flow/IR/test/BUILD b/iree/compiler/Dialect/Flow/IR/test/BUILD
index 1d152ed..0cb0f17 100644
--- a/iree/compiler/Dialect/Flow/IR/test/BUILD
+++ b/iree/compiler/Dialect/Flow/IR/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -28,7 +28,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Flow/IR/test/CMakeLists.txt b/iree/compiler/Dialect/Flow/IR/test/CMakeLists.txt
index 0a9ef9b..9cbe808 100644
--- a/iree/compiler/Dialect/Flow/IR/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/IR/test/CMakeLists.txt
@@ -22,7 +22,7 @@
     "tensor_folding.mlir"
     "tensor_ops.mlir"
     "types.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/BUILD b/iree/compiler/Dialect/Flow/Transforms/test/BUILD
index 8124541..7405dbb 100644
--- a/iree/compiler/Dialect/Flow/Transforms/test/BUILD
+++ b/iree/compiler/Dialect/Flow/Transforms/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -45,7 +45,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
index 656cd6d..ef0811b 100644
--- a/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
@@ -39,7 +39,7 @@
     "test_partitionable_loops_interface.mlir"
     "transformation.mlir"
     "verify_input_ir.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/BUILD b/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/BUILD
index ab4455e..cadeba5 100644
--- a/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -26,7 +26,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/CMakeLists.txt
index 3691baa..d2f2948 100644
--- a/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/CMakeLists.txt
@@ -20,7 +20,7 @@
     "command_buffer_ops.mlir"
     "device_ops.mlir"
     "executable_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/Conversion/StandardToHAL/test/BUILD b/iree/compiler/Dialect/HAL/Conversion/StandardToHAL/test/BUILD
index 3739310..98830ed 100644
--- a/iree/compiler/Dialect/HAL/Conversion/StandardToHAL/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Conversion/StandardToHAL/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -22,7 +22,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/Conversion/StandardToHAL/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Conversion/StandardToHAL/test/CMakeLists.txt
index 20087a3..cb81022 100644
--- a/iree/compiler/Dialect/HAL/Conversion/StandardToHAL/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Conversion/StandardToHAL/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "shape_ops.mlir"
     "structural_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD b/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD
index 8679f7f..d6b4b9f 100644
--- a/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -24,7 +24,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt
index 1904dea..c13e868 100644
--- a/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt
@@ -18,7 +18,7 @@
     "resource_ops.mlir"
     "timepoint_ops.mlir"
     "transfer_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/Conversion/UtilToHAL/test/BUILD b/iree/compiler/Dialect/HAL/Conversion/UtilToHAL/test/BUILD
index b998250..772543f 100644
--- a/iree/compiler/Dialect/HAL/Conversion/UtilToHAL/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Conversion/UtilToHAL/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -19,7 +19,7 @@
         ["global_ops.mlir"],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/Conversion/UtilToHAL/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Conversion/UtilToHAL/test/CMakeLists.txt
index e7835c8..2d95202 100644
--- a/iree/compiler/Dialect/HAL/Conversion/UtilToHAL/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Conversion/UtilToHAL/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "global_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/IR/test/BUILD b/iree/compiler/Dialect/HAL/IR/test/BUILD
index c85eaf1..d545a2a 100644
--- a/iree/compiler/Dialect/HAL/IR/test/BUILD
+++ b/iree/compiler/Dialect/HAL/IR/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -37,7 +37,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/IR/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/IR/test/CMakeLists.txt
index 8da73c6..cf00bf9 100644
--- a/iree/compiler/Dialect/HAL/IR/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/IR/test/CMakeLists.txt
@@ -31,7 +31,7 @@
     "semaphore_ops.mlir"
     "tensor_op_folding.mlir"
     "tensor_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/Target/CUDA/test/BUILD b/iree/compiler/Dialect/HAL/Target/CUDA/test/BUILD
index 244d2a4..43a44f4 100644
--- a/iree/compiler/Dialect/HAL/Target/CUDA/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Target/CUDA/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -19,7 +19,7 @@
         ["smoketest.mlir"],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/Target/CUDA/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/CUDA/test/CMakeLists.txt
index c9d1760..fa0a3a8 100644
--- a/iree/compiler/Dialect/HAL/Target/CUDA/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Target/CUDA/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "smoketest.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/test/BUILD b/iree/compiler/Dialect/HAL/Target/LLVM/test/BUILD
index 57dc522..3c27a43 100644
--- a/iree/compiler/Dialect/HAL/Target/LLVM/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Target/LLVM/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -21,7 +21,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//lld",
         "@llvm-project//llvm:FileCheck",
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/LLVM/test/CMakeLists.txt
index 57078b4..896553f 100644
--- a/iree/compiler/Dialect/HAL/Target/LLVM/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Target/LLVM/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "smoketest.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
     lld
diff --git a/iree/compiler/Dialect/HAL/Target/MetalSPIRV/test/BUILD b/iree/compiler/Dialect/HAL/Target/MetalSPIRV/test/BUILD
index 244d2a4..43a44f4 100644
--- a/iree/compiler/Dialect/HAL/Target/MetalSPIRV/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Target/MetalSPIRV/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -19,7 +19,7 @@
         ["smoketest.mlir"],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/Target/MetalSPIRV/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/MetalSPIRV/test/CMakeLists.txt
index 680c1e0..1cfab85 100644
--- a/iree/compiler/Dialect/HAL/Target/MetalSPIRV/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Target/MetalSPIRV/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "smoketest.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/Target/ROCM/test/BUILD b/iree/compiler/Dialect/HAL/Target/ROCM/test/BUILD
index 244d2a4..43a44f4 100644
--- a/iree/compiler/Dialect/HAL/Target/ROCM/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Target/ROCM/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -19,7 +19,7 @@
         ["smoketest.mlir"],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/Target/ROCM/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/ROCM/test/CMakeLists.txt
index a1ab04b..0870543 100644
--- a/iree/compiler/Dialect/HAL/Target/ROCM/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Target/ROCM/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "smoketest.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/Target/VMVX/test/BUILD b/iree/compiler/Dialect/HAL/Target/VMVX/test/BUILD
index 42f730d..d46c759 100644
--- a/iree/compiler/Dialect/HAL/Target/VMVX/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Target/VMVX/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -22,7 +22,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/Target/VMVX/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/VMVX/test/CMakeLists.txt
index d1f3280..4f8031d 100644
--- a/iree/compiler/Dialect/HAL/Target/VMVX/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Target/VMVX/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "linking.mlir"
     "smoketest.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/test/BUILD b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/test/BUILD
index 9df6a80..134cd15 100644
--- a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -22,7 +22,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/test/CMakeLists.txt
index 0501957..c89053d 100644
--- a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "linking.mlir"
     "smoketest.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/HAL/Transforms/test/BUILD b/iree/compiler/Dialect/HAL/Transforms/test/BUILD
index 173486d..ab04922 100644
--- a/iree/compiler/Dialect/HAL/Transforms/test/BUILD
+++ b/iree/compiler/Dialect/HAL/Transforms/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -31,7 +31,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt b/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
index ad6377b..2099a50 100644
--- a/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
@@ -25,7 +25,7 @@
     "pack_dispatch_operands.mlir"
     "resolve_entry_point_ordinals.mlir"
     "verify_target_environment.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Modules/Check/test/BUILD b/iree/compiler/Dialect/Modules/Check/test/BUILD
index db8964b..370f60a 100644
--- a/iree/compiler/Dialect/Modules/Check/test/BUILD
+++ b/iree/compiler/Dialect/Modules/Check/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -22,7 +22,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Modules/Check/test/CMakeLists.txt b/iree/compiler/Dialect/Modules/Check/test/CMakeLists.txt
index 5d41a83..ff3a7e2 100644
--- a/iree/compiler/Dialect/Modules/Check/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Modules/Check/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "canonicalize.mlir"
     "ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Modules/VMVX/Conversion/HALToVMVX/test/BUILD b/iree/compiler/Dialect/Modules/VMVX/Conversion/HALToVMVX/test/BUILD
index f619f75..987f2e5 100644
--- a/iree/compiler/Dialect/Modules/VMVX/Conversion/HALToVMVX/test/BUILD
+++ b/iree/compiler/Dialect/Modules/VMVX/Conversion/HALToVMVX/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -19,7 +19,7 @@
         ["interface_ops.mlir"],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Modules/VMVX/Conversion/HALToVMVX/test/CMakeLists.txt b/iree/compiler/Dialect/Modules/VMVX/Conversion/HALToVMVX/test/CMakeLists.txt
index e0500b1..a19685f 100644
--- a/iree/compiler/Dialect/Modules/VMVX/Conversion/HALToVMVX/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Modules/VMVX/Conversion/HALToVMVX/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "interface_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Modules/VMVX/Conversion/StandardToVMVX/test/BUILD b/iree/compiler/Dialect/Modules/VMVX/Conversion/StandardToVMVX/test/BUILD
index ad28cc1..1f317c9 100644
--- a/iree/compiler/Dialect/Modules/VMVX/Conversion/StandardToVMVX/test/BUILD
+++ b/iree/compiler/Dialect/Modules/VMVX/Conversion/StandardToVMVX/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -20,7 +20,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Modules/VMVX/Conversion/StandardToVMVX/test/CMakeLists.txt b/iree/compiler/Dialect/Modules/VMVX/Conversion/StandardToVMVX/test/CMakeLists.txt
index dab87e3..958d5cb 100644
--- a/iree/compiler/Dialect/Modules/VMVX/Conversion/StandardToVMVX/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Modules/VMVX/Conversion/StandardToVMVX/test/CMakeLists.txt
@@ -13,7 +13,7 @@
 iree_lit_test_suite(
   NAME
     lit
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Modules/VMVX/Conversion/VMVXToVM/test/BUILD b/iree/compiler/Dialect/Modules/VMVX/Conversion/VMVXToVM/test/BUILD
index ad28cc1..1f317c9 100644
--- a/iree/compiler/Dialect/Modules/VMVX/Conversion/VMVXToVM/test/BUILD
+++ b/iree/compiler/Dialect/Modules/VMVX/Conversion/VMVXToVM/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -20,7 +20,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Modules/VMVX/Conversion/VMVXToVM/test/CMakeLists.txt b/iree/compiler/Dialect/Modules/VMVX/Conversion/VMVXToVM/test/CMakeLists.txt
index ce8feeb..2330a58 100644
--- a/iree/compiler/Dialect/Modules/VMVX/Conversion/VMVXToVM/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Modules/VMVX/Conversion/VMVXToVM/test/CMakeLists.txt
@@ -13,7 +13,7 @@
 iree_lit_test_suite(
   NAME
     lit
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Modules/VMVX/IR/test/BUILD b/iree/compiler/Dialect/Modules/VMVX/IR/test/BUILD
index ad28cc1..1f317c9 100644
--- a/iree/compiler/Dialect/Modules/VMVX/IR/test/BUILD
+++ b/iree/compiler/Dialect/Modules/VMVX/IR/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -20,7 +20,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Modules/VMVX/IR/test/CMakeLists.txt b/iree/compiler/Dialect/Modules/VMVX/IR/test/CMakeLists.txt
index 6fd9b10..7018dcf 100644
--- a/iree/compiler/Dialect/Modules/VMVX/IR/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Modules/VMVX/IR/test/CMakeLists.txt
@@ -13,7 +13,7 @@
 iree_lit_test_suite(
   NAME
     lit
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Modules/VMVX/Transforms/test/BUILD b/iree/compiler/Dialect/Modules/VMVX/Transforms/test/BUILD
index ad28cc1..1f317c9 100644
--- a/iree/compiler/Dialect/Modules/VMVX/Transforms/test/BUILD
+++ b/iree/compiler/Dialect/Modules/VMVX/Transforms/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -20,7 +20,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Modules/VMVX/Transforms/test/CMakeLists.txt b/iree/compiler/Dialect/Modules/VMVX/Transforms/test/CMakeLists.txt
index 6ec9252..67c465b 100644
--- a/iree/compiler/Dialect/Modules/VMVX/Transforms/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Modules/VMVX/Transforms/test/CMakeLists.txt
@@ -13,7 +13,7 @@
 iree_lit_test_suite(
   NAME
     lit
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD
index bc7206a..fd6f12b 100644
--- a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -23,7 +23,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt
index b571ee6..3e65a59 100644
--- a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt
@@ -17,7 +17,7 @@
     "dispatch_ops.mlir"
     "executable_ops.mlir"
     "tensor_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/BUILD b/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/BUILD
index dadfda2..bd86a3f 100644
--- a/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/BUILD
+++ b/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -21,7 +21,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/CMakeLists.txt
index 18b8e34..e33b2c8 100644
--- a/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "abi_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/BUILD b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/BUILD
index 6b5818d..72cf45f 100644
--- a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/BUILD
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -22,7 +22,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/CMakeLists.txt
index ecbfa7d..ea18a35 100644
--- a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "constant_ops.mlir"
     "structural_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/BUILD b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/BUILD
index b998250..772543f 100644
--- a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/BUILD
+++ b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -19,7 +19,7 @@
         ["global_ops.mlir"],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/CMakeLists.txt
index db36259..672edea 100644
--- a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "global_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Stream/IR/test/BUILD b/iree/compiler/Dialect/Stream/IR/test/BUILD
index 4a928f7..f8e9232 100644
--- a/iree/compiler/Dialect/Stream/IR/test/BUILD
+++ b/iree/compiler/Dialect/Stream/IR/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -32,7 +32,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Stream/IR/test/CMakeLists.txt b/iree/compiler/Dialect/Stream/IR/test/CMakeLists.txt
index f97bfff..c12baf4 100644
--- a/iree/compiler/Dialect/Stream/IR/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Stream/IR/test/CMakeLists.txt
@@ -26,7 +26,7 @@
     "tensor_ops.mlir"
     "timepoint_folding.mlir"
     "timepoint_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Stream/Transforms/test/BUILD b/iree/compiler/Dialect/Stream/Transforms/test/BUILD
index 8bebb93..3856619 100644
--- a/iree/compiler/Dialect/Stream/Transforms/test/BUILD
+++ b/iree/compiler/Dialect/Stream/Transforms/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -42,7 +42,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt b/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
index 6cc8a93..f3ad456 100644
--- a/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
@@ -36,7 +36,7 @@
     "schedule_concurrency.mlir"
     "schedule_execution.mlir"
     "specialize_dispatches.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Util/Conversion/test/BUILD b/iree/compiler/Dialect/Util/Conversion/test/BUILD
index c5f6f6e..6f12c15 100644
--- a/iree/compiler/Dialect/Util/Conversion/test/BUILD
+++ b/iree/compiler/Dialect/Util/Conversion/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -19,7 +19,7 @@
         ["hint_ops.mlir"],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Util/Conversion/test/CMakeLists.txt b/iree/compiler/Dialect/Util/Conversion/test/CMakeLists.txt
index 49b3d0d..630c3c8 100644
--- a/iree/compiler/Dialect/Util/Conversion/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Util/Conversion/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "hint_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Util/IR/test/BUILD b/iree/compiler/Dialect/Util/IR/test/BUILD
index 8c74d7e..80b81c0 100644
--- a/iree/compiler/Dialect/Util/IR/test/BUILD
+++ b/iree/compiler/Dialect/Util/IR/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -34,7 +34,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt b/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt
index b88c8cd..3030ba4 100644
--- a/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt
@@ -28,7 +28,7 @@
     "range_ops.mlir"
     "structural_folding.mlir"
     "structural_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Util/Transforms/test/BUILD b/iree/compiler/Dialect/Util/Transforms/test/BUILD
index 3465ae7..072fa38 100644
--- a/iree/compiler/Dialect/Util/Transforms/test/BUILD
+++ b/iree/compiler/Dialect/Util/Transforms/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -31,7 +31,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Util/Transforms/test/CMakeLists.txt b/iree/compiler/Dialect/Util/Transforms/test/CMakeLists.txt
index 9d467e6..edfdbaf 100644
--- a/iree/compiler/Dialect/Util/Transforms/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Util/Transforms/test/CMakeLists.txt
@@ -25,7 +25,7 @@
     "strip_debug_ops.mlir"
     "test_float_range_analysis.mlir"
     "test_float_range_analysis_linalg.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/VM/Analysis/test/BUILD b/iree/compiler/Dialect/VM/Analysis/test/BUILD
index ac96eca..173bea6 100644
--- a/iree/compiler/Dialect/VM/Analysis/test/BUILD
+++ b/iree/compiler/Dialect/VM/Analysis/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -22,7 +22,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/VM/Analysis/test/CMakeLists.txt b/iree/compiler/Dialect/VM/Analysis/test/CMakeLists.txt
index 9e51df6..a354708 100644
--- a/iree/compiler/Dialect/VM/Analysis/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/VM/Analysis/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "register_allocation.mlir"
     "value_liveness.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/VM/Conversion/MathToVM/test/BUILD b/iree/compiler/Dialect/VM/Conversion/MathToVM/test/BUILD
index 051ab77..264880a 100644
--- a/iree/compiler/Dialect/VM/Conversion/MathToVM/test/BUILD
+++ b/iree/compiler/Dialect/VM/Conversion/MathToVM/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -21,7 +21,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/VM/Conversion/MathToVM/test/CMakeLists.txt b/iree/compiler/Dialect/VM/Conversion/MathToVM/test/CMakeLists.txt
index 60efeca..52e1cb2 100644
--- a/iree/compiler/Dialect/VM/Conversion/MathToVM/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/VM/Conversion/MathToVM/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "arithmetic_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/VM/Conversion/MemRefToVM/test/BUILD b/iree/compiler/Dialect/VM/Conversion/MemRefToVM/test/BUILD
index 80c4a814..dd16226 100644
--- a/iree/compiler/Dialect/VM/Conversion/MemRefToVM/test/BUILD
+++ b/iree/compiler/Dialect/VM/Conversion/MemRefToVM/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -21,7 +21,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/VM/Conversion/MemRefToVM/test/CMakeLists.txt b/iree/compiler/Dialect/VM/Conversion/MemRefToVM/test/CMakeLists.txt
index d0c145b..2278418 100644
--- a/iree/compiler/Dialect/VM/Conversion/MemRefToVM/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/VM/Conversion/MemRefToVM/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "load_store_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/BUILD b/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/BUILD
index 9523f40..9674a4b 100644
--- a/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/BUILD
+++ b/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -29,7 +29,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/CMakeLists.txt b/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/CMakeLists.txt
index d24ca30..233a197 100644
--- a/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/CMakeLists.txt
@@ -23,7 +23,7 @@
     "func_attrs.mlir"
     "nesting.mlir"
     "structural_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/VM/Conversion/UtilToVM/test/BUILD b/iree/compiler/Dialect/VM/Conversion/UtilToVM/test/BUILD
index 4c7f5c6..609e15f 100644
--- a/iree/compiler/Dialect/VM/Conversion/UtilToVM/test/BUILD
+++ b/iree/compiler/Dialect/VM/Conversion/UtilToVM/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -26,7 +26,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/VM/Conversion/UtilToVM/test/CMakeLists.txt b/iree/compiler/Dialect/VM/Conversion/UtilToVM/test/CMakeLists.txt
index 781efa8..6e5fe78 100644
--- a/iree/compiler/Dialect/VM/Conversion/UtilToVM/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/VM/Conversion/UtilToVM/test/CMakeLists.txt
@@ -20,7 +20,7 @@
     "hint_ops.mlir"
     "list_ops.mlir"
     "status_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/VM/Conversion/VMToEmitC/test/CMakeLists.txt b/iree/compiler/Dialect/VM/Conversion/VMToEmitC/test/CMakeLists.txt
index f4c119e..883a09b 100644
--- a/iree/compiler/Dialect/VM/Conversion/VMToEmitC/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/VM/Conversion/VMToEmitC/test/CMakeLists.txt
@@ -28,7 +28,7 @@
     "shift_ops.mlir"
     "shift_ops_i64.mlir"
     "type_conversion.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/VM/IR/test/BUILD b/iree/compiler/Dialect/VM/IR/test/BUILD
index 9b25695..de403d8 100644
--- a/iree/compiler/Dialect/VM/IR/test/BUILD
+++ b/iree/compiler/Dialect/VM/IR/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -41,7 +41,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/VM/IR/test/CMakeLists.txt b/iree/compiler/Dialect/VM/IR/test/CMakeLists.txt
index a4ecd2c..80b04ee 100644
--- a/iree/compiler/Dialect/VM/IR/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/VM/IR/test/CMakeLists.txt
@@ -35,7 +35,7 @@
     "shift_ops.mlir"
     "structural_folding.mlir"
     "structural_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/test/BUILD b/iree/compiler/Dialect/VM/Target/Bytecode/test/BUILD
index c4d96bd..85f1cc1 100644
--- a/iree/compiler/Dialect/VM/Target/Bytecode/test/BUILD
+++ b/iree/compiler/Dialect/VM/Target/Bytecode/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -23,7 +23,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-translate",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/test/CMakeLists.txt b/iree/compiler/Dialect/VM/Target/Bytecode/test/CMakeLists.txt
index 6e17248..cf69b16 100644
--- a/iree/compiler/Dialect/VM/Target/Bytecode/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/VM/Target/Bytecode/test/CMakeLists.txt
@@ -17,7 +17,7 @@
     "constant_encoding.mlir"
     "module_encoding_smoke.mlir"
     "reflection_attrs.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-translate
 )
diff --git a/iree/compiler/Dialect/VM/Target/C/test/CMakeLists.txt b/iree/compiler/Dialect/VM/Target/C/test/CMakeLists.txt
index 9b6f867..d6263eb 100644
--- a/iree/compiler/Dialect/VM/Target/C/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/VM/Target/C/test/CMakeLists.txt
@@ -12,7 +12,7 @@
     lit
   SRCS
     "${_GLOB_X_MLIR}"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-translate
 )
diff --git a/iree/compiler/Dialect/VM/Transforms/test/BUILD b/iree/compiler/Dialect/VM/Transforms/test/BUILD
index b928d63..045bb78 100644
--- a/iree/compiler/Dialect/VM/Transforms/test/BUILD
+++ b/iree/compiler/Dialect/VM/Transforms/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -25,7 +25,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/VM/Transforms/test/CMakeLists.txt b/iree/compiler/Dialect/VM/Transforms/test/CMakeLists.txt
index da04059..8de536b 100644
--- a/iree/compiler/Dialect/VM/Transforms/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/VM/Transforms/test/CMakeLists.txt
@@ -19,7 +19,7 @@
     "hoist_inlined_rodata.mlir"
     "ordinal_allocation.mlir"
     "sink_defining_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Vulkan/IR/test/BUILD b/iree/compiler/Dialect/Vulkan/IR/test/BUILD
index 9e3d553..21e7a3d 100644
--- a/iree/compiler/Dialect/Vulkan/IR/test/BUILD
+++ b/iree/compiler/Dialect/Vulkan/IR/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -19,7 +19,7 @@
         ["target_env.mlir"],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Vulkan/IR/test/CMakeLists.txt b/iree/compiler/Dialect/Vulkan/IR/test/CMakeLists.txt
index e8c3f6f..ae031f8 100644
--- a/iree/compiler/Dialect/Vulkan/IR/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Vulkan/IR/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "target_env.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Dialect/Vulkan/Utils/test/BUILD b/iree/compiler/Dialect/Vulkan/Utils/test/BUILD
index 6fe1334..5f66ef9 100644
--- a/iree/compiler/Dialect/Vulkan/Utils/test/BUILD
+++ b/iree/compiler/Dialect/Vulkan/Utils/test/BUILD
@@ -5,7 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 load("//iree:build_defs.oss.bzl", "iree_cmake_extra_content")
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -30,7 +30,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Dialect/Vulkan/Utils/test/CMakeLists.txt b/iree/compiler/Dialect/Vulkan/Utils/test/CMakeLists.txt
index 2026552..e127839 100644
--- a/iree/compiler/Dialect/Vulkan/Utils/test/CMakeLists.txt
+++ b/iree/compiler/Dialect/Vulkan/Utils/test/CMakeLists.txt
@@ -19,7 +19,7 @@
     lit
   SRCS
     "target_env_conversion.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/InputConversion/Common/test/BUILD b/iree/compiler/InputConversion/Common/test/BUILD
index 47b97f3..96610ff 100644
--- a/iree/compiler/InputConversion/Common/test/BUILD
+++ b/iree/compiler/InputConversion/Common/test/BUILD
@@ -6,7 +6,7 @@
 
 # Tests for common transforms.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -24,7 +24,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/InputConversion/Common/test/CMakeLists.txt b/iree/compiler/InputConversion/Common/test/CMakeLists.txt
index 1a51c49..3a4becf 100644
--- a/iree/compiler/InputConversion/Common/test/CMakeLists.txt
+++ b/iree/compiler/InputConversion/Common/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "iree_import_public.mlir"
     "top_level_scf_to_cfg.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/InputConversion/MHLO/test/BUILD b/iree/compiler/InputConversion/MHLO/test/BUILD
index 00b1d90..d0fbcbf 100644
--- a/iree/compiler/InputConversion/MHLO/test/BUILD
+++ b/iree/compiler/InputConversion/MHLO/test/BUILD
@@ -6,7 +6,7 @@
 
 # Tests for common transforms.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -35,7 +35,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt b/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt
index fea81a3..f95c73d 100644
--- a/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt
+++ b/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt
@@ -27,7 +27,7 @@
     "mhlo_to_mhlo_preprocessing_extract_pad_from_conv.mlir"
     "missing_legalizations.mlir"
     "verify_compiler_mhlo_input_legality.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/InputConversion/TOSA/test/BUILD b/iree/compiler/InputConversion/TOSA/test/BUILD
index acfb3f1..3e66cf3 100644
--- a/iree/compiler/InputConversion/TOSA/test/BUILD
+++ b/iree/compiler/InputConversion/TOSA/test/BUILD
@@ -6,7 +6,7 @@
 
 # Tests for common transforms.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -23,7 +23,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/InputConversion/TOSA/test/CMakeLists.txt b/iree/compiler/InputConversion/TOSA/test/CMakeLists.txt
index 3d94328..59a4c50 100644
--- a/iree/compiler/InputConversion/TOSA/test/CMakeLists.txt
+++ b/iree/compiler/InputConversion/TOSA/test/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "verify_compiler_tosa_input_legality.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
 )
diff --git a/iree/compiler/Translation/test/BUILD b/iree/compiler/Translation/test/BUILD
index b107271..77ccb64 100644
--- a/iree/compiler/Translation/test/BUILD
+++ b/iree/compiler/Translation/test/BUILD
@@ -6,7 +6,7 @@
 
 # Tests for common transforms.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -24,7 +24,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/tools:iree-translate",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/compiler/Translation/test/CMakeLists.txt b/iree/compiler/Translation/test/CMakeLists.txt
index 635e002..cdc563b 100644
--- a/iree/compiler/Translation/test/CMakeLists.txt
+++ b/iree/compiler/Translation/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "hal_executable.mlir"
     "smoketest.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-translate
 )
diff --git a/iree/hal/local/elf/BUILD b/iree/hal/local/elf/BUILD
index d9b402c..aacd5f3 100644
--- a/iree/hal/local/elf/BUILD
+++ b/iree/hal/local/elf/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//build_tools/bazel:run_binary_test.bzl", "run_binary_test")
+load("//build_tools/bazel:native_binary.bzl", "native_test")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -46,9 +46,9 @@
     ],
 )
 
-run_binary_test(
+native_test(
     name = "elf_module_test",
-    test_binary = ":elf_module_test_binary",
+    src = ":elf_module_test_binary",
 )
 
 #===------------------------------------------------------------------------===#
diff --git a/iree/hal/local/elf/CMakeLists.txt b/iree/hal/local/elf/CMakeLists.txt
index 77926b4..1c5e254 100644
--- a/iree/hal/local/elf/CMakeLists.txt
+++ b/iree/hal/local/elf/CMakeLists.txt
@@ -40,10 +40,10 @@
     iree::hal::local::executable_library
 )
 
-iree_run_binary_test(
+iree_native_test(
   NAME
     "elf_module_test"
-  TEST_BINARY
+  SRC
     ::elf_module_test_binary
 )
 
diff --git a/iree/lit.cfg.py b/iree/lit.cfg.py
new file mode 100644
index 0000000..77a0498
--- /dev/null
+++ b/iree/lit.cfg.py
@@ -0,0 +1,32 @@
+# 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
+"""Lit config for IREE."""
+
+# Lint for undefined variables is disabled as config is not defined inside this
+# file, instead config is injected by way of evaluating runlit.cfg.py from
+# runlit.site.cfg.py which in turn is evaluated by lit.py.
+# pylint: disable=undefined-variable
+
+import os
+import tempfile
+
+import lit.formats
+
+config.name = "IREE"
+config.suffixes = [".mlir", ".txt"]
+config.test_format = lit.formats.ShTest(execute_external=True)
+# Forward all IREE environment variables
+passthrough_env_vars = ["VK_ICD_FILENAMES"]
+config.environment.update({
+    k: v
+    for k, v in os.environ.items()
+    if k.startswith("IREE_") or k in passthrough_env_vars
+})
+
+# Use the most preferred temp directory.
+config.test_exec_root = (os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR") or
+                         os.environ.get("TEST_TMPDIR") or
+                         os.path.join(tempfile.gettempdir(), "lit"))
diff --git a/iree/lit_test.bzl b/iree/lit_test.bzl
deleted file mode 100644
index 5e0295f..0000000
--- a/iree/lit_test.bzl
+++ /dev/null
@@ -1,96 +0,0 @@
-# Copyright 2019 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
-
-"""Bazel macros for running lit tests."""
-
-def iree_lit_test(
-        name,
-        test_file,
-        data,
-        size = "small",
-        driver = "//iree/tools:run_lit.sh",
-        **kwargs):
-    """Creates a lit test from the specified source file.
-
-    Args:
-      name: name of the generated test suite.
-      test_file: the test file with the lit test
-      data: binaries used in the lit tests.
-      size: size of the tests.
-      driver: the shell runner for the lit tests.
-      **kwargs: Any additional arguments that will be passed to the underlying sh_test.
-    """
-    data = data if data else []
-
-    # Always include llvm-sybmolizer so we get useful stack traces. Maybe it
-    # would be better to force everyone to do this explicitly, but since
-    # forgetting wouldn't cause the test to fail, only make debugging harder
-    # when it does, I think better to hardcode it here.
-    llvm_symbolizer = "@llvm-project//llvm:llvm-symbolizer"
-    if llvm_symbolizer not in data:
-        data.append(llvm_symbolizer)
-
-    # First argument is the test. The rest are tools to add to the path.
-    data = [test_file] + data
-
-    native.sh_test(
-        name = name,
-        srcs = [driver],
-        size = size,
-        data = data,
-        env = {"FILECHECK_OPTS": "--enable-var-scope"},
-        args = ["$(location {})".format(file) for file in data],
-        **kwargs
-    )
-
-def iree_lit_test_suite(
-        name,
-        srcs,
-        data,
-        size = "small",
-        driver = "//iree/tools:run_lit.sh",
-        tags = [],
-        **kwargs):
-    """Creates one lit test per source file and a test suite that bundles them.
-
-    Args:
-      name: name of the generated test suite.
-      data: binaries used in the lit tests.
-      srcs: test file sources.
-      size: size of the tests.
-      driver: the shell runner for the lit tests.
-      tags: tags to apply to the test. Note that as in standard test suites, manual
-            is treated specially and will also apply to the test suite itself.
-      **kwargs: Any additional arguments that will be passed to the underlying tests and test_suite.
-    """
-    tests = []
-    for test_file in srcs:
-        # It's generally good practice to prefix any generated names with the
-        # macro name, but we're trying to match the style of the names that are
-        # used for LLVM internally.
-        test_name = "%s.test" % (test_file)
-        iree_lit_test(
-            name = test_name,
-            test_file = test_file,
-            size = size,
-            data = data,
-            driver = driver,
-            tags = tags,
-            **kwargs
-        )
-        tests.append(test_name)
-
-    native.test_suite(
-        name = name,
-        tests = tests,
-        # Note that only the manual tag really has any effect here. Others are
-        # used for test suite filtering, but all tests are passed the same tags.
-        tags = tags,
-        # If there are kwargs that need to be passed here which only apply to
-        # the generated tests and not to test_suite, they should be extracted
-        # into separate named arguments.
-        **kwargs
-    )
diff --git a/iree/modules/check/test/BUILD b/iree/modules/check/test/BUILD
index d3f5a6f..ce646ce 100644
--- a/iree/modules/check/test/BUILD
+++ b/iree/modules/check/test/BUILD
@@ -5,7 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 load("//build_tools/bazel:iree_check_test.bzl", "iree_check_test_suite")
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -23,12 +23,12 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tags = ["hostonly"],
+    tools = [
         "//iree/tools:iree-check-module",
         "//iree/tools:iree-translate",
         "@llvm-project//llvm:FileCheck",
     ],
-    tags = ["hostonly"],
 )
 
 iree_check_test_suite(
diff --git a/iree/modules/check/test/CMakeLists.txt b/iree/modules/check/test/CMakeLists.txt
index 33d8e62..ed1b37d 100644
--- a/iree/modules/check/test/CMakeLists.txt
+++ b/iree/modules/check/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "failure.mlir"
     "success.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-check-module
     iree::tools::iree-translate
diff --git a/iree/runtime/demo/BUILD b/iree/runtime/demo/BUILD
index 7d5861a..a5e7b0e 100644
--- a/iree/runtime/demo/BUILD
+++ b/iree/runtime/demo/BUILD
@@ -5,7 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 load("//iree:build_defs.oss.bzl", "iree_cmake_extra_content")
-load("//build_tools/bazel:run_binary_test.bzl", "run_binary_test")
+load("//build_tools/bazel:native_binary.bzl", "native_test")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -39,9 +39,9 @@
     ],
 )
 
-run_binary_test(
+native_test(
     name = "hello_world_embedded_test",
-    test_binary = ":hello_world_embedded",
+    src = ":hello_world_embedded",
 )
 
 cc_binary(
@@ -56,8 +56,8 @@
     ],
 )
 
-# TODO(benvanik): run_binary_test that passes the file as a flag. Right now we
-# can't specify data through run_binary_test, though, so this isn't possible to
+# TODO(benvanik): native_test that passes the file as a flag. Right now we
+# can't specify data through native_test, though, so this isn't possible to
 # automate.
 
 cc_binary(
@@ -69,7 +69,7 @@
     ],
 )
 
-run_binary_test(
+native_test(
     name = "hello_world_terse_test",
-    test_binary = ":hello_world_terse",
+    src = ":hello_world_terse",
 )
diff --git a/iree/runtime/demo/CMakeLists.txt b/iree/runtime/demo/CMakeLists.txt
index e598af9..4500f97 100644
--- a/iree/runtime/demo/CMakeLists.txt
+++ b/iree/runtime/demo/CMakeLists.txt
@@ -26,10 +26,10 @@
     iree::runtime::testdata::simple_mul_module_c
 )
 
-iree_run_binary_test(
+iree_native_test(
   NAME
     "hello_world_embedded_test"
-  TEST_BINARY
+  SRC
     ::hello_world_embedded
 )
 
@@ -54,10 +54,10 @@
     iree::runtime::testdata::simple_mul_module_c
 )
 
-iree_run_binary_test(
+iree_native_test(
   NAME
     "hello_world_terse_test"
-  TEST_BINARY
+  SRC
     ::hello_world_terse
 )
 
diff --git a/iree/samples/custom_modules/dialect/test/BUILD b/iree/samples/custom_modules/dialect/test/BUILD
index ef078ed..5d7f742 100644
--- a/iree/samples/custom_modules/dialect/test/BUILD
+++ b/iree/samples/custom_modules/dialect/test/BUILD
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -22,7 +22,7 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tools = [
         "//iree/samples/custom_modules/dialect:custom-opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/iree/samples/custom_modules/dialect/test/CMakeLists.txt b/iree/samples/custom_modules/dialect/test/CMakeLists.txt
index bf7af51..66b0c31 100644
--- a/iree/samples/custom_modules/dialect/test/CMakeLists.txt
+++ b/iree/samples/custom_modules/dialect/test/CMakeLists.txt
@@ -16,7 +16,7 @@
   SRCS
     "conversion.mlir"
     "custom_ops.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::samples::custom_modules::dialect::custom-opt
 )
diff --git a/iree/samples/simple_embedding/BUILD b/iree/samples/simple_embedding/BUILD
index 58b5e5f..1c38f8c 100644
--- a/iree/samples/simple_embedding/BUILD
+++ b/iree/samples/simple_embedding/BUILD
@@ -6,7 +6,7 @@
 
 load("//iree:build_defs.oss.bzl", "iree_cmake_extra_content")
 load("//build_tools/bazel:iree_bytecode_module.bzl", "iree_bytecode_module")
-load("//build_tools/bazel:run_binary_test.bzl", "run_binary_test")
+load("//build_tools/bazel:native_binary.bzl", "native_test")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -52,9 +52,9 @@
     ],
 )
 
-run_binary_test(
+native_test(
     name = "simple_embedding_vmvx_sync_test",
-    test_binary = ":simple_embedding_vmvx_sync",
+    src = ":simple_embedding_vmvx_sync",
 )
 
 iree_cmake_extra_content(
@@ -177,9 +177,9 @@
     ],
 )
 
-run_binary_test(
+native_test(
     name = "simple_embedding_embedded_sync_test",
-    test_binary = ":simple_embedding_embedded_sync",
+    src = ":simple_embedding_embedded_sync",
 )
 
 iree_cmake_extra_content(
@@ -211,9 +211,9 @@
     ],
 )
 
-run_binary_test(
+native_test(
     name = "simple_embedding_dylib_test",
-    test_binary = ":simple_embedding_dylib",
+    src = ":simple_embedding_dylib",
 )
 
 iree_cmake_extra_content(
@@ -261,9 +261,9 @@
     ],
 )
 
-run_binary_test(
+native_test(
     name = "simple_embedding_vulkan_test",
-    test_binary = ":simple_embedding_vulkan",
+    src = ":simple_embedding_vulkan",
 )
 
 iree_cmake_extra_content(
@@ -305,22 +305,22 @@
 )
 
 # Simple embedding is failing in the CI.
-# run_binary_test(
+# native_test(
 #     name = "simple_embedding_cuda_test",
 #     tags = [
 #         "driver=cuda",
 #     ],
-#     test_binary = ":simple_embedding_cuda",
+#     src = ":simple_embedding_cuda",
 # )
 
 iree_cmake_extra_content(
     content = """
-iree_run_binary_test(
+iree_native_test(
   NAME
     "simple_embedding_cuda_test"
   LABELS
     "driver=cuda"
-  TEST_BINARY
+  SRC
     ::simple_embedding_cuda
 )
 
diff --git a/iree/samples/simple_embedding/CMakeLists.txt b/iree/samples/simple_embedding/CMakeLists.txt
index 555052b..329566c 100644
--- a/iree/samples/simple_embedding/CMakeLists.txt
+++ b/iree/samples/simple_embedding/CMakeLists.txt
@@ -45,10 +45,10 @@
   PUBLIC
 )
 
-iree_run_binary_test(
+iree_native_test(
   NAME
     "simple_embedding_vmvx_sync_test"
-  TEST_BINARY
+  SRC
     ::simple_embedding_vmvx_sync
 )
 
@@ -176,10 +176,10 @@
   PUBLIC
 )
 
-iree_run_binary_test(
+iree_native_test(
   NAME
     "simple_embedding_embedded_sync_test"
-  TEST_BINARY
+  SRC
     ::simple_embedding_embedded_sync
 )
 
@@ -206,10 +206,10 @@
     iree::vm::bytecode_module
 )
 
-iree_run_binary_test(
+iree_native_test(
   NAME
     "simple_embedding_dylib_test"
-  TEST_BINARY
+  SRC
     ::simple_embedding_dylib
 )
 
@@ -250,10 +250,10 @@
   PUBLIC
 )
 
-iree_run_binary_test(
+iree_native_test(
   NAME
     "simple_embedding_vulkan_test"
-  TEST_BINARY
+  SRC
     ::simple_embedding_vulkan
 )
 
@@ -292,12 +292,12 @@
   PUBLIC
 )
 
-iree_run_binary_test(
+iree_native_test(
   NAME
     "simple_embedding_cuda_test"
   LABELS
     "driver=cuda"
-  TEST_BINARY
+  SRC
     ::simple_embedding_cuda
 )
 
diff --git a/iree/samples/static_library/CMakeLists.txt b/iree/samples/static_library/CMakeLists.txt
index d7fe192..cd46e87 100644
--- a/iree/samples/static_library/CMakeLists.txt
+++ b/iree/samples/static_library/CMakeLists.txt
@@ -87,7 +87,7 @@
     static_library_demo_test
   TEST_FILE
     "static_library_demo_test.txt"
-  DATA
+  TOOLS
     ::static_library_demo
     FileCheck
   LABELS
@@ -170,7 +170,7 @@
     static_library_demo_c_test
   TEST_FILE
     "static_library_demo_c_test.txt"
-  DATA
+  TOOLS
     ::static_library_demo_c
     FileCheck
   LABELS
diff --git a/iree/samples/vision/CMakeLists.txt b/iree/samples/vision/CMakeLists.txt
index 353f8bf..c1585fc 100644
--- a/iree/samples/vision/CMakeLists.txt
+++ b/iree/samples/vision/CMakeLists.txt
@@ -43,9 +43,11 @@
     iree_run_mnist_module_test
   TEST_FILE
     "mnist_test.txt"
-  DATA
+  TOOLS
     ::iree-run-mnist-module
     FileCheck
+  DATA
+    "mnist_test.png"
   LABELS
     "hostonly"
 )
diff --git a/iree/samples/vision/mnist_test.txt b/iree/samples/vision/mnist_test.txt
index 2c4b29d..03dcbba 100644
--- a/iree/samples/vision/mnist_test.txt
+++ b/iree/samples/vision/mnist_test.txt
@@ -1,3 +1,3 @@
-// RUN: (iree-run-mnist-module) | FileCheck %s
+// RUN: (iree-run-mnist-module %S/mnist_test.png) | FileCheck %s
 // CHECK-LABEL: Detected number
 // CHECK: 4
diff --git a/iree/test/e2e/models/BUILD b/iree/test/e2e/models/BUILD
index f1aca98..ae84e50 100644
--- a/iree/test/e2e/models/BUILD
+++ b/iree/test/e2e/models/BUILD
@@ -7,7 +7,7 @@
 # Tests for end-to-end IREE support of entire models or their close derivatives.
 
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:iree_check_test.bzl", "iree_check_single_backend_test_suite")
 
 package(
@@ -38,15 +38,15 @@
             ["*.mlir"],
         exclude = CHECK_FRAMEWORK_TESTS,
     ),
-    data = [
-        "//iree/tools:iree-run-mlir",
-        "@llvm-project//lld",
-        "@llvm-project//llvm:FileCheck",
-    ],
     tags = [
         "hostonly",
         "optonly",  # swiftshader is too slow in dbg
     ],
+    tools = [
+        "//iree/tools:iree-run-mlir",
+        "@llvm-project//lld",
+        "@llvm-project//llvm:FileCheck",
+    ],
 )
 
 iree_check_single_backend_test_suite(
diff --git a/iree/test/e2e/models/CMakeLists.txt b/iree/test/e2e/models/CMakeLists.txt
index 7641127..a15eb09 100644
--- a/iree/test/e2e/models/CMakeLists.txt
+++ b/iree/test/e2e/models/CMakeLists.txt
@@ -21,7 +21,7 @@
     "mnist_fake_weights.mlir"
     "resnet50_fake_weights.mlir"
     "unidirectional_lstm.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-run-mlir
     lld
diff --git a/iree/test/e2e/regression/BUILD b/iree/test/e2e/regression/BUILD
index d2b0e4a..0331f46 100644
--- a/iree/test/e2e/regression/BUILD
+++ b/iree/test/e2e/regression/BUILD
@@ -9,7 +9,7 @@
 # those tests in iree/tools/test/
 
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:iree_check_test.bzl", "iree_check_single_backend_test_suite")
 load("//build_tools/bazel:iree_trace_runner_test.bzl", "iree_generated_trace_runner_test")
 
@@ -56,13 +56,13 @@
             "lowering_config.mlir",
         ] + BACKEND_TESTS,
     ),
-    data = [
+    tags = ["hostonly"],
+    tools = [
         "//iree/tools:iree-opt",
         "//iree/tools:iree-run-mlir",
         "@llvm-project//lld",
         "@llvm-project//llvm:FileCheck",
     ],
-    tags = ["hostonly"],
 )
 
 iree_check_single_backend_test_suite(
diff --git a/iree/test/e2e/regression/CMakeLists.txt b/iree/test/e2e/regression/CMakeLists.txt
index 914c9ce..e54c442 100644
--- a/iree/test/e2e/regression/CMakeLists.txt
+++ b/iree/test/e2e/regression/CMakeLists.txt
@@ -20,7 +20,7 @@
     "scalar.mlir"
     "trace_dispatch_tensors.mlir"
     "unused_args.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-opt
     iree::tools::iree-run-mlir
diff --git a/iree/test/e2e/tensor_ops/BUILD b/iree/test/e2e/tensor_ops/BUILD
index dbfcf21..1c694b2 100644
--- a/iree/test/e2e/tensor_ops/BUILD
+++ b/iree/test/e2e/tensor_ops/BUILD
@@ -11,7 +11,7 @@
 # See https://github.com/google/iree/blob/main/docs/developers/developing_iree/testing_guide.md#iree-core-end-to-end-tests.
 
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:iree_check_test.bzl", "iree_check_single_backend_test_suite")
 
 package(
@@ -25,13 +25,13 @@
     srcs = [
         "tensor_cast.mlir",
     ],
-    data = [
+    tags = ["hostonly"],
+    tools = [
         "//iree/tools:iree-benchmark-module",
         "//iree/tools:iree-run-mlir",
         "//iree/tools:iree-translate",
         "@llvm-project//llvm:FileCheck",
     ],
-    tags = ["hostonly"],
 )
 
 iree_check_single_backend_test_suite(
diff --git a/iree/test/e2e/tensor_ops/CMakeLists.txt b/iree/test/e2e/tensor_ops/CMakeLists.txt
index 2536e3d..2acc931 100644
--- a/iree/test/e2e/tensor_ops/CMakeLists.txt
+++ b/iree/test/e2e/tensor_ops/CMakeLists.txt
@@ -15,7 +15,7 @@
     lit
   SRCS
     "tensor_cast.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-benchmark-module
     iree::tools::iree-run-mlir
diff --git a/iree/tools/BUILD b/iree/tools/BUILD
index d72c94e..4de58cd 100644
--- a/iree/tools/BUILD
+++ b/iree/tools/BUILD
@@ -13,10 +13,6 @@
     licenses = ["notice"],  # Apache 2.0
 )
 
-exports_files([
-    "run_lit.sh",
-])
-
 cc_binary(
     name = "iree-benchmark-module",
     srcs = ["iree-benchmark-module-main.cc"],
diff --git a/iree/tools/CMakeLists.txt b/iree/tools/CMakeLists.txt
index 8e4adf7..eca2269 100644
--- a/iree/tools/CMakeLists.txt
+++ b/iree/tools/CMakeLists.txt
@@ -459,10 +459,11 @@
     HOSTONLY
   )
 
-  add_custom_target(BundledLLVMFileCheck ALL
-    COMMAND ${CMAKE_COMMAND} -E create_symlink $<TARGET_FILE:FileCheck> FileCheck
-    DEPENDS FileCheck
-  )
+  # Ensure FileCheck gets built. Tests don't have dependencies in CMake because
+  # they aren't targets. So until we fix that, we just force this to get built.
+  # Limiting this to when IREE_BUILD_TESTS is set prevents the installation
+  # below, which we use for cross-platform testing.
+  set_target_properties(FileCheck PROPERTIES EXCLUDE_FROM_ALL OFF)
 
   # lld install - required by the compiler to link codegen executables.
   install(
@@ -471,16 +472,6 @@
     RUNTIME DESTINATION bin
   )
 
-  # General test dependency binary/scripts installed.
-  install(
-    PROGRAMS
-      run_lit.bat
-      run_lit.ps1
-      run_lit.sh
-    DESTINATION "tests/bin"
-    COMPONENT Tests
-  )
-
   # Bundle the FileCheck binary from LLVM into our tests/bin directory so
   # installed FileCheck tests are hermetic.
   install(
diff --git a/iree/tools/run_lit.bat b/iree/tools/run_lit.bat
deleted file mode 100644
index 4e6735c..0000000
--- a/iree/tools/run_lit.bat
+++ /dev/null
@@ -1,10 +0,0 @@
-@ECHO OFF
-REM Copyright 2020 The IREE Authors
-REM
-REM Licensed under the Apache License v2.0 with LLVM Exceptions.
-REM See https://llvm.org/LICENSE.txt for license information.
-REM SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-
-SET RUNNER_PATH=%~dp0
-powershell.exe -NoProfile -File "%RUNNER_PATH%\run_lit.ps1" %*
-EXIT /B %ERRORLEVEL%
diff --git a/iree/tools/run_lit.ps1 b/iree/tools/run_lit.ps1
deleted file mode 100644
index e88dcb4..0000000
--- a/iree/tools/run_lit.ps1
+++ /dev/null
@@ -1,78 +0,0 @@
-# Copyright 2020 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
-
-param(
-  [Parameter(Position=0, Mandatory)]
-  [ValidateNotNullOrEmpty()]
-  [string]
-    $test_file,
-  [Parameter(Position=1, ValueFromRemainingArguments=$true)]
-  [string[]]
-    $test_data = @()
-)
-
-# NOTE: to debug first run `$DebugPreference = 'Continue'` in your shell.
-# $DebugPreference = "Continue"
-
-trap {
-  Write-Error $_
-  exit 1
-}
-
-# Search the system path for a suitable bash.exe.
-# Note that C:\Windows\system32\bash.exe is actually WSL -- which will not
-# work. Why???
-$pathFolders = $env:Path.Split(";") 
-foreach ($_ in $pathFolders) {
-  if (-not ($_ -like "*:\Windows\*")) {
-    Write-Debug "Checking for bash.exe in: $_"
-    $possibleBashExe = "$_\bash.exe"
-    if (Test-Path $possibleBashExe -PathType leaf) {
-      $bashExe = $possibleBashExe
-      break
-    }
-  }
-}
-
-if (-not $bashExe) {
-  Write-Host -ForegroundColor Red "Could not find bash.exe on path (excluding \Windows\system32)"
-  $pathFolders -join "`r`n" |  Write-Host -ForegroundColor Red
-  exit 1
-}
-Write-Debug "Using bash.exe: $bashExe"
-
-# Get all of the directories we'll want to put on our path for the test.
-$test_dirs = [System.Collections.ArrayList]@()
-foreach ($test_path in $test_data) {
-  $test_dir = Split-Path -Path $test_path -Parent
-  $test_dirs.Add($test_dir) | Out-Null
-}
-Write-Debug "Test data directories: $test_dirs"
-$test_dirs.AddRange($env:Path.Split(";"))
-$env:Path = $test_dirs -join ";"
-Write-Debug "Test PATH:"
-Write-Debug "$env:PATH"
-
-$test_lines = Get-Content -Path $test_file
-foreach ($test_line in $test_lines) {
-  if (!$test_line.StartsWith("// RUN:")) {
-    continue
-  }
-  $test_line = $test_line.Substring("// RUN: ".Length)
-  $test_line = $test_line -replace "%s", $test_file
-  $test_line = $test_line -replace "`"", "\`""
-  Write-Host -ForegroundColor Blue "Running test command:"
-  Write-Host -ForegroundColor Yellow "$test_line"
-  & $bashExe -c $test_line | Out-Default
-  if ($LASTEXITCODE -gt 0) {
-    Write-Host -ForegroundColor Red "Test failed with $LASTEXITCODE, command:"
-    Write-Host -ForegroundColor Yellow "$test_line"
-    exit $LASTEXITCODE
-  }
-}
-
-Write-Debug "All run commands completed successfully"
-exit 0
diff --git a/iree/tools/run_lit.sh b/iree/tools/run_lit.sh
deleted file mode 100755
index f1258e9..0000000
--- a/iree/tools/run_lit.sh
+++ /dev/null
@@ -1,60 +0,0 @@
-#!/bin/bash
-
-# Copyright 2019 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
-set -e
-set -o pipefail
-
-EXPLICIT_PATH=""
-
-# First argument is the src file. Remaining arguments are tools that should
-# be on the path.
-src_file="$1"
-shift
-for tool_exe in "$@"
-do
-  EXEDIR="$(dirname $tool_exe)"
-  if ! [ -z "$cygpath" ]; then
-    EXEDIR="$($cygpath -u "$EXEDIR")"
-  fi
-  EXPLICIT_PATH="${EXEDIR}:${EXPLICIT_PATH}"
-done
-
-echo "run_lit.sh: $src_file"
-echo "PWD=$(pwd)"
-echo "EXPLICIT_PATH=$EXPLICIT_PATH"
-
-# For each "// RUN:" line, run the command.
-runline_matches="$(egrep "^// RUN: " "$src_file")"
-if [ -z "$runline_matches" ]; then
-  echo "!!! No RUN lines found in test"
-  exit 1
-fi
-
-echo "$runline_matches" | while read -r runline
-do
-  echo "RUNLINE: $runline"
-  match="${runline%%// RUN: *}"
-  command="${runline##// RUN: }"
-  if [ -z "${command}" ]; then
-    echo "ERROR: Could not extract command from runline"
-    exit 1
-  fi
-
-  # Substitute any embedded '%s' with the file name.
-  full_command="${command//\%s/$src_file}"
-
-  # Run it.
-  export PATH="$EXPLICIT_PATH:$PATH"
-  echo "RUNNING TEST: $full_command"
-  echo "----------------"
-  if eval "$full_command"; then
-    echo "--- COMPLETE ---"
-  else
-    echo "!!! ERROR EVALUATING: $full_command"
-    exit 1
-  fi
-done
diff --git a/iree/tools/test/BUILD b/iree/tools/test/BUILD
index 0ddf314..be357dd 100644
--- a/iree/tools/test/BUILD
+++ b/iree/tools/test/BUILD
@@ -6,7 +6,7 @@
 
 # Smoke tests for the execution of tool binaries.
 
-load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
@@ -29,7 +29,10 @@
         ],
         include = ["*.mlir"],
     ),
-    data = [
+    tags = [
+        "hostonly",
+    ],
+    tools = [
         "//iree/tools:iree-benchmark-module",
         "//iree/tools:iree-run-mlir",
         "//iree/tools:iree-run-module",
@@ -37,20 +40,17 @@
         "@llvm-project//lld",
         "@llvm-project//llvm:FileCheck",
     ],
-    tags = [
-        "hostonly",
-    ],
 )
 
 iree_lit_test_suite(
     name = "benchmark_flags",
     srcs = ["benchmark_flags.txt"],
-    data = [
+    tags = [
+        "hostonly",
+    ],
+    tools = [
         "//iree/tools:iree-benchmark-module",
         "//iree/tools:iree-translate",
         "@llvm-project//llvm:FileCheck",
     ],
-    tags = [
-        "hostonly",
-    ],
 )
diff --git a/iree/tools/test/CMakeLists.txt b/iree/tools/test/CMakeLists.txt
index d8003c7..c37ec3f 100644
--- a/iree/tools/test/CMakeLists.txt
+++ b/iree/tools/test/CMakeLists.txt
@@ -21,7 +21,7 @@
     "multiple_exported_functions.mlir"
     "repeated_return.mlir"
     "scalars.mlir"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-benchmark-module
     iree::tools::iree-run-mlir
@@ -37,7 +37,7 @@
     benchmark_flags
   SRCS
     "benchmark_flags.txt"
-  DATA
+  TOOLS
     FileCheck
     iree::tools::iree-benchmark-module
     iree::tools::iree-translate
diff --git a/iree/vm/BUILD b/iree/vm/BUILD
index 2bca721..d39c03a 100644
--- a/iree/vm/BUILD
+++ b/iree/vm/BUILD
@@ -6,7 +6,7 @@
 
 load("//iree:build_defs.oss.bzl", "iree_cmake_extra_content")
 load("//build_tools/bazel:iree_bytecode_module.bzl", "iree_bytecode_module")
-load("//build_tools/bazel:run_binary_test.bzl", "run_binary_test")
+load("//build_tools/bazel:native_binary.bzl", "native_test")
 # load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
 
 package(
@@ -266,10 +266,10 @@
     ],
 )
 
-run_binary_test(
+native_test(
     name = "bytecode_module_benchmark_test",
+    src = ":bytecode_module_benchmark",
     args = ["--benchmark_min_time=0"],
-    test_binary = ":bytecode_module_benchmark",
 )
 
 iree_bytecode_module(
diff --git a/iree/vm/CMakeLists.txt b/iree/vm/CMakeLists.txt
index 1274218..857b8c9 100644
--- a/iree/vm/CMakeLists.txt
+++ b/iree/vm/CMakeLists.txt
@@ -226,12 +226,12 @@
   TESTONLY
 )
 
-iree_run_binary_test(
+iree_native_test(
   NAME
     "bytecode_module_benchmark_test"
   ARGS
     "--benchmark_min_time=0"
-  TEST_BINARY
+  SRC
     ::bytecode_module_benchmark
 )