[XLA] Add ability to run the XLA unit tests against a different device (#9759)

* Add ability to run the XLA unit tests against a different device

* Allow for multiple extra backend devices

* Correct merge error

* Include options for additional tags
This commit is contained in:
David Norman 2017-06-08 03:09:13 +01:00 committed by Jonathan Hseu
parent aff4d124b2
commit 91cb809bd6
2 changed files with 31 additions and 2 deletions

View File

@ -1,12 +1,14 @@
"""Build rules for Tensorflow/XLA testing.""" """Build rules for Tensorflow/XLA testing."""
load("@local_config_cuda//cuda:build_defs.bzl", "cuda_is_configured") load("@local_config_cuda//cuda:build_defs.bzl", "cuda_is_configured")
load("//tensorflow/compiler/tests:plugin.bzl", "plugins")
def all_backends(): def all_backends():
b = ["cpu"] + plugins.keys()
if cuda_is_configured(): if cuda_is_configured():
return ["cpu", "gpu"] return b + ["gpu"]
else: else:
return ["cpu"] return b
def tf_xla_py_test(name, srcs=[], deps=[], tags=[], data=[], main=None, def tf_xla_py_test(name, srcs=[], deps=[], tags=[], data=[], main=None,
disabled_backends=None, **kwargs): disabled_backends=None, **kwargs):
@ -53,6 +55,10 @@ def tf_xla_py_test(name, srcs=[], deps=[], tags=[], data=[], main=None,
backend_args += ["--test_device=XLA_GPU", backend_args += ["--test_device=XLA_GPU",
"--types=DT_FLOAT,DT_DOUBLE,DT_INT32,DT_INT64,DT_BOOL"] "--types=DT_FLOAT,DT_DOUBLE,DT_INT32,DT_INT64,DT_BOOL"]
backend_tags += ["requires-gpu-sm35"] backend_tags += ["requires-gpu-sm35"]
elif backend in plugins:
backend_args += ["--test_device=" + plugins[backend]["device"],
"--types=" + plugins[backend]["types"]]
backend_tags += plugins[backend]["tags"]
else: else:
fail("Unknown backend {}".format(backend)) fail("Unknown backend {}".format(backend))

View File

@ -0,0 +1,23 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Additional XLA devices to be included in the unit test suite."""
# If you wish to edit this file without checking it into the repo, consider:
# git update-index --assume-unchanged tensorflow/compiler/tests/plugin.bzl
plugins = {
#"poplar": {"device":"XLA_IPU", "types":"DT_FLOAT,DT_INT32", "tags":[]},
}