test/cpp/tensorexpr/test_llvm.cpp

#ifdef TORCH_ENABLE_LLVM
#include <gtest/gtest.h>

#include <test/cpp/tensorexpr/test_base.h>

#include <c10/util/irange.h>
#include <test/cpp/tensorexpr/padded_buffer.h>
#include <test/cpp/tensorexpr/test_utils.h>
#include <torch/csrc/jit/tensorexpr/eval.h>
#include <torch/csrc/jit/tensorexpr/ir.h>
#include <torch/csrc/jit/tensorexpr/ir_printer.h>
#include <torch/csrc/jit/tensorexpr/ir_simplifier.h>
#include <torch/csrc/jit/tensorexpr/llvm_codegen.h>
#include <torch/csrc/jit/tensorexpr/loopnest.h>
#include <torch/csrc/jit/tensorexpr/tensor.h>
#include <torch/csrc/jit/testing/file_check.h>

#include <cmath>
#include <numeric>

namespace torch {
namespace jit {
using namespace torch::jit::tensorexpr;

using LLVMExprEval = ExprEval<LLVMCodeGen>;

// Typed tests, can't use gtest params here due to the way we instantiate tests.
#define TEST_LLVM_SCALAR_TYPES(_) \
  _(uint8_t, Byte, 24)            \
  _(int8_t, Char, -20)            \
  _(int16_t, Short, 3332)         \
  _(int, Int, 123456)             \
  _(int64_t, Long, 2631563121321) \
  _(float, Float, 0.122)          \
  _(double, Double, 0.21312)      \
  _(at::Half, Half, 0.128f)

#define IMM_TEST(Type, Name, Val)                  \
  TEST(LLVM, Name##ImmTest) {                      \
    auto a = Name##Imm::make(Val);                 \
    LLVMExprEval cg(a);                            \
    if (std::is_floating_point<decltype(Val)>()) { \
      ASSERT_NEAR(cg.value<Type>(), Val, 0.1);     \
    } else {                                       \
      ASSERT_EQ(cg.value<Type>(), Val);            \
    }                                              \
  }
TEST_LLVM_SCALAR_TYPES(IMM_TEST)
#undef IMM_TEST

#define ADD_TEST(Type, Name, Val)                  \
  TEST(LLVM, Name##AddTest) {                      \
    auto a = Name##Imm::make(Val);                 \
    auto b = Name##Imm::make(Val * 2);             \
    auto c = Add::make(a, b);                      \
    LLVMExprEval cg(c);                            \
    if (std::is_floating_point<decltype(Val)>()) { \
      ASSERT_NEAR(cg.value<Type>(), Val * 3, 0.1); \
    } else {                                       \
      ASSERT_EQ(cg.value<Type>(), Val * 3);        \
    }                                              \
  }
TEST_LLVM_SCALAR_TYPES(ADD_TEST)
#undef ADD_TEST

#define SUB_TEST(Type, Name, Val)                  \
  TEST(LLVM, Name##SubTest) {                      \
    auto a = Name##Imm::make(Val * 2);             \
    auto b = Name##Imm::make(Val);                 \
    auto c = Sub::make(a, b);                      \
    LLVMExprEval cg(c);                            \
    if (std::is_floating_point<decltype(Val)>()) { \
      ASSERT_NEAR(cg.value<Type>(), Val, 0.1);     \
    } else {                                       \
      ASSERT_EQ(cg.value<Type>(), Val);            \
    }                                              \
  }
TEST_LLVM_SCALAR_TYPES(SUB_TEST)
#undef SUB_TEST

#define MUL_TEST(Type, Name, Val)                  \
  TEST(LLVM, Name##MulTest) {                      \
    auto a = Name##Imm::make(Val);                 \
    auto b = Name##Imm::make((Type)4);             \
    auto c = Mul::make(a, b);                      \
    LLVMExprEval cg(c);                            \
    if (std::is_floating_point<decltype(Val)>()) { \
      ASSERT_NEAR(cg.value<Type>(), Val * 4, 0.1); \
    } else {                                       \
      ASSERT_EQ(cg.value<Type>(), Val * 4);        \
    }                                              \
  }
TEST_LLVM_SCALAR_TYPES(MUL_TEST)
#undef MUL_TEST

#define DIV_TEST(Type, Name, Val)                  \
  TEST(LLVM, Name##DivTest) {                      \
    auto a = Name##Imm::make((Type)6);             \
    auto b = Name##Imm::make((Type)3);             \
    auto c = Div::make(a, b);                      \
    LLVMExprEval cg(c);                            \
    if (std::is_floating_point<decltype(Val)>()) { \
      ASSERT_NEAR(cg.value<Type>(), 2, 0.1);       \
    } else {                                       \
      ASSERT_EQ(cg.value<Type>(), 2);              \
    }                                              \
  }
TEST_LLVM_SCALAR_TYPES(DIV_TEST)
#undef DIV_TEST

TEST(LLVM, IntToFloatCastTest) {
  auto a = IntImm::make(2);
  auto b = Cast::make(kFloat, a);
  LLVMExprEval cg(b, {});
  ASSERT_EQ(cg.value<float>(), 2.0);
}

TEST(LLVM, FloatToIntCastTest) {
  auto a = FloatImm::make(2.0);
  auto b = Cast::make(kInt, a);
  LLVMExprEval cg(b);
  ASSERT_EQ(cg.value<int>(), 2);
}

TEST(LLVM, IntToLongCastTest) {
  auto a = IntImm::make(12345);
  auto b = Cast::make(kLong, a);
  LLVMExprEval cg(b);
  ASSERT_EQ(cg.value<int64_t>(), 12345);
}

TEST(LLVM, ByteToCharCastTest) {
  auto a = ByteImm::make(250);
  auto b = Cast::make(kChar, a);
  LLVMExprEval cg(b);
  ASSERT_EQ(cg.value<int8_t>(), (int8_t)250);
}

TEST(LLVM, HalfToLongCastTest) {
  auto a = HalfImm::make(2.0);
  auto b = Cast::make(kLong, a);
  LLVMExprEval cg(b);
  ASSERT_EQ(cg.value<int64_t>(), 2);
}

TEST(LLVM, ByteToDoubleCastTest) {
  auto a = ByteImm::make(2);
  auto b = Cast::make(kDouble, a);
  LLVMExprEval cg(b);
  ASSERT_EQ(cg.value<double>(), 2);
}

TEST(LLVM, FloatToByteCastTest) {
  auto a = FloatImm::make(254.0);
  auto b = Cast::make(kByte, a);
  LLVMExprEval cg(b);
  ASSERT_EQ(cg.value<uint8_t>(), 254);
}

TEST(LLVM, FloatToCharCastTest) {
  auto a = FloatImm::make(-2.0);
  auto b = Cast::make(kChar, a);
  LLVMExprEval cg(b);
  ASSERT_EQ(cg.value<int8_t>(), -2);
}

TEST(LLVM, ByteToFloatCastTest) {
  auto a = ByteImm::make(254);
  auto b = Cast::make(kFloat, a);
  LLVMExprEval cg(b);
  ASSERT_EQ(cg.value<float>(), 254.0);
}

TEST(LLVM, CharToFloatCastTest) {
  auto a = CharImm::make(-2);
  auto b = Cast::make(kFloat, a);
  LLVMExprEval cg(b);
  ASSERT_EQ(cg.value<float>(), -2.0);
}

TEST(LLVM, BitCast) {
  /* constexpr int16_t ref16 = 1337; */
  constexpr int32_t ref32 = 1337;
  constexpr int64_t ref64 = 1337;
  constexpr float reff32 = 1337.0f;
  constexpr double reff64 = 1337.0f;

  // this is broken
  /*{
    at::Half k_;
    at::Half* k = &k_;
    *reinterpret_cast<int16_t*>(k) = ref16;
    auto a = HalfImm::make(k);
    auto b = BitCast::make(kShort, a);
    LLVMExprEval cg(b);
    ASSERT_EQ(cg.value<int16_t>(), ref16);
  }*/

  {
    float k = raw_bitcast<float>(ref32);
    auto a = FloatImm::make(k);
    auto b = BitCast::make(kInt, a);
    LLVMExprEval cg(b);
    ASSERT_EQ(cg.value<int32_t>(), ref32);
  }

  {
    double k = raw_bitcast<double>(ref64);
    auto a = DoubleImm::make(k);
    auto b = BitCast::make(kLong, a);
    LLVMExprEval cg(b);
    ASSERT_EQ(cg.value<int64_t>(), ref64);
  }

  {
    int64_t k = raw_bitcast<int64_t>(reff64);
    auto a = LongImm::make(k);
    auto b = BitCast::make(kDouble, a);
    LLVMExprEval cg(b);
    ASSERT_EQ(cg.value<double>(), reff64);
  }

  {
    int32_t k = raw_bitcast<int32_t>(reff32);
    auto a = IntImm::make(k);
    auto b = BitCast::make(kFloat, a);
    LLVMExprEval cg(b);
    ASSERT_EQ(cg.value<float>(), reff32);
  }
}

TEST(LLVM, fastLogFloat) {
  const int kTotalSize = 128 * 128;
  BufHandle a_buf("A", {ExprHandle(kTotalSize)}, kFloat);
  BufHandle b_buf("B", {ExprHandle(kTotalSize)}, kFloat);

  VarHandle index = VarHandle("index", kInt);
  ExprHandle load_a = a_buf.load(index);
  StmtPtr store_b = b_buf.store({index}, fast_log(load_a));
  StmtPtr stmt = For::make(index, 0, kTotalSize, store_b);

  PaddedBuffer<float> a_v(kTotalSize);
  PaddedBuffer<float> b_v(kTotalSize);

  for (const auto i : c10::irange(kTotalSize)) {
    a_v(i) = at::randn({1}).item().to<float>();
  }

  LLVMCodeGen ir_eval(stmt, {a_buf, b_buf});
  ir_eval.call({a_v, b_v});

  for (const auto i : c10::irange(kTotalSize)) {
    auto test = b_v(i);
    auto ref = std::log(a_v(i));
    if (std::isnan(ref)) {
      ASSERT_EQ(std::isnan(test), true);
    } else {
      ASSERT_FLOAT_EQ(test, ref);
    }
  }
}

TEST(LLVM, LetTest01) {
  BufHandle a("A", {1}, kFloat);
  std::vector<float> v = {1, 0};
  std::vector<void*> args({v.data()});
  VarHandle x("x", kFloat);
  auto block = Block::make({
      Let::make(x, 3.f),
      a.store({0}, ExprHandle(2.f) + (x * ExprHandle(3.f) + ExprHandle(4.f))),
  });

  LLVMCodeGen cg(block, {a});
  ASSERT_EQ(cg.value<int>(args), 0);
  ASSERT_EQ(v[0], 2.f + 3.f * 3.f + 4.f);
}

TEST(LLVM, LetTest02) {
  BufHandle a("A", {1}, kFloat);
  std::vector<float> v = {1, 0};
  std::vector<void*> args({v.data()});
  VarHandle x("x", kFloat);
  VarHandle y("y", kFloat);
  auto block = Block::make(
      {Let::make(x, 3.f),
       Let::make(y, 6.f),
       a.store(
           {IntImm::make(0)},
           ExprHandle(2.f) + (x * ExprHandle(3.f) + y * ExprHandle(4.f)))});

  LLVMCodeGen cg(block, {a});
  ASSERT_EQ(cg.value<int>(args), 0);
  ASSERT_EQ(v[0], 2.f + 3.f * 3.f + 6.f * 4.f);
}

TEST(LLVM, LetTestMultitype) {
  BufHandle a("A", {1}, kDouble);
  std::vector<double> v = {1, 0};
  std::vector<void*> args({v.data()});
  VarHandle x("x", kByte);
  VarHandle y("y", kHalf);
  auto block = Block::make(
      {Let::make(x, 3),
       Let::make(y, 6.f),
       a.store(
           {0},
           Cast::make(
               kDouble,
               ExprHandle(2.f) +
                   (x * ExprHandle(3.f) + y * ExprHandle(4.f))))});

  LLVMCodeGen cg(block, {a});
  ASSERT_EQ(cg.value<int>(args), 0);
  ASSERT_EQ(v[0], 2.f + 3 * 3.f + 6.f * 4.f);
}

TEST(LLVM, BufferTest) {
  BufHandle a("A", {32}, kFloat);
  std::vector<int32_t> v(5);
  std::vector<void*> args({v.data()});
  auto rv = IntImm::make(0);
  LLVMExprEval cg(rv, {a});
  ASSERT_EQ(cg.value<int>(args), 0);
}

TEST(LLVM, BlockTest) {
  BufHandle a("A", {32}, kInt);
  std::vector<int32_t> v = {1, 2};
  std::vector<void*> args({v.data()});

  auto block = Block::make({
      a.store({0}, 3),
      a.store({1}, 4),
      a.store({0}, 4),
  });

  LLVMCodeGen cg(block, {a});
  ASSERT_EQ(cg.value<int>(args), 0);
  ASSERT_EQ(v[0], 4);
  ASSERT_EQ(v[1], 4);
}

TEST(LLVM, LoadStoreTest) {
  BufHandle a("A", {1}, kInt);
  BufHandle b("B", {1}, kInt);
  std::vector<int32_t> a_buffer = {42};
  std::vector<int32_t> b_buffer = {-11};

  auto store = b.store({0}, a.load(0));
  LLVMCodeGen cg(store, {a, b});
  std::vector<void*> args({a_buffer.data(), b_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);
  ASSERT_EQ(a_buffer[0], 42);
  ASSERT_EQ(b_buffer[0], 42);
}

TEST(LLVM, IfThenElseTest) {
  BufHandle a("A", {1}, kInt);
  BufHandle b("B", {1}, kInt);
  BufHandle c("C", {1}, kInt);
  std::vector<int32_t> a_buffer = {42};
  std::vector<int32_t> b_buffer = {-11};
  std::vector<int32_t> c_buffer = {1};

  auto store = b.store({0}, IfThenElse::make(c.load(0), a.load(0), 0));
  LLVMCodeGen cg(store, {a, b, c});
  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);
  ASSERT_EQ(a_buffer[0], 42);
  ASSERT_EQ(b_buffer[0], 42);
}

// if (x < 10) x = x + 1
TEST(LLVM, CondNoFalseBlockTest) {
  BufHandle x("X", {1}, kInt);
  auto cmp = CompareSelect::make(x.load(0), 10, CompareSelectOperation::kLT);
  auto cond = Cond::make(cmp, x.store({0}, x.load(0) + 1), nullptr);

  for (int32_t x_value : {0, 10, 20}) {
    std::vector<int32_t> x_buffer = {x_value};
    std::vector<void*> args({x_buffer.data()});
    LLVMCodeGen cg(cond, {x});
    ASSERT_EQ(cg.value<int>(args), 0);
    if (x_value < 10) {
      ASSERT_EQ(x_buffer[0], x_value + 1);
    } else {
      ASSERT_EQ(x_buffer[0], x_value);
    }
  }
}

// if (x < 10) {
//   x = x + 1;
// } else {
//   x = x - 1;
// }
TEST(LLVM, CondTest) {
  BufHandle x("X", {1}, kInt);
  auto cmp = CompareSelect::make(x.load(0), 10, CompareSelectOperation::kLT);
  auto cond =
      Cond::make(cmp, x.store({0}, x.load(0) + 1), x.store({0}, x.load(0) - 1));
  auto block = Block::make({
      cond,
      x.store({0}, x.load(0) * 2),
  });

  for (int32_t x_value : {0, 10, 20}) {
    std::vector<int32_t> x_buffer = {x_value};
    std::vector<void*> args({x_buffer.data()});
    LLVMCodeGen cg(block, {x});
    ASSERT_EQ(cg.value<int>(args), 0);
    if (x_value < 10) {
      ASSERT_EQ(x_buffer[0], (x_value + 1) * 2);
    } else {
      ASSERT_EQ(x_buffer[0], (x_value - 1) * 2);
    }
  }
}

// if (x < 10) {
//   if (x > 5) {
//     x = x + 1;
//   } else {
//     x = x - 1;
//   }
// } else {
//   if (x <= 15) {
//     x = x + 2;
//   } else {
//     x = x - 2;
//   }
// }
TEST(LLVM, CondNestedTest) {
  BufHandle x("X", {1}, kInt);
  auto true_cmp =
      CompareSelect::make(x.load(0), 5, CompareSelectOperation::kGT);
  auto true_cond = Cond::make(
      true_cmp, x.store({0}, x.load(0) + 1), x.store({0}, x.load(0) - 1));
  auto false_cmp =
      CompareSelect::make(x.load(0), 15, CompareSelectOperation::kLE);
  auto false_cond = Cond::make(
      false_cmp, x.store({0}, x.load(0) + 2), x.store({0}, x.load(0) - 2));
  auto cmp = CompareSelect::make(x.load(0), 10, CompareSelectOperation::kLT);
  auto cond = Cond::make(cmp, true_cond, false_cond);

  for (int32_t x_value : {0, 8, 15, 20}) {
    std::vector<int32_t> x_buffer = {x_value};
    std::vector<void*> args({x_buffer.data()});
    LLVMCodeGen cg(cond, {x});
    ASSERT_EQ(cg.value<int>(args), 0);
    if (x_value < 10) {
      if (x_value > 5) {
        ASSERT_EQ(x_buffer[0], x_value + 1);
      } else {
        ASSERT_EQ(x_buffer[0], x_value - 1);
      }
    } else {
      if (x_value <= 15) {
        ASSERT_EQ(x_buffer[0], x_value + 2);
      } else {
        ASSERT_EQ(x_buffer[0], x_value - 2);
      }
    }
  }
}

TEST(LLVM, DirectVectorization) {
  constexpr int M = 3;
  constexpr int N = 64;
  BufHandle a("a", {M, N}, kFloat);
  BufHandle b("b", {M, N}, kFloat);
  BufHandle c("c", {M, N}, kFloat);
  VarHandle m("m", kInt);
  VarHandle n("n", kInt);
  StmtPtr s = For::make(
      m,
      0,
      M,
      Store::make(
          c,
          {Ramp::make(m * 64, 1, 64)},
          Load::make({kFloat, 64}, a, {Ramp::make(m * 64, 1, 64)}) *
              Load::make({kFloat, 64}, b, {Ramp::make(m * 64, 1, 64)})));
  LLVMCodeGen cg(s, {a, b, c});
}

TEST(LLVM, VecLoadStoreTest) {
  BufHandle a("A", {1}, kInt);
  BufHandle b("B", {1}, kInt);
  std::vector<int32_t> a_buffer = {1, 1, 1, 1};
  std::vector<int32_t> b_buffer = {2, 2, 2, 2};

  auto store = b.store({Ramp::make(0, 1, 4)}, a.load({Ramp::make(0, 1, 4)}));
  LLVMCodeGen cg(store, {a, b});
  std::vector<void*> args({a_buffer.data(), b_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);
  ASSERT_EQ(a_buffer[0], 1);
  ASSERT_EQ(a_buffer[1], 1);
  ASSERT_EQ(a_buffer[2], 1);
  ASSERT_EQ(a_buffer[3], 1);
  ASSERT_EQ(b_buffer[0], 1);
  ASSERT_EQ(b_buffer[1], 1);
  ASSERT_EQ(b_buffer[2], 1);
  ASSERT_EQ(b_buffer[3], 1);
}

#define FLOAT_INTRINSICS_TEST(Name, Lanes)                                   \
  TEST(LLVM, VecFloat_##Name##Lane##Lanes##Test) {                           \
    BufHandle a("A", {1}, kFloat);                                           \
    BufHandle b("B", {1}, kFloat);                                           \
    float val = 0.5f;                                                        \
    std::vector<float> a_buffer(Lanes, val);                                 \
    std::vector<float> b_buffer(Lanes, val);                                 \
    auto store = b.store(                                                    \
        {Ramp::make(0, 1, Lanes)}, Name(a.load({Ramp::make(0, 1, Lanes)}))); \
    LLVMCodeGen cg(store, {a, b});                                           \
    std::vector<void*> args({a_buffer.data(), b_buffer.data()});             \
    ASSERT_EQ(cg.value<int>(args), 0);                                       \
    for (const auto i : c10::irange(Lanes)) {                                \
      ASSERT_FLOAT_EQ(a_buffer[i], val);                                     \
    }                                                                        \
  } // namespace jit
FLOAT_INTRINSICS_TEST(erf, 4)
FLOAT_INTRINSICS_TEST(erfc, 4)
FLOAT_INTRINSICS_TEST(acos, 4)
FLOAT_INTRINSICS_TEST(asin, 4)
FLOAT_INTRINSICS_TEST(atan, 4)
FLOAT_INTRINSICS_TEST(cosh, 4)
FLOAT_INTRINSICS_TEST(sinh, 4)
FLOAT_INTRINSICS_TEST(tanh, 4)
FLOAT_INTRINSICS_TEST(expm1, 4)
FLOAT_INTRINSICS_TEST(lgamma, 4)
FLOAT_INTRINSICS_TEST(erf, 8)
FLOAT_INTRINSICS_TEST(erfc, 8)
FLOAT_INTRINSICS_TEST(acos, 8)
FLOAT_INTRINSICS_TEST(asin, 8)
FLOAT_INTRINSICS_TEST(atan, 8)
FLOAT_INTRINSICS_TEST(cosh, 8)
FLOAT_INTRINSICS_TEST(sinh, 8)
FLOAT_INTRINSICS_TEST(tanh, 8)
FLOAT_INTRINSICS_TEST(expm1, 8)
FLOAT_INTRINSICS_TEST(lgamma, 8)
#undef FLOAT_INTRINSICS_TEST

#define DOUBLE_INTRINSICS_TEST(Name, Lanes)                                  \
  TEST(LLVM, VecDouble_##Name##Lane##Lanes##Test) {                          \
    BufHandle a("A", {1}, kDouble);                                          \
    BufHandle b("B", {1}, kDouble);                                          \
    float val = 0.5f;                                                        \
    std::vector<double> a_buffer(Lanes, val);                                \
    std::vector<double> b_buffer(Lanes, val);                                \
    auto store = b.store(                                                    \
        {Ramp::make(0, 1, Lanes)}, Name(a.load({Ramp::make(0, 1, Lanes)}))); \
    LLVMCodeGen cg(store, {a, b});                                           \
    std::vector<void*> args({a_buffer.data(), b_buffer.data()});             \
    ASSERT_EQ(cg.value<int>(args), 0);                                       \
    for (const auto i : c10::irange(Lanes)) {                                \
      ASSERT_FLOAT_EQ(a_buffer[i], val);                                     \
    }                                                                        \
  } // namespace jit
DOUBLE_INTRINSICS_TEST(erf, 2)
DOUBLE_INTRINSICS_TEST(erfc, 2)
DOUBLE_INTRINSICS_TEST(acos, 2)
DOUBLE_INTRINSICS_TEST(asin, 2)
DOUBLE_INTRINSICS_TEST(atan, 2)
DOUBLE_INTRINSICS_TEST(cosh, 2)
DOUBLE_INTRINSICS_TEST(sinh, 2)
DOUBLE_INTRINSICS_TEST(tanh, 2)
DOUBLE_INTRINSICS_TEST(expm1, 2)
DOUBLE_INTRINSICS_TEST(lgamma, 2)
DOUBLE_INTRINSICS_TEST(erf, 4)
DOUBLE_INTRINSICS_TEST(erfc, 4)
DOUBLE_INTRINSICS_TEST(acos, 4)
DOUBLE_INTRINSICS_TEST(asin, 4)
DOUBLE_INTRINSICS_TEST(atan, 4)
DOUBLE_INTRINSICS_TEST(cosh, 4)
DOUBLE_INTRINSICS_TEST(sinh, 4)
DOUBLE_INTRINSICS_TEST(tanh, 4)
DOUBLE_INTRINSICS_TEST(expm1, 4)
DOUBLE_INTRINSICS_TEST(lgamma, 4)
#undef DOUBLE_INTRINSICS_TEST

TEST(LLVM, VectorizerLoadStoreTest) {
  BufHandle a("A", {1}, kInt);

  Tensor c = Compute("c", {4}, [&](const VarHandle& i) { return a.load(i); });

  BufHandle c_buf(c.buf());
  LoopNest l({c});
  StmtPtr s = l.root_stmt();
  ASSERT_TRUE(LoopNest::vectorize(to<For>(to<Block>(s)->front())));

  ASSERT_TRUE(to<For>(to<Block>(s)->front()) == nullptr);

  LLVMCodeGen cg(s, {a, c_buf});

  std::vector<int> a_vec(4, 21);
  std::vector<int> c_vec(4, 0);
  std::vector<void*> args({a_vec.data(), c_vec.data()});
  ASSERT_EQ(cg.value<int>(args), 0);
  assertAllEqual(c_vec, 21);
}

TEST(LLVM, VectorizeBitCast) {
  BufHandle a("A", {128}, kInt);

  Tensor c = Compute("c", {128}, [&](const VarHandle& i) {
    return bitcast<float>(a.load(i));
  });

  BufHandle c_buf(c.buf());
  LoopNest l({c});
  StmtPtr s = l.root_stmt();
  ASSERT_TRUE(LoopNest::vectorize(to<For>(to<Block>(s)->front())));
  ASSERT_TRUE(to<For>(to<Block>(s)->front()) == nullptr);

  LLVMCodeGen cg(s, {a, c_buf});

  std::vector<int> a_vec(128);
  std::vector<float> c_vec(128);
  for (const auto i : c10::irange(128)) {
    a_vec[i] = raw_bitcast<int>(1337.f);
  }
  std::vector<void*> args({a_vec.data(), c_vec.data()});
  ASSERT_EQ(cg.value<int>(args), 0);
  assertAllEqual(c_vec, 1337.f);
}

TEST(LLVM, MemcpyTest) {
  constexpr int N = 32;
  BufHandle a("A", {N}, kInt);
  BufHandle b("B", {N}, kInt);
  std::vector<int32_t> a_buffer(N, 42);
  std::vector<int32_t> b_buffer(N, 0);

  VarHandle i("i", kInt);
  auto expr = For::make(i, 0, N, b.store({i}, a.load(i)));

  LLVMCodeGen cg(expr, {a, b});

  std::vector<void*> args({a_buffer.data(), b_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  assertAllEqual(a_buffer, 42);
  assertAllEqual(b_buffer, 42);
}

TEST(LLVM, BzeroTest) {
  constexpr int N = 32;
  BufHandle b("B", {N}, kInt);
  std::vector<int32_t> b_buffer(N, 11);

  VarHandle i("i", kInt);
  auto expr = For::make(i, 0, N, b.store({i}, 0));

  LLVMCodeGen cg(expr, {b});

  std::vector<void*> args({b_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(b_buffer.size(), N);
  assertAllEqual(b_buffer, 0);
}

TEST(LLVM, ElemwiseAdd) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kInt);
  BufHandle b("B", {N}, kInt);
  BufHandle c("C", {N}, kInt);
  std::vector<int32_t> a_buffer(N, 41);
  std::vector<int32_t> b_buffer(N, 1);
  std::vector<int32_t> c_buffer(N, 1);

  VarHandle i("i", kInt);
  auto expr = For::make(i, 0, N, c.store({i}, Add::make(a.load(i), b.load(i))));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);
  assertAllEqual(a_buffer, 41);
  assertAllEqual(b_buffer, 1);
  assertAllEqual(c_buffer, 42);
}

TEST(LLVM, ElemwiseAddFloat) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kFloat);
  BufHandle b("B", {N}, kFloat);
  BufHandle c("C", {N}, kFloat);
  std::vector<float> a_buffer(N, 41);
  std::vector<float> b_buffer(N, 1);
  std::vector<float> c_buffer(N, 1);

  VarHandle i("i", kInt);
  auto expr = For::make(i, 0, N, c.store({i}, a.load(i) + b.load(i)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);
  assertAllEqual(a_buffer, 41.0f);
  assertAllEqual(b_buffer, 1.0f);
  assertAllEqual(c_buffer, 42.0f);
}

TEST(LLVM, ElemwiseLog10Float) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kFloat);
  BufHandle b("B", {N}, kFloat);
  std::vector<float> a_buffer(N, 10.0f);
  std::vector<float> b_buffer(N, 2.0f);

