# Copyright 2019 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.
# ==============================================================================
"""Tests for device placement."""

from absl.testing import parameterized

from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import remote
from tensorflow.python.eager import test
from tensorflow.python.framework import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops


class SoftDevicePlacementTest(test.TestCase, parameterized.TestCase):

  def setUp(self):
    super(SoftDevicePlacementTest, self).setUp()
    context._reset_context()
    context.ensure_initialized()
    config.set_soft_device_placement(enabled=True)
    context.context().log_device_placement = True

  @test_util.run_gpu_only
  def testDefaultPlacement(self):
    a = constant_op.constant(1)
    b = constant_op.constant(2)
    c = a + b
    with ops.device('CPU'):
      d = a + b
    self.assertIn('GPU', c.device)
    self.assertIn('CPU', d.device)

  @test_util.run_gpu_only
  def testUnsupportedDevice(self):
    a = constant_op.constant(1)
    b = constant_op.constant(2)
    s = constant_op.constant(list('hello world'))
    with ops.device('GPU:0'):
      c = a + b
      t = s[a]
    self.assertIn('GPU:0', c.device)
    self.assertIn('CPU', t.device)

  @test_util.run_gpu_only
  def testUnknownDevice(self):
    a = constant_op.constant(1)
    b = constant_op.constant(2)
    with ops.device('GPU:42'):
      c = a + b
    self.assertIn('GPU:0', c.device)

  def testNoGpu(self):
    if test_util.is_gpu_available():
      # CPU only test.
      return
    a = constant_op.constant(1)
    b = constant_op.constant(2)
    c = a + b
    with ops.device('GPU'):
      d = a + b
    self.assertIn('CPU', c.device)
    self.assertIn('CPU', d.device)

  @test_util.run_gpu_only
  def testSoftPlacedGPU(self):
    a = constant_op.constant(1)
    b = constant_op.constant(2)
    with ops.device('GPU:110'):
      c = a + b
    self.assertIn('GPU:0', c.device)

  @test_util.run_gpu_only
  def testNestedDeviceScope(self):
    a = constant_op.constant(1)
    b = constant_op.constant(2)
    with ops.device('CPU:0'):
      with ops.device('GPU:42'):
        c = a + b
    # We don't support nested device placement right now.
    self.assertIn('GPU:0', c.device)

  @parameterized.named_parameters(('float', 1.0, None),
                                  ('int32', [1], dtypes.int32),
                                  ('string', ['a'], None))
  def testSoftPlacedCPUConstant(self, value, dtype):
    if test_util.is_gpu_available():
      self.skipTest('CPU only test')
    with ops.device('GPU:0'):
      a = constant_op.constant(value, dtype=dtype)
    self.assertIn('CPU:0', a.device)
    self.assertIn('CPU:0', a.backing_device)

  @parameterized.named_parameters(
      ('float', [1.0, 2.0], None, 'GPU:0'),
      ## TODO(b/179035075) enable when int32 numeric tensors are not confused
      ## with shape tensors
      # ('int32', [1, 2], dtypes.int32, 'GPU:0'),
      ('int64', [1, 2], dtypes.int64, 'GPU:0'))
  @test_util.run_gpu_only
  def testSoftPlacedNumericTensors(self, value, dtype, expect):
    with ops.device('GPU:0'):
      a = math_ops.add(constant_op.constant(value, dtype=dtype),
                       constant_op.constant(value, dtype=dtype))
    self.assertIn(expect, a.backing_device)

  @test_util.run_gpu_only
  def testSoftPlacedShapeTensor(self):
    with ops.device('GPU:0'):
      t = constant_op.constant([[1, 2], [3, 4]])
      a = math_ops.add(array_ops.shape(t),
                       constant_op.constant([10, 20], dtype=dtypes.int32))
    # Shape computations remain on CPU
    self.assertIn('CPU:0', a.backing_device)

  def testPlacedToDeviceInFunction(self):

    @def_function.function
    def f():
      a = random_ops.random_uniform([32, 32])
      return math_ops.matmul(a, a)

    gpus = config.list_physical_devices('GPU')
    if not gpus:
      self.assertIn('CPU:0', f().device)
    else:
      self.assertIn('GPU:0', f().device)

  @test_util.disable_tfrt('b/173726713: Support properly inserting device at '
                          'tf_to_corert lowering.')
  def testUnknownDeviceInFunction(self):

    @def_function.function
    def f():
      with ops.device('GPU:42'):
        # With placer, the unknown GPU:42 will be replaced with GPU:0.
        a = constant_op.constant(1) + constant_op.constant(2)
      return a + constant_op.constant(2)

    gpus = config.list_physical_devices('GPU')
    if not gpus:
      self.assertIn('CPU:0', f().device)
    else:
      self.assertIn('GPU:0', f().device)


