# 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.
# ==============================================================================
"""Benchmark comparing eager, autograph, and official dynamic_rnn.

This code is tested on TF 1.13.
"""

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

import benchmark_base

import numpy as np
import tensorflow.compat.v1 as tf


tf.enable_eager_execution()


BATCH_SIZE = 32
MAX_SEQ_LEN = 100
FEATURE_SIZE = 50
HIDDEN_SIZE = 256


class RNNBenchmark(benchmark_base.ReportingBenchmark):
  """Runs benchmarks for eager/autograph/graph variants of dynamic_rnn."""

  def _generate_fake_rnn_inputs(self,
                                batch_size=BATCH_SIZE,
                                max_seq_len=MAX_SEQ_LEN):
    np.random.seed(17)

    input_data = np.random.random([batch_size, max_seq_len,
                                   FEATURE_SIZE]).astype(np.float32)
    # Generate some varying sequence lengths but keep max(sequence_lengths)
    # a constant, for more reproducible benchmarks.
    sequence_lengths = np.concatenate(([max_seq_len],
                                       np.random.randint(
                                           max_seq_len // 2,
                                           max_seq_len,
                                           size=[batch_size - 1]))).astype(
                                               np.int32)

    for i, seq_len in enumerate(sequence_lengths):
      input_data[i, seq_len:, :] = 0

    input_data = tf.constant(input_data)
    sequence_lengths = tf.constant(sequence_lengths)

    return input_data, sequence_lengths

  def _create_rnn_cell(self, batch_size=BATCH_SIZE):
    rnn_cell = tf.nn.rnn_cell.BasicRNNCell(HIDDEN_SIZE, dtype=tf.float32)
    rnn_cell.build(tf.TensorShape([batch_size, FEATURE_SIZE]))
    return rnn_cell, rnn_cell.zero_state(batch_size, dtype=tf.float32)

  def _benchmark_eager_dynamic_rnn(self, batch_size, max_seq_len):
    input_data, sequence_lengths = self._generate_fake_rnn_inputs(
        batch_size=batch_size, max_seq_len=max_seq_len)
    rnn_cell, initial_state = self._create_rnn_cell(batch_size=batch_size)

    def eager_dynamic_rnn(rnn_cell,
                          input_data,
                          initial_state,
                          sequence_length=None):
      """An eager version of dynamic_rnn."""
      # [batch, time, features] -> [time, batch, features]
      input_data = tf.transpose(input_data, [1, 0, 2])
      outputs = []
      state = initial_state
      if sequence_length is None:
        max_seq_len = input_data.shape[0]
      else:
        max_seq_len = tf.reduce_max(sequence_length)
      for i in range(max_seq_len):
        new_output, new_state = rnn_cell(input_data[i], state)
        output = tf.where(i < sequence_length, new_output,
                          tf.zeros(new_output.shape))
        state = tf.where(i < sequence_length, new_state, state)
        outputs.append(output)
      return tf.transpose(tf.stack(outputs), [1, 0, 2]), state

    def target():
      eager_dynamic_rnn(rnn_cell, input_data, initial_state, sequence_lengths)

    self.time_execution(
        ('Eager', batch_size, max_seq_len),
        target,
        iter_volume=batch_size,
        iter_unit='examples',
        extras={
            'max_seq_len': max_seq_len,
            'batch_size': batch_size,
        })

  def _benchmark_handwritten_dynamic_rnn(self, batch_size, max_seq_len):

    def my_dynamic_rnn(rnn_cell,
                       input_data,
                       initial_state,
                       sequence_length=None):
      """A handwritten reimplementation of dynamic_rnn."""
      input_data = tf.transpose(input_data, [1, 0, 2])
      outputs = tf.TensorArray(tf.float32, input_data.shape[0])
      if sequence_length is None:
        max_seq_len = input_data.shape[0]
      else:
        max_seq_len = tf.reduce_max(sequence_length)

      def while_body(i, state, outputs):
        new_output, new_state = rnn_cell(input_data[i], state)
        output = tf.where(i < sequence_length, new_output,
                          tf.zeros(new_output.shape))
        state = tf.where(i < sequence_length, new_state, state)
        outputs = outputs.write(i, output)
        return i + 1, state, outputs

      def while_cond(i, unused_state, unused_outputs):
        return i < max_seq_len

      _, state, outputs = tf.while_loop(
          while_cond,
          while_body,
          loop_vars=(tf.constant(0), initial_state, outputs))
      return tf.transpose(outputs.stack(), [1, 0, 2]), state

    with tf.Graph().as_default():
      input_data, sequence_lengths = self._generate_fake_rnn_inputs(
          batch_size=batch_size, max_seq_len=max_seq_len)
      rnn_cell, initial_state = self._create_rnn_cell(batch_size=batch_size)
      graph_output_t = my_dynamic_rnn(rnn_cell, input_data, initial_state,
                                      sequence_lengths)

      with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())

        def target():
          sess.run(graph_output_t)

        self.time_execution(
            ('Handwritten', batch_size, max_seq_len),
            target,
            iter_volume=batch_size,
            iter_unit='examples',
            extras={
                'max_seq_len': max_seq_len,
                'batch_size': batch_size,
            })

  def benchmark_dynamic_rnn(self):
    for batch_size in (32, 64, 128):
      for max_seq_len in (64, 128):
        self._benchmark_eager_dynamic_rnn(batch_size, max_seq_len)
        self._benchmark_handwritten_dynamic_rnn(batch_size, max_seq_len)
        self._benchmark_ag_dynamic_rnn(batch_size, max_seq_len)
        self._benchmark_official_dynamic_rnn(batch_size, max_seq_len)

  def _benchmark_ag_dynamic_rnn(self, batch_size, max_seq_len):

    def ag_dynamic_rnn(rnn_cell,
                       input_data,
                       initial_state,
                       sequence_length=None):
      """An autograph-able reimplementation of subset of dynamic_rnn."""
      # [batch, time, features] -> [time, batch, features]
      input_data = tf.transpose(input_data, [1, 0, 2])
      if sequence_length is None:
        max_seq_len = input_data.shape[0]
      else:
        max_seq_len = tf.reduce_max(sequence_length)

      outputs = tf.TensorArray(tf.float32, size=max_seq_len)
      state = initial_state
      for i in tf.range(max_seq_len):
        new_output, new_state = rnn_cell(input_data[i], state)
        output = tf.where(i < sequence_length, new_output,
                          tf.zeros(new_output.shape))
        state = tf.where(i < sequence_length, new_state, state)
        outputs = outputs.write(i, output)
      return tf.transpose(outputs.stack(), [1, 0, 2]), state

    ag_dynamic_rnn = tf.autograph.to_graph(ag_dynamic_rnn)

    with tf.Graph().as_default():
      input_data, sequence_lengths = self._generate_fake_rnn_inputs(
          batch_size=batch_size, max_seq_len=max_seq_len)
      rnn_cell, initial_state = self._create_rnn_cell(batch_size=batch_size)
      rnn_output = ag_dynamic_rnn(rnn_cell, input_data, initial_state,
                                  sequence_lengths)

      with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())

        def target():
          sess.run(rnn_output)

        self.time_execution(
            ('AutoGraph', batch_size, max_seq_len),
            target,
            iter_volume=batch_size,
            iter_unit='examples',
            extras={
                'max_seq_len': max_seq_len,
                'batch_size': batch_size,
            })

  def _benchmark_official_dynamic_rnn(self, batch_size, max_seq_len):
    with tf.Graph().as_default():
      input_data, sequence_lengths = self._generate_fake_rnn_inputs(
          batch_size=batch_size, max_seq_len=max_seq_len)
      rnn_cell, initial_state = self._create_rnn_cell(batch_size=batch_size)

      rnn_output = tf.nn.dynamic_rnn(
          rnn_cell,
          input_data,
          initial_state=initial_state,
          sequence_length=sequence_lengths)

      with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())

        def target():
          sess.run(rnn_output)

        self.time_execution(
            ('tf.nn.dynamic_rnn', batch_size, max_seq_len),
            target,
            iter_volume=batch_size,
            iter_unit='examples',
            extras={
                'max_seq_len': max_seq_len,
                'batch_size': batch_size,
            })


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