| """Build rules for Tensorflow/XLA testing.""" |
| |
| load("@local_config_cuda//cuda:build_defs.bzl", "cuda_is_configured") |
| load("//tensorflow/compiler/tests:plugin.bzl", "plugins") |
| load( |
| "//tensorflow/core:platform/default/build_config_root.bzl", |
| "tf_cuda_tests_tags", |
| ) |
| |
| def all_backends(): |
| b = ["cpu"] + plugins.keys() |
| if cuda_is_configured(): |
| return b + ["gpu"] |
| else: |
| return b |
| |
| def tf_xla_py_test( |
| name, |
| srcs = [], |
| deps = [], |
| tags = [], |
| data = [], |
| main = None, |
| disabled_backends = None, |
| **kwargs): |
| """Generates py_test targets, one per XLA backend. |
| |
| This rule generates py_test() targets named name_backend, for each backend |
| in all_backends(). The rule also generates a test suite with named `name` that |
| tests all backends for the test. |
| |
| For example, the following rule generates test cases foo_test_cpu, |
| foo_test_gpu, and a test suite name foo_test that tests both. |
| tf_xla_py_test( |
| name="foo_test", |
| srcs="foo_test.py", |
| deps=[...], |
| ) |
| |
| Args: |
| name: Name of the target. |
| srcs: Sources for the target. |
| deps: Dependencies of the target. |
| tags: Tags to apply to the generated targets. |
| data: Data dependencies of the target. |
| main: Same as py_test's main attribute. |
| disabled_backends: A list of backends that should not be tested. Supported |
| values include "cpu" and "gpu". If not specified, defaults to None. |
| **kwargs: keyword arguments passed onto the generated py_test() rules. |
| """ |
| if disabled_backends == None: |
| disabled_backends = [] |
| |
| enabled_backends = [b for b in all_backends() if b not in disabled_backends] |
| test_names = [] |
| for backend in enabled_backends: |
| test_name = "{}_{}".format(name, backend) |
| backend_tags = ["tf_xla_{}".format(backend)] |
| backend_args = [] |
| backend_deps = [] |
| backend_data = [] |
| if backend == "cpu": |
| backend_args += [ |
| "--test_device=XLA_CPU", |
| "--types=DT_HALF,DT_FLOAT,DT_DOUBLE,DT_UINT8,DT_QUINT8,DT_INT8,DT_QINT8,DT_INT32,DT_QINT32,DT_INT64,DT_BOOL,DT_COMPLEX64", |
| ] |
| elif backend == "gpu": |
| backend_args += [ |
| "--test_device=XLA_GPU", |
| "--types=DT_HALF,DT_FLOAT,DT_DOUBLE,DT_UINT8,DT_QUINT8,DT_INT8,DT_QINT8,DT_INT32,DT_QINT32,DT_INT64,DT_BOOL,DT_COMPLEX64,DT_BFLOAT16", |
| ] |
| backend_tags += tf_cuda_tests_tags() |
| elif backend in plugins: |
| backend_args += [ |
| "--test_device=" + plugins[backend]["device"], |
| "--types=" + plugins[backend]["types"], |
| ] |
| backend_tags += plugins[backend]["tags"] |
| backend_args += plugins[backend]["args"] |
| backend_deps += plugins[backend]["deps"] |
| backend_data += plugins[backend]["data"] |
| else: |
| fail("Unknown backend {}".format(backend)) |
| |
| native.py_test( |
| name = test_name, |
| srcs = srcs, |
| srcs_version = "PY2AND3", |
| args = backend_args, |
| main = "{}.py".format(name) if main == None else main, |
| data = data + backend_data, |
| deps = deps + backend_deps, |
| tags = tags + backend_tags, |
| **kwargs |
| ) |
| test_names.append(test_name) |
| native.test_suite(name = name, tests = test_names) |
| |
| def generate_backend_suites(backends = []): |
| """Generates per-backend test_suites that run all tests for a backend.""" |
| if not backends: |
| backends = all_backends() |
| for backend in backends: |
| native.test_suite(name = "%s_tests" % backend, tags = ["tf_xla_%s" % backend]) |