class HardDevicePlacementTest(test.TestCase, parameterized.TestCase):

  def setUp(self):
    super(HardDevicePlacementTest, self).setUp()
    context._reset_context()
    config.set_soft_device_placement(enabled=False)
    context.context().log_device_placement = True
    cpus = context.context().list_physical_devices('CPU')
    # Set 2 virtual CPUs
    context.context().set_logical_device_configuration(cpus[0], [
        context.LogicalDeviceConfiguration(),
        context.LogicalDeviceConfiguration()
    ])
    self.assertEqual(config.get_soft_device_placement(), False)
    self.assertEqual(context.context().soft_device_placement, False)

  @test_util.run_gpu_only
  def testIdentityCanCopy(self):
    config.set_device_policy('explicit')
    with ops.device('CPU:0'):
      x = constant_op.constant(1.0)
      self.assertIn('CPU:0', x.device)
      self.assertIn('CPU:0', x.backing_device)
    with ops.device('GPU:0'):
      y = array_ops.identity(x)
      self.assertIn('GPU:0', y.device)
      self.assertIn('GPU:0', y.backing_device)

  @test_util.run_gpu_only
  def testSimpleConstantsExplicitGPU(self):
    config.set_device_policy('explicit')
    with ops.device('GPU:0'):
      self.assertAllClose(1., array_ops.ones([]))
      self.assertAllClose(0., array_ops.zeros([]))
      self.assertAllClose([1.], array_ops.fill([1], 1.))

  def testSimpleConstantsExplicitCPU(self):
    config.set_device_policy('explicit')
    with ops.device('CPU:1'):
      self.assertAllClose(1., array_ops.ones([]))
      self.assertAllClose(0., array_ops.zeros([]))
      self.assertAllClose([1.], array_ops.fill([1], 1.))
      self.assertAllClose(2., constant_op.constant(1.) * 2.)

  @parameterized.named_parameters(
      ('float_cpu0', 'CPU:0', 1.0, None),
      ('int32_cpu0', 'CPU:0', [1], dtypes.int32),
      ('string_cpu0', 'CPU:0', ['a'], None),
      ('float_cpu1', 'CPU:1', 1.0, None),
      ('int32_cpu1', 'CPU:1', [1], dtypes.int32),
      ('string_cpu1', 'CPU:1', ['a'], None),
  )
  def testHardPlacedCPUConstant(self, device, value, dtype):
    with ops.device(device):
      a = constant_op.constant(value, dtype=dtype)
      self.assertIn(device, a.device)

  @parameterized.named_parameters(
      ('float', 'GPU:0', 1.0, None),
      ('int32', 'GPU:0', [1], dtypes.int32),
      ('string', 'GPU:0', ['a'], None),
  )
  def testHardPlacedGPUConstant(self, device, value, dtype):
    if not test_util.is_gpu_available():
      self.skipTest('Test requires a GPU')
    with ops.device(device):
      a = constant_op.constant(value, dtype=dtype)
      self.assertIn(device, a.device)
      if a.dtype == dtypes.float32:
        self.assertIn(device, a.backing_device)


class ClusterPlacementTest(test.TestCase):

  def setUp(self):
    super(ClusterPlacementTest, self).setUp()
    context._reset_context()
    config.set_soft_device_placement(enabled=True)
    context.context().log_device_placement = True
    workers, _ = test_util.create_local_cluster(2, 0)
    remote.connect_to_remote_host([workers[0].target, workers[1].target])

  @test_util.disable_tfrt('remote host not supported yet.')
  def testNotFullySpecifiedTask(self):
    a = constant_op.constant(1)
    b = constant_op.constant(2)
    with ops.device('/job:worker'):
      c = a + b
    self.assertIn('/job:worker/replica:0/task:0', c.device)

  @test_util.disable_tfrt('remote host not supported yet.')
  def testRemoteUnknownDevice(self):
    a = constant_op.constant(1)
    b = constant_op.constant(2)
    # Right now we don't support soft device place on remote worker.
    with self.assertRaises(errors.InvalidArgumentError) as cm:
      with ops.device('/job:worker/replica:0/task:0/device:GPU:42'):
        c = a + b
        del c
      self.assertIn('unknown device', cm.exception.message)

  @test_util.disable_tfrt('remote host not supported yet.')
  def testUnknownDeviceInFunctionReturnUnknowDevice(self):

    @def_function.function
    def f():
      with ops.device('GPU:42'):
        return constant_op.constant(1) + constant_op.constant(2)

    gpus = config.list_physical_devices('GPU')
    if not gpus:
      self.assertIn('CPU:0', f().device)
    else:
      self.assertIn('GPU:0', f().device)

  @test_util.disable_tfrt('remote host not supported yet.')
  def testUnknownDeviceInFunction(self):

    @def_function.function
    def f():
      with ops.device('GPU:42'):
        a = constant_op.constant(1) + constant_op.constant(2)
      return a + constant_op.constant(2)

    gpus = config.list_physical_devices('GPU')
    if not gpus:
      self.assertIn('CPU:0', f().device)
    else:
      self.assertIn('GPU:0', f().device)


if __name__ == '__main__':
  test.main()