  VarHandle i("i", kInt);
  auto expr = For::make(
      i,
      0,
      N / 4,
      b.store(
          {Ramp::make(i * 4, 1, 4)}, log10(a.load({Ramp::make(i * 4, 1, 4)}))));

  LLVMCodeGen cg(expr, {a, b});

  std::vector<void*> args({a_buffer.data(), b_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  assertAllEqual(a_buffer, 10.0f);
  assertAllEqual(b_buffer, 1.0f);
}

TEST(LLVM, ElemwiseLog1pFloat) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kFloat);
  BufHandle b("B", {N}, kFloat);
  std::vector<float> a_buffer(N, expf(3.0f) - 1);
  std::vector<float> b_buffer(N, 42.0f);

  VarHandle i("i", kInt);
  auto expr = For::make(
      i,
      0,
      N / 4,
      b.store(
          {Ramp::make(i * 4, 1, 4)}, log1p(a.load({Ramp::make(i * 4, 1, 4)}))));

  LLVMCodeGen cg(expr, {a, b});

  std::vector<void*> args({a_buffer.data(), b_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  assertAllEqual(a_buffer, expf(3.0f) - 1);
  ExpectAllNear(b_buffer, 3.0f, 1e-5f);
}

TEST(LLVM, ElemwiseMaxInt) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kInt);
  BufHandle b("B", {N}, kInt);
  BufHandle c("C", {N}, kInt);
  std::vector<int> a_buffer(N, 41);
  std::vector<int> b_buffer(N, 1);
  std::vector<int> c_buffer(N, 1);

