/* 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.
==============================================================================*/
#include "tensorflow/compiler/xla/service/hlo_evaluator.h"

#include <initializer_list>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include "tensorflow/compiler/xla/client/computation_builder.h"
#include "tensorflow/compiler/xla/literal_util.h"
#include "tensorflow/compiler/xla/reference_util.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/test.h"
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
#include "tensorflow/compiler/xla/tests/literal_test_util.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/platform/types.h"

namespace xla {
namespace {

class HloEvaluatorTest : public HloTestBase {
 protected:
  HloEvaluatorTest() { evaluator_ = MakeUnique<HloEvaluator>(); }

  std::unique_ptr<HloEvaluator> evaluator_;
};

// Verifies that HloEvaluator evaluates a HLO instruction that performs clamp
// with 3 operands.
TEST_F(HloEvaluatorTest, DoesClamp) {
  auto low = Literal::CreateR2<float>({{0.f, 2.f}, {2.f, 4.f}});
  auto high = Literal::CreateR2<float>({{2.f, 4.f}, {4.f, 4.f}});
  auto value = Literal::CreateR2<float>({{0.f, 5.f}, {0.f, 4.f}});

  Shape shape = low->shape();
  auto c1 = HloInstruction::CreateConstant(std::move(low));
  auto c2 = HloInstruction::CreateConstant(std::move(high));
  auto c3 = HloInstruction::CreateConstant(std::move(value));
  auto instruction = HloInstruction::CreateTernary(
      shape, HloOpcode::kClamp, c1.get(), c2.get(), c3.get());

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(instruction.get(), {}).ConsumeValueOrDie();

  auto expected = Literal::CreateR2<float>({{0, 4}, {2, 4}});

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

// Verifies that HloEvaluator evaluates a HLO instruction that performs select
// with 3 operands.
TEST_F(HloEvaluatorTest, DoesSelect) {
  auto pred = Literal::CreateR2<bool>({{true, false}, {false, true}});
  auto on_true = Literal::CreateR2<float>({{2.f, 4.f}, {4.f, 4.f}});
  auto on_false = Literal::CreateR2<float>({{0.f, 5.f}, {0.f, 4.f}});

  Shape shape = on_true->shape();
  auto c1 = HloInstruction::CreateConstant(std::move(pred));
  auto c2 = HloInstruction::CreateConstant(std::move(on_true));
  auto c3 = HloInstruction::CreateConstant(std::move(on_false));
  auto instruction = HloInstruction::CreateTernary(
      shape, HloOpcode::kSelect, c1.get(), c2.get(), c3.get());

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(instruction.get(), {}).ConsumeValueOrDie();

  auto expected = Literal::CreateR2<float>({{2, 5}, {0, 4}});

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

// Verifies that HloEvaluator evaluates a HLO instruction that performs
// element-wise addition with 2 operands.
TEST_F(HloEvaluatorTest, DoesAdd) {
  auto lhs = Literal::CreateR2<int64>({{1, 0}, {-100, 4}});
  auto rhs = Literal::CreateR2<int64>({{2, 4}, {4, 4}});

  Shape shape = ShapeUtil::MakeShape(S64, {2, 2});
  auto c1 = HloInstruction::CreateConstant(std::move(lhs));
  auto c2 = HloInstruction::CreateConstant(std::move(rhs));
  auto instruction =
      HloInstruction::CreateBinary(shape, HloOpcode::kAdd, c1.get(), c2.get());

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(instruction.get(), {}).ConsumeValueOrDie();

  auto expected = Literal::CreateR2<int64>({{3, 4}, {-96, 8}});

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

// Verifies that HloEvaluator evaluates a HLO instruction that performs
// element-wise divide with 2 operands.
TEST_F(HloEvaluatorTest, DoesDivide) {
  auto lhs_s64 = Literal::CreateR2<int64>({{1, 0}, {-100, 4}});
  auto rhs_s64 = Literal::CreateR2<int64>({{2, 4}, {4, 4}});

  Shape shape_s64 = ShapeUtil::MakeShape(S64, {2, 2});
  auto c1_s64 = HloInstruction::CreateConstant(std::move(lhs_s64));
  auto c2_s64 = HloInstruction::CreateConstant(std::move(rhs_s64));
  auto instruction = HloInstruction::CreateBinary(shape_s64, HloOpcode::kDivide,
                                                  c1_s64.get(), c2_s64.get());

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(instruction.get(), {}).ConsumeValueOrDie();

  auto expected = Literal::CreateR2<int64>({{0, 0}, {-25, 1}});

  LiteralTestUtil::ExpectEqual(*expected, *result);

  auto lhs_f64 = Literal::CreateR2<double>({{1.0, 0.0}, {-100.0, 4.0}});
  auto rhs_f64 = Literal::CreateR2<double>({{2.2, 4.0}, {4.0, 4.0}});

  Shape shape_f64 = ShapeUtil::MakeShape(F64, {2, 2});
  auto c1_f64 = HloInstruction::CreateConstant(std::move(lhs_f64));
  auto c2_f64 = HloInstruction::CreateConstant(std::move(rhs_f64));
  instruction = HloInstruction::CreateBinary(shape_f64, HloOpcode::kDivide,
                                             c1_f64.get(), c2_f64.get());

  result = evaluator_->Evaluate(instruction.get(), {}).ConsumeValueOrDie();

  expected = Literal::CreateR2<double>({{0.45454545454545453, 0}, {-25, 1}});

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

// Verifies that HloEvaluator evaluates a HLO instruction that performs
// element-wise abs op with 1 operand.
TEST_F(HloEvaluatorTest, DoesAbs) {
  auto operand = Literal::CreateR2<int64>({{1, -20}, {-100, 4}});
  const Shape& shape = ShapeUtil::MakeShape(S64, {2, 2});
  auto c1 = HloInstruction::CreateConstant(std::move(operand));
  auto instruction =
      HloInstruction::CreateUnary(shape, HloOpcode::kAbs, c1.get());

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(instruction.get(), {}).ConsumeValueOrDie();

  auto expected = Literal::CreateR2<int64>({{1, 20}, {100, 4}});

  LiteralTestUtil::ExpectEqual(*expected, *result);

  // For R0 literal.
  const Shape& r0 = ShapeUtil::MakeShape(F32, {});
  operand = Literal::CreateR0<float>(-1.0f);
  c1 = HloInstruction::CreateConstant(std::move(operand));
  instruction = HloInstruction::CreateUnary(r0, HloOpcode::kAbs, c1.get());
  result = evaluator_->Evaluate(instruction.get()).ConsumeValueOrDie();
  expected = Literal::CreateR0<float>(1.0f);

  LiteralTestUtil::ExpectEqual(*expected, *result);

  // For R1 literal with dimension of size 0.
  Shape empty_r1 = ShapeUtil::MakeShape(F32, {0});
  operand = Literal::CreateR1<float>({});
  c1 = HloInstruction::CreateConstant(std::move(operand));
  instruction =
      HloInstruction::CreateUnary(empty_r1, HloOpcode::kAbs, c1.get());

  result = evaluator_->Evaluate(instruction.get()).ConsumeValueOrDie();
  expected = Literal::CreateR1<float>({});

  LiteralTestUtil::ExpectEqual(*expected, *result);
}  // namespace

// Verifies that HloEvaluator evaluates a HLO Computation with non-parameter nor
// constant operands.
TEST_F(HloEvaluatorTest, DoesTraverseInstructions) {
  HloComputation::Builder builder(TestName());
  auto lhs = Literal::CreateR2<int64>({{1, 0}, {-100, 4}});
  auto rhs = Literal::CreateR2<int64>({{2, 4}, {4, 4}});
  auto rhs2 = Literal::CreateR2<int64>({{1, -20}, {-100, 4}});
  std::vector<const Literal*> args = {lhs.get(), rhs.get(), rhs2.get()};

  Shape shape = ShapeUtil::MakeShape(S64, {2, 2});

  auto param_lhs = HloInstruction::CreateParameter(0, shape, "lhs");
  auto param_rhs = HloInstruction::CreateParameter(1, shape, "rhs");
  auto lhs_instruction = HloInstruction::CreateBinary(
      shape, HloOpcode::kAdd, param_lhs.get(), param_rhs.get());

  auto param_rhs2 = HloInstruction::CreateParameter(2, shape, "rhs2");
  auto root_instruction = HloInstruction::CreateBinary(
      shape, HloOpcode::kAdd, lhs_instruction.get(), param_rhs2.get());

  builder.AddInstruction(std::move(root_instruction));
  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(builder.Build().get(), args).ConsumeValueOrDie();

  auto expected = Literal::CreateR2<int64>({{4, -16}, {-196, 12}});

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

// Verifies Reshape operation is correctly evaluated.
TEST_F(HloEvaluatorTest, DoesReshape) {
  HloComputation::Builder builder(TestName());
  const int64 dimensions[] = {11, 8, 7, 5, 9};
  TF_ASSIGN_OR_ASSERT_OK(auto literal,
                         LiteralTestUtil::CreateRandomLiteral<F32>(
                             ShapeUtil::MakeShape(F32, dimensions), 0.0, 1.0));
  auto literal_clone = literal->CloneToUnique();
  HloInstruction* literal_instruction = builder.AddInstruction(
      HloInstruction::CreateConstant(std::move(literal)));

  Shape shape = ShapeUtil::MakeShape(F32, {8, 7, 11, 9, 5});
  const int64 permutation[] = {1, 2, 0, 4, 3};
  builder.AddInstruction(
      HloInstruction::CreateTranspose(shape, literal_instruction, permutation));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(builder.Build().get(), {}).ConsumeValueOrDie();

  using NativeT = typename primitive_util::PrimitiveTypeToNative<F32>::type;
  result->EachCell<NativeT>(
      [&](tensorflow::gtl::ArraySlice<int64> indices, NativeT value) {
        std::vector<int64> rindexes = Permute(permutation, indices);
        EXPECT_TRUE(value == literal_clone->Get<NativeT>(rindexes));
      });
}

// Verifies Broadcast operation is correctly evaluated.
TEST_F(HloEvaluatorTest, DoesBroadcast) {
  HloComputation::Builder builder(TestName());
  auto input_literal = Literal::CreateR2<int32>({{1, 2}, {3, 4}, {5, 6}});
  auto output_literal = Literal::CreateR3<int32>(
      {{{1, 2}, {3, 4}, {5, 6}}, {{1, 2}, {3, 4}, {5, 6}}});
  HloInstruction* literal_instruction = builder.AddInstruction(
      HloInstruction::CreateConstant(std::move(input_literal)));

  builder.AddInstruction(HloInstruction::CreateBroadcast(
      output_literal->shape(), literal_instruction, {1, 2}));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(builder.Build().get(), {}).ConsumeValueOrDie();

  LiteralTestUtil::ExpectEqual(*result, *output_literal);
}

TEST_F(HloEvaluatorTest, ConvertWithSameLayout) {
  HloComputation::Builder builder(TestName());

  auto input_literal = Literal::CreateR2<int32>({{1, 2}, {3, 4}, {5, 6}});
  auto expected =
      Literal::CreateR2<float>({{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}});
  ASSERT_TRUE(LayoutUtil::LayoutsInShapesEqual(input_literal->shape(),
                                               expected->shape()));

  HloInstruction* constant = builder.AddInstruction(
      HloInstruction::CreateConstant(std::move(input_literal)));
  builder.AddInstruction(
      HloInstruction::CreateConvert(expected->shape(), constant));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(builder.Build().get(), {}).ConsumeValueOrDie();

  LiteralTestUtil::ExpectEqual(*result, *expected);
}

TEST_F(HloEvaluatorTest, ConvertWithDifferentLayout) {
  HloComputation::Builder builder(TestName());

  auto input_literal = Literal::CreateR2WithLayout<int32>(
      {{1, 2}, {3, 4}, {5, 6}}, LayoutUtil::MakeLayout({0, 1}));
  auto expected = Literal::CreateR2WithLayout<float>(
      {{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}}, LayoutUtil::MakeLayout({1, 0}));
  ASSERT_FALSE(LayoutUtil::LayoutsInShapesEqual(input_literal->shape(),
                                                expected->shape()));

  HloInstruction* constant = builder.AddInstruction(
      HloInstruction::CreateConstant(std::move(input_literal)));
  builder.AddInstruction(
      HloInstruction::CreateConvert(expected->shape(), constant));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(builder.Build().get(), {}).ConsumeValueOrDie();

  LiteralTestUtil::ExpectEqual(*result, *expected);
}

PaddingConfig CreatePaddingConfig(
    std::initializer_list<std::array<int64, 3>> padding_dimensions) {
  PaddingConfig padding_config;

  for (auto& paddings_per_dim : padding_dimensions) {
    auto dimension = padding_config.add_dimensions();
    dimension->set_edge_padding_low(paddings_per_dim[0]);
    dimension->set_edge_padding_high(paddings_per_dim[1]);
    dimension->set_interior_padding(paddings_per_dim[2]);
  }
  return padding_config;
}

TEST_F(HloEvaluatorTest, Pad2DIntegerArrayWithZeroDimension) {
  auto operand = Literal::CreateR2<int32>({{}, {}});
  auto operand_instruction = HloInstruction::CreateConstant(std::move(operand));

  constexpr int32 kPadValue = 10;
  auto pad_value = Literal::CreateR0<int32>(kPadValue);
  auto padding_value_instruction =
      HloInstruction::CreateConstant(std::move(pad_value));

  auto padding_config = CreatePaddingConfig({{{1, 0, 2}}, {{0, 2, 1}}});
  Shape shape = ShapeUtil::MakeShape(S32, {5, 2});
  auto pad_instruction = HloInstruction::CreatePad(
      shape, operand_instruction.get(), padding_value_instruction.get(),
      padding_config);

  auto result = evaluator_->Evaluate(pad_instruction.get()).ConsumeValueOrDie();

  auto expected = Literal::CreateR2<int32>(
      {{10, 10}, {10, 10}, {10, 10}, {10, 10}, {10, 10}});

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, Pad4DFloatArrayWithInteriorPadding) {
  HloComputation::Builder b(TestName());

  Array4D<float> input_array(3, 2, 1, 1, {1, 2, 3, 4, 5, 6});
  auto input = Literal::CreateR4FromArray4D<float>(input_array);
  HloInstruction* input_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(input)));
  constexpr float kPadValue = 1.5;
  auto pad_value = Literal::CreateR0<float>(kPadValue);
  HloInstruction* pad_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(pad_value)));

  Shape shape = ShapeUtil::MakeShape(F32, {8, 5, 1, 1});
  auto r4_padding_on_dim0_dim1 =
      CreatePaddingConfig({{{1, 0, 2}}, {{0, 2, 1}}, {{0, 0, 0}}, {{0, 0, 0}}});
  b.AddInstruction(HloInstruction::CreatePad(
      shape, input_instruction, pad_instruction, r4_padding_on_dim0_dim1));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  auto expected_array = MakeUnique<Array4D<float>>(8, 5, 1, 1);
  expected_array->Fill(kPadValue);
  (*expected_array)(1, 0, 0, 0) = 1.0f;
  (*expected_array)(1, 2, 0, 0) = 2.0f;
  (*expected_array)(4, 0, 0, 0) = 3.0f;
  (*expected_array)(4, 2, 0, 0) = 4.0f;
  (*expected_array)(7, 0, 0, 0) = 5.0f;
  (*expected_array)(7, 2, 0, 0) = 6.0f;

  auto expected = Literal::CreateR4FromArray4D<float>(*expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, NegativePadding2D) {
  HloComputation::Builder b(TestName());

  // input_array:
  // f32[4,3] {
  //  { 1, 2, 3 },
  //  { 5, 6, 7 },
  //  { 9, 10, 11 },
  //  { 13, 14, 15 },
  // }
  auto input_array = MakeUnique<Array2D<float>>(4, 3);
  input_array->FillUnique(1.0f);
  auto input = Literal::CreateR2FromArray2D<float>(*input_array);
  HloInstruction* input_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(input)));

  auto pad_value_instruction = b.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateR0<float>(2.718f)));

