#if defined(USE_CUDA)
#include <gtest/gtest.h>

#include <torch/csrc/jit/codegen/cuda/arith.h>
#include <torch/csrc/jit/codegen/cuda/codegen.h>
#include <torch/csrc/jit/codegen/cuda/disjoint_set.h>
#include <torch/csrc/jit/codegen/cuda/executor.h>
#include <torch/csrc/jit/codegen/cuda/executor_launch_params.h>
#include <torch/csrc/jit/codegen/cuda/expr_evaluator.h>
#include <torch/csrc/jit/codegen/cuda/fusion.h>
#include <torch/csrc/jit/codegen/cuda/fusion_segmenter.h>
#include <torch/csrc/jit/codegen/cuda/interface.h>
#include <torch/csrc/jit/codegen/cuda/ir_all_nodes.h>
#include <torch/csrc/jit/codegen/cuda/ir_builder.h>
#include <torch/csrc/jit/codegen/cuda/ir_graphviz.h>
#include <torch/csrc/jit/codegen/cuda/ir_iostream.h>
#include <torch/csrc/jit/codegen/cuda/ir_utils.h>
#include <torch/csrc/jit/codegen/cuda/iter_visitor.h>
#include <torch/csrc/jit/codegen/cuda/kernel_cache.h>
#include <torch/csrc/jit/codegen/cuda/kernel_expr_evaluator.h>
#include <torch/csrc/jit/codegen/cuda/kernel_ir.h>
#include <torch/csrc/jit/codegen/cuda/lower2device.h>
#include <torch/csrc/jit/codegen/cuda/mutator.h>
#include <torch/csrc/jit/codegen/cuda/root_domain_map.h>
#include <torch/csrc/jit/codegen/cuda/scheduler/all_schedulers.h>
#include <torch/csrc/jit/codegen/cuda/scheduler/utils.h>
#include <torch/csrc/jit/codegen/cuda/transform_replay.h>
#include <torch/csrc/jit/codegen/cuda/transform_rfactor.h>

// fuser and IR parser
#include "test_gpu_validator.h"

#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/Exceptions.h>
#include <c10/cuda/CUDAStream.h>

#include <algorithm>
#include <iostream>

// Tests go in torch::jit
namespace torch {
namespace jit {

using namespace torch::jit::fuser::cuda;
using namespace at::indexing;

namespace {

// Make a tensor that is known to be fully contiguous of dimensionality=ndims,
// but unknown sizes
TensorView* makeContigTensor(size_t ndims, DataType dtype = DataType::Float) {
  return TensorViewBuilder()
      .ndims(ndims)
      .dtype(dtype)
      .contiguity(std::vector<bool>(ndims, true))
      .build();
}

// Make a tensor that is known to be non-contiguous of dimensionality=ndims,
// but unknown sizes
TensorView* makeSymbolicTensor(size_t ndims, DataType dtype = DataType::Float) {
  return TensorViewBuilder().ndims(ndims).dtype(dtype).build();
}

// Make a non-contiguous tensor of compile-time known sizes
TensorView* makeConcreteTensor(
    std::vector<int64_t> shape,
    DataType dtype = DataType::Float) {
  return TensorViewBuilder().shape(shape).dtype(dtype).build();
}

void checkIntValue(
    ExpressionEvaluator& evaluator,
    Val* val,
    Int::ScalarType expected_value) {
  TORCH_CHECK(val->isAnInt());
  const auto actual_value = evaluator.evaluate(val);
  TORCH_CHECK(actual_value.has_value());
  TORCH_CHECK(actual_value.value() == expected_value);
}

void checkIntValue(
    kir::ExpressionEvaluator& evaluator,
    const Val* val,
    Int::ScalarType expected_value) {
  const auto actual_value = evaluator.evaluate(val);
  TORCH_CHECK(actual_value.has_value());
  TORCH_CHECK(actual_value.value() == expected_value);
}

// Used to signify invalid ranges, i.e., values at offset 0 to
// start_offset, and values at offset stop_offset to the end of the
// domain.
static constexpr int invalid_marker = 1;

// ATen version of tensor shifting
auto shift(
    at::Tensor tensor,
    const std::vector<int>& offsets,
    std::vector<int> padding = {}) {
  TORCH_INTERNAL_ASSERT(tensor.ndimension() == offsets.size());
  if (padding.empty()) {
    padding = offsets;
    for (auto& p : padding) {
      p = std::abs(p);
    }
  }
  at::Tensor t = tensor;
  for (size_t i = 0; i < offsets.size(); ++i) {
    auto offset = offsets[i];
    t = t.roll(offsets[i], i);
    if (offset == 0) {
      continue;
    }
    // Zero padding
    std::vector<at::indexing::TensorIndex> indices(
        tensor.ndimension(), at::indexing::Slice(0, at::indexing::None));
    if (offset > 0) {
      indices[i] = at::indexing::Slice(0, offset);
    } else {
      indices[i] = at::indexing::Slice(offset, at::indexing::None);
    }
    t.index(indices) = 0;
    // Fill the outside range by the special marker value.
    const auto pad = padding[i];
    if (offset > 0) {
      indices[i] = at::indexing::Slice(0, offset - pad);
    } else {
      offset += pad;
      TORCH_INTERNAL_ASSERT(offset <= 0);
      if (offset == 0) {
        continue;
      }
      indices[i] = at::indexing::Slice(offset, at::indexing::None);
    }
    t.index(indices) = invalid_marker;
  }
  return t;
}

// ATen version of tensor gather
auto gather(
    at::Tensor tensor,
    const std::vector<int>& window_shape,
    const std::vector<std::vector<int>>& pad_width,
    std::vector<int> strides = {}) {
  TORCH_CHECK(
      tensor.ndimension() == window_shape.size(),
      "Invalid window shape: ",
      window_shape,
      ". Size of the window shape is different from the tensor dimension.");
  TORCH_CHECK(
      tensor.ndimension() == pad_width.size(),
      "Invalid pad width: ",
      pad_width,
      ". Size of the pad width is different from the tensor dimension.");
  if (strides.empty()) {
    strides = std::vector<int>(tensor.ndimension(), 1);
  } else {
    TORCH_CHECK(
        tensor.ndimension() == strides.size(),
        "Invalid strides: ",
        strides,
        ". Size of strides is different from the tensor dimension.");
  }
  at::Tensor t = tensor;
  for (size_t i = 0; i < window_shape.size(); ++i) {
    const auto w_size = window_shape[i];
    TORCH_CHECK(w_size != 0);
    const auto& pad = pad_width[i];
    TORCH_CHECK(pad.size() == 2);
    const auto out_extent_adj = -w_size + 1 + pad[0] + pad[1];
    TORCH_INTERNAL_ASSERT(out_extent_adj <= 0);
    const auto stride = strides[i];
    TORCH_CHECK(stride >= 1);

    at::Tensor concat_tensor;

    for (int w = 0; w < w_size; ++w) {
      std::vector<int> shift_offsets(t.ndimension(), 0);
      shift_offsets[i] = pad[0] - w;
      auto shifted = shift(t, shift_offsets);
      // Apply stride
      if (stride != 1) {
        std::vector<at::indexing::TensorIndex> indices(
            shifted.ndimension(), at::indexing::Slice(0, at::indexing::None));
        if (out_extent_adj == 0) {
          indices[i] = at::indexing::Slice(0, at::indexing::None, strides[i]);
        } else {
          indices[i] = at::indexing::Slice(0, out_extent_adj, strides[i]);
        }
        shifted = shifted.index(indices);
      }
      shifted = shifted.unsqueeze(-1);
      if (w == 0) {
        concat_tensor = shifted;
      } else {
        concat_tensor = at::cat({concat_tensor, shifted}, -1);
      }
    }
    t = concat_tensor;
  }

  // Fill invalid regions with the marker. Note that when non-unit
  // stride is used, it trims invalid regions, so no marking is
  // necessary.
  for (size_t i = 0; i < window_shape.size(); ++i) {
    if (strides[i] != 1) {
      continue;
    }

    const auto out_extent_adj =
        -window_shape[i] + 1 + pad_width[i][0] + pad_width[i][1];
    if (out_extent_adj < 0) {
      std::vector<at::indexing::TensorIndex> indices(
          t.ndimension(), at::indexing::Slice(0, at::indexing::None));
      indices[i] = at::indexing::Slice(out_extent_adj, at::indexing::None);
      t.index(indices) = invalid_marker;
    }
  }

  return t;
}

} // namespace

// Shift an input tensor
TEST_F(NVFuserTest, FusionShift1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = shift(tv0, {-1, 0});
  fusion.addOutput(tv1);

  auto tv2 = shift(tv0, {0, 1});
  fusion.addOutput(tv2);

  auto tv3 = shift(tv0, {2, 2});
  fusion.addOutput(tv3);

  auto tv4 = shift(tv0, {-2, -2});
  fusion.addOutput(tv4);

  int numel_x = 9;
  int numel_y = 11;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = shift(t0, {-1, 0});
  TORCH_CHECK(t1.equal(outputs[0]));

  auto t2 = shift(t0, {0, 1});
  TORCH_CHECK(t2.equal(outputs[1]));

  auto t3 = shift(t0, {2, 2});
  TORCH_CHECK(t3.equal(outputs[2]));

  auto t4 = shift(t0, {-2, -2});
  TORCH_CHECK(t4.equal(outputs[3]));
}

// Shifts an intermediate tensor
TEST_F(NVFuserTest, FusionShift2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {-1, 0});
  fusion.addOutput(tv2);

  // make it a little more complex
  auto tv3 = add(tv0, IrBuilder::create<Double>(3));
  auto tv4 = add(tv3, IrBuilder::create<Double>(4));
  auto tv5 = shift(tv4, {-1, 0});
  auto tv6 = shift(tv4, {0, -1});
  auto tv7 = shift(tv4, {1, 0});
  auto tv8 = shift(tv4, {0, 0});
  auto tv9 = add(tv5, tv6);
  auto tv10 = add(tv9, tv7);
  auto tv11 = add(tv10, tv8);
  fusion.addOutput(tv11);

  for (auto tv : {tv1, tv2, tv3, tv4, tv5, tv6, tv7, tv8, tv9, tv10, tv11}) {
    tv->setMemoryType(MemoryType::Global);
  }

  // t1 allocation: (t1.size[0] + 1) * (t1.size[1])
  // t3 allocation: (t3.size[0] + 2) * (t3.size[1] + 1)
  // t4 allocation: (t3.size[0] + 2) * (t3.size[1] + 1)
  GpuLower gpulw(&fusion);

  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1 || tensor_name == 3 || tensor_name == 4) {
        TORCH_CHECK(alloc->shape().size() == 2);
        for (int i = 0; i < 2; ++i) {
          if (tensor_name == 1 && i == 1) {
            TORCH_CHECK(alloc->shape().at(i)->isA<NamedScalar>());
            continue;
          }
          auto def =
              dynamic_cast<BinaryOp*>(alloc->shape().at(i)->definition());
          TORCH_CHECK(
              def != nullptr && def->getBinaryOpType() == BinaryOpType::Add);
          TORCH_CHECK(def->as<BinaryOp>()->lhs()->isA<NamedScalar>());
          auto rhs = dynamic_cast<Int*>(def->as<BinaryOp>()->rhs());
          TORCH_CHECK(rhs != nullptr && rhs->isConst());
          int rhs_value = *rhs->value();
          if (tensor_name == 1) {
            TORCH_CHECK(i == 0);
            TORCH_CHECK(rhs_value == 1);
          } else {
            if (i == 0) {
              TORCH_CHECK(rhs_value == 2);
            } else {
              TORCH_CHECK(rhs_value == 1);
            }
          }
        }
      }
    }
  }

  int numel_x = 9;
  int numel_y = 11;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {-1, 0});

  auto t3 = t0 + 3;
  auto t4 = t3 + 4;
  auto t5 = shift(t4, {-1, 0});
  auto t6 = shift(t4, {0, -1});
  auto t7 = shift(t4, {1, 0});
  auto t8 = shift(t4, {0, 0});
  auto t9 = t5 + t6;
  auto t10 = t9 + t7;
  auto t11 = t10 + t8;

  testValidate(&fusion, outputs, inputs, {t2, t11}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftRightOfCA_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {0, 1});
  fusion.addOutput(tv2);

  tv0->computeAt(tv2, -2);

  tv1->setMemoryType(MemoryType::Global);

  int numel_x = 100;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {0, 1});

  TORCH_CHECK(t2.allclose(outputs[0]));
}

TEST_F(NVFuserTest, FusionShiftLeftOfCA_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = add(tv1, IrBuilder::create<Double>(1));
  auto tv3 = shift(tv2, {-1, 0});
  auto tv4 = add(tv3, IrBuilder::create<Double>(1));
  fusion.addOutput(tv4);

  tv0->computeAt(tv4, -1);

  // Lowering should trigger an assertion failure as a shifted axis is
  // found inside an allocation position.
  ASSERT_ANY_THROW(fusion.printKernel());
}