  VarHandle i("i", kInt);
  auto expr =
      For::make(i, 0, N, c.store({i}, Max::make(a.load(i), b.load(i), false)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);
  assertAllEqual(a_buffer, 41);
  assertAllEqual(b_buffer, 1);
  assertAllEqual(c_buffer, 41);
}

TEST(LLVM, ElemwiseMinInt) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kInt);
  BufHandle b("B", {N}, kInt);
  BufHandle c("C", {N}, kInt);
  std::vector<int> a_buffer(N, 41);
  std::vector<int> b_buffer(N, 1);
  std::vector<int> c_buffer(N, 1);

  VarHandle i("i", kInt);
  auto expr =
      For::make(i, 0, N, c.store({i}, Min::make(a.load(i), b.load(i), false)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);
  assertAllEqual(a_buffer, 41);
  assertAllEqual(b_buffer, 1);
  assertAllEqual(c_buffer, 1);
}

TEST(LLVM, ElemwiseMaxFloat) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kFloat);
  BufHandle b("B", {N}, kFloat);
  BufHandle c("C", {N}, kFloat);
  std::vector<float> a_buffer(N, 41);
  std::vector<float> b_buffer(N, 1);
  std::vector<float> c_buffer(N, 1);

  VarHandle i("i", kInt);
  auto expr =
      For::make(i, 0, N, c.store({i}, Max::make(a.load(i), b.load(i), false)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);
  assertAllEqual(a_buffer, 41.0f);
  assertAllEqual(b_buffer, 1.0f);
  assertAllEqual(c_buffer, 41.0f);
}

