/* Copyright 2015 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.
==============================================================================*/

#include <functional>
#include <memory>
#include <vector>

#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/test_benchmark.h"

namespace tensorflow {
namespace {

class QuantizeAndDequantizeTest : public OpsTestBase {};

// Convert a simple scalar tensor.
TEST_F(QuantizeAndDequantizeTest, Convert_scalar_tensor) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV2")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Attr("signed_input", true)
          .Attr("num_bits", 8)
          .Attr("range_given", false)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({1}), {-3.5});
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Max

  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({1}));
  test::FillValues<float>(&expected, {-3.5});
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));

  // Ensure that the inputs haven't been changed.
  EXPECT_EQ(inputs_[1]->scalar<float>()(), 0.0);
  EXPECT_EQ(inputs_[2]->scalar<float>()(), 0.0);
}

TEST_F(QuantizeAndDequantizeTest, Convert_scalar_tensor_V3) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV3")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_INT32))
          .Attr("signed_input", true)
          .Attr("range_given", false)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({1}), {-3.5});
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Max
  AddInputFromArray<int32>(TensorShape({}), {8});    // num_bits

  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({1}));
  test::FillValues<float>(&expected, {-3.5});
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));

  // Ensure that the inputs haven't been changed.
  EXPECT_EQ(inputs_[1]->scalar<float>()(), 0.0);
  EXPECT_EQ(inputs_[2]->scalar<float>()(), 0.0);
}

// Convert a 1D tensor with signed 8 bits.
TEST_F(QuantizeAndDequantizeTest, Convert_1D_tensor_with_int8) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV2")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Attr("signed_input", true)
          .Attr("num_bits", 8)
          .Attr("range_given", false)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({6}), {-1, -0.5, 0, 0.3, 0.8, 0.555});
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Max

  // With int8, the tensor is quantized to {-128, -64, 0, 38, 102, 71}.
  // Scale is: 1/127
  // Then it is dequantized to {-1, -0.5, 0, 38.0/128, 102.0/128, 71.0/128}
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({6}));
  test::FillValues<float>(&expected,
                          {-1, -0.5, 0, 38.0 / 128, 102.0 / 128, 71.0 / 128});
  test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-5);

  // Ensure that the inputs haven't been changed.
  EXPECT_EQ(inputs_[1]->scalar<float>()(), 0.0);
  EXPECT_EQ(inputs_[2]->scalar<float>()(), 0.0);
}

// Convert a 1D tensor with signed 8 bits.
TEST_F(QuantizeAndDequantizeTest, Convert_1D_tensor_with_int8_V3) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV3")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_INT32))
          .Attr("signed_input", true)
          .Attr("range_given", false)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({6}), {-1, -0.5, 0, 0.3, 0.8, 0.555});
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Max
  AddInputFromArray<int32>(TensorShape({}), {8});    // num_bits

  // With int8, the tensor is quantized to {-128, -64, 0, 38, 102, 71}.
  // Scale is: 1/128
  // Then it is dequantized to {-1, -64.0/128, 0, 38.0/128, 102.0/128, 71.0/128}
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({6}));
  test::FillValues<float>(&expected,
                          {-1, -0.5, 0, 38.0 / 128, 102.0 / 128, 71.0 / 128});
  test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-5);

  // Ensure that the inputs haven't been changed.
  EXPECT_EQ(inputs_[1]->scalar<float>()(), 0.0);
  EXPECT_EQ(inputs_[2]->scalar<float>()(), 0.0);
}

// Convert a 1D tensor with signed 4 bits.
TEST_F(QuantizeAndDequantizeTest, Convert_1D_tensor_with_int4) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV2")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Attr("signed_input", true)
          .Attr("num_bits", 4)
          .Attr("range_given", false)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({6}), {-1, -0.5, 0, 0.3, 0.8, 0.555});
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Max

  // With int4, the tensor is quantized to {-8, -4, 0, 2, 6, 4}.
  // Scale is: 1/8
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({6}));
  test::FillValues<float>(&expected, {-1, -0.5, 0, 0.25, 0.75, 0.5});
  test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-5);

  // Ensure that the inputs haven't been changed.
  EXPECT_EQ(inputs_[1]->scalar<float>()(), 0.0);
  EXPECT_EQ(inputs_[2]->scalar<float>()(), 0.0);
}

