# 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.
# ==============================================================================
"""Tests for the behavior of the auto-compilation pass."""

import numpy as np

from tensorflow.compiler.tests import xla_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import googletest

CPU_DEVICE = "/job:localhost/replica:0/task:0/cpu:0"


class ClusteringTest(xla_test.XLATestCase):

  def testAdd(self):
    val1 = np.array([4, 3, 2, 1], dtype=np.float32)
    val2 = np.array([5, 6, 7, 8], dtype=np.float32)
    expected = val1 + val2
    with self.session():
      with self.test_scope():
        input1 = constant_op.constant(val1, name="const1")
        input2 = constant_op.constant(val2, name="const2")
        output = math_ops.add(input1, input2)
      result = self.evaluate(output)
    self.assertAllClose(result, expected, rtol=1e-3)

  def testAddFromCpuMultiple(self):
    val1 = np.array([4, 3, 2, 1]).astype(np.float32)
    val2 = np.array([5, 6, 7, 8]).astype(np.float32)
    expected = val1 + val2
    with self.session():
      with ops.device(CPU_DEVICE):
        input1 = constant_op.constant(val1, name="const1")
        input2 = constant_op.constant(val2, name="const2")
      with self.test_scope():
        output = math_ops.add(input1, input2)
      for _ in range(10):
        result = self.evaluate(output)
        self.assertAllClose(result, expected, rtol=1e-3)

  def testDeadlock(self):
    # Builds a graph of the form:
    #  x -> y
    #       | \
    #       z -> w
    # where x and z are placed on the CPU and y and w are placed on the XLA
    # device. If y and w are clustered for compilation, then the graph will
    # deadlock since the clustered graph will contain a self-loop.
    with self.session() as sess:
      with ops.device(CPU_DEVICE):
        x = array_ops.placeholder(dtypes.float32, [2])
      with self.test_scope():
        y = x * 2
      with ops.device(CPU_DEVICE):
        z = y * y
      with self.test_scope():
        w = y + z
      result = sess.run(w, {x: [1.5, 0.5]})
    self.assertAllClose(result, [12., 2.], rtol=1e-3)

  def testHostMemory(self):
    with self.session() as sess:
      x = array_ops.placeholder(dtypes.int32)
      with self.test_scope():
        y = x + 1
      with ops.device(CPU_DEVICE):
        # Place a computation on the CPU, so y and w cannot be merged into the
        # same JIT compilation.
        z = y * 2
      with self.test_scope():
        # Argument 'y' is a non-constant output of a previous cluster. Make sure
        # it is properly copied to host memory so it can be used as a
        # compile-time constant input for this cluster.
        w = array_ops.reshape(z, y)
      result = sess.run(w, {x: [1, 0]})
      expected = np.array([[4], [2]], dtype=np.int32)
      self.assertAllClose(expected, result, rtol=1e-3)


if __name__ == "__main__":
  googletest.main()