TEST(LLVM, ElemwiseMaxNaNFloat) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kFloat);
  BufHandle b("B", {N}, kFloat);
  BufHandle c("C", {N}, kFloat);
  std::vector<float> a_buffer(N, NAN);
  std::vector<float> b_buffer(N, 1);
  std::vector<float> c_buffer(N, 1);

  VarHandle i("i", kInt);
  auto expr =
      For::make(i, 0, N, c.store({i}, Max::make(a.load(i), b.load(i), false)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);
  assertAllEqual(b_buffer, 1.0f);
  for (auto const& elt : c_buffer) {
    ASSERT_TRUE(std::isnan(elt));
  }
}

TEST(LLVM, ElemwiseMinFloat) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kFloat);
  BufHandle b("B", {N}, kFloat);
  BufHandle c("C", {N}, kFloat);
  std::vector<float> a_buffer(N, 41);
  std::vector<float> b_buffer(N, 1);
  std::vector<float> c_buffer(N, 1);

  VarHandle i("i", kInt);
  auto expr =
      For::make(i, 0, N, c.store({i}, Min::make(a.load(i), b.load(i), false)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);
  assertAllEqual(a_buffer, 41.0f);
  assertAllEqual(b_buffer, 1.0f);
  assertAllEqual(c_buffer, 1.0f);
}