TEST_F(NVFuserTest, FusionShiftSplit1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {0, 1});
  auto tv3 = shift(tv1, {0, -2});
  fusion.addOutput(tv2);
  fusion.addOutput(tv3);

  int split_factor = 4;
  tv2->split(-1, split_factor);
  tv3->split(-1, split_factor);

  tv0->computeAt(tv2, -2);
  tv0->computeAt(tv3, -2);

  // t1 allocation: 7
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1) {
        TORCH_CHECK(alloc->shape().size() == 1);
        auto size = dynamic_cast<Int*>(alloc->shape().at(0));
        TORCH_CHECK(
            size != nullptr && size->isConst() && size->value().value() == 7);
      }
    }
  }

  int numel_x = 9;
  int numel_y = 11;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {0, 1});
  auto t3 = shift(t1, {0, -2});

  testValidate(&fusion, outputs, inputs, {t2, t3}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftSplit2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = add(tv1, IrBuilder::create<Double>(1));
  auto tv3 = shift(tv2, {0, -1});
  auto tv4 = shift(tv2, {0, 1});
  auto tv5 = add(tv3, tv4);
  fusion.addOutput(tv5);

  auto tv6 = add(tv0, IrBuilder::create<Double>(1));
  auto tv7 = shift(tv6, {0, 0});
  auto tv8 = add(tv7, IrBuilder::create<Double>(1));
  fusion.addOutput(tv8);

  int split_factor = 4;

  tv5->split(-1, split_factor);
  tv8->split(-1, split_factor);

  tv0->computeAt(tv5, -2);
  tv0->computeAt(tv8, -2);

  // t1 and t2 allocation: 6
  // t4 allocation: 4
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1 || tensor_name == 2) {
        TORCH_CHECK(alloc->shape().size() == 1);
        auto size = dynamic_cast<Int*>(alloc->shape().at(0));
        TORCH_CHECK(
            size != nullptr && size->isConst() && size->value().value() == 6);
      } else if (tensor_name == 4) {
        TORCH_CHECK(alloc->shape().size() == 1);
        auto size = dynamic_cast<Int*>(alloc->shape().at(0));
        TORCH_CHECK(size != nullptr && size->isConst());
        int size_value = *size->value();
        TORCH_CHECK(size_value == split_factor);
      }
    }
  }

  int numel_x = 9;
  int numel_y = 11;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 2;
  auto t3 = shift(t1, {0, -1});
  auto t4 = shift(t1, {0, 1});
  auto t5 = t3 + t4;

  auto t6 = t0 + 1;
  auto t7 = t6;
  auto t8 = t7 + 1;

  testValidate(&fusion, outputs, inputs, {t5, t8}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftDoubleSplit_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = add(tv1, IrBuilder::create<Double>(2));
  auto tv3 = shift(tv2, {0, 1});
  fusion.addOutput(tv3);

  int split_factor1 = 8;
  int split_factor2 = 4;

  tv3->split(-1, split_factor1);

  tv0->computeAt(tv3, -2);

  tv1->split(-1, split_factor2);

  // t1: [i1, i2/8, 8/4, 4]
  // t2: [i1, i2/8, 8]
  // t3: [i1, i2/8, 8]

  // t1 and t2 allocation: (split_factor1 + 1) = 9
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1 || tensor_name == 2) {
        TORCH_CHECK(alloc->shape().size() == 1);
        auto size = dynamic_cast<Int*>(alloc->shape().at(0));
        TORCH_CHECK(
            size != nullptr && size->isConst() && size->value().value() == 9);
      }
    }
  }

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 3;
  auto ref = shift(t1, {0, 1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShift3ptStencil_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // 3-pt stencil
  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  std::vector<std::vector<int>> offsets = {{-1}, {1}};

  std::vector<TensorView*> tvs;
  for (const auto& offset : offsets) {
    tvs.push_back(shift(tv0, offset));
  }

  auto tv_out = tv0;

  for (auto tv : tvs) {
    tv_out = add(tv_out, tv);
  }

  tv_out = div(tv_out, IrBuilder::create<Double>(tvs.size() + 1));

  fusion.addOutput(tv_out);

  int split_factor = 4;

  tv_out->split(0, split_factor);

  // This seems fine but not verified yet
  // tv_out->axis(-1)->parallelize(ParallelType::Unswitch);

  auto cache = tv0->cache_after();

  tv0->computeAt(tv_out, 1);

  // Inline completely except for the cache
  for (auto tv : tvs) {
    tv->computeAt(tv_out, -1);
  }

  // cache allocation: (split_factor + 2)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == cache->name()) {
        TORCH_CHECK(alloc->shape().size() == 1);
        auto size = dynamic_cast<Int*>(alloc->shape().at(0));
        TORCH_CHECK(
            size != nullptr && size->isConst() &&
            size->value().value() == split_factor + 2);
      }
    }
  }

  cache->doubleBuffer();

  int numel_x = 99;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = (t0 + shift(t0, {-1}) + shift(t0, {1})) / 3;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShift5ptStencil_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // 5-pt stencil
  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  std::vector<std::vector<int>> offsets = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}};

  std::vector<TensorView*> tvs;
  for (const auto& offset : offsets) {
    tvs.push_back(shift(tv0, offset));
  }

  auto tv_out = tv0;

  for (auto tv : tvs) {
    tv_out = add(tv_out, tv);
  }

  tv_out = div(tv_out, IrBuilder::create<Double>(tvs.size() + 1));

  fusion.addOutput(tv_out);

  std::vector<int> split_factor({4, 8});

  tv_out->split(-1, split_factor[1]);
  tv_out->split(0, split_factor[0]);
  tv_out->reorder({{1, 2}, {2, 1}});

  auto cache = tv0->cache_after();

  tv0->computeAt(tv_out, 2);

  // Inline completely except for the cache
  for (auto tv : tvs) {
    tv->computeAt(tv_out, -1);
  }

  // cache allocation: (split_factor + 2) * (split_factor + 2)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == cache->name()) {
        TORCH_CHECK(alloc->shape().size() == 2);
        for (int i = 0; i < 2; ++i) {
          auto size = dynamic_cast<Int*>(alloc->shape().at(i));
          TORCH_CHECK(
              size != nullptr && size->isConst() &&
              size->value().value() == split_factor[i] + 2);
        }
      }
    }
  }

  cache->doubleBuffer();

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = t0;
  for (const auto& offset : offsets) {
    ref = ref + shift(t0, offset);
  }
  ref = ref / int(offsets.size() + 1);

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShift9ptStencil_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // 9-pt stencil
  std::vector<std::vector<int>> offsets;
  for (int i = -1; i < 2; ++i) {
    for (int j = -1; j < 2; ++j) {
      if (i == 0 && j == 0) {
        continue;
      }
      offsets.push_back({i, j});
    }
  }

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  std::vector<TensorView*> tvs;
  for (const auto& offset : offsets) {
    tvs.push_back(shift(tv0, offset));
  }

  auto tv_out = tv0;

  for (auto tv : tvs) {
    tv_out = add(tv_out, tv);
  }

  tv_out = div(tv_out, IrBuilder::create<Double>(tvs.size() + 1));

  fusion.addOutput(tv_out);

  std::vector<int> split_factor({4, 8});
  tv_out->split(-1, split_factor[1]);
  tv_out->split(0, split_factor[0]);
  tv_out->reorder({{1, 2}, {2, 1}});

  auto cache = tv0->cache_after();

  tv0->computeAt(tv_out, 2);

  // Inline completely except for the cache
  for (auto tv : tvs) {
    tv->computeAt(tv_out, -1);
  }

  // This seems fine but not yet verified
  // tv_out->axis(-1)->parallelize(ParallelType::Unswitch);

  // cache allocation: (split_factor + 2) * (split_factor + 2)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == cache->name()) {
        TORCH_CHECK(alloc->shape().size() == 2);
        for (int i = 0; i < 2; ++i) {
          auto size = dynamic_cast<Int*>(alloc->shape().at(i));
          TORCH_CHECK(
              size != nullptr && size->isConst() &&
              size->value().value() == split_factor[i] + 2);
        }
      }
    }
  }

  cache->doubleBuffer();

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = t0;
  for (const auto& offset : offsets) {
    ref = ref + shift(t0, offset);
  }
  ref = ref / int(offsets.size() + 1);

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftSmemBlocking_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {0, 1});
  fusion.addOutput(tv2);

  int smem_block_factor = 32;

  tv2->split(-1, smem_block_factor);

  tv0->computeAt(tv2, -2);

  tv1->axis(-1)->parallelize(ParallelType::TIDx);
  tv2->axis(-1)->parallelize(ParallelType::TIDx);

  tv1->setMemoryType(MemoryType::Shared);

  // tv1 allocation: (split_factor + 1)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == tv1->name()) {
        TORCH_CHECK(alloc->shape().size() == 1);
        for (int i = 0; i < 1; ++i) {
          auto size = dynamic_cast<Int*>(alloc->shape().at(i));
          TORCH_CHECK(
              size != nullptr && size->isConst() &&
              size->value().value() == smem_block_factor + 1);
        }
      }
    }
  }

  int numel_x = 100;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {0, 1});
  auto ref = t2;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShift3ptStencilParallel_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // 3-pt stencil
  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);
  std::vector<TensorView*> tvs;
  tvs.push_back(shift(tv0, {-1}));
  tvs.push_back(shift(tv0, {1}));

  auto tv_out = tv0;

  for (auto tv : tvs) {
    tv_out = add(tv_out, tv);
  }

  tv_out = div(tv_out, IrBuilder::create<Double>(tvs.size() + 1));

  fusion.addOutput(tv_out);

  int smem_block_factor = 32;

  tv_out->split(0, smem_block_factor);
  // tv_out->axis(-1)->parallelize(ParallelType::Unswitch);

  auto tv0_cache = tv0->cache_after();

  tv0->computeAt(tv_out, 1);

  for (auto tv : tvs) {
    tv->computeAt(tv_out, -1);
  }

  tv0_cache->setMemoryType(MemoryType::Shared);
  tv_out->axis(-1)->parallelize(ParallelType::TIDx);
  tv0_cache->axis(-1)->parallelize(ParallelType::TIDx);

  tv0_cache->doubleBuffer();

  int numel_x = 99;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = (t0 + shift(t0, {-1}) + shift(t0, {1})) / 3;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShift5ptStencilParallel_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // 5-pt stencil
  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  std::vector<std::vector<int>> offsets = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}};

  std::vector<TensorView*> tvs;
  for (const auto& offset : offsets) {
    tvs.push_back(shift(tv0, offset));
  }

  auto tv_out = tv0;

  for (auto tv : tvs) {
    tv_out = add(tv_out, tv);
  }

  tv_out = div(tv_out, IrBuilder::create<Double>(tvs.size() + 1));

  fusion.addOutput(tv_out);

  int smem_block_factor = 32;

  tv_out->split(-1, smem_block_factor);
  tv_out->split(0, smem_block_factor);

  tv_out->reorder({{1, 2}, {2, 1}});

  auto tv0_cache = tv0->cache_after();

  tv0->computeAt(tv_out, 2);

  for (auto tv : tvs) {
    tv->computeAt(tv_out, -1);
  }

  tv_out->axis(-1)->parallelize(ParallelType::TIDx);
  tv_out->axis(-2)->parallelize(ParallelType::TIDy);
  tv_out->axis(-3)->parallelize(ParallelType::BIDx);
  tv_out->axis(-4)->parallelize(ParallelType::BIDy);

  tv0_cache->setMemoryType(MemoryType::Shared);
  tv0_cache->axis(-1)->parallelize(ParallelType::TIDx);
  tv0_cache->axis(-2)->parallelize(ParallelType::TIDy);

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = t0;
  for (const auto& offset : offsets) {
    ref = ref + shift(t0, offset);
  }
  ref = ref / int(offsets.size() + 1);

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftMerge1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {-1, 1});
  fusion.addOutput(tv2);

  int split_factor = 4;

  tv2->split(-1, split_factor);
  tv2->split(0, split_factor);
  tv2->reorder({{1, 2}, {2, 1}});
  tv2->merge(2, 3);

  tv0->computeAt(tv2, 2);

  // t1 allocation: (split_factor + 1) * (split_factor + 1)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1) {
        TORCH_CHECK(alloc->shape().size() == 2);
        for (int i = 0; i < 2; ++i) {
          auto size = dynamic_cast<Int*>(alloc->shape().at(i));
          TORCH_CHECK(
              size != nullptr && size->isConst() &&
              size->value().value() == split_factor + 1);
        }
      }
    }
  }

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {-1, 1});
  auto ref = t2;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftMerge2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {1, -1});
  auto tv3 = shift(tv1, {-1, 1});
  auto tv4 = add(tv2, tv3);
  fusion.addOutput(tv4);

  int split_factor = 4;

  tv4->split(-1, split_factor);
  tv4->split(0, split_factor);
  tv4->reorder({{1, 2}, {2, 1}});
  tv4->merge(2, 3);

  tv0->computeAt(tv4, -2);

  // t1 allocation: (split_factor + 2) * (split_factor + 2)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1) {
        TORCH_CHECK(alloc->shape().size() == 2);
        for (int i = 0; i < 2; ++i) {
          auto size = dynamic_cast<Int*>(alloc->shape().at(i));
          TORCH_CHECK(
              size != nullptr && size->isConst() &&
              size->value().value() == split_factor + 2);
        }
      }
    }
  }

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {1, -1});
  auto t3 = shift(t1, {-1, 1});
  auto t4 = t2 + t3;

  TORCH_CHECK(t4.allclose(outputs[0]));
}

TEST_F(NVFuserTest, FusionShiftGlobal_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {0, 1});
  auto tv3 = shift(tv1, {-1, 0});
  auto tv4 = add(tv2, tv3);
  fusion.addOutput(tv4);

  tv1->split(-1, 4);
  tv2->split(-1, 8);
  tv3->split(-1, 2);
  tv4->split(-1, 3);

  tv1->merge(-2, -1);

  tv1->setMemoryType(MemoryType::Global);
  tv2->setMemoryType(MemoryType::Global);
  tv3->setMemoryType(MemoryType::Global);

  // t1 allocation: (t1.size[0] + 1) * (t1.size[1] + 1)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1) {
        TORCH_CHECK(alloc->shape().size() == 2);
        for (int i = 0; i < 2; ++i) {
          auto def =
              dynamic_cast<BinaryOp*>(alloc->shape().at(i)->definition());
          TORCH_CHECK(
              def != nullptr && def->getBinaryOpType() == BinaryOpType::Add);
          TORCH_CHECK(def->as<BinaryOp>()->lhs()->isA<NamedScalar>());
          auto rhs = dynamic_cast<Int*>(def->as<BinaryOp>()->rhs());
          TORCH_CHECK(rhs != nullptr && rhs->isConst());
          int rhs_value = *rhs->value();
          TORCH_CHECK(rhs_value == 1);
        }
      }
    }
  }

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {0, 1});
  auto t3 = shift(t1, {-1, 0});
  auto t4 = t2 + t3;
  auto ref = t4;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftDoubleSplitMerge1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = add(tv1, IrBuilder::create<Double>(2));
  auto tv3 = shift(tv2, {0, 1});
  fusion.addOutput(tv3);

  int split_factor1 = 8;
  int split_factor2 = 4;

  tv3->split(-1, split_factor1);

  tv0->computeAt(tv3, -2);

  tv1->split(-1, split_factor2);
  tv1->merge(-2, -1);

  // t1 and t2 allocation: (split_factor1 + 1)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1 || tensor_name == 2) {
        auto size = dynamic_cast<Int*>(alloc->shape().at(0));
        TORCH_CHECK(
            size != nullptr && size->isConst() &&
            size->value().value() == split_factor1 + 1);
      }
    }
  }

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 3;
  auto ref = shift(t1, {0, 1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftDoubleSplitMerge2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = add(tv1, IrBuilder::create<Double>(2));
  auto tv3 = shift(tv2, {1, 1});
  fusion.addOutput(tv3);

  auto out = tv3;

  int split_factor1 = 32;
  int split_factor2 = 4;

  out->split(-1, split_factor1);
  out->split(-1, split_factor2);
  out->split(0, split_factor1);
  out->split(1, split_factor2);
  out->reorder({{3, 1}, {1, 2}, {4, 3}, {2, 4}});
  out->merge(2, 3);
  out->merge(2, 3);
  out->merge(2, 3);
  out->merge(0, 1);

  TransformPropagator::from(out);

  tv0->computeAt(out, 1);

  out->axis(0)->parallelize(ParallelType::BIDx);
  out->axis(1)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out, {tv1, tv2});

  for (auto tv : {tv1, tv2}) {
    tv->setMemoryType(MemoryType::Shared);
  }

  // t1 and t2 allocation: (split_factor1 + 1) * (split_factor1 + 1)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1 || tensor_name == 2) {
        TORCH_CHECK(alloc->shape().size() == 2);
        for (int i = 0; i < 2; ++i) {
          auto size = dynamic_cast<Int*>(alloc->shape().at(i));
          TORCH_CHECK(
              size != nullptr && size->isConst() &&
              size->value().value() == split_factor1 + 1);
        }
      }
    }
  }

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = shift(t0 + 1 + 2, {1, 1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShift5ptStencilParallel1DThreadBlock_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // 5-pt stencil
  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  std::vector<std::vector<int>> offsets = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}};

  std::vector<TensorView*> tvs;
  for (const auto& offset : offsets) {
    tvs.push_back(shift(tv0, offset));
  }

  auto tv_out = tv0;

  for (auto tv : tvs) {
    tv_out = add(tv_out, tv);
  }

  tv_out = div(tv_out, IrBuilder::create<Double>(tvs.size() + 1));

  fusion.addOutput(tv_out);

  std::vector<int> split_factor({4, 32});

  tv_out->split(-1, split_factor[1]);
  tv_out->split(0, split_factor[0]);
  tv_out->reorder({{1, 2}, {2, 1}});

  auto tv0_cache = tv0->cache_after();

  // Merge the inner-most two axes and create
  // a 1D thread block of split_factor1*split_factor2 threads
  tv_out->merge(-2, -1);

  tv0->computeAt(tv_out, 2);

  // Inline completely except for the cache
  for (auto tv : tvs) {
    tv->computeAt(tv_out, -1);
  }

  tv0_cache->merge(-2, -1);

  tv_out->axis(-1)->parallelize(ParallelType::TIDx);
  tv_out->axis(1)->parallelize(ParallelType::BIDx);
  tv_out->axis(0)->parallelize(ParallelType::BIDy);

  tv0_cache->setMemoryType(MemoryType::Shared);
  tv0_cache->axis(-1)->parallelize(ParallelType::TIDx);

  // cache allocation: (split_factor1 + 2) * (split_factor2 + 2)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == tv0_cache->name()) {
        TORCH_CHECK(alloc->shape().size() == 2);
        for (int i = 0; i < 2; ++i) {
          auto size = dynamic_cast<Int*>(alloc->shape().at(i));
          TORCH_CHECK(
              size != nullptr && size->isConst() &&
              size->value().value() == split_factor[i] + 2);
        }
      }
    }
  }

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = t0;
  for (const auto& offset : offsets) {
    ref = ref + shift(t0, offset);
  }
  ref = ref / int(offsets.size() + 1);

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftChain1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = shift(tv0, {0, 1});
  auto tv2 = shift(tv1, {0, 1});
  fusion.addOutput(tv2);

  int split_factor = 4;
  tv2->split(-1, split_factor);

  tv0->computeAt(tv2, -2);

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = shift(shift(t0, {0, 1}), {0, 1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftChain2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = shift(tv0, {0, 1});
  auto tv2 = shift(tv1, {0, -1});
  fusion.addOutput(tv2);

  tv2->split(-1, 4);

  tv0->computeAt(tv2, -2);

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = shift(shift(t0, {0, 1}), {0, -1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftChain3_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {0, 1});
  auto tv3 = shift(tv2, {0, 1});
  fusion.addOutput(tv3);

  int split_factor = 4;
  tv3->split(-1, split_factor);

  tv0->computeAt(tv3, -2);

  // Halo size of tv1 is 2 as it needs to account for both of the two
  // shift operations , while that of tv2 is still just 1

  // tv1: (split_factor + 2)
  // tv2: (split_factor + 1)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1 || tensor_name == 2) {
        TORCH_CHECK(alloc->shape().size() == 1);
        for (int i = 0; i < 1; ++i) {
          auto size = dynamic_cast<Int*>(alloc->shape().at(i));
          TORCH_CHECK(size != nullptr && size->isConst());
          if (tensor_name == 1) {
            TORCH_CHECK(size->value().value() == split_factor + 2);
          } else if (tensor_name == 2) {
            TORCH_CHECK(size->value().value() == split_factor + 1);
          }
        }
      }
    }
  }

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {0, 1});
  auto t3 = shift(t2, {0, 1});
  auto ref = t3;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftChain4_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = shift(tv0, {1, -1});
  auto tv2 = shift(tv1, {2, -2});
  auto tv3 = shift(tv2, {3, -3});
  auto tv4 = shift(tv3, {4, -4});
  auto tv_out = tv4;

  fusion.addOutput(tv_out);

  int split_factor = 4;

  tv_out->split(-1, split_factor);
  tv_out->split(0, split_factor);
  tv_out->reorder({{1, 2}, {2, 1}});

  tv0->computeAt(tv_out, 2);

  tv1->merge(-2, -1);
  tv2->merge(-2, -1);
  tv3->merge(-2, -1);

  // tv1: (split_factor + 9) * (split_factor + 9)
  // tv2: (split_factor + 7) * (split_factor + 7)
  // tv3: (split_factor + 4) * (split_factor + 4)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == 1 || tensor_name == 2) {
        TORCH_CHECK(alloc->shape().size() == 2);
        for (int i = 0; i < 2; ++i) {
          auto size = dynamic_cast<Int*>(alloc->shape().at(i));
          TORCH_CHECK(size != nullptr && size->isConst());
          auto size_val = size->value().value();
          if (tensor_name == 1) {
            TORCH_CHECK(size_val == split_factor + 9);
          } else if (tensor_name == 2) {
            TORCH_CHECK(size_val == split_factor + 7);
          } else if (tensor_name == 3) {
            TORCH_CHECK(size_val == split_factor + 4);
          }
        }
      }
    }
  }

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = shift(t0, {1, -1});
  auto t2 = shift(t1, {2, -2});
  auto t3 = shift(t2, {3, -3});
  auto t4 = shift(t3, {4, -4});
  auto ref = t4;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShift5ptStencilChain_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  std::vector<std::vector<int>> offsets = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}};

  // First stencil: 5pt stencil
  // stencil1 = (tv0 + tv0[+1][0] + tv0[-1][0] + tv0[0][+1] + tv0[0][-1]) / 5
  std::vector<TensorView*> tv_stencil1_shifts;
  for (const auto& offset : offsets) {
    tv_stencil1_shifts.push_back(shift(tv0, offset));
  }

  auto tv_stencil1 = tv0;
  for (auto tv : tv_stencil1_shifts) {
    tv_stencil1 = add(tv_stencil1, tv);
  }

  tv_stencil1 = div(
      tv_stencil1, IrBuilder::create<Double>(tv_stencil1_shifts.size() + 1));

  // Second stencil: Same 5pt stencil
  std::vector<TensorView*> tv_stencil2_shifts;
  for (const auto& offset : offsets) {
    tv_stencil2_shifts.push_back(shift(tv_stencil1, offset));
  }

  auto tv_stencil2 = tv_stencil1;
  for (auto tv : tv_stencil2_shifts) {
    tv_stencil2 = add(tv_stencil2, tv);
  }

  tv_stencil2 = div(
      tv_stencil2, IrBuilder::create<Double>(tv_stencil2_shifts.size() + 1));

  auto tv_out = tv_stencil2;

  fusion.addOutput(tv_out);

  auto tv0_cache = tv0->cache_after();

  std::vector<int> split_factor({16, 16});

  tv_out->split(-1, split_factor[1]);
  tv_out->split(0, split_factor[0]);
  tv_out->reorder({{1, 2}, {2, 1}});

  tv0->computeAt(tv_out, 2);

  // Inline completely all inputs to the first stencil output, except for the
  // tv0 cache
  for (auto tv : tv_stencil1_shifts) {
    tv->computeAt(tv_stencil1, -1);
  }

  // Inline completely all inputs to the second stencil output, except
  // for the first stencil output
  for (auto tv : tv_stencil2_shifts) {
    tv->computeAt(tv_stencil2, -1);
  }

  tv_out->axis(1)->parallelize(ParallelType::BIDx);
  tv_out->axis(0)->parallelize(ParallelType::BIDy);

  auto all_values = DependencyCheck::getAllValsBetween(
      {fusion.inputs().begin(), fusion.inputs().end()}, fusion.outputs());
  for (auto tv : ir_utils::filterByType<TensorView>(all_values)) {
    tv->axis(-1)->parallelize(ParallelType::TIDx);
    tv->axis(-2)->parallelize(ParallelType::TIDy);
  }

  tv0_cache->setMemoryType(MemoryType::Shared);
  tv_stencil1->setMemoryType(MemoryType::Shared);

  // tv0_cache: (split_factor + 4) * (split_factor + 4)
  // tv_stencil1: (split_factor + 2) * (split_factor + 2)
  GpuLower gpulw(&fusion);
  for (const auto expr : gpulw.kernel()->unordered_exprs()) {
    if (auto alloc = dynamic_cast<kir::Allocate*>(expr)) {
      auto tensor_name = alloc->buffer()->name();
      if (tensor_name == tv0_cache->name() ||
          tensor_name == tv_stencil1->name()) {
        TORCH_CHECK(alloc->shape().size() == 2);
        for (int i = 0; i < 2; ++i) {
          auto size = dynamic_cast<Int*>(alloc->shape().at(i));
          TORCH_CHECK(size != nullptr && size->isConst());
          if (tensor_name == tv0_cache->name()) {
            TORCH_CHECK(size->value().value() == split_factor[i] + 4);
          } else if (tensor_name == tv_stencil1->name()) {
            TORCH_CHECK(size->value().value() == split_factor[i] + 2);
          }
        }
      }
    }
  }

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto stencil1 = t0;
  for (const auto& offset : offsets) {
    stencil1 = stencil1 + shift(t0, offset);
  }
  stencil1 = stencil1 / int(offsets.size() + 1);
  auto stencil2 = stencil1;
  for (const auto& offset : offsets) {
    stencil2 = stencil2 + shift(stencil1, offset);
  }
  stencil2 = stencil2 / int(offsets.size() + 1);
  auto ref = stencil2;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// Shift a reduced tensor
