tensorflow/compiler/tests/xla_test.py - platform/external/tensorflow - Git at Google

 # 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.
 # ==============================================================================
 """Definition of XLA test case."""

 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function

 import contextlib
 import os
 import random
 import re

 import numpy as np

 from tensorflow.contrib.compiler import jit
 from tensorflow.core.framework import types_pb2
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.python.client import session
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import random_seed
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import flags
 from tensorflow.python.platform import test
 from tensorflow.python.platform import tf_logging as logging

 FLAGS = flags.FLAGS

 flags.DEFINE_string('test_device', None,
                     'Tensorflow device on which to place operators under test')
 flags.DEFINE_string('types', None, 'Types to test. Comma-separated list.')
 flags.DEFINE_string('disabled_manifest', None,
                     'Path to a file with a list of tests that should not run.')
 flags.DEFINE_string('tf_xla_flags', None,
                     'Value to set the TF_XLA_FLAGS environment variable to')


 def parse_disabled_manifest(manifest_content):
   comments_re = re.compile('#.*$')
   disabled_tests = []
   disabled_method_types = []
   for l in manifest_content.splitlines():
     stripped = comments_re.sub('', l).strip()
     if not stripped:
       continue
     entry = stripped.split(' ')
     if len(entry) == 1:
       disabled_tests.append(entry[0])
     elif len(entry) == 2:
       disabled_method_types.append((entry[0], entry[1].strip().split(',')))
     else:
       raise ValueError('Bad entry in manifest file.')

   disabled_regex = '|'.join(disabled_tests)
   method_types_filter = dict()
   for method, types in disabled_method_types:
     method_types_filter[method] = set([
         dtypes.as_dtype(types_pb2.DataType.Value(name)).as_numpy_dtype
         for name in types
     ])
   return disabled_regex, method_types_filter


 class XLATestCase(test.TestCase):
   """XLA test cases are parameterized test cases."""

   def __init__(self, method_name='runTest'):
     super(XLATestCase, self).__init__(method_name)
     self.device = FLAGS.test_device
     self.has_custom_call = (self.device == 'XLA_CPU')
     self._all_tf_types = set([
         dtypes.as_dtype(types_pb2.DataType.Value(name))
         for name in FLAGS.types.split(',')
     ])
     self.int_tf_types = set([
         dtype for dtype in self._all_tf_types if dtype.is_integer
     ])
     self._float_tf_types = set([
         dtype for dtype in self._all_tf_types if dtype.is_floating
     ])
     self.complex_tf_types = set([
         dtype for dtype in self._all_tf_types if dtype.is_complex
     ])
     self._numeric_tf_types = set(
         self.int_tf_types | self._float_tf_types | self.complex_tf_types)
     self.quantized_tf_types = set(
         dtype for dtype in self._all_tf_types if dtype.is_quantized)

     # Quantized types don't have a numpy equivalent, include them in
     # all_tf_types but not in all_types.
     # TODO(b/115960798): Parametrize tests on TF types instead of numpy types
     # and remove all_types.
     self._all_types = set(dtype.as_numpy_dtype
                           for dtype in self._all_tf_types
                           if not dtype.is_quantized)
     self._int_types = set([dtype.as_numpy_dtype for dtype in self.int_tf_types])
     self.signed_int_types = set(dtype.as_numpy_dtype
                                 for dtype in self.int_tf_types
                                 if not dtype.is_unsigned)
     self.unsigned_int_types = set(dtype.as_numpy_dtype
                                   for dtype in self.int_tf_types
                                   if dtype.is_unsigned)
     self._float_types = set(
         [dtype.as_numpy_dtype for dtype in self._float_tf_types])
     self.complex_types = set([
         dtype.as_numpy_dtype for dtype in self.complex_tf_types
     ])
     self._numeric_types = set(self._int_types | self._float_types
                               | self.complex_types)

     # Parse the manifest file, if any, into a regex identifying tests to
     # disable
     # TODO(xpan): Make it text proto if it doesn't scale.
     # Each line of the manifest file specifies an entry. The entry can be
     # 1) TestNameRegex  // E.g. CumprodTest.* Or
     # 2) TestName TypeName  // E.g. AdamOptimizerTest.testSharing DT_BFLOAT16
     # The 1) disables the entire test. While 2) only filter some numeric types
     # so that they are not used in those tests.
     self.disabled_regex = None
     self._method_types_filter = {}

     if FLAGS.disabled_manifest is not None:
       with open(FLAGS.disabled_manifest, 'r') as manifest_file:
         disabled_regex, self._method_types_filter = (
             parse_disabled_manifest(manifest_file.read()))
         if disabled_regex:
           self.disabled_regex = re.compile(disabled_regex)

     if FLAGS.tf_xla_flags is not None:
       os.environ['TF_XLA_FLAGS'] = FLAGS.tf_xla_flags

   @property
   def all_tf_types(self):
     name = '{}.{}'.format(type(self).__name__, self._testMethodName)
     tf_types = set([dtypes.as_dtype(t)
                     for t in self._method_types_filter.get(name, set())])
     return self._all_tf_types - tf_types

   @property
   def float_types(self):
     name = '{}.{}'.format(type(self).__name__, self._testMethodName)
     return self._float_types - self._method_types_filter.get(name, set())

   @property
   def float_tf_types(self):
     name = '{}.{}'.format(type(self).__name__, self._testMethodName)
     return self._float_tf_types - self._method_types_filter.get(name, set())

   @property
   def int_types(self):
     name = '{}.{}'.format(type(self).__name__, self._testMethodName)
     return self._int_types - self._method_types_filter.get(name, set())

   @property
   def numeric_tf_types(self):
     name = '{}.{}'.format(type(self).__name__, self._testMethodName)
     tf_types = set([dtypes.as_dtype(t)
                     for t in self._method_types_filter.get(name, set())])
     return self._numeric_tf_types - tf_types

   @property
   def numeric_types(self):
     name = '{}.{}'.format(type(self).__name__, self._testMethodName)
     return self._numeric_types - self._method_types_filter.get(name, set())

   @property
   def all_types(self):
     name = '{}.{}'.format(type(self).__name__, self._testMethodName)
     return self._all_types - self._method_types_filter.get(name, set())

   def setUp(self):
     super(XLATestCase, self).setUp()
     name = '{}.{}'.format(type(self).__name__, self._testMethodName)
     if self.disabled_regex is not None and self.disabled_regex.match(name):
       logging.info('Disabled test case: %s', name)
       self.skipTest('{} is disabled by manifest.'.format(name))
       return
     logging.info('Start test case: %s', name)

     random.seed(random_seed.DEFAULT_GRAPH_SEED)
     np.random.seed(random_seed.DEFAULT_GRAPH_SEED)

   def tearDown(self):
     super(XLATestCase, self).tearDown()
     logging.info('End test case: %s', self._testMethodName)

   @contextlib.contextmanager
   def test_session(self):
     """Custom implementation of test_session() for XLA tests.

     We override the standard Tensorflow test_session() since it is too
     specific to CPU and GPU tests. In particular, we want to disable soft
     placement and explicitly assign ops to devices under test.

     Yields:
       A session to use when running a test case.
     """
     graph = ops.Graph()
     with session.Session(graph=graph) as sess, graph.as_default():
       yield sess

   @contextlib.contextmanager
   def test_scope(self):
     """Test scope that runs tests on a Tensorflow/XLA device.

     Uses a compilation_scope() to mark operators to compile.

     Yields:
       A scope to apply to the operators under test.
     """
     with ops.device('device:{}:0'.format(self.device)):
       yield


 def Benchmark(tf_bench,
               builder_fn,
               use_xla_jit,
               device,
               separate_compiled_gradients=False):
   """Build a graph and run benchmarks against it, with or without XLA.

   Args:
     tf_bench: An instance of tf.test.Benchmark, used to run the benchmark.
     builder_fn: A function that builds a graph when invoked, and returns
         (name, fetches), where name is the name of the test, and fetches
         is a list of tensors to fetch as output.
     use_xla_jit: If true compile with the XLA JIT, otherwise use regular TF.
     device: The tensorflow device to run on, e.g. "cpu", "gpu".
     separate_compiled_gradients: If true put each gradient subgraph into a
       separate compilation scope. This gives fine-grained control over which
       portions of the graph will be compiled as a single unit. Compiling
       gradients separately may yield better performance for some graphs.
       The scope is named based on the scope of the forward computation as well
       as the name of the gradients. As a result, the gradients will be compiled
       in a scope that is separate from both the forward computation, and from
       other gradients.
   """

   with ops.Graph().as_default():
     name = None
     targets = []
     with ops.device(device):
       fetches = []
       jit_scope = jit.experimental_jit_scope
       with jit_scope(
           compile_ops=use_xla_jit,
           separate_compiled_gradients=separate_compiled_gradients):
         name, fetches = builder_fn()

       # We only want to benchmark the operations themselves, and not the data
       # transfer of the result(s).  Non-compiled identity ops ensure XLA
       # doesn't know we're dropping the results, otherwise it might compile
       # away the entire computation.
       for fetch in fetches:
         targets.append(array_ops.identity(fetch).op)

     config = config_pb2.ConfigProto(allow_soft_placement=True)
     with session.Session(config=config) as sess:
       sess.run(variables.global_variables_initializer())
       xla = 'xla_' if use_xla_jit else ''
       tf_bench.run_op_benchmark(
           sess, targets, name='%s_%s%s' % (name, xla, device))
	# 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.
	# ==============================================================================
	"""Definition of XLA test case."""

	from __future__ import absolute_import
	from __future__ import division
	from __future__ import print_function

	import contextlib
	import os
	import random
	import re

	import numpy as np

	from tensorflow.contrib.compiler import jit
	from tensorflow.core.framework import types_pb2
	from tensorflow.core.protobuf import config_pb2
	from tensorflow.python.client import session
	from tensorflow.python.framework import dtypes
	from tensorflow.python.framework import ops
	from tensorflow.python.framework import random_seed
	from tensorflow.python.ops import array_ops
	from tensorflow.python.ops import variables
	from tensorflow.python.platform import flags
	from tensorflow.python.platform import test
	from tensorflow.python.platform import tf_logging as logging

	FLAGS = flags.FLAGS

	flags.DEFINE_string('test_device', None,
	'Tensorflow device on which to place operators under test')
	flags.DEFINE_string('types', None, 'Types to test. Comma-separated list.')
	flags.DEFINE_string('disabled_manifest', None,
	'Path to a file with a list of tests that should not run.')
	flags.DEFINE_string('tf_xla_flags', None,
	'Value to set the TF_XLA_FLAGS environment variable to')


	def parse_disabled_manifest(manifest_content):
	comments_re = re.compile('#.*$')
	disabled_tests = []
	disabled_method_types = []
	for l in manifest_content.splitlines():
	stripped = comments_re.sub('', l).strip()
	if not stripped:
	continue
	entry = stripped.split(' ')
	if len(entry) == 1:
	disabled_tests.append(entry[0])
	elif len(entry) == 2:
	disabled_method_types.append((entry[0], entry[1].strip().split(',')))
	else:
	raise ValueError('Bad entry in manifest file.')

	disabled_regex = '\|'.join(disabled_tests)
	method_types_filter = dict()
	for method, types in disabled_method_types:
	method_types_filter[method] = set([
	dtypes.as_dtype(types_pb2.DataType.Value(name)).as_numpy_dtype
	for name in types
	])
	return disabled_regex, method_types_filter


	class XLATestCase(test.TestCase):
	"""XLA test cases are parameterized test cases."""

	def __init__(self, method_name='runTest'):
	super(XLATestCase, self).__init__(method_name)
	self.device = FLAGS.test_device
	self.has_custom_call = (self.device == 'XLA_CPU')
	self._all_tf_types = set([
	dtypes.as_dtype(types_pb2.DataType.Value(name))
	for name in FLAGS.types.split(',')
	])
	self.int_tf_types = set([
	dtype for dtype in self._all_tf_types if dtype.is_integer
	])
	self._float_tf_types = set([
	dtype for dtype in self._all_tf_types if dtype.is_floating
	])
	self.complex_tf_types = set([
	dtype for dtype in self._all_tf_types if dtype.is_complex
	])
	self._numeric_tf_types = set(
	self.int_tf_types \| self._float_tf_types \| self.complex_tf_types)
	self.quantized_tf_types = set(
	dtype for dtype in self._all_tf_types if dtype.is_quantized)

	# Quantized types don't have a numpy equivalent, include them in
	# all_tf_types but not in all_types.
	# TODO(b/115960798): Parametrize tests on TF types instead of numpy types
	# and remove all_types.
	self._all_types = set(dtype.as_numpy_dtype
	for dtype in self._all_tf_types
	if not dtype.is_quantized)
	self._int_types = set([dtype.as_numpy_dtype for dtype in self.int_tf_types])
	self.signed_int_types = set(dtype.as_numpy_dtype
	for dtype in self.int_tf_types
	if not dtype.is_unsigned)
	self.unsigned_int_types = set(dtype.as_numpy_dtype
	for dtype in self.int_tf_types
	if dtype.is_unsigned)
	self._float_types = set(
	[dtype.as_numpy_dtype for dtype in self._float_tf_types])
	self.complex_types = set([
	dtype.as_numpy_dtype for dtype in self.complex_tf_types
	])
	self._numeric_types = set(self._int_types \| self._float_types
	\| self.complex_types)

	# Parse the manifest file, if any, into a regex identifying tests to
	# disable
	# TODO(xpan): Make it text proto if it doesn't scale.
	# Each line of the manifest file specifies an entry. The entry can be
	# 1) TestNameRegex // E.g. CumprodTest.* Or
	# 2) TestName TypeName // E.g. AdamOptimizerTest.testSharing DT_BFLOAT16
	# The 1) disables the entire test. While 2) only filter some numeric types
	# so that they are not used in those tests.
	self.disabled_regex = None
	self._method_types_filter = {}

	if FLAGS.disabled_manifest is not None:
	with open(FLAGS.disabled_manifest, 'r') as manifest_file:
	disabled_regex, self._method_types_filter = (
	parse_disabled_manifest(manifest_file.read()))
	if disabled_regex:
	self.disabled_regex = re.compile(disabled_regex)

	if FLAGS.tf_xla_flags is not None:
	os.environ['TF_XLA_FLAGS'] = FLAGS.tf_xla_flags

	@property
	def all_tf_types(self):
	name = '{}.{}'.format(type(self).__name__, self._testMethodName)
	tf_types = set([dtypes.as_dtype(t)
	for t in self._method_types_filter.get(name, set())])
	return self._all_tf_types - tf_types

	@property
	def float_types(self):
	name = '{}.{}'.format(type(self).__name__, self._testMethodName)
	return self._float_types - self._method_types_filter.get(name, set())

	@property
	def float_tf_types(self):
	name = '{}.{}'.format(type(self).__name__, self._testMethodName)
	return self._float_tf_types - self._method_types_filter.get(name, set())

	@property
	def int_types(self):
	name = '{}.{}'.format(type(self).__name__, self._testMethodName)
	return self._int_types - self._method_types_filter.get(name, set())

	@property
	def numeric_tf_types(self):
	name = '{}.{}'.format(type(self).__name__, self._testMethodName)
	tf_types = set([dtypes.as_dtype(t)
	for t in self._method_types_filter.get(name, set())])
	return self._numeric_tf_types - tf_types

	@property
	def numeric_types(self):
	name = '{}.{}'.format(type(self).__name__, self._testMethodName)
	return self._numeric_types - self._method_types_filter.get(name, set())

	@property
	def all_types(self):
	name = '{}.{}'.format(type(self).__name__, self._testMethodName)
	return self._all_types - self._method_types_filter.get(name, set())

	def setUp(self):
	super(XLATestCase, self).setUp()
	name = '{}.{}'.format(type(self).__name__, self._testMethodName)
	if self.disabled_regex is not None and self.disabled_regex.match(name):
	logging.info('Disabled test case: %s', name)
	self.skipTest('{} is disabled by manifest.'.format(name))
	return
	logging.info('Start test case: %s', name)

	random.seed(random_seed.DEFAULT_GRAPH_SEED)
	np.random.seed(random_seed.DEFAULT_GRAPH_SEED)

	def tearDown(self):
	super(XLATestCase, self).tearDown()
	logging.info('End test case: %s', self._testMethodName)

	@contextlib.contextmanager
	def test_session(self):
	"""Custom implementation of test_session() for XLA tests.

	We override the standard Tensorflow test_session() since it is too
	specific to CPU and GPU tests. In particular, we want to disable soft
	placement and explicitly assign ops to devices under test.

	Yields:
	A session to use when running a test case.
	"""
	graph = ops.Graph()
	with session.Session(graph=graph) as sess, graph.as_default():
	yield sess

	@contextlib.contextmanager
	def test_scope(self):
	"""Test scope that runs tests on a Tensorflow/XLA device.

	Uses a compilation_scope() to mark operators to compile.

	Yields:
	A scope to apply to the operators under test.
	"""
	with ops.device('device:{}:0'.format(self.device)):
	yield


	def Benchmark(tf_bench,
	builder_fn,
	use_xla_jit,
	device,
	separate_compiled_gradients=False):
	"""Build a graph and run benchmarks against it, with or without XLA.

	Args:
	tf_bench: An instance of tf.test.Benchmark, used to run the benchmark.
	builder_fn: A function that builds a graph when invoked, and returns
	(name, fetches), where name is the name of the test, and fetches
	is a list of tensors to fetch as output.
	use_xla_jit: If true compile with the XLA JIT, otherwise use regular TF.
	device: The tensorflow device to run on, e.g. "cpu", "gpu".
	separate_compiled_gradients: If true put each gradient subgraph into a
	separate compilation scope. This gives fine-grained control over which
	portions of the graph will be compiled as a single unit. Compiling
	gradients separately may yield better performance for some graphs.
	The scope is named based on the scope of the forward computation as well
	as the name of the gradients. As a result, the gradients will be compiled
	in a scope that is separate from both the forward computation, and from
	other gradients.
	"""

	with ops.Graph().as_default():
	name = None
	targets = []
	with ops.device(device):
	fetches = []
	jit_scope = jit.experimental_jit_scope
	with jit_scope(
	compile_ops=use_xla_jit,
	separate_compiled_gradients=separate_compiled_gradients):
	name, fetches = builder_fn()

	# We only want to benchmark the operations themselves, and not the data
	# transfer of the result(s). Non-compiled identity ops ensure XLA
	# doesn't know we're dropping the results, otherwise it might compile
	# away the entire computation.
	for fetch in fetches:
	targets.append(array_ops.identity(fetch).op)

	config = config_pb2.ConfigProto(allow_soft_placement=True)
	with session.Session(config=config) as sess:
	sess.run(variables.global_variables_initializer())
	xla = 'xla_' if use_xla_jit else ''
	tf_bench.run_op_benchmark(
	sess, targets, name='%s_%s%s' % (name, xla, device))