TEST(LLVM, ElemwiseMinNaNFloat) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kFloat);
  BufHandle b("B", {N}, kFloat);
  BufHandle c("C", {N}, kFloat);
  std::vector<float> a_buffer(N, NAN);
  std::vector<float> b_buffer(N, 1);
  std::vector<float> c_buffer(N, 1);

  VarHandle i("i", kInt);
  auto expr =
      For::make(i, 0, N, c.store({i}, Min::make(a.load(i), b.load(i), false)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);
  assertAllEqual(b_buffer, 1.0f);
  for (auto const& elt : c_buffer) {
    ASSERT_TRUE(std::isnan(elt));
  }
}

TEST(LLVM, ElemwiseMod) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kInt);
  BufHandle b("B", {N}, kInt);
  BufHandle c("C", {N}, kInt);
  std::vector<int32_t> a_buffer(N, 41);
  std::vector<int32_t> b_buffer(N, 23);
  std::vector<int32_t> c_buffer(N, 18);

  VarHandle i("i", kInt);
  auto expr = For::make(i, 0, N, c.store({i}, Mod::make(a.load(i), b.load(i))));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);
  assertAllEqual(a_buffer, 41);
  assertAllEqual(b_buffer, 23);
  assertAllEqual(c_buffer, 18);
}

TEST(LLVM, CompareSelectIntEQ) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kInt);
  BufHandle b("B", {N}, kInt);
  BufHandle c("C", {N}, kInt);
  std::vector<int> a_buffer(N, 1);
  std::vector<int> b_buffer(N, 1);
  std::vector<int> c_buffer(N, 0);
  std::vector<int> c_ref(N, 1);

  for (int i = 0; i < N / 2; i++) {
    b_buffer[i] = 0;
    c_ref[i] = 0;
  }

  VarHandle i("i", kInt);
  auto expr = For::make(
      i,
      0,
      N,
      c.store(
          {i},
          CompareSelect::make(
              a.load(i), b.load(i), CompareSelectOperation::kEQ)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);

  assertAllEqual(a_buffer, 1);
  for (const auto i : c10::irange(N)) {
    ASSERT_EQ(c_ref[i], c_buffer[i]);
  }
}

TEST(LLVM, CompareSelectFloatEQ) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kFloat);
  BufHandle b("B", {N}, kFloat);
  BufHandle c("C", {N}, kInt);
  std::vector<float> a_buffer(N, 1.0f);
  std::vector<float> b_buffer(N, 1.0f);
  std::vector<int> c_buffer(N, 0);

  VarHandle i("i", kInt);
  auto expr = For::make(
      i,
      0,
      N,
      c.store(
          {i},
          CompareSelect::make(
              a.load(i), b.load(i), CompareSelectOperation::kEQ)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);

  assertAllEqual(a_buffer, 1.0f);
  assertAllEqual(b_buffer, 1.0f);
  assertAllEqual(c_buffer, 1);
}

TEST(LLVM, CompareSelectByteGT) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kByte);
  BufHandle b("B", {N}, kByte);
  BufHandle c("C", {N}, kInt);
  std::vector<uint8_t> a_buffer(N, 0);
  std::vector<uint8_t> b_buffer(N, 0);
  std::vector<int> c_buffer(N, 0);
  std::vector<int> c_ref(N, 0);

  for (int i = 0; i < N / 2; i++) {
    a_buffer[i] = 128;
    c_ref[i] = 1;
  }

  VarHandle i("i", kInt);
  auto expr = For::make(
      i,
      0,
      N,
      c.store(
          {i},
          CompareSelect::make(
              a.load(i), b.load(i), CompareSelectOperation::kGT)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);

  assertAllEqual(b_buffer, uint8_t(0));
  for (const auto i : c10::irange(N)) {
    ASSERT_EQ(c_ref[i], c_buffer[i]);
  }
}

TEST(LLVM, CompareSelectByteGE) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kByte);
  BufHandle b("B", {N}, kByte);
  BufHandle c("C", {N}, kInt);
  std::vector<uint8_t> a_buffer(N, 0);
  std::vector<uint8_t> b_buffer(N, 0);
  std::vector<int> c_buffer(N, 0);
  std::vector<int> c_ref(N, 1);

  VarHandle i("i", kInt);
  auto expr = For::make(
      i,
      0,
      N,
      c.store(
          {i},
          CompareSelect::make(
              a.load(i), b.load(i), CompareSelectOperation::kGE)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);

  assertAllEqual(b_buffer, uint8_t(0));
  for (const auto i : c10::irange(N)) {
    ASSERT_EQ(c_ref[i], c_buffer[i]);
  }
}

TEST(LLVM, CompareSelectByteLT) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kByte);
  BufHandle b("B", {N}, kByte);
  BufHandle c("C", {N}, kInt);
  std::vector<uint8_t> a_buffer(N, 0);
  std::vector<uint8_t> b_buffer(N, 128);
  std::vector<int> c_buffer(N, 0);
  std::vector<int> c_ref(N, 1);

  for (int i = 0; i < N / 2; i++) {
    a_buffer[i] = 128;
    c_ref[i] = 0;
  }

  VarHandle i("i", kInt);
  auto expr = For::make(
      i,
      0,
      N,
      c.store(
          {i},
          CompareSelect::make(
              a.load(i), b.load(i), CompareSelectOperation::kLT)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);

  assertAllEqual(b_buffer, uint8_t(128));
  for (const auto i : c10::irange(N)) {
    ASSERT_EQ(c_ref[i], c_buffer[i]);
  }
}

TEST(LLVM, CompareSelectByteLE) {
  constexpr int N = 1024;
  BufHandle a("A", {N}, kByte);
  BufHandle b("B", {N}, kByte);
  BufHandle c("C", {N}, kInt);
  std::vector<uint8_t> a_buffer(N, 0);
  std::vector<uint8_t> b_buffer(N, 128);
  std::vector<int> c_buffer(N, 0);
  std::vector<int> c_ref(N, 1);

  VarHandle i("i", kInt);
  auto expr = For::make(
      i,
      0,
      N,
      c.store(
          {i},
          CompareSelect::make(
              a.load(i), b.load(i), CompareSelectOperation::kLE)));

  LLVMCodeGen cg(expr, {a, b, c});

  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  ASSERT_EQ(a_buffer.size(), N);
  ASSERT_EQ(b_buffer.size(), N);
  ASSERT_EQ(c_buffer.size(), N);

  assertAllEqual(b_buffer, uint8_t(128));
  for (const auto i : c10::irange(N)) {
    ASSERT_EQ(c_ref[i], c_buffer[i]);
  }
}