// Convert a 1D tensor with signed 4 bits.
TEST_F(QuantizeAndDequantizeTest, Convert_1D_tensor_with_int4_V3) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV3")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_INT32))
          .Attr("signed_input", true)
          .Attr("range_given", false)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({6}), {-1, -0.5, 0, 0.3, 0.8, 0.555});
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Max
  AddInputFromArray<int32>(TensorShape({}), {4});    // num_bits

  // With int4, the tensor is quantized to {-8, -4, 0, 2, 6, 4}.
  // Scale is: 1/8
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({6}));
  test::FillValues<float>(&expected, {-1, -0.5, 0, 0.25, 0.75, 0.5});
  test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-5);

  // Ensure that the inputs haven't been changed.
  EXPECT_EQ(inputs_[1]->scalar<float>()(), 0.0);
  EXPECT_EQ(inputs_[2]->scalar<float>()(), 0.0);
}

// Convert a 2D tensor with signed 8 bits with given range.
TEST_F(QuantizeAndDequantizeTest, Convert_2D_tensor_with_int8_range_given) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV2")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Attr("signed_input", true)
          .Attr("num_bits", 8)
          .Attr("range_given", true)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  // Note that the last two values are saturated.
  AddInputFromArray<float>(TensorShape({2, 4}),
                           {-0.8, -0.5, 0, 0.3, 0.8, 0.555, -2, 33});
  AddInputFromArray<float>(TensorShape({}), {-1.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {1.0});   // Max

  // Note that the range is given as [-1, 1].
  // With int8, the tensor is quantized to {-102, -64, 0, 38, 102, 70, -128,
  // 127}.
  // Scale is: 1/127
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 4}));
  test::FillValues<float>(
      &expected, {-102.0 / 127, -64.0 / 127, 0, 38.0 / 127, 102.0 / 127,
                  70.0 / 127, -128.0 / 127, 1});
  test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-5);
}

// Convert a 2D tensor with signed 8 bits with given range.
TEST_F(QuantizeAndDequantizeTest, Convert_2D_tensor_with_int8_range_given_V3) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV3")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_INT32))
          .Attr("signed_input", true)
          .Attr("range_given", true)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  // Note that the last two values are saturated.
  AddInputFromArray<float>(TensorShape({2, 4}),
                           {-0.8, -0.5, 0, 0.3, 0.8, 0.555, -2, 33});
  AddInputFromArray<float>(TensorShape({}), {-1.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {1.0});   // Max
  AddInputFromArray<int32>(TensorShape({}), {8});     // num_bits

  // Note that the range is given as [-1, 1].
  // With int8, the tensor is quantized to {-102, -64, 0, 38, 102, 70, -128,
  // 127}.
  // Scale is: 1/127
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 4}));
  test::FillValues<float>(
      &expected, {-102.0 / 127, -64.0 / 127, 0, 38.0 / 127, 102.0 / 127,
                  70.0 / 127, -128.0 / 127, 1});
  test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-5);
}

// Convert a 4D tensor with unsigned 8 bits with given range.
TEST_F(QuantizeAndDequantizeTest, Convert_4D_tensor_with_uint8_range_given) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV2")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Attr("signed_input", false)
          .Attr("num_bits", 8)
          .Attr("range_given", true)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({2, 2, 1, 1}), {-0.5, 0, 0.3, 0.8});
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {1.0});  // Max

  // Note that the range is given as [0, 1].
  // With int8, the tensor is quantized to {0, 0, 76, 204}
  // Scale is: 1/255
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 2, 1, 1}));
  test::FillValues<float>(&expected, {0, 0, 76.0 / 255, 204.0 / 255});
  test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-5);
}