TEST_F(NVFuserTest, FusionShiftReduction1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = sum(tv1, {1});
  auto tv3 = shift(tv2, {1});
  fusion.addOutput(tv3);

  tv3->split(0, 4);
  tv0->computeAt(tv3, 1);
  tv0->computeAt(tv2, -1);

  const int numel_x = 9;
  const int numel_y = 11;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = sum(t1, {1});
  auto t3 = shift(t2, {1});
  auto ref = t3;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// Parallelized version of FusionShiftReduction1
TEST_F(NVFuserTest, FusionShiftReduction2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = sum(tv1, {1});
  auto tv3 = shift(tv2, {1});
  fusion.addOutput(tv3);

  tv3->split(0, 4);
  tv0->computeAt(tv3, 1);

  tv2->split(-1, 32);
  tv0->computeAt(tv2, -1);

  tv2->axis(-1)->parallelize(ParallelType::TIDx);

  tv2->setMemoryType(MemoryType::Shared);

  const int numel_x = 201;
  const int numel_y = 301;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = sum(t1, {1});
  auto t3 = shift(t2, {1});
  auto ref = t3;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftRfactor1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = sum(tv1, {1});
  auto tv3 = shift(tv2, {1});
  fusion.addOutput(tv3);

  tv3->split(0, 4);
  tv0->computeAt(tv3, 1);

  tv2->split(-1, 32);
  auto rf = tv2->rFactor({-2});
  tv0->computeAt(tv2, -1);
  tv0->computeAt(rf, -1);

  tv2->axis(-1)->parallelize(ParallelType::TIDx);

  tv2->setMemoryType(MemoryType::Shared);

  const int numel_x = 201;
  const int numel_y = 301;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = sum(t1, {1});
  auto t3 = shift(t2, {1});
  auto ref = t3;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftBcast1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);
  auto tv1 = makeSymbolicTensor(2);
  fusion.addInput(tv1);
  auto tv2 = broadcast(tv0, {false, true});
  auto tv3 = shift(tv2, {0, 1});
  auto tv4 = add(tv3, tv1);
  fusion.addOutput(tv4);

  tv0->computeAt(tv4, -1);
  tv1->computeAt(tv4, -1);

  const int numel_x = 9;
  const int numel_y = 11;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x}, options);
  at::Tensor t1 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0, t1};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t4 = t0.unsqueeze(-1).expand({numel_x, numel_y}) + t1;
  auto ref = t4;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftBcast2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);
  auto tv1 = makeSymbolicTensor(2);
  fusion.addInput(tv1);
  auto tv2 = broadcast(tv0, {false, true});
  auto tv3 = shift(tv2, {1, 0});
  auto tv4 = add(tv3, tv1);
  fusion.addOutput(tv4);

  tv4->split(0, 4);
  tv0->computeAt(tv4, 1);

  const int numel_x = 9;
  const int numel_y = 11;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x}, options);
  at::Tensor t1 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0, t1};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t2 = t0.unsqueeze(-1).expand({numel_x, numel_y});
  auto t3 = shift(t2, {1, 0});
  auto ref = t3 + t1;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// Combine ShiftBcast1 and ShiftBcast2 with parallelization
TEST_F(NVFuserTest, FusionShiftBcast3_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);
  auto tv1 = makeSymbolicTensor(2);
  fusion.addInput(tv1);
  auto tv2 = broadcast(tv0, {false, true});
  auto tv3 = shift(tv2, {1, 0});
  auto tv4 = shift(tv2, {0, 1});
  auto tv5 = shift(tv2, {-1, -1});
  auto tv6 = add(tv3, tv4);
  auto tv7 = add(tv6, tv5);
  auto tv8 = add(tv7, tv1);
  fusion.addOutput(tv8);

  tv8->split(0, 4);
  tv8->split(-1, 4);
  tv0->computeAt(tv8, 1);

  tv8->axis(-1)->parallelize(ParallelType::TIDx);
  for (auto tv : {tv8, tv7, tv6, tv5, tv4, tv3, tv2}) {
    tv->axis(1)->parallelize(ParallelType::TIDy);
  }

  tv2->setMemoryType(MemoryType::Shared);

  const int numel_x = 101;
  const int numel_y = 201;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x}, options);
  at::Tensor t1 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0, t1};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t2 = t0.unsqueeze(-1).expand({numel_x, numel_y});
  auto t3 = shift(t2, {1, 0});
  auto t4 = t2;
  auto t5 = shift(t2, {-1, 0});
  auto ref = t3 + t4 + t5 + t1;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// See issue #893
TEST_F(NVFuserTest, FusionShiftSyncPlacement1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = add(tv0, IrBuilder::create<Double>(2));
  auto tv3 = add(tv1, tv2);
  auto tv4 = shift(tv3, {0, 1});
  fusion.addOutput(tv4);

  tv4->split(1, 8);
  tv0->computeAt(tv4, 2);

  tv2->computeAt(tv3, -1);

  tv1->setMemoryType(MemoryType::Shared);
  tv3->setMemoryType(MemoryType::Shared);

  tv1->axis(-1)->parallelize(ParallelType::TIDx);
  tv3->axis(-1)->parallelize(ParallelType::TIDx);
  tv4->axis(-1)->parallelize(ParallelType::TIDx);

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = t0 + 2;
  auto t3 = add(t1, t2);
  auto t4 = shift(t3, {0, 1});

  testValidate(&fusion, outputs, inputs, {t4}, __LINE__, __FILE__);
}

// See issue #893. Top-level placement.
TEST_F(NVFuserTest, FusionShiftSyncPlacement2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = add(tv0, IrBuilder::create<Double>(2));
  auto tv3 = add(tv1, tv2);
  auto tv4 = shift(tv3, {1});
  fusion.addOutput(tv4);

  tv2->computeAt(tv3, -1);

  tv1->setMemoryType(MemoryType::Shared);
  tv3->setMemoryType(MemoryType::Shared);

  tv1->axis(-1)->parallelize(ParallelType::TIDx);
  tv3->axis(-1)->parallelize(ParallelType::TIDx);
  tv4->axis(-1)->parallelize(ParallelType::TIDx);

  int numel_x = 99;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = t0 + 2;
  auto t3 = add(t1, t2);
  auto t4 = shift(t3, {1});

  testValidate(&fusion, outputs, inputs, {t4}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftSyncPlacement3_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = add(tv1, IrBuilder::create<Double>(2));
  auto tv3 = shift(tv2, {1});
  fusion.addOutput(tv3);

  // This doesn't work. syncthreads is needed between tv1 and tv2, but
  // both the loop extent of both tv1 and tv2 has halo, so the loop is
  // not eliminated even though it is parallelized. Moving syncthreads
  // out of the loop would make it placed before tv1, which would make
  // it meaningless.
  // Ideally, an exception should be thrown at this computeAt, but at
  // this point, the fusion is not yet parallelized, nor memory type
  // is set, so this computeAt itself is not an error yet.
  tv1->computeAt(tv2, -1);

  tv1->setMemoryType(MemoryType::Shared);
  tv2->setMemoryType(MemoryType::Shared);

  tv1->axis(-1)->parallelize(ParallelType::TIDx);
  tv2->axis(-1)->parallelize(ParallelType::TIDx);
  tv3->axis(-1)->parallelize(ParallelType::TIDx);

  // The error should be detected when the fusion is lowered.
  ASSERT_ANY_THROW(fusion.printKernel());
}

// Based on original CUDA provided by Vishal Mehta.
// Major differences with the original version:
// - The original version uses additional 2 warps to load the halos
//   along the Y dimension. The other 10 warps are used to load a 32x10
//   tile, and all warps will do coalesced loads. No such optimization
//   is done in the fuser version.
TEST_F(NVFuserTest, FusionHdiff_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);
  auto coeff = makeSymbolicTensor(3);
  fusion.addInput(coeff);

  std::vector<std::vector<int>> offsets{
      {0, 1, 0}, {0, -1, 0}, {0, 0, 1}, {0, 0, -1}};

  // T2, T3, T4, T5
  std::vector<TensorView*> inp_neighbors;
  for (const auto& offset : offsets) {
    inp_neighbors.push_back(shift(inp, offset, false));
  }

  // T8
  TensorView* sum_of_neighbors = nullptr;
  for (auto inp_neighbor : inp_neighbors) {
    if (sum_of_neighbors == nullptr) {
      sum_of_neighbors = inp_neighbor;
    } else {
      sum_of_neighbors = add(sum_of_neighbors, inp_neighbor);
    }
  }

  // T9 = T0 * 4
  // T10 = T9 - T8
  auto lap = sub(mul(inp, IrBuilder::create<Double>(4)), sum_of_neighbors);

  // T11 = shift(T10)
  // T12 = T11 - T10
  auto flx = sub(shift(lap, {0, 0, -1}, false), lap);
  // T14 = T13 - T0
  // T15 = T12 * T14
  // T16 = T15 > 0
  // T17 = T16 ? 0 : T12
  auto flx_cond =
      gt(mul(flx, sub(shift(inp, {0, 0, -1}, false), inp)),
         IrBuilder::create<Double>(0));
  auto flx0 = where(flx_cond, IrBuilder::create<Double>(0), flx);

  // T18 = shift(T10)
  // T19 = T18 - T10
  auto fly = sub(shift(lap, {0, -1, 0}, false), lap);
  // T20 = shift(T0)
  // T21 = T20 - T0
  // T22 = T19 * T21
  // T23 = T22 > 0
  auto fly_cond =
      gt(mul(fly, sub(shift(inp, {0, -1, 0}, false), inp)),
         IrBuilder::create<Double>(0));
  // T24 = T23 ? 0 : T19
  auto fly0 = where(fly_cond, IrBuilder::create<Double>(0), fly);

  // T25 = shift(flx0)
  // T26 = T17 - T25
  // T27 = shift(fly0)
  // T28 = T24 - T27
  // T29 = T26 + T28
  // T30 = T1 * T29
  // T31 = T0 - T30
  auto out =
      sub(inp,
          mul(coeff,
              add(sub(flx0, shift(flx0, {0, 0, 1}, false)),
                  sub(fly0, shift(fly0, {0, 1, 0}, false)))));

  fusion.addOutput(out);

  /////////////////////////////////
  // Scheduling
  /////////////////////////////////

  out->setContiguity(false);

  // Step 1: 2D Tiling

  const int tile_x = 32;
  const int tile_y = 8;

  out->split(-1, tile_x);
  out->split(-3, tile_y);
  out->reorder({{-2, -3}});
  inp->computeAt(out, -3);
  coeff->computeAt(out, -3);

  // Step 2: Inlining

  // Inline inputs to lap
  auto lap_vals = DependencyCheck::getAllValsBetween({inp}, {lap});
  for (auto val : ir_utils::filterByType<TensorView>(lap_vals)) {
    if (val != lap && val != inp) {
      val->computeAt(lap, -1);
    }
  }

  // Inline inputs to flx0
  auto flx0_vals = DependencyCheck::getAllValsBetween({lap, inp}, {flx0});
  for (auto val : ir_utils::filterByType<TensorView>(flx0_vals)) {
    if (val != lap && val != flx0 && val != inp) {
      val->computeAt(flx0, -1);
    }
  }

  // Inline inputs to fly0
  auto flxy_vals = DependencyCheck::getAllValsBetween({lap, inp}, {fly0});
  for (auto val : ir_utils::filterByType<TensorView>(flxy_vals)) {
    if (val != lap && val != fly0 && val != inp) {
      val->computeAt(fly0, -1);
    }
  }

  // Inline inputs to out
  auto out_vals = DependencyCheck::getAllValsBetween({flx0, fly0}, {out});
  for (auto val : ir_utils::filterByType<TensorView>(out_vals)) {
    if (val != flx0 && val != fly0 && val != out) {
      val->computeAt(out, -1);
    }
  }

  // Step 3: Parallelization

  // Block parallelization
  out->axis(0)->parallelize(ParallelType::BIDz);
  out->axis(1)->parallelize(ParallelType::BIDy);
  out->axis(2)->parallelize(ParallelType::BIDx);
  // Thread parallelization
  out->axis(3)->parallelize(ParallelType::TIDy);
  out->axis(4)->parallelize(ParallelType::TIDx);
  // Apply the same parallelization to all other tensors
  scheduler_utils::parallelizeAllLike(out, ir_utils::allTvs(&fusion));

  // Store intermediate stencil results on smem so that they can be
  // accessed by threads
  for (auto tv : {flx0, fly0, lap}) {
    tv->setMemoryType(MemoryType::Shared);
  }

  /////////////////////////////////
  int numel_x = 101;
  int numel_y = 99;
  int numel_z = 10;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor inp_at = at::randn({numel_z, numel_y, numel_x}, options);
  at::Tensor coeff_at = at::randn({numel_z, numel_y, numel_x}, options);
  std::vector<IValue> inputs = {inp_at, coeff_at};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto fuser_output = fe.runFusion(inputs)[0];

  // Trim the outer rim
  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(0, at::indexing::None),
      at::indexing::Slice(2, -2),
      at::indexing::Slice(2, -2)};
  fuser_output = fuser_output.index(indices);

  {
    at::Tensor zeros = at::zeros({numel_z, numel_y, numel_x}, options);
    auto lap = inp_at * 4 -
        (shift(inp_at, {0, 1, 0}) + shift(inp_at, {0, -1, 0}) +
         shift(inp_at, {0, 0, 1}) + shift(inp_at, {0, 0, -1}));
    auto flx = shift(lap, {0, 0, -1}) - lap;
    auto flx_cond = (flx * (shift(inp_at, {0, 0, -1}) - inp_at)) > 0;
    auto flx0 = at::where(flx_cond, zeros, flx);
    auto fly = shift(lap, {0, -1, 0}) - lap;
    auto fly_cond = (fly * (shift(inp_at, {0, -1, 0}) - inp_at)) > 0;
    auto fly0 = at::where(fly_cond, zeros, fly);

    auto ref = inp_at -
        coeff_at *
            ((flx0 - shift(flx0, {0, 0, 1})) + (fly0 - shift(fly0, {0, 1, 0})));
    ref = ref.index(indices);

    testValidate(&fusion, {fuser_output}, inputs, {ref}, __LINE__, __FILE__);
  }
}

