mirror of
https://github.com/zebrajr/tensorflow.git
synced 2025-12-07 12:20:24 +01:00
[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:
parent
aff4d124b2
commit
91cb809bd6
|
|
@ -1,12 +1,14 @@
|
|||
"""Build rules for Tensorflow/XLA testing."""
|
||||
|
||||
load("@local_config_cuda//cuda:build_defs.bzl", "cuda_is_configured")
|
||||
load("//tensorflow/compiler/tests:plugin.bzl", "plugins")
|
||||
|
||||
def all_backends():
|
||||
b = ["cpu"] + plugins.keys()
|
||||
if cuda_is_configured():
|
||||
return ["cpu", "gpu"]
|
||||
return b + ["gpu"]
|
||||
else:
|
||||
return ["cpu"]
|
||||
return b
|
||||
|
||||
def tf_xla_py_test(name, srcs=[], deps=[], tags=[], data=[], main=None,
|
||||
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",
|
||||
"--types=DT_FLOAT,DT_DOUBLE,DT_INT32,DT_INT64,DT_BOOL"]
|
||||
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:
|
||||
fail("Unknown backend {}".format(backend))
|
||||
|
||||
|
|
|
|||
23
tensorflow/compiler/tests/plugin.bzl
Normal file
23
tensorflow/compiler/tests/plugin.bzl
Normal 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":[]},
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user