TEST(LLVM, StoreFloat) {
  BufHandle result("result", {1}, kFloat);
  std::vector<float> result_buffer = {0.0f};
  auto expr = result.store({0}, FloatImm::make(3.14f));
  LLVMCodeGen cg(expr, {result});
  std::vector<void*> args({result_buffer.data()});
  ASSERT_EQ(cg.value<int>(args), 0);
  ASSERT_EQ(result_buffer[0], 3.14f);
}

TEST(LLVM, SimpleMath01) {
  const int N = 1024;
  Tensor tensor = Compute(
      "f", {N}, [](const VarHandle& i) { return cast<float>(i * i + 1); });
  LoopNest l({tensor});
  StmtPtr stmt = l.root_stmt();
  BufHandle f_buf(tensor.buf());
  LLVMCodeGen cg(stmt, {f_buf});

  PaddedBuffer<float> f_v(N, "f_v");
  std::vector<void*> args({f_v.data()});
  int value = cg.value<int>(args);
  ASSERT_EQ(value, 0);
  PaddedBuffer<float> f_ref(N, "f_ref");
  for (const auto i : c10::irange(N)) {
    f_ref(i) = i * i + 1;
  }
  ExpectAllNear(f_v, f_ref, 1e-5);
}

TEST(LLVM, ComputeMul) {
  const int N = 1024;
  BufHandle a("a", {N}, kFloat);
  BufHandle b("b", {N}, kFloat);
  Tensor c = Compute(
      "c", {N}, [&](const VarHandle& i) { return a.load(i) * b.load(i); });

  BufHandle c_buf(c.buf());
  LoopNest l({c});
  StmtPtr s = l.root_stmt();

  LLVMCodeGen cg(s, {a, b, c_buf});

  std::vector<float> a_vec(N, 21.0f);
  std::vector<float> b_vec(N, 2.0f);
  std::vector<float> c_vec(N, 0.0f);
  std::vector<void*> args({a_vec.data(), b_vec.data(), c_vec.data()});
  ASSERT_EQ(cg.value<int>(args), 0);
  assertAllEqual(c_vec, 42.0f);
}

TEST(LLVM, BroadcastAdd) {
  const int M = 32;
  const int N = 1024;
  BufHandle a("a", {M, N}, kFloat);
  BufHandle b("b", {N}, kFloat);
  Tensor c = Compute("c", {M, N}, [&](const VarHandle& i, const VarHandle& j) {
    return a.load(i, j) + b.load(j);
  });

  BufHandle c_buf(c.buf());
  LoopNest l({c});
  l.prepareForCodegen();
  StmtPtr s = l.root_stmt();

  LLVMCodeGen cg(s, {a, b, c_buf});

  std::vector<float> av(M * N);
  std::iota(av.begin(), av.end(), 0);
  std::vector<float> bv(N);
  std::iota(bv.begin(), bv.end(), 0);
  std::vector<float> cv(M * N, 0);
  std::vector<void*> args({av.data(), bv.data(), cv.data()});
  ASSERT_EQ(cg.value<int>(args), 0);

  for (const auto i : c10::irange(M)) {
    for (const auto j : c10::irange(N)) {
      ASSERT_EQ(cv[i * N + j], av[i * N + j] + bv[j]);
    }
  }
}

TEST(LLVM, BitwiseOps) {
  auto a = IntImm::make(59);
  auto b = IntImm::make(11);
  auto c = IntImm::make(101);
  auto d = IntImm::make(2);

  ExprHandle f = (((a ^ (b << 1)) & c) >> 2) | d;
  LLVMExprEval cg(f);

  ASSERT_EQ(cg.value<int>(), 11);
}

TEST(LLVM, ArithmeticRightShift) {
  auto a = CharImm::make(-4);
  auto b = CharImm::make(1);
  ExprHandle f = a >> b;
  LLVMExprEval cg(f);
  ASSERT_EQ(cg.value<int8_t>(), -2);
}

TEST(LLVM, LogicalRightShift) {
  auto a = ByteImm::make(0xfc);
  auto b = ByteImm::make(1);
  ExprHandle f = a >> b;
  LLVMExprEval cg(f);
  ASSERT_EQ(cg.value<uint8_t>(), 0x7e);
}

TEST(LLVM, DynamicShapeAdd) {
  auto testWithSize = [](int32_t size) {
    VarHandle n("n", kInt);
    BufHandle a("a", {n}, kFloat);
    BufHandle b("b", {n}, kFloat);
    BufHandle c("c", {n}, kFloat);
    VarHandle i("i", kInt);
    StmtPtr s = For::make(i, 0, n, c.store({i}, a.load(i) + b.load(i)));
    std::vector<float> aData(size, 1.0f);
    std::vector<float> bData(size, 2.0f);
    std::vector<float> cData(size, 0.0f);
    LLVMCodeGen cg(s, {a, b, c, n});
    std::vector<void*> args({aData.data(), bData.data(), cData.data(), &size});
    cg.value<float>(args);
    ExpectAllNear(cData, std::vector<float>(size, 3.0f), 1e-7);
  };
  testWithSize(1);
  testWithSize(16);
  testWithSize(37);
}

TEST(LLVM, BindDynamicShapeAdd) {
  auto testWithSize = [](int32_t size) {
    VarHandle n("n", kInt);
    BufHandle a("a", {n}, kFloat);
    BufHandle b("b", {n}, kFloat);
    BufHandle c("c", {n}, kFloat);
    VarHandle i("i", kInt);
    StmtPtr s = For::make(i, 0, n, c.store({i}, a.load(i) + b.load(i)));
    std::vector<float> aData(size, 1.0f);
    std::vector<float> bData(size, 2.0f);
    std::vector<float> cData(size, 0.0f);
    LLVMCodeGen cg(s, {a, b, c, n});
    cg.call({aData, bData, cData, size});
    ExpectAllNear(cData, std::vector<float>(size, 3.0f), 1e-7);
  };
  testWithSize(1);
  testWithSize(16);
  testWithSize(37);
}

TEST(LLVM, TensorDynamicShapeAdd) {
  auto testWithSize = [](int32_t size) {
    VarHandle n("n", kInt);
    BufHandle a("a", {n}, kFloat);
    BufHandle b("b", {n}, kFloat);
    Tensor c = Compute(
        "c", {n}, [&](const VarHandle& i) { return a.load(i) + b.load(i); });
    LoopNest l({c});
    StmtPtr s = l.root_stmt();
    LLVMCodeGen cg(s, {a, b, c, n});
    std::vector<float> aData(size, 1.0f);
    std::vector<float> bData(size, 2.0f);
    std::vector<float> cData(size, 0.0f);
    cg.call({aData, bData, cData, size});
    ExpectAllNear(cData, std::vector<float>(size, 3.0f), 1e-7);
  };
  testWithSize(1);
  testWithSize(16);
  testWithSize(37);
}