TEST_F(NVFuserTest, FusionHdiffPartialSplitUnswitch_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);
  auto coeff = makeSymbolicTensor(3);
  fusion.addInput(coeff);

  std::vector<std::vector<int>> offsets{
      {0, 1, 0}, {0, -1, 0}, {0, 0, 1}, {0, 0, -1}};

  // T2, T3, T4, T5
  std::vector<TensorView*> inp_neighbors;
  for (const auto& offset : offsets) {
    inp_neighbors.push_back(shift(inp, offset, false));
  }

  // T8
  TensorView* sum_of_neighbors = nullptr;
  for (auto inp_neighbor : inp_neighbors) {
    if (sum_of_neighbors == nullptr) {
      sum_of_neighbors = inp_neighbor;
    } else {
      sum_of_neighbors = add(sum_of_neighbors, inp_neighbor);
    }
  }

  // T9 = T0 * 4
  // T10 = T9 - T8
  auto lap = sub(mul(inp, IrBuilder::create<Double>(4)), sum_of_neighbors);

  // T11 = shift(T10)
  // T12 = T11 - T10
  auto flx = sub(shift(lap, {0, 0, -1}, false), lap);
  // T14 = T13 - T0
  // T15 = T12 * T14
  // T16 = T15 > 0
  // T17 = T16 ? 0 : T12
  auto flx_cond =
      gt(mul(flx, sub(shift(inp, {0, 0, -1}, false), inp)),
         IrBuilder::create<Double>(0));
  auto flx0 = where(flx_cond, IrBuilder::create<Double>(0), flx);

  // T18 = shift(T10)
  // T19 = T18 - T10
  auto fly = sub(shift(lap, {0, -1, 0}, false), lap);
  // T20 = shift(T0)
  // T21 = T20 - T0
  // T22 = T19 * T21
  // T23 = T22 > 0
  auto fly_cond =
      gt(mul(fly, sub(shift(inp, {0, -1, 0}, false), inp)),
         IrBuilder::create<Double>(0));
  // T24 = T23 ? 0 : T19
  auto fly0 = where(fly_cond, IrBuilder::create<Double>(0), fly);

  // T25 = shift(flx0)
  // T26 = T17 - T25
  // T27 = shift(fly0)
  // T28 = T24 - T27
  // T29 = T26 + T28
  // T30 = T1 * T29
  // T31 = T0 - T30
  auto out =
      sub(inp,
          mul(coeff,
              add(sub(flx0, shift(flx0, {0, 0, 1}, false)),
                  sub(fly0, shift(fly0, {0, 1, 0}, false)))));

  fusion.addOutput(out);

  out->setContiguity(false);

  /////////////////////////////////
  // Scheduling
  /////////////////////////////////

  const auto all_vals = fusion.usedMathVals();
  const std::vector<TensorView*> all_tensors(
      {ir_utils::filterByType<TensorView>(all_vals).begin(),
       ir_utils::filterByType<TensorView>(all_vals).end()});

  // Step 1: Blocking
  // - Thread block size: (tile_x, tile_y)
  // - Each thread computes a vertical column of length tile_z along the Z
  // axis.
  // - Grid dize: (NX / block_x, NY / block_y, NZ / tile_z)

  const int tile_x = 32;
  const int tile_y = 8;
  const int tile_z = 16;

  out->split(0, tile_z);
  out->split(-1, tile_x, true, true);
  out->split(-3, tile_y, true, true);
  // out: [NZ/tz, tz, NY/by, by, NX/bx, bx]
  out->reorder({{1, 3}, {2, 1}, {3, 4}, {4, 2}});
  // out: [NZ/tz, NY/by, NX/bx, tz, by, bx]

  TransformPropagator::from(out);

  inp->computeAt(out, 4);

  // Step 2: Inlining

  // Inline inputs to lap
  auto lap_vals = DependencyCheck::getAllValsBetween({inp}, {lap});
  for (auto val : ir_utils::filterByType<TensorView>(lap_vals)) {
    if (val != lap && val != inp) {
      val->computeAt(lap, -1);
    }
  }

  // Inline inputs to flx0
  auto flx0_vals = DependencyCheck::getAllValsBetween({lap, inp}, {flx0});
  for (auto val : ir_utils::filterByType<TensorView>(flx0_vals)) {
    if (val != lap && val != flx0 && val != inp) {
      val->computeAt(flx0, -1);
    }
  }

  // Inline inputs to fly0
  auto flxy_vals = DependencyCheck::getAllValsBetween({lap, inp}, {fly0});
  for (auto val : ir_utils::filterByType<TensorView>(flxy_vals)) {
    if (val != lap && val != fly0 && val != inp) {
      val->computeAt(fly0, -1);
    }
  }

  // Inline inputs to out
  auto out_vals = DependencyCheck::getAllValsBetween({flx0, fly0}, {out});
  for (auto val : ir_utils::filterByType<TensorView>(out_vals)) {
    if (val != flx0 && val != fly0 && val != out) {
      val->computeAt(out, -1);
    }
  }

  // Step 3: Parallelization

  // Block parallelization
  out->axis(0)->parallelize(ParallelType::BIDz);
  out->axis(1)->parallelize(ParallelType::BIDy);
  out->axis(2)->parallelize(ParallelType::BIDx);
  out->axis(4)->parallelize(ParallelType::TIDy);
  out->axis(5)->parallelize(ParallelType::TIDx);
  // Unswitch at the tz axis
  out->axis(3)->parallelize(ParallelType::Unswitch);

  scheduler_utils::parallelizeAllLike(out, all_tensors);

  // These need to be on smem
  for (auto tv : {flx0, fly0, lap}) {
    tv->setMemoryType(MemoryType::Shared);
  }

  /////////////////////////////////
  const int halo_extent = 2;
  const int numel_x = 64 + halo_extent * 2;
  const int numel_y = 64 + halo_extent * 2;
  const int numel_z = 32;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor inp_at = at::randn({numel_z, numel_y, numel_x}, options);
  at::Tensor coeff_at = at::randn({numel_z, numel_y, numel_x}, options);
  std::vector<IValue> inputs = {inp_at, coeff_at};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto fuser_output = fe.runFusion(inputs)[0];

  // Trim the outer rim
  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(0, at::indexing::None),
      at::indexing::Slice(2, -2),
      at::indexing::Slice(2, -2)};
  fuser_output = fuser_output.index(indices);

  {
    at::Tensor zeros = at::zeros({numel_z, numel_y, numel_x}, options);
    auto lap = inp_at * 4 -
        (shift(inp_at, {0, 1, 0}) + shift(inp_at, {0, -1, 0}) +
         shift(inp_at, {0, 0, 1}) + shift(inp_at, {0, 0, -1}));
    auto flx = shift(lap, {0, 0, -1}) - lap;
    auto flx_cond = (flx * (shift(inp_at, {0, 0, -1}) - inp_at)) > 0;
    auto flx0 = at::where(flx_cond, zeros, flx);
    auto fly = shift(lap, {0, -1, 0}) - lap;
    auto fly_cond = (fly * (shift(inp_at, {0, -1, 0}) - inp_at)) > 0;
    auto fly0 = at::where(fly_cond, zeros, fly);

    auto ref = inp_at -
        coeff_at *
            ((flx0 - shift(flx0, {0, 0, 1})) + (fly0 - shift(fly0, {0, 1, 0})));
    ref = ref.index(indices);

    testValidate(&fusion, {fuser_output}, inputs, {ref}, __LINE__, __FILE__);
  }
}

