text: # 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.
  # ==============================================================================
  """Benchmarks for low-level eager execution primitives.
  
  Packaged as a test to ensure that this code is exercised by continuous
  integration tests. To get numbers:
  
    bazel build -c opt :benchmarks_test &&
    ./bazel-bin/tensorflow/python/eager/benchmarks_test --iters=0
  """
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import argparse
  import contextlib
  import sys
  import time
  
  import numpy as np
  from six.moves import xrange  # pylint: disable=redefined-builtin
  
  from tensorflow.python import pywrap_tensorflow
  from tensorflow.python.eager import backprop  # pylint: disable=unused-import
  from tensorflow.python.eager import context
  from tensorflow.python.eager import function
  from tensorflow.python.eager import test
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import gen_math_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import random_ops
  
  FLAGS = None
  
  
  @contextlib.contextmanager
  def timer(label, iters=30000):
    start = time.time()
    yield xrange(iters)
    end = time.time()
    t = (end - start) * 1e6 / iters
    print("%-40s took %.2fus (%d iterations)" % (label, t, iters))
  
  
  def benchmark_create_tensor(n):
    """Benchmark overheads of creating a Tensor object."""
  
    def label(s):
      return "{:20s}".format(s)
  
    with timer(label("np.array([[3.0]])"), iters=n) as iters:
      for _ in iters:
        np.array([[3.0]])
  
    ctx = context.context()
    handle = ctx._handle
    device = ctx.device_name
    # May be warmup GPU.
    ops.EagerTensor([[3.0]], context=handle, device=device)
  
    # float32
    dtype = dtypes.float32.as_datatype_enum
    three = [[3.0]]
    with timer(label("EagerTensor([[3.0]])"), iters=n) as iters:
      for _ in iters:
        ops.EagerTensor(three, context=handle, device=device, dtype=dtype)
  
    np_3 = np.array([[3.0]], dtype=np.float32)
    with timer(label("EagerTensor(np.array([[3.0]]))"), iters=n) as iters:
      for _ in iters:
        ops.EagerTensor(np_3, context=handle, device=device, dtype=dtype)
  
    # int32.
    # This is interesting since int32 will be kept on host memory for the GPU
    # case.
    dtype = dtypes.int32.as_datatype_enum
    three = [[3]]
    with timer(label("EagerTensor([[3]])"), iters=n) as iters:
      for _ in iters:
        ops.EagerTensor(three, context=handle, device=device, dtype=dtype)
  
    np_3 = np.array([[3]], dtype=np.int32)
    with timer(label("EagerTensor(np.array([[3]]))"), iters=n) as iters:
      for _ in iters:
        ops.EagerTensor(np_3, context=handle, device=device, dtype=dtype)
  
  
  def benchmark_matmul(shape, n, use_gpu=False):
    """Benchmark for matrix multiplication using tf.matmul."""
    transpose_b = (shape[0] != shape[1])
    m = random_ops.random_uniform(shape)
    if use_gpu:
      m = m.as_gpu_tensor()
      # Warm up the GPU - the very first kernel invocation
      # seems to require a bunch of setup.
      math_ops.matmul(m, m, transpose_b=transpose_b)
  
    def label(s):
      return "MatMul {}: {:30s}".format(shape, s)
  
    if not use_gpu:
      a = m.as_cpu_tensor().numpy()
      b = a.T if transpose_b else a
      with timer(label("np.dot"), iters=n) as iters:
        for _ in iters:
          np.dot(a, b)
  
    with timer(label("tf.matmul"), iters=n) as iters:
      for _ in iters:
        math_ops.matmul(m, m, transpose_b=transpose_b)
  
    with timer(label("gen_math_ops.mat_mul"), iters=n) as iters:
      for _ in iters:
        gen_math_ops._mat_mul(m, m, transpose_b=transpose_b)
  
    inputs = [m, m]
    # pylint: disable=protected-access
    ctx_handle = context.context()._handle
    # pylint: enable=protected-access
    attrs = ("transpose_a", False, "transpose_b", transpose_b, "T",
             m.dtype.as_datatype_enum)
    with timer(label("TFE_Py_Execute"), iters=n) as iters:
      for _ in iters:
        pywrap_tensorflow.TFE_Py_Execute(ctx_handle, None, "MatMul",
                                         inputs, attrs, 1)
  
    f = function.defun(math_ops.matmul)
    with timer(label("defun(tf.matmul)"), iters=n) as iters:
      for _ in iters:
        f(m, m, transpose_b=transpose_b)
  
  
  class BenchmarksTest(test_util.TensorFlowTestCase):
  
    def testBenchmarks(self):
      # This isn't actually a test, but benchmarks packaged as a test
      # so that continuous integration runs catch any breakages.
      print(context.context())
      benchmark_create_tensor(FLAGS.iters or 30000)
      benchmark_matmul([2, 2], FLAGS.iters or 30000)
      benchmark_matmul([100, 28 * 28], FLAGS.iters or 1000)
  
      if context.context().num_gpus() > 0:
        print("---- RUNNING ON GPU NOW ----")
        with context.device("/device:GPU:0"):
          benchmark_create_tensor(FLAGS.iters or 30000)
        benchmark_matmul([2, 2], FLAGS.iters or 30000, use_gpu=True)
        benchmark_matmul([100, 28 * 28], FLAGS.iters or 1000, use_gpu=True)
  
  
  if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # Default iterations to 1 to keep continuos integration test times low.
    parser.add_argument(
        "--iters",
        type=int,
        default=1,
        help="Number of iterators for each test. None or 0 for auto-selection")
    FLAGS, unparsed = parser.parse_known_args()
    sys.argv = [sys.argv[0]] + unparsed
    test.main()
  