TEST(LLVM, DynamicShape2D) {
  auto testWithSize = [](int32_t M, int32_t N) {
    VarHandle m("m", kInt);
    VarHandle n("n", kInt);
    BufHandle a("a", {m, n}, kFloat);
    BufHandle b("b", {m, n}, kFloat);
    Tensor c =
        Compute("c", {m, n}, [&](const VarHandle& i, const VarHandle& j) {
          return a.load(i, j) + b.load(i, j);
        });
    LoopNest l({c});
    l.prepareForCodegen();
    StmtPtr s = l.root_stmt();
    LLVMCodeGen cg(s, {a, b, c, m, n});
    std::vector<float> aData(M * N, 1.0f);
    std::vector<float> bData(M * N, 2.0f);
    std::vector<float> cData(M * N, 0.0f);
    cg.call({aData, bData, cData, M, N});
    ExpectAllNear(cData, std::vector<float>(M * N, 3.0f), 1e-7);
  };
  testWithSize(1, 8);
  testWithSize(16, 32);
  testWithSize(37, 11);
}

TEST(LLVM, EmptyStmt) {
  StmtPtr s = alloc<Block>(std::vector<StmtPtr>({}));

  LLVMCodeGen cg(s, {});
  cg.call({});
  // Just don't crash.
}

TEST(LLVM, EliminatedStmt) {
  BufHandle a("a", {1}, kFloat);

  Tensor c = Compute("c", {0}, [&](const VarHandle& m) { return m; });

  LoopNest l({c});
  l.prepareForCodegen();
  StmtPtr s = l.root_stmt();
  s = IRSimplifier::simplify(s);
  LLVMCodeGen cg(s, {a, c});
  std::vector<float> aData(1, 1.0f);
  std::vector<float> cData(0, 0.0f);
  cg.call({aData, cData});
}

TEST(LLVM, SimpleReduction) {
  int M = 128;
  int N = 64;

  BufHandle a("a", {1, M, N}, kFloat);

  Tensor b = Reduce("sum", {1}, Sum(), a, {M, N});
  LoopNest loop({b});

  loop.prepareForCodegen();
  StmtPtr s = loop.root_stmt();
  s = IRSimplifier::simplify(s);

  LLVMCodeGen cg(s, {a, b});

  PaddedBuffer<float> a_v(1, M, N, "a_v");
  PaddedBuffer<float> b_v(1, "b_v");
  PaddedBuffer<float> b_ref(1, "b_ref");

  b_ref(0) = 0;
  for (const auto i : c10::irange(M)) {
    for (const auto j : c10::irange(N)) {
      int v = i + j;
      a_v(0, i, j) = v;
      b_ref(0) += v;
    }
  }

  cg.call({a_v, b_v});

  ExpectAllNear(b_v, b_ref, 1e-5);
}

TEST(LLVM, RFactorReduction) {
  int M = 128;
  int N = 64;

  BufHandle a("a", {1, M, N}, kFloat);

  Tensor b = Reduce("sum", {1}, Sum(), a, {M, N});
  LoopNest loop({b});

  std::vector<ForPtr> loops = loop.getLoopStmtsFor(b);
  ForPtr loop_m = loops.at(1);
  ForPtr loop_n = loops.at(2);
  loop.reorderAxis(loop_m, loop_n);

  loops = loop.getLoopStmtsFor(b);
  loop_m = loops.at(2);
  loop_n = loops.at(1);
  auto b_body = loop.getAllWritesToBuf(b.buf())[1];
  ASSERT_TRUE(loop.rfactor(b_body, loop_n));

  loop.prepareForCodegen();
  StmtPtr s = loop.root_stmt();
  s = IRSimplifier::simplify(s);

  LLVMCodeGen cg(s, {a, b});

  PaddedBuffer<float> a_v(1, M, N, "a_v");
  PaddedBuffer<float> b_v(1, "b_v");
  PaddedBuffer<float> b_ref(1, "b_ref");

  b_ref(0) = 0;
  for (const auto i : c10::irange(M)) {
    for (const auto j : c10::irange(N)) {
      int v = i + j;
      a_v(0, i, j) = v;
      b_ref(0) += v;
    }
  }

  cg.call({a_v, b_v});

  ExpectAllNear(b_v, b_ref, 1e-5);
}

TEST(LLVM, RFactorVectorizedReduction) {
  int M = 128;
  int N = 64;

  BufHandle a("a", {1, M, N}, kFloat);

  Tensor b = Reduce("sum", {1}, Sum(), a, {M, N});
  LoopNest loopnest({b});
  std::vector<ForPtr> loops = loopnest.getLoopStmtsFor(b);
  // Reorder n and m loops
  loopnest.reorderAxis(loops.at(1), loops.at(2));
  auto b_body = loopnest.getAllWritesToBuf(b.buf()).at(1);
  auto all_loops = loopnest.getAllLoopNestsWritingToBuf(b.buf());
  ASSERT_TRUE(all_loops.size() == 2 && all_loops[1].size() == 3);
  ASSERT_TRUE(loopnest.rfactor(b_body, all_loops[1][1]));
  auto distributed_loops = loopnest.distributeLoop(all_loops[1][1]);

  // Vectorize initializer of rfac_buf
  ASSERT_TRUE(LoopNest::vectorize(distributed_loops[0]));
  // Vectorize producer of rfac_buf
  ASSERT_TRUE(LoopNest::vectorize(distributed_loops[1]));
  loopnest.simplify();

  loopnest.prepareForCodegen();

  StmtPtr s = IRSimplifier::simplify(loopnest.root_stmt());
  LLVMCodeGen cg(s, {a, b});

  PaddedBuffer<float> a_v(1, M, N, "a_v");
  PaddedBuffer<float> b_v(1, "b_v");
  PaddedBuffer<float> b_ref(1, "b_ref");

  b_ref(0) = 0;
  for (const auto i : c10::irange(M)) {
    for (const auto j : c10::irange(N)) {
      int v = i + j;
      a_v(0, i, j) = v;
      b_ref(0) += v;
    }
  }

  cg.call({a_v, b_v});

  ExpectAllNear(b_v, b_ref, 1e-5);
}

template <bool outer, bool inner>
static void testSimpleParallel() {
  // Compute a simple operation, and try all loop-axis combination to be
  // parallel or sequential.
  const int M = 4;
  const int N = 6;
  Tensor f = Compute("f", {M, N}, [](const VarHandle& m, const VarHandle& n) {
    return cast<float>(m + n);
  });
  LoopNest loop_nest({f});
  auto const& loops = loop_nest.getLoopStmtsFor(f);
  ForPtr m = loops[0];
  ForPtr n = loops[1];
  if (outer) {
    m->set_parallel();
  }
  if (inner) {
    n->set_parallel();
  }
  loop_nest.prepareForCodegen();
  StmtPtr stmt = loop_nest.root_stmt();
  LLVMCodeGen cg(stmt, {f});

  PaddedBuffer<float> f_v(M, N, "f_v");
  std::vector<void*> args({f_v.data()});
  int value = cg.value<int>(args);
  ASSERT_EQ(value, 0);
  PaddedBuffer<float> f_ref(M, N, "f_ref");
  for (const auto m : c10::irange(M)) {
    for (const auto n : c10::irange(N)) {
      f_ref(m, n) = m + n;
    }
  }
  ExpectAllNear(f_v, f_ref, 1e-5);
}

TEST(LLVM, SimpleParallelSS) {
  testSimpleParallel<false, false>();
}
TEST(LLVM, SimpleParallelSP) {
  testSimpleParallel<false, true>();
}
TEST(LLVM, SimpleParallelPS) {
  testSimpleParallel<true, false>();
}
TEST(LLVM, SimpleParallelPP) {
  testSimpleParallel<true, true>();
}