// 3x3 max pooling
TEST_F(NVFuserTest, FusionMaxPooling_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Format: CHW
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // 3x3 pooling of the HW spatial domain
  std::vector<std::vector<int>> offsets;
  for (int i = -1; i <= 1; ++i) {
    for (int j = -1; j <= 1; ++j) {
      if (i == 0 && j == 0) {
        continue;
      }
      offsets.push_back({i, j});
    }
  }

  std::vector<TensorView*> inp_tile({inp});
  for (auto offset : offsets) {
    offset.insert(offset.begin(), 0);
    inp_tile.push_back(shift(inp, offset));
  }

  TensorView* max_tensor = nullptr;
  for (auto tv : inp_tile) {
    if (max_tensor == nullptr) {
      max_tensor = tv;
    } else {
      max_tensor = binaryOp(BinaryOpType::Max, max_tensor, tv);
    }
  }

  fusion.addOutput(max_tensor);

  ////////////////////////////////////

  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Tiling the spatial domain
  const int tile_x = 32;
  const int tile_y = 8;

  max_tensor->split(-2, tile_y);
  max_tensor->axis(-2)->parallelize(ParallelType::TIDy);
  max_tensor->split(-1, tile_x);
  max_tensor->axis(-1)->parallelize(ParallelType::TIDx);
  max_tensor->reorder({{-3, -2}});

  inp_cache->computeAt(max_tensor, 3);
  inp_cache->axis(-2)->parallelize(ParallelType::TIDy);
  inp_cache->axis(-1)->parallelize(ParallelType::TIDx);
  inp_cache->setMemoryType(MemoryType::Shared);

  auto max_tensor_dep =
      DependencyCheck::getAllValsBetween({inp_cache}, {max_tensor});
  for (auto tv : ir_utils::filterByType<TensorView>(max_tensor_dep)) {
    if (tv == inp_cache || tv == max_tensor) {
      continue;
    }
    tv->computeAt(max_tensor, -1);
  }

  max_tensor->axis(0)->parallelize(ParallelType::BIDx);

  const int hw = 50;
  const int num_channels = 20;
  const int pooling_window = 3;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor aten_inp = at::randn({num_channels, hw, hw}, options);
  // shift always pads by zero, so if all surrounding values are
  // negative, max pooling would pick a padded value, which isn't the
  // correct behavior. We need to be able to choose the value of
  // padding. In this case, padding by the minimum value would not
  // have this problem. For now, avoid the problem by making sure all
  // values are not negative.
  aten_inp = at::abs(aten_inp);
  std::vector<IValue> inputs = {aten_inp};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = at::max_pool2d(
      aten_inp, {pooling_window, pooling_window}, {1, 1}, {1, 1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionGather1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  const std::vector<int> window_shape = {1, 3};
  const std::vector<std::vector<int>> padding_width = {{0, 0}, {1, 1}};

  auto tv1 = gather(tv0, window_shape, padding_width);

  fusion.addOutput(tv1);

  const int s1 = 11;
  const int s2 = 13;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({s1, s2}, options);

  auto ref = gather(t0, window_shape, padding_width);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto outputs = fe.runFusion({t0});

  TORCH_CHECK(ref.equal(outputs[0]));
}

TEST_F(NVFuserTest, FusionGather2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int> window_shape = {1, 3};
  const std::vector<std::vector<int>> padding_width = {{0, 0}, {1, 1}};

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));

  auto tv2 = gather(tv1, window_shape, padding_width);

  auto tv3 = sum(tv2, {-1});

  fusion.addOutput(tv3);

  tv3->split(1, 32);
  tv0->computeAt(tv3, 2);
  tv2->computeAt(tv3, -1);

  tv3->axis(0)->parallelize(ParallelType::BIDy);
  tv3->axis(1)->parallelize(ParallelType::BIDx);
  tv3->axis(2)->parallelize(ParallelType::TIDx);
  tv1->axis(2)->parallelize(ParallelType::TIDx);

  tv1->setMemoryType(MemoryType::Shared);

  const int s1 = 99;
  const int s2 = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({s1, s2}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = gather(t1, window_shape, padding_width);
  auto ref = sum(t2, {-1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionGather3_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  const std::vector<int> window_shape = {1, 3};
  const std::vector<std::vector<int>> padding_width = {{0, 0}, {0, 0}};

  auto tv1 = gather(tv0, window_shape, padding_width);

  fusion.addOutput(tv1);

  const int s1 = 11;
  const int s2 = 13;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  std::vector<int64_t> size({s1, s2});
  at::Tensor t0 = at::randn(size, options);
  size.insert(size.end(), window_shape.begin(), window_shape.end());
  // Use a pre-allocated output tensor filled with 1 so that invalid
  // writes to outside valid ranges can be detected
  at::Tensor output = at::ones(size, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto outputs = fe.runFusion({t0}, {output});

  auto ref = gather(t0, window_shape, padding_width);
  TORCH_CHECK(ref.equal(outputs[0]));
}

TEST_F(NVFuserTest, FusionGather4_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  const std::vector<int> window_shape = {3, 3};
  const std::vector<std::vector<int>> padding_width = {{0, 0}, {0, 0}};

  auto tv1 = gather(tv0, window_shape, padding_width);

  fusion.addOutput(tv1);

  const int s1 = 11;
  const int s2 = 13;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  std::vector<int64_t> size({s1, s2});
  at::Tensor t0 = at::randn(size, options);
  size.insert(size.end(), window_shape.begin(), window_shape.end());
  // Use a pre-allocated output tensor filled with 1 so that invalid
  // writes to outside valid ranges can be detected
  at::Tensor output = at::ones(size, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto outputs = fe.runFusion({t0}, {output});

  auto ref = gather(t0, window_shape, padding_width);

  TORCH_CHECK(ref.equal(outputs[0]));
}

TEST_F(NVFuserTest, FusionGather5_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  const std::vector<int> window_shape = {3, 3};
  const std::vector<std::vector<int>> padding_width = {{1, 0}, {0, 1}};

  auto tv1 = gather(tv0, window_shape, padding_width);

  fusion.addOutput(tv1);

  const int s1 = 11;
  const int s2 = 13;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  std::vector<int64_t> size({s1, s2});
  at::Tensor t0 = at::randn(size, options);
  size.insert(size.end(), window_shape.begin(), window_shape.end());
  // Use a pre-allocated output tensor filled with 1 so that invalid
  // writes to outside valid ranges can be detected
  at::Tensor output = at::ones(size, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto outputs = fe.runFusion({t0}, {output});

  auto ref = gather(t0, window_shape, padding_width);

  TORCH_CHECK(ref.equal(outputs[0]));
}

// Conv-like pattern with no padding
TEST_F(NVFuserTest, FusionGather6_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  const std::vector<int> window_shape = {3, 4};
  const std::vector<std::vector<int>> padding_width = {{0, 0}, {0, 0}};

  auto tv1 = gather(tv0, window_shape, padding_width);

  fusion.addOutput(tv1);

  // Blocking the spatial dimensions
  const int block_x = 16;
  const int block_y = 8;

  auto tv0_cache = tv0->cache_after();
  auto out = tv1;
  auto out_cache = out->cache_before();

  out->split(1, block_x);
  out->split(0, block_y);
  out->reorder({{1, 2}, {2, 1}});

  TransformPropagator::from(out);

  tv0->computeAt(out, 2);

  tv0_cache->setMemoryType(MemoryType::Shared);

  out->axis(0)->parallelize(ParallelType::BIDy);
  out->axis(1)->parallelize(ParallelType::BIDx);
  out->axis(2)->parallelize(ParallelType::TIDy);
  out->axis(3)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(out, ir_utils::allTvs(&fusion));

  const int s1 = 101;
  const int s2 = 99;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  std::vector<int64_t> size({s1, s2});
  at::Tensor t0 = at::randn(size, options);
  size.insert(size.end(), window_shape.begin(), window_shape.end());
  // Use a pre-allocated output tensor filled with 1 so that invalid
  // writes to outside valid ranges can be detected
  at::Tensor output = at::ones(size, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto outputs = fe.runFusion({t0}, {output});

  auto ref = gather(t0, window_shape, padding_width);

  TORCH_CHECK(ref.equal(outputs[0]));
}

// Conv-like pattern with irregular padding
TEST_F(NVFuserTest, FusionGather7_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  const std::vector<int> window_shape = {3, 4};
  const std::vector<std::vector<int>> padding_width = {{0, 2}, {2, 1}};

  auto tv1 = gather(tv0, window_shape, padding_width);

  fusion.addOutput(tv1);

  // Blocking the spatial dimensions
  const int block_x = 16;
  const int block_y = 8;

  auto tv0_cache = tv0->cache_after();
  auto out = tv1;
  auto out_cache = out->cache_before();

  out->split(1, block_x);
  out->split(0, block_y);
  out->reorder({{1, 2}, {2, 1}});

  TransformPropagator::from(out);

  tv0->computeAt(out, 2);

  tv0_cache->setMemoryType(MemoryType::Shared);

  out->axis(0)->parallelize(ParallelType::BIDy);
  out->axis(1)->parallelize(ParallelType::BIDx);
  out->axis(2)->parallelize(ParallelType::TIDy);
  out->axis(3)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(out, ir_utils::allTvs(&fusion));

  const int s1 = 101;
  const int s2 = 99;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  std::vector<int64_t> size({s1, s2});
  at::Tensor t0 = at::randn(size, options);
  size.insert(size.end(), window_shape.begin(), window_shape.end());
  at::Tensor output = at::ones(size, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto outputs = fe.runFusion({t0}, {output});

  auto ref = gather(t0, window_shape, padding_width);

  TORCH_CHECK(ref.equal(outputs[0]));
}

// With no padding but with striding
TEST_F(NVFuserTest, FusionGather8_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  const std::vector<int> window_shape = {2, 3};
  const std::vector<std::vector<int>> padding_width = {{0, 0}, {0, 0}};
  const std::vector<int> strides = {3, 3};

  auto tv1 = gather(tv0, window_shape, padding_width, strides);

  fusion.addOutput(tv1);

  const int s1 = 11;
  const int s2 = 13;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  std::vector<int64_t> size({s1, s2});
  at::Tensor t0 = at::randn(size, options);
  for (const auto i : c10::irange(size.size())) {
    size[i] = ceilDiv(
        size[i] - window_shape[i] + 1 + padding_width[i][0] +
            padding_width[i][1],
        strides[i]);
  }
  size.insert(size.end(), window_shape.begin(), window_shape.end());
  // Use a pre-allocated output tensor filled with 1 so that invalid
  // writes to outside valid ranges can be detected
  at::Tensor output = at::ones(size, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto outputs = fe.runFusion({t0}, {output});

  auto ref = gather(t0, window_shape, padding_width, strides);

  TORCH_CHECK(ref.equal(outputs[0]));
}

// Similar to Gather8 but with splitting and parallelization
TEST_F(NVFuserTest, FusionGather9_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  const std::vector<int> window_shape = {3, 4};
  const std::vector<std::vector<int>> padding_width = {{0, 0}, {0, 0}};
  const std::vector<int> strides = {2, 2};

  auto tv1 = gather(tv0, window_shape, padding_width, strides);

  fusion.addOutput(tv1);

  // Blocking the spatial dimensions
  const int block_x = 16;
  const int block_y = 8;

  auto tv0_cache = tv0->cache_after();
  auto out = tv1;
  auto out_cache = out->cache_before();

  out->split(1, block_x);
  out->split(0, block_y);
  out->reorder({{1, 2}, {2, 1}});

  TransformPropagator::from(out);

  tv0->computeAt(out, 2);

  tv0_cache->setMemoryType(MemoryType::Shared);

  out->axis(0)->parallelize(ParallelType::BIDy);
  out->axis(1)->parallelize(ParallelType::BIDx);
  out->axis(2)->parallelize(ParallelType::TIDy);
  out->axis(3)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(out, ir_utils::allTvs(&fusion));

  const int s1 = 101;
  const int s2 = 99;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  std::vector<int64_t> size({s1, s2});
  at::Tensor t0 = at::randn(size, options);
  for (const auto i : c10::irange(size.size())) {
    size[i] = ceilDiv(
        size[i] - window_shape[i] + 1 + padding_width[i][0] +
            padding_width[i][1],
        strides[i]);
  }
  size.insert(size.end(), window_shape.begin(), window_shape.end());
  // Use a pre-allocated output tensor filled with 1 so that invalid
  // writes to outside valid ranges can be detected
  at::Tensor output = at::ones(size, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto outputs = fe.runFusion({t0}, {output});

  auto ref = gather(t0, window_shape, padding_width, strides);

  TORCH_CHECK(ref.equal(outputs[0]));
}

TEST_F(NVFuserTest, FusionConv2D_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input: [C, H, W]
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // Weights: [K, C, 3, 3]
  auto w = makeSymbolicTensor(4);
  fusion.addInput(w);

  // Gather a neighbor tile of [3, 3] with padding size of 1 for each
  // side of the spatial dimensions
  auto inp_tile = gather(inp, {1, 3, 3}, {{0, 0}, {1, 1}, {1, 1}});
  // inp_tile: [C, H, W, 1, 3, 3]

  auto inp_bc =
      broadcast(inp_tile, {true, false, false, false, false, false, false});
  auto w_bc = broadcast(w, {false, false, true, true, true, false, false});

  auto inp_times_w = mul(inp_bc, w_bc);

  // Reduce the channel and neighbor tile dimensions
  auto out = sum(inp_times_w, {1, 4, 5, 6});

  fusion.addOutput(out);

  ////////////////////////////////////

  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Blocking the spatial dimensions
  const int block_w = 16;
  const int block_h = 4;
  // Blocking the channel dimension
  const int block_c = 8;

  out->split(2, block_h);
  out->split(4, block_w);
  out->reorder({{3, 4}});
  // out: [K, C, Ho, Wo, Hi, Wi, 1, 3, 3]

  out->split(1, block_c);
  // out: [K, Co, Ci, Ho, Wo, Hi, Wi, 1, 3, 3]

  auto out_rf = out->rFactor({1, -3, -2, -1});
  // out_rf: [K, rCo, Ci, Ho, Wo, Hi, Wi, 1, 3, 3]
  // out_rf: [K, Ci, Ho, Wo, Hi, Wi]

  // Create a [block_x, block_y] tile on smem
  inp_cache->computeAt(out, 4);
  // inp_cache: [Co, Ho, Wo, Ci, Hi, Wi]
  inp_cache->setMemoryType(MemoryType::Shared);

  // Move Ci forward
  out_rf->reorder({{-4, -6}, {-5, -4}, {-6, -5}});
  inp_cache->computeAt(out_rf, 5);

  inp_tile->computeAt(out_rf, -1);
  w->computeAt(out_rf, -1);

  out->axis(0)->parallelize(ParallelType::BIDx);
  out->axis(1)->parallelize(ParallelType::TIDz);
  out->axis(4)->parallelize(ParallelType::TIDy);
  out->axis(5)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out, {inp_cache, out_rf});

  const int dim_h = 99;
  const int dim_w = 101;
  const int dim_c = 10;
  const int dim_f = 20;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor at_inp = at::randn({dim_c, dim_h, dim_w}, options);
  at::Tensor at_w = at::randn({dim_f, dim_c, 3, 3}, options);
  std::vector<IValue> inputs = {at_inp, at_w};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto cg_outputs = fe.runFusion(inputs);

  at_inp = at_inp.unsqueeze(0); // at::conv2d needs the N axis
  auto at_out = at::conv2d(at_inp, at_w, {}, 1, 1);
  at_out = at_out.squeeze(0); // drop the N axis

  testValidate(&fusion, cg_outputs, inputs, {at_out}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionConv2DNoPadding_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input: [C, H, W]
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // Weights: [K, C, 3, 3]
  auto w = makeSymbolicTensor(4);
  fusion.addInput(w);

  // Gather a neighbor tile of [3, 3] with no padding
  auto inp_tile =
      gather(inp, {1, 3, 3}, {{0, 0}, {0, 0}, {0, 0}}, {1, 1, 1}, true);
  // inp_tile: [C, H-2, W-2, 1, 3, 3]

  auto inp_bc =
      broadcast(inp_tile, {true, false, false, false, false, false, false});
  auto w_bc = broadcast(w, {false, false, true, true, true, false, false});

  auto inp_times_w = mul(inp_bc, w_bc);

  // Reduce the channel and neighbor tile dimensions
  auto out = sum(inp_times_w, {1, 4, 5, 6});

  fusion.addOutput(out);

  ////////////////////////////////////

  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Blocking the spatial dimensions
  const int block_w = 16;
  const int block_h = 4;
  // Blocking the channel dimension
  const int block_c = 8;

  out->split(2, block_h);
  out->split(4, block_w);
  out->reorder({{3, 4}});
  // out: [K, C, Ho, Wo, Hi, Wi, 1, 3, 3]

  out->split(1, block_c);
  // out: [K, Co, Ci, Ho, Wo, Hi, Wi, 1, 3, 3]

  auto out_rf = out->rFactor({1, -3, -2, -1});
  // out_rf: [K, rCo, Ci, Ho, Wo, Hi, Wi, 1, 3, 3]
  // out_rf: [K, Ci, Ho, Wo, Hi, Wi]

  // Create a [block_x, block_y] tile on smem
  inp_cache->computeAt(out, 4);
  // inp_cache: [Co, Ho, Wo, Ci, Hi, Wi]
  inp_cache->setMemoryType(MemoryType::Shared);

  // Move Ci forward
  out_rf->reorder({{-4, -6}, {-5, -4}, {-6, -5}});
  inp_cache->computeAt(out_rf, 5);

  inp_tile->computeAt(out_rf, -1);
  w->computeAt(out_rf, -1);

  out->axis(0)->parallelize(ParallelType::BIDx);
  out->axis(1)->parallelize(ParallelType::TIDz);
  out->axis(4)->parallelize(ParallelType::TIDy);
  out->axis(5)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out, {inp_cache, out_rf});

  const int dim_h = 99;
  const int dim_w = 101;
  const int dim_c = 10;
  const int dim_f = 20;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor at_inp = at::randn({dim_c, dim_h, dim_w}, options);
  at::Tensor at_w = at::randn({dim_f, dim_c, 3, 3}, options);
  std::vector<IValue> inputs = {at_inp, at_w};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto cg_outputs = fe.runFusion(inputs);

  at_inp = at_inp.unsqueeze(0); // at::conv2d needs the N axis
  std::vector<int64_t> stride = {1, 1};
  std::vector<int64_t> padding = {0, 0};
  auto at_out = at::conv2d(at_inp, at_w, {}, stride, padding);
  at_out = at_out.squeeze(0); // drop the N axis

  testValidate(&fusion, cg_outputs, inputs, {at_out}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionConv2DNoPaddingStrided_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input: [C, H, W]
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // Weights: [K, C, 3, 3]
  auto w = makeSymbolicTensor(4);
  fusion.addInput(w);

  // Gather a neighbor tile of [2, 2] with no padding and strides of
  // [2, 2]
  auto inp_tile = gather(inp, {1, 2, 2}, {{0, 0}, {0, 0}, {0, 0}}, {1, 2, 2});
  // inp_tile: [C, H/2, W/2, 1, 2, 2]

  auto inp_bc =
      broadcast(inp_tile, {true, false, false, false, false, false, false});
  auto w_bc = broadcast(w, {false, false, true, true, true, false, false});

  auto inp_times_w = mul(inp_bc, w_bc);

  // Reduce the channel and neighbor tile dimensions
  auto out = sum(inp_times_w, {1, 4, 5, 6});

  fusion.addOutput(out);

  ////////////////////////////////////

  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Blocking the spatial dimensions
  const int block_w = 16;
  const int block_h = 4;
  // Blocking the channel dimension
  const int block_c = 8;

  out->split(2, block_h);
  out->split(4, block_w);
  out->reorder({{3, 4}});
  // out: [K, C, Ho, Wo, Hi, Wi, 1, 3, 3]

  out->split(1, block_c);
  // out: [K, Co, Ci, Ho, Wo, Hi, Wi, 1, 3, 3]

  auto out_rf = out->rFactor({1, -3, -2, -1});
  // out_rf: [K, rCo, Ci, Ho, Wo, Hi, Wi, 1, 3, 3]
  // out_rf: [K, Ci, Ho, Wo, Hi, Wi]

  // Create a [block_x, block_y] tile on smem
  inp_cache->computeAt(out, 4);
  // inp_cache: [Co, Ho, Wo, Ci, Hi, Wi]
  inp_cache->setMemoryType(MemoryType::Shared);

  // Move Ci forward
  out_rf->reorder({{-4, -6}, {-5, -4}, {-6, -5}});
  inp_cache->computeAt(out_rf, 5);

  inp_tile->computeAt(out_rf, -1);
  w->computeAt(out_rf, -1);

  out->axis(0)->parallelize(ParallelType::BIDx);
  out->axis(1)->parallelize(ParallelType::TIDz);
  out->axis(4)->parallelize(ParallelType::TIDy);
  out->axis(5)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out, {inp_cache, out_rf});

  const int dim_h = 99;
  const int dim_w = 101;
  const int dim_c = 10;
  const int dim_f = 20;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor at_inp = at::randn({dim_c, dim_h, dim_w}, options);
  at::Tensor at_w = at::randn({dim_f, dim_c, 2, 2}, options);
  std::vector<IValue> inputs = {at_inp, at_w};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto cg_outputs = fe.runFusion(inputs);

  at_inp = at_inp.unsqueeze(0); // at::conv2d needs the N axis
  std::vector<int64_t> stride = {2, 2};
  std::vector<int64_t> padding = {0, 0};
  auto at_out = at::conv2d(at_inp, at_w, {}, stride, padding);
  at_out = at_out.squeeze(0); // drop the N axis

  testValidate(&fusion, cg_outputs, inputs, {at_out}, __LINE__, __FILE__);
}

// 5x5 followed by 3x3
TEST_F(NVFuserTest, FusionConv2DChain_CUDA) {
  const int dim_w1_h = 5;
  const int dim_w1_w = 5;
  const int dim_pad1_h = (dim_w1_h - 1) / 2;
  const int dim_pad1_w = (dim_w1_w - 1) / 2;
  const int dim_w2_h = 3;
  const int dim_w2_w = 3;
  const int dim_pad2_h = (dim_w2_h - 1) / 2;
  const int dim_pad2_w = (dim_w2_w - 1) / 2;

  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input: [K1, H, W]
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // Weights: [K2, K1, S1, T1]
  auto w1 = makeSymbolicTensor(4);
  fusion.addInput(w1);

  // Weights: [K3, K2, S2, T2]
  auto w2 = makeSymbolicTensor(4);
  fusion.addInput(w2);

  // Gather a neighbor tile of [w1_h, w1_w] with padding
  auto inp_tile = gather(
      inp,
      {1, dim_w1_h, dim_w1_w},
      {{0, 0}, {dim_pad1_h, dim_pad1_h}, {dim_pad1_w, dim_pad1_w}});
  // inp_tile: [C, 1, H - w1_h + 1, W - w1_w + 1, w1_h, w1_w]

  auto inp_bc =
      broadcast(inp_tile, {true, false, false, false, false, false, false});
  auto w1_bc = broadcast(w1, {false, false, true, true, true, false, false});

  auto inp_times_w1 = mul(inp_bc, w1_bc);

  // Reduce the channel and neighbor tile dimensions
  auto out1 = sum(inp_times_w1, {1, 4, 5, 6});

  // Second conv
  auto out1_tile = gather(
      out1,
      {1, dim_w2_h, dim_w2_w},
      {{0, 0}, {dim_pad2_h, dim_pad2_h}, {dim_pad2_w, dim_pad2_w}});

  auto out1_bc =
      broadcast(out1_tile, {true, false, false, false, false, false, false});
  auto w2_bc = broadcast(w2, {false, false, true, true, true, false, false});

  auto out1_times_w2 = mul(out1_bc, w2_bc);

  auto out2 = sum(out1_times_w2, {1, 4, 5, 6});

  fusion.addOutput(out2);

  ////////////////////////////////////
  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Blocking the spatial dimensions
  const int block_w = 16;
  const int block_h = 4;

  out2->split(2, block_h);
  out2->split(4, block_w);
  out2->reorder({{3, 4}});
  // out2: [K3, K2, Ho, Wo, Hi, Wi, 1, 3, 3]

  // Create a [block_x, block_y] tile on smem
  inp_cache->computeAt(out2, 4);
  // inp_cache: [Co, Ho, Wo, Ci, Hi, Wi]
  inp_cache->setMemoryType(MemoryType::Shared);

  // Move Ci forward
  out1->reorder({{5, 3}, {3, 4}, {4, 5}});
  out1->setMemoryType(MemoryType::Shared);

  inp_cache->computeAt(out1, 4);

  inp_tile->computeAt(out1, -1);
  w1->computeAt(out1, -1);

  out1_tile->computeAt(out2, -1);
  w2->computeAt(out2, -1);

  out2->axis(0)->parallelize(ParallelType::BIDx);
  out2->axis(4)->parallelize(ParallelType::TIDy);
  out2->axis(5)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out2, {inp_cache, out1});

  const int dim_h = 99;
  const int dim_w = 101;
  const int dim_k1 = 3;
  const int dim_k2 = 5;
  const int dim_k3 = 7;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor at_inp = at::randn({dim_k1, dim_h, dim_w}, options);
  at::Tensor at_w1 = at::randn({dim_k2, dim_k1, dim_w1_h, dim_w1_w}, options);
  at::Tensor at_w2 = at::randn({dim_k3, dim_k2, dim_w2_h, dim_w2_w}, options);
  std::vector<IValue> inputs = {at_inp, at_w1, at_w2};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto cg_outputs = fe.runFusion(inputs);

  at_inp = at_inp.unsqueeze(0); // at::conv2d needs the N axis
  auto at_out1 = at::conv2d(at_inp, at_w1, {}, 1, 2);
  auto at_out2 = at::conv2d(at_out1, at_w2, {}, 1, 1);
  at_out2 = at_out2.squeeze(0); // drop the N axis

  testValidate(&fusion, cg_outputs, inputs, {at_out2}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionConv2DStaticEvenSizedWindow_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input: [C, H, W]
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // Weights: [K, C, 2, 2]
  auto w = makeSymbolicTensor(4);
  fusion.addInput(w);

  // Gather a neighbor tile of [2, 2] with padding size of 1 only for
  // the right side of the spatial dimensions. The left padding is
  // zero so that the output axis stays the same.
  auto inp_tile = gather(inp, {1, 2, 2}, {{0, 0}, {0, 1}, {0, 1}});
  // inp_tile: [C, H, W, 1, 2, 2]

  auto inp_bc =
      broadcast(inp_tile, {true, false, false, false, false, false, false});
  auto w_bc = broadcast(w, {false, false, true, true, true, false, false});

  auto inp_times_w = mul(inp_bc, w_bc);

  // Reduce the channel and neighbor tile dimensions
  auto out = sum(inp_times_w, {1, 4, 5, 6});

  fusion.addOutput(out);

  ////////////////////////////////////

  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Blocking the spatial dimensions
  const int block_w = 16;
  const int block_h = 4;
  // Blocking the channel dimension
  const int block_c = 8;

  out->split(2, block_h);
  out->split(4, block_w);
  out->reorder({{3, 4}});
  // out: [K, C, Ho, Wo, Hi, Wi, 1, 2, 2]

  out->split(1, block_c);
  // out: [K, Co, Ci, Ho, Wo, Hi, Wi, 1, 2, 2]

  auto out_rf = out->rFactor({1, -3, -2, -1});
  // out_rf: [K, rCo, Ci, Ho, Wo, Hi, Wi, 1, 2, 2]
  // out_rf: [K, Ci, Ho, Wo, Hi, Wi]

  // Create a [block_x, block_y] tile on smem
  inp_cache->computeAt(out, 4);
  // inp_cache: [Co, Ho, Wo, Ci, Hi, Wi]
  inp_cache->setMemoryType(MemoryType::Shared);

  // Move Ci forward
  out_rf->reorder({{-4, -6}, {-5, -4}, {-6, -5}});
  inp_cache->computeAt(out_rf, 5);

  inp_tile->computeAt(out_rf, -1);
  w->computeAt(out_rf, -1);

  out->axis(0)->parallelize(ParallelType::BIDx);
  out->axis(1)->parallelize(ParallelType::TIDz);
  out->axis(4)->parallelize(ParallelType::TIDy);
  out->axis(5)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out, {inp_cache, out_rf});

  const int dim_h = 99;
  const int dim_w = 101;
  const int dim_c = 10;
  const int dim_f = 20;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor at_inp = at::randn({dim_c, dim_h, dim_w}, options);
  at::Tensor at_w = at::randn({dim_f, dim_c, 2, 2}, options);
  std::vector<IValue> inputs = {at_inp, at_w};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto cg_outputs = fe.runFusion(inputs);

  at_inp = at_inp.unsqueeze(0); // at::conv2d needs the N axis
  auto at_out = at::conv2d(at_inp, at_w, {}, 1, 1);
  at_out = at_out.squeeze(0); // drop the N axis
  // The shape of the spatial domain is (dim_h+1)x(dim_w+1), whereas
  // the fuser output has dim_h*dim_w. Drop the first elements to make
  // it match with the fuser output.
  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(0, at::indexing::None),
      at::indexing::Slice(1, at::indexing::None),
      at::indexing::Slice(1, at::indexing::None)};
  at_out = at_out.index(indices);

  testValidate(&fusion, cg_outputs, inputs, {at_out}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionConv4x4Pad1x1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input: [C, H, W]
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // Weights: [K, C, 4, 4]
  auto w = makeSymbolicTensor(4);
  fusion.addInput(w);

  // Gather a neighbor tile of [4, 4] with padding size of 1 for both
  // sides of the spatial dimensions. The resulting extent is
  // decreased by one.
  auto inp_tile =
      gather(inp, {1, 4, 4}, {{0, 0}, {1, 1}, {1, 1}}, {1, 1, 1}, true);
  // inp_tile: [C, H-1, W-1, 1, 4, 4]

  auto inp_bc =
      broadcast(inp_tile, {true, false, false, false, false, false, false});
  auto w_bc = broadcast(w, {false, false, true, true, true, false, false});

  auto inp_times_w = mul(inp_bc, w_bc);

  // Reduce the channel and neighbor tile dimensions
  auto out = sum(inp_times_w, {1, 4, 5, 6});

  fusion.addOutput(out);

  ////////////////////////////////////

  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Blocking the spatial dimensions
  const int block_w = 16;
  const int block_h = 4;
  // Blocking the channel dimension
  const int block_c = 8;

  out->split(2, block_h);
  out->split(4, block_w);
  out->reorder({{3, 4}});
  // out: [K, C, Ho, Wo, Hi, Wi, 1, 4, 4]

  out->split(1, block_c);
  // out: [K, Co, Ci, Ho, Wo, Hi, Wi, 1, 4, 4]

  auto out_rf = out->rFactor({1, -3, -2, -1});
  // out_rf: [K, rCo, Ci, Ho, Wo, Hi, Wi, 1, 4, 4]
  // out_rf: [K, Ci, Ho, Wo, Hi, Wi]

  // Create a [block_x, block_y] tile on smem
  inp_cache->computeAt(out, 4);
  // inp_cache: [Co, Ho, Wo, Ci, Hi, Wi]
  inp_cache->setMemoryType(MemoryType::Shared);

  // Move Ci forward
  out_rf->reorder({{-4, -6}, {-5, -4}, {-6, -5}});
  inp_cache->computeAt(out_rf, 5);

  inp_tile->computeAt(out_rf, -1);
  w->computeAt(out_rf, -1);

  out->axis(0)->parallelize(ParallelType::BIDx);
  out->axis(1)->parallelize(ParallelType::TIDz);
  out->axis(4)->parallelize(ParallelType::TIDy);
  out->axis(5)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out, {inp_cache, out_rf});

  const int dim_h = 99;
  const int dim_w = 101;
  const int dim_c = 10;
  const int dim_f = 20;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor at_inp = at::randn({dim_c, dim_h, dim_w}, options);
  at::Tensor at_w = at::randn({dim_f, dim_c, 4, 4}, options);
  std::vector<IValue> inputs = {at_inp, at_w};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto cg_outputs = fe.runFusion(inputs);

  at_inp = at_inp.unsqueeze(0); // at::conv2d needs the N axis
  auto at_out =
      at::conv2d(at_inp.to(at::kDouble), at_w.to(at::kDouble), {}, 1, 1);
  at_out = at_out.squeeze(0); // drop the N axis

  testValidate(&fusion, cg_outputs, inputs, {at_out}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionConv4x5Pad1x2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input: [C, H, W]
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // Weights: [K, C, 4, 4]
  auto w = makeSymbolicTensor(4);
  fusion.addInput(w);

  // Gather a neighbor tile of [4, 5] with padding size of 1 and 2 for
  // each side of the spatial dimensions.
  auto inp_tile =
      gather(inp, {1, 4, 5}, {{0, 0}, {1, 1}, {2, 2}}, {1, 1, 1}, true);
  // inp_tile: [C, H-1, W, 1, 4, 5]

  auto inp_bc =
      broadcast(inp_tile, {true, false, false, false, false, false, false});
  auto w_bc = broadcast(w, {false, false, true, true, true, false, false});

  auto inp_times_w = mul(inp_bc, w_bc);

  // Reduce the channel and neighbor tile dimensions
  auto out = sum(inp_times_w, {1, 4, 5, 6});

  fusion.addOutput(out);

  ////////////////////////////////////

  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Blocking the spatial dimensions
  const int block_w = 16;
  const int block_h = 4;
  // Blocking the channel dimension
  const int block_c = 8;

  out->split(2, block_h);
  out->split(4, block_w);
  out->reorder({{3, 4}});
  // out: [K, C, Ho, Wo, Hi, Wi, 1, 4, 5]

  out->split(1, block_c);
  // out: [K, Co, Ci, Ho, Wo, Hi, Wi, 1, 4, 5]

  auto out_rf = out->rFactor({1, -3, -2, -1});
  // out_rf: [K, rCo, Ci, Ho, Wo, Hi, Wi, 1, 4, 5]
  // out_rf: [K, Ci, Ho, Wo, Hi, Wi]

  // Create a [block_x, block_y] tile on smem
  inp_cache->computeAt(out, 4);
  // inp_cache: [Co, Ho, Wo, Ci, Hi, Wi]
  inp_cache->setMemoryType(MemoryType::Shared);

  // Move Ci forward
  out_rf->reorder({{-4, -6}, {-5, -4}, {-6, -5}});
  inp_cache->computeAt(out_rf, 5);

  inp_tile->computeAt(out_rf, -1);
  w->computeAt(out_rf, -1);

  out->axis(0)->parallelize(ParallelType::BIDx);
  out->axis(1)->parallelize(ParallelType::TIDz);
  out->axis(4)->parallelize(ParallelType::TIDy);
  out->axis(5)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out, {inp_cache, out_rf});

  const int dim_h = 99;
  const int dim_w = 101;
  const int dim_c = 10;
  const int dim_f = 20;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor at_inp = at::randn({dim_c, dim_h, dim_w}, options);
  at::Tensor at_w = at::randn({dim_f, dim_c, 4, 5}, options);
  std::vector<IValue> inputs = {at_inp, at_w};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto cg_outputs = fe.runFusion(inputs);

  at_inp = at_inp.unsqueeze(0); // at::conv2d needs the N axis
  auto at_out =
      at::conv2d(at_inp.to(at::kDouble), at_w.to(at::kDouble), {}, 1, {1, 2});
  at_out = at_out.squeeze(0); // drop the N axis

  testValidate(&fusion, cg_outputs, inputs, {at_out}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionConv4x4Pad1x1Stride4_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input: [C, H, W]
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // Weights: [K, C, 3, 3]
  auto w = makeSymbolicTensor(4);
  fusion.addInput(w);

  // Gather a neighbor tile of [4, 4] with padding size of 1 for both
  // sides of the spatial dimensions. Set the stride width as 4.
  auto inp_tile = gather(inp, {1, 4, 4}, {{0, 0}, {1, 1}, {1, 1}}, {1, 4, 4});
  // inp_tile: [C, H/4, s4, W/4, s4, 1, 4, 4]

  auto inp_bc =
      broadcast(inp_tile, {true, false, false, false, false, false, false});
  auto w_bc = broadcast(w, {false, false, true, true, true, false, false});

  auto inp_times_w = mul(inp_bc, w_bc);

  // Reduce the channel and neighbor tile dimensions
  auto out = sum(inp_times_w, {1, 4, 5, 6});

  fusion.addOutput(out);

  ////////////////////////////////////

  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Blocking the spatial dimensions
  const int block_w = 16;
  const int block_h = 4;
  const int block_c = 2;

  // [K, C, H/s, W/s, 1, 4, 4]
  out->split(2, block_h);
  // [K, C, H/s/block_h, block_h, W/s, 1, 4, 4]
  out->split(4, block_w);
  // [K, C, H/s/block_h, block_h, W/s/block_w, block_w, 1, 4, 4]
  out->reorder({{3, 4}});
  // [K, C, H/s/block_h, W/s/block_w, block_h, block_w, 1, 4, 4]
  out->split(1, block_c);
  // [K, C/block_c, block_c, H/s/block_h, W/s/block_w, block_h, block_w, 1, 4,
  // 4]
  out->split(4, 1);
  // [K, C/block_c, block_c, H/s/block_h, W/s/block_w, 1, block_h, block_w, 1,
  // 4, 4]

  auto out_rf = out->rFactor({1, -3, -2, -1});
  // [K, C/block_c, block_c, H/s/block_h, W/s/block_w, 1, block_h, block_w, 1,
  // 4, 4]

  // out: [K, block_c, H/s/block_h, W/s/block_w, 1, block_h, block_w]

  inp_cache->computeAt(out, 5);
  inp_cache->setMemoryType(MemoryType::Shared);
  // [K, block_c, H/s/block_h, W/s/block_w, 1, block_h, block_w, C/block_c, 1,
  // 4, 4]

  // Move C/block_c before block_h/2 and share the domain from
  // inp_cache to out_rf
  out_rf->reorder({{7, 5}, {5, 6}, {6, 7}});
  inp_cache->computeAt(out_rf, 6);

  inp_tile->computeAt(out_rf, -1);
  w->computeAt(out_rf, -1);

  out->axis(0)->parallelize(ParallelType::BIDx);
  out->axis(1)->parallelize(ParallelType::TIDz);
  out->axis(4)->parallelize(ParallelType::Unswitch);
  out->axis(5)->parallelize(ParallelType::TIDy);
  out->axis(6)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out, {inp_cache, out_rf});

  const int dim_h = 99;
  const int dim_w = 101;
  const int dim_c = 10;
  const int dim_f = 20;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor at_inp = at::randn({dim_c, dim_h, dim_w}, options);
  at::Tensor at_w = at::randn({dim_f, dim_c, 4, 4}, options);
  std::vector<IValue> inputs = {at_inp, at_w};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto cg_outputs = fe.runFusion(inputs);

  at_inp = at_inp.unsqueeze(0); // at::conv2d needs the N axis
  auto at_out =
      at::conv2d(at_inp.to(at::kDouble), at_w.to(at::kDouble), {}, 4, {1, 1});
  at_out = at_out.squeeze(0); // drop the N axis

  testValidate(&fusion, cg_outputs, inputs, {at_out}, __LINE__, __FILE__);
}

// POC implementation of im2col for 3-by-3 kernels
TEST_F(NVFuserTest, FusionIm2Col_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input: [N, C, H, W]
  auto inp = makeSymbolicTensor(4);
  fusion.addInput(inp);

  // Gather a neighbor tile of [3, 3] with padding size of 1 for each
  // side of the spatial dimensions
  auto inp_tile = gather(inp, {1, 1, 3, 3}, {{0, 0}, {0, 0}, {1, 1}, {1, 1}});
  // inp_tile: [N, C, H, W, 1, 1, 3, 3]

  auto inp_col = transpose(inp_tile, {{1, 3}, {2, 1}, {3, 2}});
  // inp_col: [N, H, W, C, 1, 1, 3, 3]

  fusion.addOutput(inp_col);

  ////////////////////////////////////

  // Cache the input tensor
  auto inp_cache = inp->cache_after();

  // Blocking the spatial dimensions
  const int block_w = 16;
  const int block_h = 4;

  auto out = inp_col;

  out->split(1, block_h);
  out->split(3, block_w);
  out->reorder({{2, 3}});
  // out: [N, Ho, Wo, Hi, Wi, C, 1, 1, 3, 3]
  // Move the C axis out of Hi*Wi
  out->reorder({{5, 3}, {3, 4}, {4, 5}});
  // out: [N, Ho, Wo, C, Hi, Wi, 1, 1, 3, 3]

  // Create a [block_x, block_y] tile on smem
  inp_cache->computeAt(out, 4);
  inp_cache->setMemoryType(MemoryType::Shared);
  // Fully inline inp_tile
  inp_tile->computeAt(out, -1);

  out->axis(0)->parallelize(ParallelType::BIDz);
  out->axis(1)->parallelize(ParallelType::BIDy);
  out->axis(2)->parallelize(ParallelType::BIDx);
  out->axis(4)->parallelize(ParallelType::TIDy);
  out->axis(5)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out, {inp_cache, inp_tile});

  const int dim_h = 31;
  const int dim_w = 33;
  const int dim_c = 5;
  const int dim_n = 3;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor at_inp = at::randn({dim_n, dim_c, dim_h, dim_w}, options);
  std::vector<IValue> inputs = {at_inp};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto cg_outputs = fe.runFusion(inputs);

  auto at_out = at::im2col(at_inp, {3, 3}, {1, 1}, {1, 1}, {1, 1});

  // at::im2col outputs [N, C*3*3, N*H]
  at_out = at::transpose(at_out, 1, 2);
  at_out = at::reshape(at_out, {dim_n, dim_h, dim_w, dim_c, 1, 1, 3, 3});

  testValidate(&fusion, cg_outputs, inputs, {at_out}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftNoPadding1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {1, -1}, false);
  auto tv3 = shift(tv1, {-1, 1}, false);
  auto tv4 = add(tv2, tv3);
  auto tv5 = sum(tv4, {0, 1});

  fusion.addOutput(tv5);

  tv1->setMemoryType(MemoryType::Shared);

  tv5->split(0, 4);
  tv5->split(-1, 8);
  tv5->reorder({{1, 2}});

  TransformPropagator::from(tv5);

  tv2->computeAt(tv5, -1);
  tv3->computeAt(tv5, -1);

  tv5->axis(-1)->parallelize(ParallelType::TIDx);
  tv5->axis(-2)->parallelize(ParallelType::TIDy);
  scheduler_utils::parallelizeAllLike(tv5, ir_utils::allTvs(&fusion));

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {1, -1});
  auto t3 = shift(t1, {-1, 1});
  auto t4 = t2 + t3;
  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(1, -1), at::indexing::Slice(1, -1)};
  t4 = t4.index(indices);
  auto ref = t4.sum(at::ArrayRef<int64_t>{0, 1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// Split and merge
TEST_F(NVFuserTest, FusionShiftNoPadding2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {1, -1}, false);
  auto tv3 = shift(tv1, {-1, 1}, false);
  auto tv4 = add(tv2, tv3);
  auto tv5 = sum(tv4, {0, 1});

  fusion.addOutput(tv5);

  tv1->setMemoryType(MemoryType::Shared);

  tv5->split(0, 4);
  tv5->split(-1, 8);
  tv5->reorder({{1, 2}});
  tv5->merge(-2, -1);

  TransformPropagator::from(tv5);

  tv2->computeAt(tv5, -1);
  tv3->computeAt(tv5, -1);

  tv5->axis(-1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv5, ir_utils::allTvs(&fusion));

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {1, -1});
  auto t3 = shift(t1, {-1, 1});
  auto t4 = t2 + t3;
  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(1, -1), at::indexing::Slice(1, -1)};
  t4 = t4.index(indices);
  auto ref = t4.sum(at::ArrayRef<int64_t>{0, 1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// Split and merge, then welford
TEST_F(NVFuserTest, FusionShiftNoPadding3_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {1, -1}, false);
  auto tv3 = shift(tv1, {-1, 1}, false);
  auto tv4 = add(tv2, tv3);
  auto tvs = Welford(tv4, {0, 1});
  auto tv_avg = tvs.avg;
  auto tv_M2 = tvs.var_sum;
  auto tv_N = tvs.n;

  fusion.addOutput(tv_avg);
  fusion.addOutput(tv_M2);
  fusion.addOutput(tv_N);

  tv1->setMemoryType(MemoryType::Shared);

  tv_avg->split(0, 4);
  tv_avg->split(-1, 8);
  tv_avg->reorder({{1, 2}});
  tv_avg->merge(-2, -1);

  TransformPropagator::from(tv_avg);

  tv2->computeAt(tv_avg, -1);
  tv3->computeAt(tv_avg, -1);

  tv_avg->axis(-1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv_avg, ir_utils::allTvs(&fusion));

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto options_int = at::TensorOptions().dtype(at::kLong).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  outputs[1] /= (numel_x - 2) * (numel_y - 2);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {1, -1});
  auto t3 = shift(t1, {-1, 1});
  auto t4 = t2 + t3;
  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(1, -1), at::indexing::Slice(1, -1)};
  t4 = t4.index(indices);
  auto ref_avg = t4.mean(at::ArrayRef<int64_t>{0, 1});
  auto ref_M2 = t4.var(at::ArrayRef<int64_t>{0, 1}, false);
  auto ref_N = at::ones({}, options_int) * (numel_x - 2) * (numel_y - 2);

  testValidate(
      &fusion, outputs, inputs, {ref_avg, ref_M2, ref_N}, __LINE__, __FILE__);
}

// Shift indexing and predication with contiguous merge
TEST_F(NVFuserTest, FusionShiftNoPaddingContigMerge_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {1, -1}, true);
  auto tv3 = shift(tv1, {-1, 1}, false);
  auto tv4 = add(tv2, tv3);
  fusion.addOutput(tv4);

  tv2->merge(0);
  tv3->merge(0);
  tv4->merge(0);

  tv1->setMemoryType(MemoryType::Global);
  tv2->setMemoryType(MemoryType::Global);
  tv3->setMemoryType(MemoryType::Global);

  int numel_x = 9;
  int numel_y = 11;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(1, -1), at::indexing::Slice(1, -1)};

  auto fuser_out = outputs[0].index(indices);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {1, -1});
  auto t3 = shift(t1, {-1, 1});
  auto ref = t2 + t3;

  ref = ref.index(indices);

  testValidate(&fusion, {fuser_out}, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftNoPaddingChain_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {1, -1}, false);
  auto tv3 = shift(tv2, {1, -1}, false);
  auto tv4 = sum(tv3, {0, 1});
  fusion.addOutput(tv4);

  tv1->setMemoryType(MemoryType::Shared);
  tv2->setMemoryType(MemoryType::Shared);

  tv4->split(0, 4);
  tv4->split(-1, 8);
  tv4->reorder({{1, 2}});

  tv1->computeAt(tv4, 2);

  tv4->axis(-1)->parallelize(ParallelType::TIDx);
  tv4->axis(-2)->parallelize(ParallelType::TIDy);

  tv4->axis(0)->parallelize(ParallelType::BIDy);
  tv4->axis(1)->parallelize(ParallelType::BIDx);

  scheduler_utils::parallelizeAllLike(tv4, {tv1, tv2, tv3});

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto options_int = at::TensorOptions().dtype(at::kLong).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {1, -1});
  auto t3 = shift(t2, {1, -1});
  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(2, at::indexing::None), at::indexing::Slice(0, -2)};
  t3 = t3.index(indices);
  auto ref = t3.sum(at::ArrayRef<int64_t>{0, 1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// Rfactor is not allowed with partial domains
TEST_F(NVFuserTest, FusionShiftNoPaddingRfactor_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {1, -1}, false);
  auto tv3 = sum(tv2, {0, 1});
  fusion.addOutput(tv3);

  tv3->split(0, 4);
  tv3->split(-1, 8);
  tv3->reorder({{1, 2}});

  ASSERT_ANY_THROW(tv3->rFactor({-2}));
}

TEST_F(NVFuserTest, FusionShiftPadding1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {2, -2}, {1, 1});
  auto tv3 = shift(tv1, {-3, 2}, {2, 2});
  auto tv4 = add(tv2, tv3);
  auto tv5 = sum(tv4, {0, 1});

  fusion.addOutput(tv5);

  tv1->setMemoryType(MemoryType::Shared);

  tv5->split(0, 4);
  tv5->split(-1, 8);
  tv5->reorder({{1, 2}});

  TransformPropagator::from(tv5);

  tv2->computeAt(tv5, -1);
  tv3->computeAt(tv5, -1);

  tv5->axis(-1)->parallelize(ParallelType::TIDx);
  tv5->axis(-2)->parallelize(ParallelType::TIDy);
  scheduler_utils::parallelizeAllLike(tv5, ir_utils::allTvs(&fusion));

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = shift(t1, {2, -2});
  auto t3 = shift(t1, {-3, 2});
  auto t4 = t2 + t3;
  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(1, -1), at::indexing::Slice(0, -1)};
  t4 = t4.index(indices);
  auto ref = t4.sum(at::ArrayRef<int64_t>{0, 1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionPartialSplit1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1);
  // [I]
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(0));
  // [I]
  auto tv2 = shift(tv1, {1}, false);
  // [1:I]
  auto tv3 = shift(tv1, {-1}, false);
  // [0:I-1]
  auto tv4 = add(tv2, tv3);
  // [1:I-1]
  fusion.addOutput(tv4);

  // Partial split of tv4. Split only the valid range, which is
  // [1:-1].
  tv4->split(0, 8, true, true);
  // [(I-2)/8, 8]

  // Propagates the partial split back to tv1. This means that all of
  // the other tensors are also shaped as [(I-2)/8, 8], which appears
  // to mean only the sub region of ((I-2)/8 * 8) is
  // computed for tv1, tv2 and tv3. It's fine for the tv2 and tv3
  // tensors as only that sub region is used by tv4. It's also fine
  // for tv1 since it has halo of size one at each side, so the whole
  // region is actually calculated for tv1.
  tv1->computeAt(tv4, 1);

  tv4->axis(-1)->parallelize(ParallelType::TIDx);
  tv4->axis(-2)->parallelize(ParallelType::BIDx);
  scheduler_utils::parallelizeAllLike(tv4, {tv1, tv2, tv3});

  tv1->setMemoryType(MemoryType::Shared);

  // gridDim.x is ceilDiv(numel_x - 2, 8), not ceilDiv(numel_x, 8),
  // so it's going to be just 2 rather than 3.
  const int numel_x = 18;

  ExpressionEvaluator evaluator(&fusion);
  auto root_extent = tv4->getRootDomain()[0]->extent();
  evaluator.bind(root_extent, numel_x);
  auto extent_eval = evaluator.evaluate(tv4->axis(0)->extent());
  TORCH_CHECK(
      extent_eval.has_value(),
      "Invalid evaluation of outer domain extent of partial split");
  TORCH_CHECK(
      extent_eval.value() == (numel_x - 2) / 8,
      "Invalid extent of outer domain of partial split");

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto options_int = at::TensorOptions().dtype(at::kLong).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor t0 = at::randn({numel_x}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  std::vector<at::indexing::TensorIndex> indices{at::indexing::Slice(1, -1)};

  outputs[0] = outputs[0].index(indices);

  auto ref = (shift(t0, {1}) + shift(t0, {-1})).index(indices);

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionPartialSplit2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(0));
  auto tv2 = shift(tv1, {1}, false);
  auto tv3 = shift(tv1, {-1}, false);
  auto tv4 = add(tv2, tv3);
  fusion.addOutput(tv4);

  auto tv5 = add(tv1, IrBuilder::create<Double>(1));
  auto tv6 = add(tv5, IrBuilder::create<Double>(1));
  fusion.addOutput(tv6);

  tv4->split(0, 4, true, true);

  // This causes tv5 and tv6 also to be split with the same partial
  // offsets, however, since they need to be calculated entirely, the
  // resulting code would be invalid. It should be detected as part of
  // initial fusion validation during lowering.
  tv1->computeAt(tv4, 1);

  // Validation should throw an error due to tv5 and tv6.
  ASSERT_ANY_THROW(fusion.printKernel());
}