// Convert a 4D tensor with unsigned 8 bits with given range.
TEST_F(QuantizeAndDequantizeTest, Convert_4D_tensor_with_uint8_range_given_V3) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV3")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_INT32))
          .Attr("signed_input", false)
          .Attr("range_given", true)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({2, 2, 1, 1}), {-0.5, 0, 0.3, 0.8});
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {1.0});  // Max
  AddInputFromArray<int32>(TensorShape({}), {8});    // num_bits

  // Note that the range is given as [0, 1].
  // With int8, the tensor is quantized to {0, 0, 76, 204}
  // Scale is: 1/255
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 2, 1, 1}));
  test::FillValues<float>(&expected, {0, 0, 76.0 / 255, 204.0 / 255});
  test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-5);
}

// Convert a tensor with all 0.
TEST_F(QuantizeAndDequantizeTest, Convert_tensor_with_all_0) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV2")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Attr("signed_input", false)
          .Attr("num_bits", 8)
          .Attr("range_given", false)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({2, 2, 1, 1}), {0, 0, 0, 0});
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Max

  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 2, 1, 1}));
  test::FillValues<float>(&expected, {0, 0, 0, 0});
  test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-5);
}

// Convert a tensor with all 0.
TEST_F(QuantizeAndDequantizeTest, Convert_tensor_with_all_0_V3) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_op", "QuantizeAndDequantizeV3")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_INT32))
          .Attr("signed_input", false)
          .Attr("range_given", false)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({2, 2, 1, 1}), {0, 0, 0, 0});
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Max
  AddInputFromArray<int32>(TensorShape({}), {8});    // num_bits

  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 2, 1, 1}));
  test::FillValues<float>(&expected, {0, 0, 0, 0});
  test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-5);
}

// Range is invalid
TEST_F(QuantizeAndDequantizeTest, Invalid_range_given) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_Op", "QuantizeAndDequantizeV2")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Attr("num_bits", 8)
          .Attr("range_given", true)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({2, 2, 1, 1}), {-0.5, 0, 0.3, 0.8});
  AddInputFromArray<float>(TensorShape({}), {1.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Max

  Status s = RunOpKernel();
  EXPECT_TRUE(str_util::StrContains(s.ToString(),
                                    "Invalid range: input_min 1 > input_max 0"))
      << s;
}

// Range is invalid
TEST_F(QuantizeAndDequantizeTest, Invalid_range_given_V3) {
  TF_ASSERT_OK(
      NodeDefBuilder("quantize_and_dequantize_Op", "QuantizeAndDequantizeV3")
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_FLOAT))
          .Input(FakeInput(DT_INT32))
          .Attr("range_given", true)
          .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  AddInputFromArray<float>(TensorShape({2, 2, 1, 1}), {-0.5, 0, 0.3, 0.8});
  AddInputFromArray<float>(TensorShape({}), {1.0});  // Min
  AddInputFromArray<float>(TensorShape({}), {0.0});  // Max
  AddInputFromArray<int32>(TensorShape({}), {8});    // num_bits

  Status s = RunOpKernel();
  EXPECT_TRUE(str_util::StrContains(s.ToString(),
                                    "Invalid range: input_min 1 > input_max 0"))
      << s;
}

#define BM_SIMPLE_QUAN_DEQUAN(DEVICE)                     \
  static void BM_SIMPLE_QUAN_DEQUAN_##DEVICE(int iters) { \
    auto root = Scope::NewRootScope().ExitOnError();      \
    ops::QuantizeAndDequantizeV2(root, -3.5, -3.5, -3.5); \
    TF_CHECK_OK(root.status());                           \
    Graph* g = new Graph(OpRegistry::Global());           \
    TF_CHECK_OK(root.ToGraph(g));                         \
    test::Benchmark(#DEVICE, g).Run(iters);               \
  }                                                       \
  BENCHMARK(BM_SIMPLE_QUAN_DEQUAN_##DEVICE);

BM_SIMPLE_QUAN_DEQUAN(cpu);
BM_SIMPLE_QUAN_DEQUAN(gpu);

}  // namespace
}  // namespace tensorflow