TEST(LLVM, CompositeParallel) {
  int loop_count = 6;
  int test_count = 1 << loop_count;
  // Compute a composite operation, and try all loop-axis combination to be
  // parallel or sequential.
  for (const auto test_cfg : c10::irange(test_count)) {
    int M = 5;
    int N = 7;
    Tensor t1 = Compute("t1", {M}, [](const VarHandle& m) { return m + 1.f; });
    Tensor t2 = Compute("t2", {N}, [](const VarHandle& n) { return n + 2.f; });
    Tensor t3 =
        Compute("t3", {M, N}, [=](const VarHandle& m, const VarHandle& n) {
          return t1.load(m) * t2.load(n);
        });
    Tensor t4 =
        Compute("t4", {M, N}, [=](const VarHandle& m, const VarHandle& n) {
          return t3.load(m, n) + m + n;
        });
    LoopNest loop_nest({t4}, {t1, t2, t3, t4});
    std::vector<ForPtr> loop_list;
    {
      auto const& loops = loop_nest.getLoopStmtsFor(t1);
      loop_list.push_back(loops[0]);
    }
    {
      auto const& loops = loop_nest.getLoopStmtsFor(t2);
      loop_list.push_back(loops[0]);
    }
    {
      auto const& loops = loop_nest.getLoopStmtsFor(t3);
      loop_list.push_back(loops[0]);
      loop_list.push_back(loops[1]);
    }
    {
      auto const& loops = loop_nest.getLoopStmtsFor(t4);
      loop_list.push_back(loops[0]);
      loop_list.push_back(loops[1]);
    }
    ASSERT_EQ(loop_list.size(), loop_count);
    for (const auto i : c10::irange(loop_count)) {
      if (test_cfg & (1 << i)) {
        loop_list[i]->set_parallel();
      }
    }
    loop_nest.prepareForCodegen();
    StmtPtr stmt = loop_nest.root_stmt();
    LLVMCodeGen cg(stmt, {t4});

    PaddedBuffer<float> t4_v(M, N, "t4_v");
    std::vector<void*> args({t4_v.data()});
    int value = cg.value<int>(args);
    ASSERT_EQ(value, 0);
    PaddedBuffer<float> t4_ref(M, N, "t4_ref");
    for (const auto m : c10::irange(M)) {
      for (const auto n : c10::irange(N)) {
        t4_ref(m, n) = (m + 1) * (n + 2) + m + n;
      }
    }
    ExpectAllNear(t4_v, t4_ref, 1e-5);
  }
}

TEST(LLVM, VectorizedGEMM) {
  int M = 32;
  int N = 32;
  int K = 48;

  BufHandle AP("A", {M, K}, kFloat);
  BufHandle BP("B", {K, N}, kFloat);
  Tensor CT = Reduce(
      "gemm",
      {M, N},
      Sum(),
      [&](const ExprHandle& m, const ExprHandle& n, const ExprHandle& k) {
        return AP.load(m, k) * BP.load(k, n);
      },
      {K});
  LoopNest loop({CT});

  {
    auto const& loops = loop.getLoopStmtsFor(CT);
    ForPtr m = loops[0];
    loop.splitWithMask(m, 16);
  }
  {
    auto const& loops = loop.getLoopStmtsFor(CT);
    ForPtr n = loops[2];
    loop.splitWithMask(n, 16);
  }
  // mo, mi, no, ni, k ->
  // mo, no, mi, ni, k
  {
    auto const& loops = loop.getLoopStmtsFor(CT);
    ForPtr mi = loops[1];
    ForPtr no = loops[2];
    loop.reorderAxis(mi, no);
  }
  // mo, no, mi, ni, k ->
  // mo, no, mi, k, ni
  {
    auto const& loops = loop.getLoopStmtsFor(CT);
    ForPtr ni = loops[3];
    ForPtr k = loops[4];
    loop.reorderAxis(ni, k);
  }
  // mo, no, mi, k, ni ->
  // mo, no, k, mi, ni
  {
    auto const& loops = loop.getLoopStmtsFor(CT);
    ForPtr mi = loops[2];
    ForPtr k = loops[3];
    loop.reorderAxis(mi, k);
  }
  {
    auto loops = NodeFinder<For>::find(loop.root_stmt());
    ASSERT_TRUE(LoopNest::vectorize(loops[3]));
    ASSERT_TRUE(LoopNest::vectorize(loops.back()));
  }

  loop.prepareForCodegen();

  StmtPtr s = loop.root_stmt();
  s = IRSimplifier::simplify(s);
  LLVMCodeGen cg(s, {AP, BP, CT});

  PaddedBuffer<float> a_v(M, K, "a_v");
  PaddedBuffer<float> b_v(K, N, "b_v");
  PaddedBuffer<float> c_v(M, N, "c_v");
  PaddedBuffer<float> c_ref(M, N, "c_ref");

  for (const auto m : c10::irange(M)) {
    for (const auto n : c10::irange(N)) {
      c_ref(m, n) = 0.f;
      for (const auto k : c10::irange(K)) {
        c_ref(m, n) += a_v(m, k) * b_v(k, n);
      }
    }
  }

  cg.call({a_v, b_v, c_v});

  ExpectAllNear(c_v, c_ref, 1e-5);
}

TEST(LLVM, CallRaw) {
  const int M = 32;
  VarHandle N("N", kInt);
  BufHandle a("a", {M, N}, kFloat);
  BufHandle b("b", {N}, kFloat);
  Tensor c = Compute("c", {M, N}, [&](const VarHandle& i, const VarHandle& j) {
    return a.load(i, j) + b.load(j);
  });

  LoopNest l({c});
  l.prepareForCodegen();
  StmtPtr s = l.root_stmt();

  int32_t N_value = 1024;
  std::vector<float> av(M * N_value);
  std::iota(av.begin(), av.end(), 0);
  std::vector<float> bv(N_value);
  std::iota(bv.begin(), bv.end(), 0);
  std::vector<float> cv(M * N_value, 0);
  std::vector<void*> args({av.data(), bv.data(), cv.data(), &N_value});

  LLVMCodeGen cg(s, {a, b, BufHandle(c.buf()), N});
  cg.call_raw(args);

  for (const auto i : c10::irange(M)) {
    for (const auto j : c10::irange(N_value)) {
      ASSERT_EQ(cv[i * N_value + j], av[i * N_value + j] + bv[j]);
    }
  }

  SimpleIREvaluator eval(s, {a, b, BufHandle(c.buf()), N});
  eval.call_raw(args);

  for (const auto i : c10::irange(M)) {
    for (const auto j : c10::irange(N_value)) {
      ASSERT_EQ(cv[i * N_value + j], av[i * N_value + j] + bv[j]);
    }
  }
}

TEST(LLVM, CustomTarget) {
  constexpr int M = 16;
  BufHandle a("a", {M}, kFloat);
  BufHandle b("b", {M}, kFloat);
  BufHandle c("c", {M}, kFloat);
  Tensor d = Compute("d", {M}, [&](const VarHandle& m) {
    return a.load(m) * b.load(m) + c.load(m);
  });
  LoopNest nest({d});
  nest.prepareForCodegen();
  auto cg = LLVMCodeGenBuilder(nest.root_stmt(), {a, b, c, d})
                .triple("i686-elf")
                .cpu("i386")
                .build();
  std::ostringstream ss;
  ss << cg->getCodeText("asm");
  torch::jit::testing::FileCheck()
      .check("fadds")
      ->check("fmuls")
      ->check_not("vfmadd")
      ->run(ss.str());
}

TEST(LLVM, CodeGenKernelFuncName) {
  BufHandle a("A", {1}, kInt);
  BufHandle b("B", {1}, kInt);
  std::vector<int32_t> a_buffer = {42};
  std::vector<int32_t> b_buffer = {-11};
  auto store = b.store({0}, a.load(0));

  {
    LLVMCodeGen cg(store, {a, b});
    // Check that the kernel function name used by LLVMCodeGen
    // is not empty.
    ASSERT_NE(cg.kernel_func_name(), "");
  }

  {
    LLVMCodeGen cg(store, {a, b}, at::kCPU, "new_func");
    // Check that the kernel function name used by LLVMCodeGen
    // is the one that was given above.
    ASSERT_EQ(cg.kernel_func_name(), "new_func");
  }
}

} // namespace jit
} // namespace torch

#endif // TORCH_ENABLE_LLVM