// 2D version of PartialSplit1
TEST_F(NVFuserTest, FusionPartialSplit3_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(0));
  auto tv2 = shift(tv1, {1, 2}, false);
  auto tv3 = shift(tv1, {-2, -1}, false);
  auto tv4 = add(tv2, tv3);
  fusion.addOutput(tv4);

  tv4->split(1, 8, true, true);
  tv4->split(0, 4, true, true);
  tv4->reorder({{1, 2}, {2, 1}});

  tv1->computeAt(tv4, 2);

  tv4->axis(0)->parallelize(ParallelType::BIDy);
  tv4->axis(1)->parallelize(ParallelType::BIDx);
  tv4->axis(2)->parallelize(ParallelType::TIDy);
  tv4->axis(3)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv4, {tv1, tv2, tv3});

  tv1->setMemoryType(MemoryType::Shared);

  const int numel_x = 32 + 3;
  const int numel_y = 32 + 3;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto options_int = at::TensorOptions().dtype(at::kLong).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(1, -2), at::indexing::Slice(2, -1)};

  outputs[0] = outputs[0].index(indices);

  auto ref = (shift(t0, {1, 2}) + shift(t0, {-2, -1})).index(indices);

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// Almost same fusion with Shift5ptStencilChain but non-padded shift
// and partial split.
TEST_F(NVFuserTest, FusionPartialSplit4_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  std::vector<std::vector<int>> offsets = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}};

  // First stencil: 5pt stencil
  // stencil1 = (tv0 + tv0[+1][0] + tv0[-1][0] + tv0[0][+1] + tv0[0][-1]) / 5
  std::vector<TensorView*> tv_stencil1_shifts;
  for (const auto& offset : offsets) {
    tv_stencil1_shifts.push_back(shift(tv0, offset, false));
  }

  auto tv_stencil1 = tv0;
  for (auto tv : tv_stencil1_shifts) {
    tv_stencil1 = add(tv_stencil1, tv);
  }

  tv_stencil1 = div(
      tv_stencil1, IrBuilder::create<Double>(tv_stencil1_shifts.size() + 1));

  // Second stencil: Same 5pt stencil
  std::vector<TensorView*> tv_stencil2_shifts;
  for (const auto& offset : offsets) {
    tv_stencil2_shifts.push_back(shift(tv_stencil1, offset, false));
  }

  auto tv_stencil2 = tv_stencil1;
  for (auto tv : tv_stencil2_shifts) {
    tv_stencil2 = add(tv_stencil2, tv);
  }

  tv_stencil2 = div(
      tv_stencil2, IrBuilder::create<Double>(tv_stencil2_shifts.size() + 1));

  auto tv_out = tv_stencil2;

  fusion.addOutput(tv_out);

  auto tv0_cache = tv0->cache_after();

  std::vector<int> split_factor({16, 16});

  tv_out->split(-1, split_factor[1], true, true);
  tv_out->split(0, split_factor[0], true, true);
  tv_out->reorder({{1, 2}, {2, 1}});

  tv0->computeAt(tv_out, 2);

  // Inline completely all inputs to the first stencil output, except for the
  // tv0 cache
  for (auto tv : tv_stencil1_shifts) {
    tv->computeAt(tv_stencil1, -1);
  }

  // Inline completely all inputs to the second stencil output, except
  // for the first stencil output
  for (auto tv : tv_stencil2_shifts) {
    tv->computeAt(tv_stencil2, -1);
  }

  tv_out->axis(0)->parallelize(ParallelType::BIDy);
  tv_out->axis(1)->parallelize(ParallelType::BIDx);
  tv_out->axis(2)->parallelize(ParallelType::TIDy);
  tv_out->axis(3)->parallelize(ParallelType::TIDx);

  auto all_values = DependencyCheck::getAllValsBetween(
      {fusion.inputs().begin(), fusion.inputs().end()}, fusion.outputs());
  for (auto tv : ir_utils::filterByType<TensorView>(all_values)) {
    scheduler_utils::parallelizeAllLike(tv_out, {tv});
  }

  tv0_cache->setMemoryType(MemoryType::Shared);
  tv_stencil1->setMemoryType(MemoryType::Shared);

  // Input matrix size is 68x68, and the output is 64x64. Both
  // gridDim.x and gridim.y should be ceilDiv(numel - 4,
  // split_factor), which is 4. If full split is used, the grid
  // dimension would be 5.
  const int numel_x = 64 + 4;
  const int numel_y = 64 + 4;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(2, -2), at::indexing::Slice(2, -2)};

  outputs[0] = outputs[0].index(indices);

  auto stencil1 = t0;
  for (const auto& offset : offsets) {
    stencil1 = stencil1 + shift(t0, offset);
  }
  stencil1 = stencil1 / int(offsets.size() + 1);
  auto stencil2 = stencil1;
  for (const auto& offset : offsets) {
    stencil2 = stencil2 + shift(stencil1, offset);
  }
  stencil2 = stencil2 / int(offsets.size() + 1);
  auto ref = stencil2.index(indices);

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionPartialSplit5_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const int numel_x = 10;
  const int numel_y = 11;

  // auto tv0 = makeSymbolicTensor(2);
  auto tv0 = makeConcreteTensor({numel_x, numel_y});
  fusion.addInput(tv0);

  auto tv1 = shift(tv0, {0, 1}, false);
  auto tv2 = add(tv1, IrBuilder::create<Double>(1));

  fusion.addOutput(tv2);

  // Partially split tv2 but not tv1. Producer indexing with tv2 as a consumer
  // requires adjustment of the index to account for the difference of split
  // offsets.
  tv2->split(1, 4, true, true);
  tv1->split(1, 4);

  tv1->computeAt(tv2, 1);

  tv2->axis(1)->parallelize(ParallelType::TIDx);
  tv1->axis(1)->parallelize(ParallelType::TIDx);

  tv1->setMemoryType(MemoryType::Shared);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(0, at::indexing::None),
      at::indexing::Slice(1, at::indexing::None)};

  outputs[0] = outputs[0].index(indices);

  auto ref = (shift(t0, {0, 1}) + 1).index(indices);

  testValidate(&fusion, outputs, {t0}, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionPartialSplit6_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const int numel_x = 9;

  auto tv0 = makeConcreteTensor({numel_x});
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {1}, false);
  auto tv3 = add(tv2, IrBuilder::create<Double>(1));

  fusion.addOutput(tv3);

  // Another mix of partial and non-partial split
  tv1->split(0, 4);
  tv2->split(0, 4, true, true);
  tv3->split(0, 4);

  // Just make it easier for compute-sanitizer to flag invalid memory accesses
  tv1->setMemoryType(MemoryType::Shared);
  tv2->setMemoryType(MemoryType::Shared);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  std::vector<at::indexing::TensorIndex> indices{
      at::indexing::Slice(1, at::indexing::None)};

  outputs[0] = outputs[0].index(indices);

  auto ref = (shift(t0 + 1, {1}) + 1).index(indices);

  testValidate(&fusion, outputs, {t0}, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionShiftUnswitch1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = shift(tv0, {-1, 0});
  fusion.addOutput(tv1);

  auto tv2 = shift(tv0, {0, 1});
  fusion.addOutput(tv2);

  auto tv3 = shift(tv0, {2, 2});
  fusion.addOutput(tv3);

  auto tv4 = shift(tv0, {-2, -2});
  fusion.addOutput(tv4);

  auto tv5 = add(tv0, IrBuilder::create<Double>(1));
  auto tv6 = shift(tv5, {0, -1});
  fusion.addOutput(tv6);

  tv1->axis(1)->parallelize(ParallelType::Unswitch);
  tv2->axis(1)->parallelize(ParallelType::Unswitch);
  tv3->axis(0)->parallelize(ParallelType::Unswitch);
  tv4->axis(0)->parallelize(ParallelType::Unswitch);

  tv5->axis(1)->parallelize(ParallelType::TIDx);
  tv6->axis(1)->parallelize(ParallelType::TIDx);
  tv5->axis(0)->parallelize(ParallelType::Unswitch);
  tv5->setMemoryType(MemoryType::Shared);

  int numel_x = 9;
  int numel_y = 11;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = shift(t0, {-1, 0});
  TORCH_CHECK(t1.equal(outputs[0]));

  auto t2 = shift(t0, {0, 1});
  TORCH_CHECK(t2.equal(outputs[1]));

  auto t3 = shift(t0, {2, 2});
  TORCH_CHECK(t3.equal(outputs[2]));

  auto t4 = shift(t0, {-2, -2});
  TORCH_CHECK(t4.equal(outputs[3]));

  auto t6 = shift(t0 + 1, {0, -1});
  TORCH_CHECK(t6.equal(outputs[4]));
}