  auto r2_padding_on_dim0_dim1 =
      CreatePaddingConfig({{{-1, -2, 0}}, {{-2, 4, 0}}});
  Shape shape = ShapeUtil::MakeShape(F32, {1, 5});
  b.AddInstruction(HloInstruction::CreatePad(shape, input_instruction,
                                             pad_value_instruction,
                                             r2_padding_on_dim0_dim1));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  // f32[1,5] { 7.0, 2.718, 2.718, 2.718, 2.718 }
  auto expected_array = MakeUnique<Array2D<float>>(1, 5);
  (*expected_array)(0, 0) = 7.0f;
  (*expected_array)(0, 1) = 2.718f;
  (*expected_array)(0, 2) = 2.718f;
  (*expected_array)(0, 3) = 2.718f;
  (*expected_array)(0, 4) = 2.718f;
  auto expected = Literal::CreateR2FromArray2D<float>(*expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, NegativeAndInteriorPadding2D) {
  HloComputation::Builder b(TestName());

  // f32[4,3] {
  //  { 1, 2, 3 },
  //  { 5, 6, 7 },
  //  { 9, 10, 11 },
  //  { 13, 14, 15 },
  // }
  auto input_array = MakeUnique<Array2D<float>>(4, 3);
  input_array->FillUnique(1.0f);
  auto input = Literal::CreateR2FromArray2D<float>(*input_array);
  HloInstruction* input_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(input)));

  auto pad_value_instruction = b.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateR0<float>(2.718f)));

  PaddingConfig padding_config = MakeNoPaddingConfig(2);

  // Negative padding that results in zero dimensions.
  auto r2_padding_on_dim0_dim1 =
      CreatePaddingConfig({{{-2, -5, 1}}, {{-2, 4, 2}}});

  Shape shape = ShapeUtil::MakeShape(F32, {0, 9});
  b.AddInstruction(HloInstruction::CreatePad(shape, input_instruction,
                                             pad_value_instruction,
                                             r2_padding_on_dim0_dim1));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  auto expected_array = MakeUnique<Array2D<float>>(0, 9);
  auto expected = Literal::CreateR2FromArray2D<float>(*expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, DotRank2AndRank1) {
  HloComputation::Builder b(TestName());

  // lhs:
  // f32[4,1] {
  //  { 1 },
  //  { 2 },
  //  { 3 },
  //  { 4 },
  // }
  auto lhs_array = MakeUnique<Array2D<float>>(4, 1);
  lhs_array->FillUnique(1.0f);
  auto lhs_literal = Literal::CreateR2FromArray2D<float>(*lhs_array);
  HloInstruction* lhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal)));

  // rhs:
  // f32[2] { 1, 2 },
  auto rhs_literal = Literal::CreateR2<float>({{1, 2}});
  HloInstruction* rhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal)));

  Shape shape = ShapeUtil::MakeShape(F32, {4, 2});
  b.AddInstruction(HloInstruction::CreateBinary(
      shape, HloOpcode::kDot, lhs_instruction, rhs_instruction));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  // clang-format off
  auto expected_array = Array2D<float>({
      {1.f, 2.f},
      {2.f, 4.f},
      {3.f, 6.f},
      {4.f, 8.f},
  });
  // clang-format on
  auto expected = Literal::CreateR2FromArray2D<float>(expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, DotRank1AndRank2) {
  HloComputation::Builder b(TestName());

  // lhs:
  // f32[3]
  //  { 1, 2, 3 },
  auto lhs_literal = Literal::CreateR1<float>({1, 2, 3});
  HloInstruction* lhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal)));

  // rhs:
  // f32[3,2] {
  //  { 1, 2 },
  //  { 3, 4 },
  //  { 5, 6 },
  // }
  auto rhs_array = MakeUnique<Array2D<float>>(3, 2);
  rhs_array->FillUnique(1.0f);
  auto rhs_literal = Literal::CreateR2FromArray2D<float>(*rhs_array);
  HloInstruction* rhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal)));

  Shape shape = ShapeUtil::MakeShape(F32, {2});
  b.AddInstruction(HloInstruction::CreateBinary(
      shape, HloOpcode::kDot, lhs_instruction, rhs_instruction));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  auto expected = Literal::CreateR1<float>({22.f, 28.f});

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, DotRank2AndRank2) {
  HloComputation::Builder b(TestName());

  // lhs:
  // f32[4,3] {
  //  { 1, 2, 3 },
  //  { 5, 6, 7 },
  //  { 9, 10, 11 },
  //  { 13, 14, 15 },
  // }
  auto lhs_array = MakeUnique<Array2D<float>>(4, 3);
  lhs_array->FillUnique(1.0f);
  auto lhs_literal = Literal::CreateR2FromArray2D<float>(*lhs_array);
  HloInstruction* lhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal)));

  // rhs:
  // f32[3,2] {
  //  { 1, 2 },
  //  { 3, 4 },
  //  { 5, 6 },
  // }
  auto rhs_array = MakeUnique<Array2D<float>>(3, 2);
  rhs_array->FillUnique(1.0f);
  auto rhs_literal = Literal::CreateR2FromArray2D<float>(*rhs_array);
  HloInstruction* rhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal)));

  Shape shape = ShapeUtil::MakeShape(F32, {4, 2});
  b.AddInstruction(HloInstruction::CreateBinary(
      shape, HloOpcode::kDot, lhs_instruction, rhs_instruction));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  auto expected_array = Array2D<float>({
      {22.f, 28.f}, {58.f, 76.f}, {94.f, 124.f}, {130.f, 172.f},
  });
  auto expected = Literal::CreateR2FromArray2D<float>(expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, SimpleConv1D) {
  HloComputation::Builder b(TestName());

  Array3D<float> lhs_array = {{{1, 2, 3}}};
  auto lhs_literal = Literal::CreateR3FromArray3D<float>(lhs_array);
  HloInstruction* lhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal)));

  Array3D<float> rhs_array = {{{3.f, 4.f}}};
  auto rhs_literal = Literal::CreateR3FromArray3D<float>(rhs_array);
  HloInstruction* rhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal)));

  Window window;
  WindowDimension dim;
  dim.set_size(2);
  dim.set_stride(1);
  dim.set_padding_low(0);
  dim.set_padding_high(1);
  dim.set_window_dilation(1);
  dim.set_base_dilation(1);
  *window.add_dimensions() = dim;

  ConvolutionDimensionNumbers dnums;
  dnums.set_batch_dimension(0);
  dnums.set_feature_dimension(1);
  dnums.add_spatial_dimensions(2);

  dnums.set_kernel_output_feature_dimension(0);
  dnums.set_kernel_input_feature_dimension(1);
  dnums.add_kernel_spatial_dimensions(2);

  const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 3});
  b.AddInstruction(HloInstruction::CreateConvolve(
      shape, lhs_instruction, rhs_instruction, window, dnums));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  Array3D<float> expected_array = {{{11.f, 18.f, 9.f}}};
  auto expected = Literal::CreateR3FromArray3D<float>(expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, Simple4x4Conv2DWith2x2Kernel) {
  HloComputation::Builder b(TestName());

  Array4D<float> lhs_array(1, 1, 4, 4);
  // clang-format off
  lhs_array.FillWithYX(Array2D<float>({
    {1,  2,  3,  4 },
    {5,  6,  7,  8 },
    {9,  10, 11, 12},
    {13, 14, 15, 16},
  }));
  // clang-format on
  auto lhs_literal = Literal::CreateR4FromArray4D<float>(lhs_array);
  HloInstruction* lhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal)));

  Array4D<float> rhs_array(1, 1, 2, 2);
  // clang-format off
  rhs_array.FillWithYX(Array2D<float>({
    {5, 6},
    {7, 8},
  }));
  // clang-format on
  auto rhs_literal = Literal::CreateR4FromArray4D<float>(rhs_array);
  HloInstruction* rhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal)));

  Window window;
  WindowDimension dim;
  dim.set_size(2);
  dim.set_stride(1);
  dim.set_padding_low(0);
  dim.set_padding_high(1);
  dim.set_window_dilation(1);
  dim.set_base_dilation(1);
  *window.add_dimensions() = dim;
  *window.add_dimensions() = dim;

  ConvolutionDimensionNumbers dnums =
      ComputationBuilder::CreateDefaultConvDimensionNumbers(2);

  const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 4, 4});
  b.AddInstruction(HloInstruction::CreateConvolve(
      shape, lhs_instruction, rhs_instruction, window, dnums));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  Array4D<float> expected_array(1, 1, 4, 4);
  // clang-format off
  expected_array.FillWithYX(Array2D<float>({
    {100, 126, 152,  76},
    {204, 230, 256, 124},
    {308, 334, 360, 172},
    {149, 160, 171,  80},
  }));
  // clang-format on
  auto expected = Literal::CreateR4FromArray4D<float>(expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, Conv2DGeneralDimensions) {
  HloComputation::Builder b(TestName());

  // clang-format off
  // Input dimensions: [feature=2, height=3, batch=1, width=4]
  Array4D<float> input({
    {{{1, 2, 3, 4}},
     {{5, 6, 7, 8}},
     {{9, 10, 11, 12}}},
    {{{13, 14, 15, 16}},
     {{17, 18, 19, 20}},
     {{21, 22, 23, 24}}}
  });
  // Weight dimensions:
  // [kernel_output_feature=1, width=3, kernel_input_feature=2, height=3]
  Array4D<float> weight({{
    {{1, 7, 13},
     {4, 10, 16}},
    {{2, 8, 14},
     {5, 11, 17}},
    {{3, 9, 15},
     {6, 12, 18}}
  }});
  // clang-format on

  auto lhs_literal = Literal::CreateR4FromArray4D<float>(input);
  HloInstruction* lhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal)));

  auto rhs_literal = Literal::CreateR4FromArray4D<float>(weight);
  HloInstruction* rhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal)));

  Window window;
  WindowDimension dim;
  dim.set_size(3);
  dim.set_stride(1);
  dim.set_padding_low(0);
  dim.set_padding_high(0);
  dim.set_window_dilation(1);
  dim.set_base_dilation(1);
  *window.add_dimensions() = dim;
  *window.add_dimensions() = dim;

  ConvolutionDimensionNumbers dnums;
  dnums.set_batch_dimension(2);
  dnums.set_feature_dimension(0);
  dnums.add_spatial_dimensions(1);
  dnums.add_spatial_dimensions(3);

  dnums.set_kernel_output_feature_dimension(0);
  dnums.set_kernel_input_feature_dimension(2);
  dnums.add_kernel_spatial_dimensions(3);
  dnums.add_kernel_spatial_dimensions(1);

  const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 1, 2});
  b.AddInstruction(HloInstruction::CreateConvolve(
      shape, lhs_instruction, rhs_instruction, window, dnums));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  // clang-format off
  // Result dimensions: [feature=1, height=1, batch=1, width=2]
  Array4D<float> expected_array({{{{2514, 2685}}}});
  // clang-format on
  auto expected = Literal::CreateR4FromArray4D<float>(expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, DilatedBaseConv2DWithHighPadding) {
  HloComputation::Builder b(TestName());

  Array4D<float> lhs_array(1, 1, 4, 4);
  // clang-format off
  lhs_array.FillWithYX(Array2D<float>({
    {1,  2,  3,  4 },
    {5,  6,  7,  8 },
    {9,  10, 11, 12},
    {13, 14, 15, 16},
  }));
  // clang-format on
  auto lhs_literal = Literal::CreateR4FromArray4D<float>(lhs_array);
  HloInstruction* lhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal)));

  Array4D<float> rhs_array(1, 1, 2, 2);
  // clang-format off
  rhs_array.FillWithYX(Array2D<float>({
    {5, 6},
    {7, 8},
  }));
  // clang-format on
  auto rhs_literal = Literal::CreateR4FromArray4D<float>(rhs_array);
  HloInstruction* rhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal)));

  Window window;
  WindowDimension dim;
  dim.set_size(2);
  dim.set_stride(1);
  dim.set_padding_low(0);
  dim.set_padding_high(1);
  dim.set_window_dilation(1);
  dim.set_base_dilation(2);
  *window.add_dimensions() = dim;
  *window.add_dimensions() = dim;

  ConvolutionDimensionNumbers dnums =
      ComputationBuilder::CreateDefaultConvDimensionNumbers(2);

  const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 7, 7});
  b.AddInstruction(HloInstruction::CreateConvolve(
      shape, lhs_instruction, rhs_instruction, window, dnums));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  Array4D<float> expected_array(1, 1, 7, 7);
  expected_array.FillWithYX(Array2D<float>({
      {5, 12, 10, 18, 15, 24, 20},
      {35, 48, 42, 56, 49, 64, 56},
      {25, 36, 30, 42, 35, 48, 40},
      {63, 80, 70, 88, 77, 96, 84},
      {45, 60, 50, 66, 55, 72, 60},
      {91, 112, 98, 120, 105, 128, 112},
      {65, 84, 70, 90, 75, 96, 80},
  }));
  auto expected = Literal::CreateR4FromArray4D<float>(expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest, DilatedBaseConv2DWithLowAndHighPadding) {
  HloComputation::Builder b(TestName());

  Array4D<float> lhs_array(1, 1, 4, 4);
  // clang-format off
  lhs_array.FillWithYX(Array2D<float>({
    {1,  2,  3,  4 },
    {5,  6,  7,  8 },
    {9,  10, 11, 12},
    {13, 14, 15, 16},
  }));
  // clang-format on
  auto lhs_literal = Literal::CreateR4FromArray4D<float>(lhs_array);
  HloInstruction* lhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal)));

  Array4D<float> rhs_array(1, 1, 2, 2);
  // clang-format off
  rhs_array.FillWithYX(Array2D<float>({
    {5, 6},
    {7, 8},
  }));
  // clang-format on
  auto rhs_literal = Literal::CreateR4FromArray4D<float>(rhs_array);
  HloInstruction* rhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal)));

  Window window;
  WindowDimension dim;
  dim.set_size(2);
  dim.set_stride(1);
  dim.set_padding_low(1);
  dim.set_padding_high(1);
  dim.set_window_dilation(1);
  dim.set_base_dilation(2);
  *window.add_dimensions() = dim;
  *window.add_dimensions() = dim;

  ConvolutionDimensionNumbers dnums =
      ComputationBuilder::CreateDefaultConvDimensionNumbers(2);

  const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 8, 8});
  b.AddInstruction(HloInstruction::CreateConvolve(
      shape, lhs_instruction, rhs_instruction, window, dnums));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  Array4D<float> expected_array(1, 1, 8, 8);
  expected_array.FillWithYX(Array2D<float>({
      {8, 7, 16, 14, 24, 21, 32, 28},
      {6, 5, 12, 10, 18, 15, 24, 20},
      {40, 35, 48, 42, 56, 49, 64, 56},
      {30, 25, 36, 30, 42, 35, 48, 40},
      {72, 63, 80, 70, 88, 77, 96, 84},
      {54, 45, 60, 50, 66, 55, 72, 60},
      {104, 91, 112, 98, 120, 105, 128, 112},
      {78, 65, 84, 70, 90, 75, 96, 80},
  }));
  auto expected = Literal::CreateR4FromArray4D<float>(expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

TEST_F(HloEvaluatorTest,
       DilatedWindowAndBaseConv2DWithDifferentLowAndHighPaddingAndStrides) {
  HloComputation::Builder b(TestName());

  Array4D<float> lhs_array(1, 1, 4, 4);
  // clang-format off
  lhs_array.FillWithYX(Array2D<float>({
    {1,  2,  3,  4 },
    {5,  6,  7,  8 },
    {9,  10, 11, 12},
    {13, 14, 15, 16},
  }));
  // clang-format on
  auto lhs_literal = Literal::CreateR4FromArray4D<float>(lhs_array);
  HloInstruction* lhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal)));

  Array4D<float> rhs_array(1, 1, 2, 3);
  // clang-format off
  rhs_array.FillWithYX(Array2D<float>({
    {5, 6, 7},
    {8, 9, 10},
  }));
  // clang-format on
  auto rhs_literal = Literal::CreateR4FromArray4D<float>(rhs_array);
  HloInstruction* rhs_instruction =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal)));

  Window window;
  WindowDimension dim;
  dim.set_size(2);
  dim.set_stride(1);
  dim.set_padding_low(2);
  dim.set_padding_high(2);
  dim.set_window_dilation(2);
  dim.set_base_dilation(2);
  *window.add_dimensions() = dim;
  dim.set_size(3);
  dim.set_stride(3);
  dim.set_padding_low(2);
  dim.set_padding_high(-1);
  dim.set_window_dilation(1);
  dim.set_base_dilation(3);
  *window.add_dimensions() = dim;

  ConvolutionDimensionNumbers dnums =
      ComputationBuilder::CreateDefaultConvDimensionNumbers(2);

  const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 9, 3});
  b.AddInstruction(HloInstruction::CreateConvolve(
      shape, lhs_instruction, rhs_instruction, window, dnums));

  std::unique_ptr<Literal> result =
      evaluator_->Evaluate(b.Build().get(), {}).ConsumeValueOrDie();

  Array4D<float> expected_array(1, 1, 9, 3);
  expected_array.FillWithYX(Array2D<float>({
      {10, 20, 30},
      {0, 0, 0},
      {57, 74, 91},
      {0, 0, 0},
      {125, 142, 159},
      {0, 0, 0},
      {193, 210, 227},
      {0, 0, 0},
      {91, 98, 105},
  }));
  auto expected = Literal::CreateR4FromArray4D<float>(expected_array);

  LiteralTestUtil::ExpectEqual(*expected, *result);
}

}  // namespace
}  // namespace xla