TEST_F(NVFuserTest, FusionGatherUnswitch1_CUDA) {
  const int tv1_gather = 3;
  const int tv1_gather_pad = 1;
  const int tv2_gather = 5;
  const int tv2_gather_pad = 2;

  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = gather(tv0, {tv1_gather}, {{tv1_gather_pad, tv1_gather_pad}});
  fusion.addOutput(tv1);

  auto tv2 = gather(tv0, {tv2_gather}, {{tv2_gather_pad, tv2_gather_pad}});
  fusion.addOutput(tv2);

  // Static gather
  auto tv3 = gather(tv0, {3}, {{1, 1}});
  fusion.addOutput(tv3);

  // Static gather
  auto tv4 = gather(tv0, {5}, {{2, 2}});
  fusion.addOutput(tv4);

  auto tv0_cache = tv0->cache_after();
  tv0_cache->setMemoryType(MemoryType::Shared);

  tv4->split(0, 32);

  tv0->computeAt(tv4, 1);

  tv4->axis(0)->parallelize(ParallelType::Unswitch);
  tv4->axis(1)->parallelize(ParallelType::TIDx);

  const int numel_x = 100;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = gather(t0, {tv1_gather}, {{tv1_gather_pad, tv1_gather_pad}});
  TORCH_CHECK(t1.equal(outputs[0]));

  auto t2 = gather(t0, {tv2_gather}, {{tv2_gather_pad, tv2_gather_pad}});
  TORCH_CHECK(t2.equal(outputs[1]));

  auto t3 = gather(t0, {3}, {{1, 1}});
  TORCH_CHECK(t3.equal(outputs[2]));

  auto t4 = gather(t0, {5}, {{2, 2}});
  TORCH_CHECK(t4.equal(outputs[3]));
}

TEST_F(NVFuserTest, FusionGatherStrided1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  const std::vector<int> window_shape = {1, 3};
  const std::vector<std::vector<int>> padding_width = {{0, 0}, {1, 1}};

  const std::vector<int> strides = {1, 3};

  auto tv1 = gather(tv0, window_shape, padding_width, strides);

  fusion.addOutput(tv1);

  const int s1 = 11;
  const int s2 = 13;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({s1, s2}, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto outputs = fe.runFusion({t0});

  // tv1 has a stride dimension, so its number of dimensions should be
  // input_ndims + window_ndims + stride.
  TORCH_CHECK(tv1->nDims() == tv0->nDims() * 2 + 1);

  // However, the number of dimensions of the Aten tensor should still
  // be just the twice of the number of dimensions of the input
  // tensor.
  auto fuser_out = outputs[0];
  TORCH_CHECK(
      fuser_out.ndimension() == tv0->nDims() * 2,
      "Invalid dimensionality of output tensor: ",
      fuser_out.ndimension());

  // Each output dimension should be: ceilDiv(input_size + padding_width -
  // window, stride).
  for (const auto i : c10::irange(window_shape.size())) {
    auto valid_dim = ceilDiv(
        t0.size(i) + padding_width[i][0] + padding_width[i][1] -
            window_shape[i] + 1,
        strides[i]);
    auto actual_dim = outputs[0].size(i);
    TORCH_CHECK(
        valid_dim == actual_dim,
        "Invalid output size at dimension ",
        i,
        ". Expected: ",
        valid_dim,
        ", actual: ",
        actual_dim);
  }

  auto ref = gather(t0, window_shape, padding_width, strides);

  TORCH_CHECK(ref.equal(outputs[0]));
}

// Split strided domain
TEST_F(NVFuserTest, FusionGatherStrided2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int> window_shape = {3};
  const std::vector<std::vector<int>> padding_width = {{1, 1}};
  const std::vector<int> strides = {3};

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));

  auto tv2 = gather(tv1, window_shape, padding_width, strides);

  auto tv3 = sum(tv2, {-1});

  fusion.addOutput(tv3);

  // Split the strided domain
  tv3->split(0, 4);

  // Propagate the split by 4 of the tv3 domain to pre-stride domains,
  // making them split by 4 * 3
  tv0->computeAt(tv3, 1);

  tv2->computeAt(tv3, -1);

  tv3->axis(0)->parallelize(ParallelType::BIDx);
  tv3->axis(1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv3, {tv1, tv2});

  tv1->setMemoryType(MemoryType::Shared);

  const int s1 = 100;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({s1}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = gather(t1, window_shape, padding_width, strides);
  auto ref = sum(t2, {-1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// Outer split
TEST_F(NVFuserTest, FusionGatherStrided3_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int> window_shape = {3};
  const std::vector<std::vector<int>> padding_width = {{1, 1}};
  const std::vector<int> strides = {3};

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));

  auto tv2 = gather(tv1, window_shape, padding_width, strides);

  auto tv3 = sum(tv2, {-1});
  fusion.addOutput(tv3);

  // Outer split
  tv3->split(0, 2, false);

  tv0->computeAt(tv3, 1);

  tv3->axis(0)->parallelize(ParallelType::BIDx);
  tv3->axis(1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv3, {tv1, tv2});

  tv1->setMemoryType(MemoryType::Shared);
  tv2->setMemoryType(MemoryType::Shared);

  const int s1 = 100;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({s1}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = gather(t1, window_shape, padding_width, strides);
  auto ref = sum(t2, {-1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionGatherStrided4_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int> window_shape = {3};
  const std::vector<std::vector<int>> padding_width = {{1, 1}};
  const std::vector<int> strides = {3};

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));

  // Test propagation of split from one gather output to another
  auto tv2 = gather(tv1, window_shape, padding_width, strides);
  auto tv3 = gather(tv1, window_shape, padding_width, strides);

  auto tv4 = sum(tv2, {-1});
  fusion.addOutput(tv4);

  auto tv5 = sum(tv3, {-1});
  fusion.addOutput(tv5);

  tv4->split(0, 2);

  // Test forward computeAt propagation from tv1 to tv3
  tv0->computeAt(tv4, 1);

  const int s1 = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({s1}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = gather(t1, window_shape, padding_width, strides);
  auto ref = sum(t2, {-1});

  testValidate(&fusion, outputs, inputs, {ref, ref}, __LINE__, __FILE__);
}

// Same as GatherStrided1 but with stride != window
TEST_F(NVFuserTest, FusionGatherStrided5_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  const std::vector<int> window_shape = {1, 3};
  const std::vector<std::vector<int>> padding_width = {{0, 0}, {1, 1}};

  const std::vector<int> strides = {1, 2};

  auto tv1 = gather(tv0, window_shape, padding_width, strides);

  fusion.addOutput(tv1);

  const int s1 = 11;
  const int s2 = 13;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({s1, s2}, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto outputs = fe.runFusion({t0});

  auto ref = gather(t0, window_shape, padding_width, strides);

  TORCH_CHECK(ref.equal(outputs[0]));
}

// Same as GatherStrided2 but with stride != window
TEST_F(NVFuserTest, FusionGatherStrided6_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int> window_shape = {3};
  const std::vector<std::vector<int>> padding_width = {{1, 1}};
  const std::vector<int> strides = {2};

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));

  auto tv2 = gather(tv1, window_shape, padding_width, strides);

  auto tv3 = sum(tv2, {-1});

  fusion.addOutput(tv3);

  // Split the strided domain
  tv3->split(0, 4);

  // Propagate the split by 4 of the tv3 domain to pre-stride domains,
  // making them split by 4 * 2
  tv0->computeAt(tv3, 1);

  tv2->computeAt(tv3, -1);

  tv3->axis(0)->parallelize(ParallelType::BIDx);
  tv3->axis(1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv3, {tv1, tv2});

  tv1->setMemoryType(MemoryType::Shared);

  const int s1 = 100;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({s1}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = gather(t1, window_shape, padding_width, strides);
  auto ref = sum(t2, {-1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// Same as GatherStrided4 but different strides
TEST_F(NVFuserTest, FusionGatherStrided7_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int> window_shape = {3};
  const std::vector<std::vector<int>> padding_width = {{1, 1}};

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));

  // Use different strides
  auto tv2 = gather(tv1, window_shape, padding_width, {3});
  auto tv3 = gather(tv1, window_shape, padding_width, {2});

  auto tv4 = sum(tv2, {-1});
  fusion.addOutput(tv4);

  auto tv5 = sum(tv3, {-1});
  fusion.addOutput(tv5);

  tv4->split(0, 2);

  // Since tv3 has a different stride factor, this should fail.
  ASSERT_ANY_THROW(tv0->computeAt(tv4, 1));
}

// Same as GatherStrided2 but with unswitch
TEST_F(NVFuserTest, FusionGatherStrided8_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int> window_shape = {3};
  const std::vector<std::vector<int>> padding_width = {{1, 1}};
  const std::vector<int> strides = {3};

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));

  auto tv2 = gather(tv1, window_shape, padding_width, strides);

  auto tv3 = sum(tv2, {-1});

  fusion.addOutput(tv3);

  const int tidx = 32;

  // Split the strided domain
  tv3->split(0, tidx);

  // Split for unswitch
  tv3->split(0, 1);

  tv0->computeAt(tv3, 2);

  tv2->computeAt(tv3, -1);

  tv3->axis(0)->parallelize(ParallelType::BIDx);
  tv3->axis(1)->parallelize(ParallelType::Unswitch);
  tv3->axis(2)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv3, {tv1, tv2});

  tv1->setMemoryType(MemoryType::Shared);

  const int s1 = 1023;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({s1}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = t0 + 1;
  auto t2 = gather(t1, window_shape, padding_width, strides);
  auto ref = sum(t2, {-1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

// Chained strided gather. Not supported yet.
TEST_F(NVFuserTest, FusionGatherStridedChain_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int> window_shape = {3};
  const std::vector<std::vector<int>> padding_width = {{1, 1}};
  const std::vector<int> strides = {3};
  // const std::vector<int> strides = {1};

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));

  auto tv2 = gather(tv1, window_shape, padding_width, strides);
  // Reduce gathered window
  auto tv3 = sum(tv2, {-1});

  // Repeat
  auto tv4 = gather(tv3, window_shape, padding_width, strides);
  auto tv5 = sum(tv4, {-1});
  auto out = tv5;

  fusion.addOutput(out);

  // This should throw an error at HaloInfo::build.
  ASSERT_ANY_THROW(GpuLower gpulw(&fusion));
}

TEST_F(NVFuserTest, FusionMaxPoolingStrided_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input:  CHW
  // Pooling window: 3x3
  // Strides: 3
  // Padding: 1 at each end of the inner 2 dimensions

  // [C, H, W]
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // [C, H/3, W/3, 1, 3, 3]
  auto inp_tile = gather(inp, {1, 3, 3}, {{0, 0}, {1, 1}, {1, 1}}, {1, 3, 3});

  // [C, H/3, W/3]
  auto max_tensor = reductionOp(
      BinaryOpType::Max,
      {-3, -2, -1},
      IrBuilder::create<Double>(std::numeric_limits<float>::lowest()),
      inp_tile);
  fusion.addOutput(max_tensor);

  ////////////////////////////////////

  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Tiling the spatial domain
  const int tile_x = 32;
  const int tile_y = 8;

  max_tensor->split(1, tile_y);
  max_tensor->split(3, tile_x);
  max_tensor->reorder({{2, 3}});
  // [C, H/tile_y, W/tile_x, tile_y, tile_x]
  max_tensor->split(2, 1);
  // [C, H/tile_y, W/tile_x, 1, tile_y, tile_x]

  inp->computeAt(max_tensor, 4);

  max_tensor->axis(0)->parallelize(ParallelType::BIDx);
  max_tensor->axis(3)->parallelize(ParallelType::Unswitch);
  max_tensor->axis(4)->parallelize(ParallelType::TIDy);
  max_tensor->axis(5)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(max_tensor, ir_utils::allTvs(&fusion));

  inp_cache->setMemoryType(MemoryType::Shared);

  const int hw = 50;
  const int num_channels = 20;
  const int pooling_window = 3;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor aten_inp = at::randn({num_channels, hw, hw}, options);
  // We always pad inputs by zero, so if all surrounding values are
  // negative, max pooling would pick a padded value, which isn't the
  // correct behavior. We need to be able to choose the value of
  // padding. In this case, padding by the minimum value would not
  // have this problem. For now, avoid the problem by making sure all
  // values are not negative.
  aten_inp = at::abs(aten_inp);
  std::vector<IValue> inputs = {aten_inp};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto ref = at::max_pool2d(
      aten_inp, {pooling_window, pooling_window}, {3, 3}, {1, 1});

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionConv2DStaticStrided_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  // Input: [C, H, W]
  auto inp = makeSymbolicTensor(3);
  fusion.addInput(inp);

  // Weights: [K, C, 3, 3]
  auto w = makeSymbolicTensor(4);
  fusion.addInput(w);

  // Gather a neighbor tile of [3, 3] with padding size of 1 for each
  // side of the spatial dimensions
  auto inp_tile = gather(inp, {1, 3, 3}, {{0, 0}, {1, 1}, {1, 1}}, {1, 3, 3});
  // inp_tile: [C, H/3, s3, W/3, s3, 1, 3, 3]

  auto inp_bc =
      broadcast(inp_tile, {true, false, false, false, false, false, false});
  auto w_bc = broadcast(w, {false, false, true, true, true, false, false});

  auto inp_times_w = mul(inp_bc, w_bc);

  // Reduce the channel and neighbor tile dimensions
  auto out = sum(inp_times_w, {1, 4, 5, 6});

  fusion.addOutput(out);

  ////////////////////////////////////

  // Cache the input and weight tensors
  auto inp_cache = inp->cache_after();

  // Blocking the spatial dimensions
  const int block_w = 16;
  const int block_h = 4;
  const int block_c = 2;

  // [K, C, H/s, W/s, 1, 3, 3]
  out->split(2, block_h);
  // [K, C, H/s/block_h, block_h, W/s, 1, 3, 3]
  out->split(4, block_w);
  // [K, C, H/s/block_h, block_h, W/s/block_w, block_w, 1, 3, 3]
  out->reorder({{3, 4}});
  // [K, C, H/s/block_h, W/s/block_w, block_h, block_w, 1, 3, 3]
  out->split(1, block_c);
  // [K, C/block_c, block_c, H/s/block_h, W/s/block_w, block_h, block_w, 1, 3,
  // 3]
  out->split(4, 1);
  // [K, C/block_c, block_c, H/s/block_h, W/s/block_w, 1, block_h, block_w, 1,
  // 3, 3]

  auto out_rf = out->rFactor({1, -3, -2, -1});
  // [K, C/block_c, block_c, H/s/block_h, W/s/block_w, 1, block_h, block_w, 1,
  // 3, 3]

  // out: [K, block_c, H/s/block_h, W/s/block_w, 1, block_h, block_w]

  inp_cache->computeAt(out, 5);
  inp_cache->setMemoryType(MemoryType::Shared);
  // [K, block_c, H/s/block_h, W/s/block_w, 1, block_h, block_w, C/block_c, 1,
  // 3, 3]

  // Move C/block_c before block_h/2 and share the domain from
  // inp_cache to out_rf
  out_rf->reorder({{7, 5}, {5, 6}, {6, 7}});
  inp_cache->computeAt(out_rf, 6);

  inp_tile->computeAt(out_rf, -1);
  w->computeAt(out_rf, -1);

  out->axis(0)->parallelize(ParallelType::BIDx);
  out->axis(1)->parallelize(ParallelType::TIDz);
  out->axis(4)->parallelize(ParallelType::Unswitch);
  out->axis(5)->parallelize(ParallelType::TIDy);
  out->axis(6)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(out, {inp_cache, out_rf});

  const int dim_h = 99;
  const int dim_w = 101;
  const int dim_c = 10;
  const int dim_f = 20;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::manual_seed(0);
  at::Tensor at_inp = at::randn({dim_c, dim_h, dim_w}, options);
  at::Tensor at_w = at::randn({dim_f, dim_c, 3, 3}, options);
  std::vector<IValue> inputs = {at_inp, at_w};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto cg_outputs = fe.runFusion(inputs);

  at_inp = at_inp.unsqueeze(0); // at::conv2d needs the N axis
  auto at_out = at::conv2d(at_inp, at_w, {}, 3, 1);
  at_out = at_out.squeeze(0); // drop the N axis

  testValidate(&fusion, cg_outputs, inputs, {at_out}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionNonDivisibleHalo1_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));
  auto tv2 = shift(tv1, {-1});
  fusion.addOutput(tv2);

  // [I]
  tv2->split(0, 8);
  // [I/8, 8]
  tv2->split(1, 3);
  // [I/8, 3, 3]

  tv0->computeAt(tv2, -2);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({24}, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto cg_outputs = fe.runFusion({t0});

  auto ref = shift((t0 + 1), {-1});

  testValidate(&fusion, cg_outputs, {t0}, {ref}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionNonDivisibleHalo2_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = gather(tv0, {3, 3}, {{1, 1}, {1, 1}});
  auto tv2 = sum(tv1, {-2, -1});
  auto tv3 = add(tv0, tv2);
  auto tv4 = sum(tv3, {0, 1});
  fusion.addOutput(tv4);

  const int gy = 50;
  const int gx = 50;
  const int by = 8;
  const int bx = 16;

  auto tv5 = tv0->cache_after();

  // [I, J]
  tv4->split(0, gy);
  // [I/gy, gy, J]
  tv4->split(1, by);
  // [I/gy, gy/by, by, J]
  tv4->split(-1, gx);
  // [I/gy, gy/by, by, J/gx, gx]
  tv4->split(-1, bx);
  // [I/gy, gy/by, by, J/gx, gx/bx, bx]
  tv4->reorder({{3, 1}, {1, 2}, {4, 3}, {2, 4}});
  // [I/gy, J/gx, gy/by, gx/bx, by, bx]

  auto tv6 = tv4->rFactor({2, 3});

  tv0->computeAt(tv6, 4);

  tv4->axis(0)->parallelize(ParallelType::BIDy);
  tv4->axis(1)->parallelize(ParallelType::BIDx);
  tv4->axis(2)->parallelize(ParallelType::TIDy);
  tv4->axis(3)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(tv4, {tv1, tv2, tv3, tv5, tv6});

  tv5->setMemoryType(MemoryType::Shared);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({111, 222}, options);

  FusionExecutor fe;
  fe.compileFusion(&fusion, {t0});
  auto cg_outputs = fe.runFusion({t0});

  auto t1 = gather(t0, {3, 3}, {{1, 1}, {1, 1}});
  auto t2 = t1.sum({-2, -1});
  auto t3 = t0 + t2;
  auto t4 = t3.sum({-2, -1});

  testValidate(&fusion, cg_outputs, {t0}, {t4}, __LINE__, __FILE__);
}

TEST_F(NVFuserTest, FusionGather9ptStencilDoubleBuffering_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = gather(tv0, {3, 3}, {{1, 1}, {1, 1}});
  auto tv2 = sum(tv1, {-2, -1});
  auto tv3 = div(tv2, IrBuilder::create<Double>(9));

  auto out = tv3;

  fusion.addOutput(out);

  auto tv0_cache = tv0->cache_after();

  tv0_cache->setMemoryType(MemoryType::Shared);

  out->split(-2, 4);
  out->split(-1, 32);
  out->reorder({{1, 2}, {2, 1}});
  TransformPropagator::from(out);

  tv0->computeAt(out, 2);

  out->axis(3)->parallelize(ParallelType::TIDx);
  out->axis(2)->parallelize(ParallelType::TIDy);
  out->axis(0)->parallelize(ParallelType::BIDx);

  scheduler_utils::parallelizeAllLike(out, ir_utils::allTvs(&fusion));

  tv0_cache->doubleBuffer();

  int numel_x = 99;
  int numel_y = 101;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({numel_x, numel_y}, options);
  std::vector<IValue> inputs = {t0};

  FusionExecutor fe;
  fe.compileFusion(&fusion, inputs);
  auto outputs = fe.runFusion(inputs);

  auto t1 = gather(t0, {3, 3}, {{1, 1}, {1, 1}});
  auto t2 = sum(t1, {-2, -1});
  auto t3 = t2 / 9;
  auto ref = t3;

  testValidate(&fusion, outputs, inputs, {ref}, __LINE__, __FILE__);
}

} // namespace jit
} // namespace torch
#endif // #if defined(USE_CUDA)