test_reduction.cpp

// clang-format off
/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
 * All rights reserved.
 * SPDX-License-Identifier: BSD-3-Clause
 */
// clang-format on
#include <algorithm>
#include <iostream>

// fuser and IR parser
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/Exceptions.h>

#include <gtest/gtest.h>

#include <c10/cuda/CUDAStream.h>

#include "codegen.h"
#include "device_lower/lower2device.h"
#include "disjoint_set.h"
#include "exceptions.h"
#include "expr_evaluator.h"
#include "fusion.h"
#include "fusion_segmenter.h"
#include "grouped_reduction.h"
#include "ir/all_nodes.h"
#include "ir/builder.h"
#include "ir/graphviz.h"
#include "ir/iostream.h"
#include "ir/utils.h"
#include "iter_visitor.h"
#include "kernel_ir.h"
#include "logical_domain_map.h"
#include "ops/all_ops.h"
#include "runtime/executor.h"
#include "runtime/executor_params.h"
#include "runtime/fusion_executor_cache.h"
#include "scheduler/all_schedulers.h"
#include "scheduler/reduction_outer_tma.h"
#include "scheduler/reduction_tma.h"
#include "scheduler/reduction_utils.h"
#include "scheduler/tools/inlining.h"
#include "scheduler/utils.h"
#include "tests/cpp/utils.h"
#include "transform_replay.h"
#include "transform_rfactor.h"
#include "validator_utils.h"

namespace nvfuser {

using namespace at::indexing;

namespace {

void validateNoParallelBroadcastExist(kir::Kernel* kernel) {
  for (auto expr : KernelExprVisitor::getAllExprs(kernel)) {
    BroadcastOp* bc = dynamic_cast<BroadcastOp*>(expr);
    if (bc == nullptr) {
      auto grid_bc = dynamic_cast<kir::GridBroadcast*>(expr);
      if (grid_bc != nullptr) {
        std::cerr << "Grid broadcast: " << grid_bc->toString();
        bc = grid_bc->broadcast_op();
      }
    }
    if (bc == nullptr) {
      continue;
    }
    NVF_CHECK(
        kernel->summary().broadcast_parallel_types.at(bc).none(),
        "Parallel broadcast should not exist but was found: ",
        bc->toString());
  }
}

} // namespace

using ReductionTest = NVFuserTest;

TEST_F(ReductionTest, GridAllreduce1) {
  const int nx = 999;
  const int tidx = 128;
  const int bidx = 4;

  if (ceilDiv(nx, tidx) > deviceSMCount()) {
    GTEST_SKIP() << "Not enough SMs to run this test";
  }

  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {0});
  auto tv2 = broadcast(tv1, {true});
  auto tv3 = add(tv0, tv2);

  fusion.addOutput(tv3);

  tv3->split(0, tidx);
  tv3->split(0, bidx);
  tv3->split(0, 1); // unswitch
  TransformPropagator propagator(tv3);
  MaxLogicalDomainInfoSpanningTree(tv3).traverse(&propagator);

  tv3->axis(0)->parallelize(ParallelType::BIDy);
  tv3->axis(2)->parallelize(ParallelType::BIDx);
  tv3->axis(3)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv3);

  // Just to make sure fused_reduction and work buffers are allocated
  // uniquely
  tv1->axis(1)->parallelize(ParallelType::Unswitch);

  GpuLower gpulw(&fusion);
  validateNoParallelBroadcastExist(gpulw.run());

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

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto cg_outputs = ke.run({t0});

  auto ref = sum(t0).unsqueeze(0) + t0;

  testValidate(&fusion, cg_outputs, {t0}, {ref}, __LINE__, __FILE__);
}

TEST_F(ReductionTest, GridAllreduce2) {
  const int nx = 99;
  const int tidx = 32;

  if (ceilDiv(nx, tidx) > deviceSMCount()) {
    GTEST_SKIP() << "Not enough SMs to run this test";
  }

  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {0});
  auto tv2 = broadcast(tv1, {true});
  auto tv3 = add(tv0, tv2);

  fusion.addOutput(tv3);

  tv3->split(0, tidx);
  TransformPropagator propagator(tv3);
  MaxLogicalDomainInfoSpanningTree(tv3).traverse(&propagator);

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

  // Broadcast on TIDy instead of TIDx. This still uses the fused
  // reduction as it's broadcast on BIDx as well. Since TIDy is not
  // predicated, the broadcast becomes a set op.
  tv1->axis(0)->parallelize(ParallelType::BIDx);
  tv1->axis(1)->parallelize(ParallelType::TIDy);

  GpuLower gpulw(&fusion);
  validateNoParallelBroadcastExist(gpulw.run());

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

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto cg_outputs = ke.run({t0});

  auto ref = sum(t0).unsqueeze(0) + t0;

  testValidate(&fusion, cg_outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Grid reduction with serial non-reduction axis. The global work
// buffer is circular buffered.
TEST_F(ReductionTest, GridAllreduce3) {
  const int nx = 100;
  const int ny = 5000;
  const int tidx = 128;

  if (ceilDiv(ny, tidx) > deviceSMCount()) {
    GTEST_SKIP() << "Not enough SMs to run this test";
  }

  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {1});
  auto tv2 = broadcast(tv1, {false, true});
  auto tv3 = add(tv0, tv2);

  fusion.addOutput(tv3);

  tv3->split(1, tidx);
  TransformPropagator propagator(tv3);
  MaxLogicalDomainInfoSpanningTree(tv3).traverse(&propagator);

  tv0->computeAt(tv3, 1);

  tv3->axis(1)->parallelize(ParallelType::BIDx);
  tv3->axis(2)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv3);

  GpuLower gpulw(&fusion);
  validateNoParallelBroadcastExist(gpulw.run());

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

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto cg_outputs = ke.run({t0});

  auto ref = sum(t0, {1}).unsqueeze(-1) + t0;

  testValidate(&fusion, cg_outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Indirect reduction and broadcast
TEST_F(ReductionTest, GridAllreduce4) {
  const int nx = 999;
  const int tidx = 128;

  if (ceilDiv(nx, tidx) > deviceSMCount()) {
    GTEST_SKIP() << "Not enough SMs to run this test";
  }

  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {0});
  auto tv2 = add(tv1, IrBuilder::create<Val>(1.0));
  auto tv3 = broadcast(tv2, {true});
  auto tv4 = add(tv0, tv3);

  fusion.addOutput(tv4);

  tv4->split(0, tidx);
  TransformPropagator propagator(tv4);
  MaxLogicalDomainInfoSpanningTree(tv4).traverse(&propagator);

  tv4->axis(0)->parallelize(ParallelType::BIDx);
  tv4->axis(1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv4);

  GpuLower gpulw(&fusion);
  validateNoParallelBroadcastExist(gpulw.run());

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

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto cg_outputs = ke.run({t0});

  auto ref = (sum(t0) + 1).unsqueeze(0) + t0;

  testValidate(&fusion, cg_outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Unused block dimension in the kernel
TEST_F(ReductionTest, GridAllreduce5) {
  const int nx = 999;
  const int tidx = 128;
  const int iter = 2;
  const int bdimx = 9; // One more than required by the reduction
  const int bdimy = 3; // Want an unused dimension

  // Going to bump the bdimx count for this test, ignor
  if (bdimx * bdimy > deviceSMCount()) {
    GTEST_SKIP() << "Not enough SMs to run this test";
  }

  Fusion fusion;
  FusionGuard fg(&fusion);

  // Didn't setup this test with inlining for register usage, so just leave the
  // iter dimension concrete
  auto tv0 = makeConcreteTensor({iter, -1});
  fusion.addInput(tv0);

  auto tv1 = sum(tv0, {1});
  auto tv2 = add(tv1, IrBuilder::create<Val>(1.0));
  auto tv3 = broadcast(tv2, {false, true});
  auto tv4 = add(tv0, tv3);

  fusion.addOutput(tv4);

  // Dummy op to mess with parallelization
  auto tv5 = makeSymbolicTensor(2);
  fusion.addInput(tv5);
  auto tv6 = set(tv5);
  fusion.addOutput(tv6);

  // Setup the reduction
  tv4->split(1, tidx);
  TransformPropagator propagator(tv4);
  MaxLogicalDomainInfoSpanningTree(tv4).traverse(&propagator);

  tv4->axis(1)->parallelize(ParallelType::BIDx);
  tv4->axis(2)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv4);

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

  GpuLower gpulw(&fusion);
  validateNoParallelBroadcastExist(gpulw.run());

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn({iter, nx}, options);
  auto t5 = at::randn({bdimy, bdimx}, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0, t5});
  auto cg_outputs = ke.run({t0, t5});

  auto ref = (sum(t0, {1}) + 1).unsqueeze(-1) + t0;

  testValidate(&fusion, cg_outputs, {t0, t5}, {ref, t5}, __LINE__, __FILE__);
}

TEST_F(ReductionTest, GridAllreduce6) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  std::vector<int64_t> shape({99, 200});

  const int vec = 4;
  const int tidx = 32;
  const int tidy = 8;
  const int bdimx = ceilDiv(shape[1], vec * tidx);
  const int bdimy = ceilDiv(shape[0], tidy);

  if (bdimx * bdimy > deviceSMCount()) {
    GTEST_SKIP() << "Not enough SMs to run this test";
  }

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

  auto tv1 = set(tv0);
  auto tv2 = sum(tv1, {0});
  auto tv3 = broadcast(tv2, {true, false});
  auto tv4 = add(tv0, tv3);
  fusion.addOutput(tv4);

  tv1->split(1, vec);
  tv1->split(1, tidx);
  tv1->split(0, tidy);
  TransformPropagator propagator(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator);

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

  scheduler_utils::parallelizeAllLike(tv1);

  tv1->axis(4)->parallelize(ParallelType::Vectorize);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t0_double = t0.to(at::kDouble);
  auto ref = t0_double + t0_double.sum({0}).unsqueeze(0);

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

TEST_F(ReductionTest, GridAllreduceWelford1) {
  const int nx = 999;
  const int tidx = 128;

  if (ceilDiv(nx, tidx) > deviceSMCount()) {
    GTEST_SKIP() << "Not enough SMs to run this test";
  }

  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tvs = Welford(tv0, {0});
  auto tv2 = broadcast(tvs.avg, {true});
  auto tv3 = broadcast(tvs.var_sum, {true});
  auto tv4 = add(tv0, tv2);
  auto tv5 = add(tv4, tv3);

  fusion.addOutput(tv5);

  tv5->split(0, tidx);
  TransformPropagator propagator(tv5);
  MaxLogicalDomainInfoSpanningTree(tv5).traverse(&propagator);

  tv5->axis(0)->parallelize(ParallelType::BIDx);
  tv5->axis(1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv5);

  GpuLower gpulw(&fusion);
  validateNoParallelBroadcastExist(gpulw.run());

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

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto cg_outputs = ke.run({t0});

  auto ref =
      (t0.mean({0}).unsqueeze(0) + t0) + t0.var({0}, false).unsqueeze(0) * nx;

  testValidate(&fusion, cg_outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Grid welford reduction with serial non-reduction axis. The global
// work buffer is circular buffered.
TEST_F(ReductionTest, GridAllreduceWelford2) {
  const int nx = 100;
  const int ny = 5000;
  const int tidx = 128;

  if (ceilDiv(ny, tidx) > deviceSMCount()) {
    GTEST_SKIP() << "Not enough SMs to run this test";
  }

  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tvs = Welford(tv0, {1});
  auto tv2 = broadcast(tvs.avg, {false, true});
  auto tv3 = add(tv0, tv2);

  fusion.addOutput(tv3);

  tv3->split(1, tidx);
  TransformPropagator propagator(tv3);
  MaxLogicalDomainInfoSpanningTree(tv3).traverse(&propagator);

  tv0->computeAt(tv3, 1);

  tv3->axis(1)->parallelize(ParallelType::BIDx);
  tv3->axis(2)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv3);

  // There must be no parallel broadcast
  GpuLower gpulw(&fusion);
  validateNoParallelBroadcastExist(gpulw.run());

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

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto cg_outputs = ke.run({t0});

  auto ref = (sum(t0, {1}) / ny).unsqueeze(-1) + t0;

  testValidate(&fusion, cg_outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Persistent batchnorm. Uses the fused reduction for grid welford and
// broadcast.
TEST_F(ReductionTest, FusedReductionBatchnorm) {
  const std::vector<int64_t> input_shape{256, 2048, 14, 14};

  std::unique_ptr<Fusion> fusion_ptr = std::make_unique<Fusion>();
  Fusion& fusion = *fusion_ptr.get();
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(4, DataType::Half);
  fusion.addInput(tv0);
  auto tv1 = makeSymbolicTensor(1, DataType::Half);
  fusion.addInput(tv1);
  auto tv2 = makeSymbolicTensor(1, DataType::Half);
  fusion.addInput(tv2);
  auto tv3 = makeSymbolicTensor(1, DataType::Float);
  fusion.addInput(tv3);
  auto tv4 = makeSymbolicTensor(1, DataType::Float);
  fusion.addInput(tv4);

  auto d34 = IrBuilder::create<Val>(1.0);
  auto tv5 = castOp(DataType::Float, tv0);
  auto tv6 = castOp(DataType::Float, tv1);
  auto tv7 = castOp(DataType::Float, tv2);
  auto tvs = Welford(tv5, {0, 2, 3});
  auto tv8 = tvs.avg;
  auto tv9 = tvs.var_sum;
  auto tv11 = mul(tv8, IrBuilder::create<Val>(0.1));
  auto tv12 = mul(tv3, d34);
  auto tv13 = add(tv12, tv11);
  auto d43 = IrBuilder::create<Val>(0.5);
  auto tv14 = mul(tv9, d43);
  auto tv15 = mul(tv14, IrBuilder::create<Val>(0.1));
  auto tv16 = mul(tv4, d34);
  auto tv17 = add(tv16, tv15);
  auto tv18 = broadcast(tv8, {true, false, true, true});
  auto tv19 = sub(tv5, tv18);
  auto tv20 = mul(tv9, d43);
  auto tv21 = add(tv20, IrBuilder::create<Val>(0.0001));
  auto tv22 = rsqrt(tv21);
  auto tv23 = broadcast(tv22, {true, false, true, true});
  auto tv24 = mul(tv19, tv23);
  auto tv25 = broadcast(tv6, {true, false, true, true});
  auto tv26 = mul(tv24, tv25);
  auto tv27 = broadcast(tv7, {true, false, true, true});
  auto tv28 = add(tv26, tv27);
  auto tv29 = castOp(DataType::Half, tv28);
  fusion.addOutput(tv13);
  fusion.addOutput(tv17);
  fusion.addOutput(tv29);

  tv0->cacheAfter();
  tv1->cacheAfter();
  tv2->cacheAfter();
  tv3->cacheAfter();
  tv4->cacheAfter();

  tv13->cacheBefore();
  tv17->cacheBefore();
  tv29->cacheBefore();

  tv0->split(1, NamedScalar::getParallelDim(ParallelType::BIDx), false);
  tv0->split(0, NamedScalar::getParallelDim(ParallelType::BIDy), false);
  tv0->split(1, 8, false);
  tv0->split(2, 8, false);
  tv0->merge(-2, -1);
  tv0->split(-1, 2);
  tv0->split(-2, 1, false);
  tv0->split(-2, 1, false);
  tv0->reorder(
      {{4, 0},
       {5, 1},
       {0, 2},
       {3, 3},
       {8, 4},
       {1, 5},
       {7, 6},
       {2, 7},
       {9, 8},
       {6, 9}});

  TransformPropagator propagator(tv0);
  MaxLogicalDomainInfoSpanningTree(tv0).traverse(&propagator);

  ir_utils::rFactorHelper(tvs.avg, {-5, -4, -3, -2, -1});

  tv0->computeAt(tv29, 2);
  tv1->computeAt(tv29, 2);
  tv2->computeAt(tv29, 2);
  tv3->computeAt(tv13, 2);
  tv4->computeAt(tv17, 2);

  tv29->axis(0)->parallelize(ParallelType::BIDx);
  tv29->axis(2)->parallelize(ParallelType::BIDy);
  tv29->axis(3)->parallelize(ParallelType::TIDz);
  tv29->axis(4)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv29);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto options_half = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  auto t0 = at::randn(input_shape, options_half);
  auto t1 = at::randn(input_shape[1], options_half);
  auto t2 = at::randn(input_shape[1], options_half);
  auto t3 = at::randn(input_shape[1], options);
  auto t4 = at::randn(input_shape[1], options);
  KernelArgumentHolder inputs({t0, t1, t2, t3, t4});

  GpuLower gpulw(&fusion);
  validateNoParallelBroadcastExist(gpulw.run());

  KernelExecutor ke;
  LaunchParams launch_params(2, 2, -1, -1, -1, -1);
  ke.compile(&fusion, inputs, launch_params);
  auto cg_outputs = ke.run(inputs, {}, launch_params);

  auto t5 = t0.to(at::kFloat);
  auto t6 = t1.to(at::kFloat);
  auto t7 = t2.to(at::kFloat);
  auto t8 = t5.mean({0, 2, 3});
  auto t9 = t5.var({0, 2, 3}, false) * input_shape[0] * input_shape[2] *
      input_shape[3];
  auto t11 = t8 * 0.1;
  auto t12 = t3 * 1;
  auto t13 = t12 + t11;
  auto t14 = t9 * 0.5;
  auto t15 = t14 * 0.1;
  auto t16 = t4 * 1;
  auto t17 = t16 + t15;
  auto t18 = t8.unsqueeze(0).unsqueeze(-1).unsqueeze(-1);
  auto t19 = t5 - t18;
  auto t20 = t9 * 0.5;
  auto t21 = t20 + 0.0001;
  auto t22 = rsqrt(t21);
  auto t23 = t22.unsqueeze(0).unsqueeze(-1).unsqueeze(-1);
  auto t24 = t19 * t23;
  auto t25 = t6.unsqueeze(0).unsqueeze(-1).unsqueeze(-1);
  auto t26 = t24 * t25;
  auto t27 = t7.unsqueeze(0).unsqueeze(-1).unsqueeze(-1);
  auto t28 = t26 + t27;
  auto t29 = t28.to(at::kHalf);

  testValidate(
      &fusion,
      cg_outputs,
      inputs,
      {t13, t17, t29},
      __LINE__,
      __FILE__,
      "",
      launch_params);
}

// Simple grouped reduction
TEST_F(ReductionTest, GroupedReduction1) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {1});
  auto tv2 = sum(tv0, {1});
  auto tv3 = add(tv1, tv2);
  fusion.addOutput(tv3);

  groupReductions({tv1, tv2});

  tv2->axis(0)->parallelize(ParallelType::BIDx);
  tv2->axis(1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv2);

  std::vector<int64_t> shape({99, 999});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto ref = t0.sum({1}) * 2;

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Grouping reductions with different ops
TEST_F(ReductionTest, GroupedReduction2) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = add(tv0, IrBuilder::create<Val>(1.0));
  auto tv2 = sum(tv1, {1});

  auto tv3 = add(tv0, IrBuilder::create<Val>(2.0));
  auto tv4 = max(tv3, {1});

  auto tv5 = add(tv2, tv4);
  fusion.addOutput(tv5);

  groupReductions({tv2, tv4});

  tv2->split(1, 128);
  TransformPropagator propagator(tv2);
  MaxLogicalDomainInfoSpanningTree(tv2).traverse(&propagator);

  tv0->computeAt(tv4, -1, ComputeAtMode::MostInlined);

  // tv4 is automatically parallelized in the same way
  tv2->axis(0)->parallelize(ParallelType::BIDy);
  tv2->axis(1)->parallelize(ParallelType::BIDx);
  tv2->axis(2)->parallelize(ParallelType::TIDx);

  std::vector<int64_t> shape({99, 999});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto ref = (t0 + 1).sum({1}) + std::get<0>((t0 + 2).max(1));

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Grouped reduction with different types
TEST_F(ReductionTest, GroupedReduction3) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {1});

  auto tv2 = castOp(DataType::Double, tv0);
  auto tv3 = sum(tv2, {1});
  auto tv4 = castOp(DataType::Float, tv3);

  auto tv5 = add(tv1, tv4);
  fusion.addOutput(tv5);

  groupReductions({tv1, tv3});
  tv1->split(1, 128);
  TransformPropagator propagator(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator);

  tv0->computeAt(tv5, -1, ComputeAtMode::MostInlined);

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

  std::vector<int64_t> shape({99, 999});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto ref = t0.sum({1}) + t0.to(c10::kDouble).sum({1}).to(c10::kFloat);

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Testing validation
TEST_F(ReductionTest, GroupedReduction4) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv2 = sum(tv0, {1});
  auto tv3 = sum(tv1, {1});
  auto tv4 = add(tv2, tv3);
  fusion.addOutput(tv4);

  // Invalid grouping as tv2 and tv3 are not guaranteed to have the
  // same shape
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-goto,hicpp-avoid-goto)
  ASSERT_ANY_THROW(groupReductions({tv2, tv3}));
}

// Testing validation
TEST_F(ReductionTest, GroupedReduction5) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {1});
  auto tv2 = sum(tv0, {1});
  auto tv3 = add(tv1, tv2);
  fusion.addOutput(tv3);

  tv1->split(1, 128);
  tv2->split(1, 64);

  // Invalid grouping as tv1 and tv2 don't have the same
  // transformations
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-goto,hicpp-avoid-goto)
  ASSERT_ANY_THROW(groupReductions({tv1, tv2}));
}

// Grouping 3 reductions
TEST_F(ReductionTest, GroupedReduction6) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = add(tv0, IrBuilder::create<Val>(1.0));
  auto tv2 = sum(tv1, {1});

  auto tv3 = add(tv0, IrBuilder::create<Val>(2.0));
  auto tv4 = sum(tv3, {1});

  auto tv5 = add(tv0, IrBuilder::create<Val>(3.0));
  auto tv6 = sum(tv5, {1});

  auto tv7 = add(add(tv2, tv4), tv6);

  fusion.addOutput(tv7);

  groupReductions({tv2, tv4, tv6});

  // There's no runtime grid reduction function that can take more
  // than 2 inputs, yet.
  tv2->axis(0)->parallelize(ParallelType::BIDx);
  tv2->axis(1)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(tv2);

  std::vector<int64_t> shape({99, 999});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, __LINE__, __FILE__);
}

TEST_F(ReductionTest, GroupedReduction7) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {1});
  auto tv2 = broadcast(tv1, {false, true});
  auto tv3 = add(tv0, tv2);
  auto tv4 = sum(tv3, {1});
  fusion.addOutput(tv4);

  // Invalid grouping as tv3 depends on tv1
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-goto,hicpp-avoid-goto)
  ASSERT_ANY_THROW(groupReductions({tv1, tv4}));
}

// Grouping rfactor'ed reductions
TEST_F(ReductionTest, GroupedReductionRfactor1) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {0});
  auto tv2 = sum(tv0, {0});
  auto tv3 = add(tv1, tv2);
  fusion.addOutput(tv3);

  const int64_t gdimx = 10;
  const int64_t bdimx = 128;

  tv1->split(0, gdimx, false);
  tv1->split(1, bdimx);
  auto tv1_rf = tv1->rFactor({1});

  tv2->split(0, gdimx, false);
  tv2->split(1, bdimx);
  auto tv2_rf = tv2->rFactor({1});

  groupReductions({tv1_rf, tv2_rf});
  groupReductions({tv1, tv2});

  tv1_rf->axis(0)->parallelize(ParallelType::BIDx);
  tv1_rf->axis(2)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(tv1_rf);

  std::vector<int64_t> shape({12345});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto ref = t0.sum({0}) * 2;

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Rfactoring grouped reductions
TEST_F(ReductionTest, GroupedReductionRfactor2) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {0});
  auto tv2 = sum(tv0, {0});
  auto tv3 = add(tv1, tv2);
  fusion.addOutput(tv3);

  groupReductions({tv1, tv2});

  const int64_t gdimx = 10;
  const int64_t bdimx = 128;

  tv1->split(0, gdimx, false);
  tv1->split(1, bdimx);

  // This should rfactor tv2 as well
  auto rf_tvs = tv1->rFactor({1}, {tv1, tv2});
  auto tv1_rf = rf_tvs.at(0);

  tv1_rf->axis(0)->parallelize(ParallelType::BIDx);
  tv1_rf->axis(2)->parallelize(ParallelType::TIDx);

  scheduler_utils::parallelizeAllLike(tv1_rf);

  std::vector<int64_t> shape({12345});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto ref = t0.sum({0}) * 2;

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Group reductions of tensors that have computeAt positions set
TEST_F(ReductionTest, GroupedReductionAfterComputeAt) {
  Fusion fusion;
  FusionGuard fg(&fusion);
  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Val>(1.0));
  auto tv2 = sum(tv1, {1});
  auto tv3 = sum(tv1, {1});
  auto tv4 = add(tv2, tv3);
  fusion.addOutput(tv4);

  const int64_t bdimx = 128;

  tv2->split(1, bdimx);
  auto tv2_rf = tv2->rFactor({1});
  tv2_rf->reorder({{1, 2}});

  tv3->split(1, bdimx);
  auto tv3_rf = tv3->rFactor({1});
  tv3_rf->reorder({{1, 2}});

  tv0->computeAt(tv4, -1, ComputeAtMode::MostInlined);

  groupReductions({tv2_rf, tv3_rf});
  groupReductions({tv2, tv3});

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

  std::vector<int64_t> shape({3, 1234});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto ref = (t0 + 1).sum({1}) * 2;

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

TEST_F(ReductionTest, GroupAllreduce1) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {0});
  auto tv2 = broadcast(tv1, {true});
  auto tv3 = sum(tv0, {0});
  auto tv4 = broadcast(tv3, {true});
  auto tv5 = add(tv0, tv2);
  auto tv6 = add(tv5, tv4);
  fusion.addOutput(tv6);

  groupReductions({tv1, tv3});

  tv2->split(0, 128);
  TransformPropagator propagator(tv2);
  MaxLogicalDomainInfoSpanningTree(tv2).traverse(&propagator);

  tv2->axis(0)->parallelize(ParallelType::BIDx);
  tv2->axis(1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv2);

  std::vector<int64_t> shape({999});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t3 = t0.sum({0}).unsqueeze(-1);
  auto ref = t0 + t3 + t3;

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Grid reductionso of different types
TEST_F(ReductionTest, GroupAllreduce2) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {1});
  auto tv2 = broadcast(tv1, {false, true});

  auto tv3 = castOp(DataType::Double, tv0);
  auto tv4 = sum(tv3, {1});
  auto tv5 = broadcast(tv4, {false, true});
  auto tv6 = castOp(DataType::Float, tv5);

  auto tv7 = add(tv0, tv2);
  auto tv8 = add(tv7, tv6);
  fusion.addOutput(tv8);

  const int tidx = 512;
  groupReductions({tv1, tv4});
  tv1->split(1, tidx);
  TransformPropagator propagator(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator);

  tv0->computeAt(tv8, -1, ComputeAtMode::MostInlined);

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

  std::vector<int64_t> shape({10, 999});

  if (shape.at(0) * ceilDiv(shape.at(1), tidx) > deviceSMCount()) {
    GTEST_SKIP() << "Not enough SMs to run this test";
  }

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t2 = t0.sum({1}).unsqueeze(-1);
  auto t6 = t0.to(c10::kDouble).sum({1}).unsqueeze(-1).to(c10::kFloat);
  auto ref = t0 + t2 + t6;

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Grouping 3 grid allreduces
TEST_F(ReductionTest, GroupAllreduce3) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = sum(tv0, {0});
  auto tv2 = broadcast(tv1, {true});
  auto tv3 = div(tv0, tv2);
  auto tv4 = max(tv0, {0});
  auto tv5 = broadcast(tv4, {true});
  auto tv6 = div(tv0, tv5);
  auto tv7 = min(tv0, {0});
  auto tv8 = broadcast(tv7, {true});
  auto tv9 = sub(tv0, tv8);
  fusion.addOutput(tv3);
  fusion.addOutput(tv6);
  fusion.addOutput(tv9);

  groupReductions({tv1, tv4, tv7});

  tv1->split(0, 128);
  TransformPropagator propagator(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator);

  tv1->axis(0)->parallelize(ParallelType::BIDx);
  tv1->axis(1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv1);

  std::vector<int64_t> shape({999});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t3 = t0 / t0.sum({0}).unsqueeze(0);
  auto t6 = t0 / std::get<0>(t0.max(0)).unsqueeze(0);
  auto t9 = t0 - std::get<0>(t0.min(0)).unsqueeze(0);

  testValidate(
      ke.compiledKernel()->kernel(),
      outputs,
      {t0},
      {t3, t6, t9},
      __LINE__,
      __FILE__);
}

// Grouping 8 grid allreduces
TEST_F(ReductionTest, GroupAllreduce4) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const int num_reductions = 8;

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

  auto tv_sum = tv0;
  std::vector<TensorView*> reduction_tvs;

  for (int i = 0; i < num_reductions; ++i) {
    auto reduction =
        sum(add(tv0, IrBuilder::create<Val>(static_cast<double>(i))), {0});
    reduction_tvs.push_back(reduction);
    auto avg = div(reduction, tv0->axis(0)->extent());
    auto bc = broadcast(avg, {true});
    tv_sum = add(tv_sum, bc);
  }

  fusion.addOutput(tv_sum);

  groupReductions(reduction_tvs);

  auto reduction_tv = reduction_tvs.at(0);

  reduction_tv->split(0, 128);
  TransformPropagator propagator(reduction_tv);
  MaxLogicalDomainInfoSpanningTree(reduction_tv).traverse(&propagator);

  reduction_tv->axis(0)->parallelize(ParallelType::BIDx);
  reduction_tv->axis(1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(reduction_tv);

  std::vector<int64_t> shape({999});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  at::Tensor ref = t0;
  for (int i = 0; i < num_reductions; ++i) {
    auto reduction = sum(add(t0, i), {0});
    auto avg = reduction / t0.sizes()[0];
    auto bc = avg.unsqueeze(0);
    ref = add(ref, bc);
  }

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Variation of GroupAllreduce5 but with different
// types. Exercise grouped allreduces with different types.
TEST_F(ReductionTest, GroupAllreduce5) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1, DataType::Float);
  fusion.addInput(tv0);
  auto tv1 = sum(tv0, {0});
  auto tv2 = broadcast(tv1, {true});
  auto tv3 = div(tv0, tv2);

  auto tv4 = makeSymbolicTensor(1, DataType::Double);
  fusion.addInput(tv4);
  auto tv5 = sum(tv4, {0});
  auto tv6 = broadcast(tv5, {true});
  auto tv7 = div(tv4, tv6);

  auto tv8 = makeSymbolicTensor(1, DataType::Int);
  fusion.addInput(tv8);
  auto tv9 = sum(tv8, {0});
  auto tv10 = broadcast(tv9, {true});
  auto tv11 = div(tv8, tv10);

  auto tv12 = makeSymbolicTensor(1, DataType::ComplexFloat);
  fusion.addInput(tv12);
  auto tv13 = sum(tv12, {0});
  auto tv14 = broadcast(tv13, {true});
  auto tv15 = div(tv12, tv14);

  auto tv16 = makeSymbolicTensor(1, DataType::ComplexDouble);
  fusion.addInput(tv16);
  auto tv17 = sum(tv16, {0});
  auto tv18 = broadcast(tv17, {true});
  auto tv19 = div(tv16, tv18);

  auto outFDI =
      add(add(castOp(DataType::ComplexDouble, tv3), tv7),
          castOp(DataType::ComplexDouble, tv11));
  auto out = add(add(tv15, outFDI), tv19);
  fusion.addOutput(out);

  groupReductions({tv1, tv5, tv9, tv13, tv17});

  tv1->split(0, 128);
  TransformPropagator propagator(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator);

  tv1->axis(0)->parallelize(ParallelType::BIDx);
  tv1->axis(1)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv1);

  std::vector<int64_t> shape({999});

  auto options_float =
      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto options_double =
      at::TensorOptions().dtype(at::kDouble).device(at::kCUDA, 0);
  auto options_complex_float =
      at::TensorOptions().dtype(at::kComplexFloat).device(at::kCUDA, 0);
  auto options_complex_double =
      at::TensorOptions().dtype(at::kComplexDouble).device(at::kCUDA, 0);
  auto options_long = at::TensorOptions().dtype(at::kLong).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options_float);
  auto t4 = at::randn(shape, options_double);
  auto t8 = at::randint(0, 1000, shape, options_long);
  auto t12 = at::randn(shape, options_complex_float);
  auto t16 = at::randn(shape, options_complex_double);
  KernelArgumentHolder inputs({t0, t4, t8, t12, t16});

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

  KernelExecutor ke;
  ke.compile(&fusion, inputs);
  auto outputs = ke.run(inputs);

  auto t3 = t0 / t0.sum({0}).unsqueeze(0).to(at::kComplexDouble);
  auto t7 = t4 / t4.sum({0}).unsqueeze(0).to(at::kComplexDouble);
  auto t11 = at::div(t8, t8.sum({0}).unsqueeze(0), "trunc");
  auto t15 = t12 / t12.sum({0}).unsqueeze(0).to(at::kComplexDouble);
  auto t19 = t16 / t16.sum({0}).unsqueeze(0).to(at::kComplexDouble);
  auto ref = t3 + t7 + t11 + t15 + t19;
  testValidate(
      ke.compiledKernel()->kernel(),
      outputs,
      inputs,
      {ref},
      __LINE__,
      __FILE__);
}

// Persistent batchnorm backward with grouped allreduce
TEST_F(ReductionTest, PersistentBNBackwardAllreduce) {
  const std::vector<int64_t> shape({64, 1024, 14, 14});

  Fusion fusion;
  FusionGuard fg(&fusion);

  auto input = makeContigTensor(4);
  fusion.addInput(input);
  auto grad_output = makeContigTensor(4);
  fusion.addInput(grad_output);
  auto weight = makeContigTensor(1);
  fusion.addInput(weight);
  auto save_mean = makeContigTensor(1);
  fusion.addInput(save_mean);
  auto save_invstd = makeContigTensor(1);
  fusion.addInput(save_invstd);

  const bool kTraining = true;
  const bool channels_last = false;

  const int64_t kNumberOfDims =
      TensorDomain::noReductions(input->getLogicalDomain()).size();
  int64_t c_axis = channels_last ? kNumberOfDims - 1 : 1;

  std::vector<int64_t> reduction_axes;
  std::vector<bool> broadcast_mask(kNumberOfDims, false);
  Val* num_features = nullptr;
  for (const auto axis : arange(kNumberOfDims)) {
    if (axis != c_axis) {
      reduction_axes.push_back(axis);
      broadcast_mask[axis] = true;
      if (num_features == nullptr) {
        num_features =
            castOp(DataType::Double, input->getLoopDomain()[axis]->extent());
      } else {
        num_features =
            mul(num_features, input->getLoopDomain()[axis]->extent());
      }
    }
  }

  auto mean = save_mean;
  auto invstd = save_invstd;

  mean = broadcast(mean, broadcast_mask);

  auto norm = reciprocal(num_features);

  auto grad_output_sum = sum(grad_output, reduction_axes);
  auto dot_p = sum(mul(grad_output, sub(input, mean)), reduction_axes);

  auto grad_mean = broadcast(mul(grad_output_sum, norm), broadcast_mask);

  auto proj_scale =
      broadcast(mul(mul(dot_p, norm), mul(invstd, invstd)), broadcast_mask);

  TensorView* grad_scale = nullptr;

  if (weight == nullptr) {
    grad_scale =
        mul(broadcast(invstd, broadcast_mask), IrBuilder::create<Val>(1.0));
  } else {
    grad_scale = mul(
        broadcast(invstd, broadcast_mask), broadcast(weight, broadcast_mask));
  }

  TensorView* grad_input = nullptr;
  if (kTraining) {
    auto proj = mul(sub(input, mean), proj_scale);
    grad_input = mul(sub(sub(grad_output, proj), grad_mean), grad_scale);
  } else {
    grad_input = mul(grad_output, grad_scale);
  }

  fusion.addOutput(grad_input);

  // Scheduling strategy
  // 1. Cache inputs
  // 2. Group the reductions (automatically fused with broadcasts)
  // 3. Merge HW and vectorize with the outer parallelized by TIDx
  // 4. Split N by TIDy with the outer parallelized by BIDx and
  // inner by TIDy
  // 5. Split C by BIDy and let the outer be the serial outermost loop

  auto input_cache = input->cacheAfter();
  auto grad_output_cache = grad_output->cacheAfter();
  weight->cacheAfter();
  save_mean->cacheAfter();
  save_invstd->cacheAfter();

  // Group the two reductions
  groupReductions({grad_output_sum, dot_p});

  // Transform grad_input to: [C/bidy, N/tidy, tidy, bidy, HW/vec_width,
  // vec_width]
  const int tidy = 8;
  const int bidy = 4;
  const int bidx = ceilDiv(shape[0], (int64_t)tidy);
  const int vec_width = 4;
  NVF_CHECK(
      (shape[2] * shape[3]) % vec_width == 0,
      "Invalid vector width: ",
      vec_width);

  grad_input->merge(-2, -1);
  grad_input->split(-1, vec_width);

  grad_input->split(0, tidy);
  grad_input->split(2, bidy);
  NVF_CHECK(
      grad_input->nDims() == 6,
      "Unexpected number of dimensions: ",
      grad_input->toString());

  grad_input->reorder({{2, 0}, {0, 1}, {1, 2}});

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

  TransformPropagator propagator(grad_input);
  MaxLogicalDomainInfoSpanningTree(grad_input).traverse(&propagator);

  auto rf_tensors = grad_output_sum->rFactor(
      {-1}, std::vector<TensorView*>({grad_output_sum, dot_p}));

  for (auto fusion_input :
       ir_utils::filterByType<TensorView>(fusion.inputs())) {
    fusion_input->computeAt(grad_input, 1);
  }

  // Parallelization
  scheduler_utils::parallelizeAllLike(grad_input);
  input_cache->axis(-1)->parallelize(ParallelType::Vectorize);
  grad_output_cache->axis(-1)->parallelize(ParallelType::Vectorize);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto at_input = at::randn(shape, options);
  auto at_grad_output = at::randn(shape, options);
  auto at_weight = at::randn({shape[c_axis]}, options);
  auto at_save_mean = at::randn({shape[c_axis]}, options);
  auto at_save_invstd = at::randn({shape[c_axis]}, options);
  KernelArgumentHolder inputs(
      {at_input, at_grad_output, at_weight, at_save_mean, at_save_invstd});

  GpuLower gpulw(&fusion);
  validateNoParallelBroadcastExist(gpulw.run());

  KernelExecutor ke;
  ke.compile(&fusion, inputs);

  if (bidx * bidy > deviceSMCount()) {
    GTEST_SKIP() << "Not enough SMs to run this test";
  }

  auto outputs = ke.run(inputs);

  std::vector<int64_t> at_reduction_axes;
  std::copy(
      reduction_axes.begin(),
      reduction_axes.end(),
      std::back_inserter(at_reduction_axes));

  // MSVC bug on lambda non-capture of const integral type
  // https://developercommunity.visualstudio.com/t/lambda-fails-to-implicitly-capture-constexpr-value/610504
  auto at_bcast = [=](const auto& tensor) {
    if (channels_last) {
      tensor.unsqueeze(0).unsqueeze(0).unsqueeze(0);
    } else {
      return tensor.unsqueeze(0).unsqueeze(-1).unsqueeze(-1);
    }
  };

  auto at_mean = at_save_mean;
  const auto& at_invstd = at_save_invstd;
  at_mean = at_bcast(at_mean);
  auto at_norm = 1.0f / static_cast<float>(shape[0] * shape[2] * shape[3]);

  auto at_grad_output_sum = sum(at_grad_output, at_reduction_axes);
  auto at_dot_p =
      sum(mul(at_grad_output, sub(at_input, at_mean)), at_reduction_axes);

  auto at_grad_mean = at_bcast(at_grad_output_sum * at_norm);

  auto at_proj_scale = at_bcast((at_dot_p * at_norm) * (at_invstd * at_invstd));

  at::Tensor at_grad_scale;

  if (weight == nullptr) {
    at_grad_scale = at_bcast(at_invstd);
  } else {
    at_grad_scale = at_bcast(at_invstd) * at_bcast(at_weight);
  }

  at::Tensor at_grad_input;
  if (kTraining) {
    auto at_proj = (at_input - at_mean) * at_proj_scale;
    at_grad_input = (at_grad_output - at_proj - at_grad_mean) * at_grad_scale;
  } else {
    at_grad_input = at_grad_output * at_grad_scale;
  }

  testValidate(
      ke.compiledKernel()->kernel(),
      outputs,
      inputs,
      {at_grad_input},
      __LINE__,
      __FILE__);
}

TEST_F(ReductionTest, GroupedReductionReEntrant1) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = add(tv0, IrBuilder::create<Val>(1.0));
  auto tv2 = sum(tv1, {0});

  auto tv3 = add(tv0, IrBuilder::create<Val>(2.0));
  auto tv4 = sum(tv3, {0});

  auto tv5 = add(tv2, tv4);
  fusion.addOutput(tv5);

  groupReductions({tv2, tv4});

  auto tv0_cache = tv0->cacheAfter();

  const int vec = 2;
  const int tidx = 64;
  const int tidy = 8;

  tv2->split(1, vec);
  tv2->split(1, tidx);

  tv2->split(0, tidy);
  TransformPropagator propagator(tv2);
  MaxLogicalDomainInfoSpanningTree(tv2).traverse(&propagator);

  tv0_cache->axis(-1)->parallelize(ParallelType::Vectorize);

  tv0->computeAt(tv4, -1, ComputeAtMode::MostInlined);

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

  scheduler_utils::parallelizeAllLike(tv2);

  std::vector<int64_t> shape({99, 999});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t0_double = t0.to(at::kDouble);
  auto ref = (t0_double + 1).sum({0}) + (t0_double + 2).sum({0});

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Channels-last batch norm with vectorization. Relies on re-entrant
// GroupedGridReduction
TEST_F(ReductionTest, GroupedReductionChannelsLastBatchNormLike) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int64_t> shape({64, 14, 14, 32});

  auto tv0 = makeContigTensor(4, DataType::Half);
  fusion.addInput(tv0);
  auto tv1 = makeContigTensor(4, DataType::Half);
  fusion.addInput(tv1);
  auto tv2 = makeContigTensor(1);
  fusion.addInput(tv2);

  std::vector<int64_t> reduction_axes({0, 1, 2});
  std::vector<bool> broadcast_mask({true, true, true, false});

  auto tv3 = castOp(DataType::Float, tv0);
  auto tv4 = castOp(DataType::Float, tv1);

  auto tv5 = sum(tv3, reduction_axes);

  auto tv6 = broadcast(tv2, broadcast_mask);
  auto tv7 = sub(tv4, tv6);
  auto tv8 = mul(tv3, tv7);
  auto tv9 = sum(tv8, reduction_axes);

  auto tv10 = castOp(DataType::Half, tv5);
  auto tv11 = castOp(DataType::Half, tv9);

  fusion.addOutput(tv10);
  fusion.addOutput(tv11);

  groupReductions({tv5, tv9});

  // Applies the outer-reduction schedule
  const int64_t num_channels = shape.back();
  const int64_t vector = 2;
  NVF_CHECK(num_channels % vector == 0);
  // Use at most 32 TIDx threads
  const int64_t tidx = std::min<int64_t>(32l, num_channels / vector);
  const auto bidx = ceilDiv(num_channels, tidx * vector);

  const int64_t tidy = 8;
  const auto bidy = ceilDiv(
      at::cuda::getCurrentDeviceProperties()->multiProcessorCount * 4, bidx);

  auto tv0_cache = tv0->cacheAfter();
  auto tv1_cache = tv1->cacheAfter();

  auto ref = tv5;

  // Move the reduction domains inner positions
  ref->reorder({{0, 1}, {1, 2}, {2, 3}, {3, 0}});

  // Parallelizing the reduction domains
  ref->merge(2, 3);
  ref->merge(1, 2);
  ref->split(1, tidy);
  ref->split(1, bidy, false);

  // Parallelizing the iteration domains
  ref->split(0, vector);
  ref->split(0, tidx);

  // Move the vector axis to the innermost position
  ref->reorder({{2, 5}, {3, 2}, {4, 3}, {5, 4}});
  // Move the serial reduction to the right of the vector axis
  ref->reorder({{3, 4}, {4, 3}});

  TransformPropagator propagator(ref);
  MaxLogicalDomainInfoSpanningTree(ref).traverse(&propagator);

  auto rf_tvs = tv5->rFactor({-2}, {tv5, tv9});
  auto tv5_rf = rf_tvs.at(0);
  auto tv9_rf = rf_tvs.at(1);

  tv0->computeAt(tv5_rf, -2, ComputeAtMode::BestEffort);
  tv1->computeAt(tv9_rf, -2, ComputeAtMode::BestEffort);
  tv3->computeAt(tv5_rf, -1, ComputeAtMode::BestEffort);
  tv4->computeAt(tv9_rf, -1, ComputeAtMode::BestEffort);

  ref = tv5_rf;

  ref->axis(0)->parallelize(ParallelType::BIDx);
  ref->axis(1)->parallelize(ParallelType::TIDx);
  ref->axis(2)->parallelize(ParallelType::BIDy);
  ref->axis(3)->parallelize(ParallelType::TIDy);
  ref->axis(4)->parallelize(ParallelType::Serial);
  ref->axis(5)->parallelize(ParallelType::Serial);

  scheduler_utils::parallelizeAllLike(ref);

  tv0_cache->axis(-1)->parallelize(ParallelType::Vectorize);
  tv1_cache->axis(-1)->parallelize(ParallelType::Vectorize);

  auto options_half = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  auto options_float =
      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape, options_half);
  auto t1 = at::randn(shape, options_half);
  auto t2 = at::randn({shape.back()}, options_float);

  KernelExecutor ke;
  ke.compile(&fusion, {t0, t1, t2});
  auto outputs = ke.run({t0, t1, t2});

  auto t0_double = t0.to(at::kDouble);
  auto t1_double = t1.to(at::kDouble);
  auto t2_double = t2.to(at::kDouble);

  std::vector<int64_t> at_reduction_axes(
      {reduction_axes.begin(), reduction_axes.end()});
  auto t5 = t0_double.sum(at_reduction_axes);
  auto t8 = t0_double *
      (t1_double - t2_double.unsqueeze(0).unsqueeze(0).unsqueeze(0));
  auto t9 = t8.sum(at_reduction_axes);

  testValidate(
      ke.compiledKernel()->kernel(),
      outputs,
      {t0, t1, t2},
      {t5, t9},
      __LINE__,
      __FILE__);
}

// Test the grouped grid allreduce with BN-like outer reductions
TEST_F(ReductionTest, GroupedReductionPersistentChannelsLastBatchNormLike) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int64_t> shape({64, 14, 14, 32});

  auto tv0 = makeContigTensor(4, DataType::Half);
  fusion.addInput(tv0);
  auto tv1 = makeContigTensor(4, DataType::Half);
  fusion.addInput(tv1);
  auto tv2 = makeContigTensor(1);
  fusion.addInput(tv2);

  std::vector<int64_t> reduction_axes({0, 1, 2});
  std::vector<bool> broadcast_mask({true, true, true, false});

  auto tv3 = castOp(DataType::Float, tv0);
  auto tv4 = castOp(DataType::Float, tv1);

  auto tv5 = sum(tv3, reduction_axes);

  auto tv6 = broadcast(tv2, broadcast_mask);
  auto tv7 = sub(tv4, tv6);
  auto tv8 = mul(tv3, tv7);
  auto tv9 = sum(tv8, reduction_axes);

  auto tv10 = broadcast(tv5, broadcast_mask);
  auto tv11 = add(tv3, tv10);

  auto tv12 = broadcast(tv9, broadcast_mask);
  auto tv13 = add(tv4, tv12);

  auto tv14 = castOp(DataType::Half, tv11);
  auto tv15 = castOp(DataType::Half, tv13);

  fusion.addOutput(tv14);
  fusion.addOutput(tv15);

  groupReductions({tv5, tv9});

  // Applies the outer-reduction schedule
  const int64_t num_channels = shape.back();
  const int64_t vector = 2;
  NVF_CHECK(num_channels % vector == 0);
  // Use at most 32 TIDx threads
  const int64_t tidx = std::min<int64_t>(32l, num_channels / vector);
  ceilDiv(num_channels, tidx * vector);

  const int64_t tidy = 8;
  const int64_t reduction_work_per_thread = 8;

  auto tv0_cache = tv0->cacheAfter();
  auto tv1_cache = tv1->cacheAfter();

  auto ref = tv5;

  // Move the reduction domains inner positions
  ref->reorder({{0, 1}, {1, 2}, {2, 3}, {3, 0}});

  // Parallelizing the reduction domains
  ref->merge(2, 3);
  ref->merge(1, 2);
  ref->split(1, tidy);
  ref->split(1, reduction_work_per_thread);

  // Parallelizing the iteration domains
  ref->split(0, vector);
  ref->split(0, tidx);

  // Move the vector axis to the innermost position
  ref->reorder({{2, 5}, {3, 2}, {4, 3}, {5, 4}});
  // Move the serial reduction to the right of the vector axis
  ref->reorder({{3, 4}, {4, 3}});

  TransformPropagator propagator(ref);
  MaxLogicalDomainInfoSpanningTree(ref).traverse(&propagator);

  auto rf_tvs = tv5->rFactor({-2}, {tv5, tv9});
  auto tv5_rf = rf_tvs.at(0);
  auto tv9_rf = rf_tvs.at(1);

  tv0->computeAt(tv5_rf, -2, ComputeAtMode::BestEffort);
  tv1->computeAt(tv9_rf, -2, ComputeAtMode::BestEffort);
  tv3->computeAt(tv5_rf, -1, ComputeAtMode::BestEffort);
  tv4->computeAt(tv9_rf, -1, ComputeAtMode::BestEffort);

  ref = tv5_rf;

  ref->axis(0)->parallelize(ParallelType::BIDx);
  ref->axis(1)->parallelize(ParallelType::TIDx);
  ref->axis(2)->parallelize(ParallelType::BIDy);
  ref->axis(3)->parallelize(ParallelType::TIDy);
  ref->axis(4)->parallelize(ParallelType::Serial);
  ref->axis(5)->parallelize(ParallelType::Serial);

  scheduler_utils::parallelizeAllLike(ref);

  tv0_cache->axis(-1)->parallelize(ParallelType::Vectorize);
  tv1_cache->axis(-1)->parallelize(ParallelType::Vectorize);

  tv5->axis(-1)->parallelize(ParallelType::Group);

  auto options_half = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  auto options_float =
      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape, options_half);
  auto t1 = at::randn(shape, options_half);
  auto t2 = at::randn({shape.back()}, options_float);

  KernelExecutor ke;
  ke.compile(&fusion, {t0, t1, t2});
  auto outputs = ke.run({t0, t1, t2});

  auto t0_double = t0.to(at::kDouble);
  auto t1_double = t1.to(at::kDouble);
  auto t2_double = t2.to(at::kDouble);

  std::vector<int64_t> at_reduction_axes(
      {reduction_axes.begin(), reduction_axes.end()});
  auto t5 = t0_double.sum(at_reduction_axes);
  auto t8 = t0_double *
      (t1_double - t2_double.unsqueeze(0).unsqueeze(0).unsqueeze(0));
  auto t9 = t8.sum(at_reduction_axes);

  auto t10 = t5.unsqueeze(0).unsqueeze(0).unsqueeze(0);
  auto t11 = t0_double + t10;
  auto t12 = t9.unsqueeze(0).unsqueeze(0).unsqueeze(0);
  auto t13 = t1_double + t12;

  testValidate(
      ke.compiledKernel()->kernel(),
      outputs,
      {t0, t1, t2},
      {t11, t13},
      __LINE__,
      __FILE__);
}

TEST_F(ReductionTest, CrossIterationGroupedGridAllreduce1) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = set(tv0);
  auto tv2 = sum(tv1, {0});
  auto tv3 = broadcast(tv2, {true, false});
  auto tv4 = add(tv0, tv3);
  fusion.addOutput(tv4);

  const int vec = 2;
  const int tidx = 32;
  const int tidy = 8;

  tv1->split(1, vec);
  tv1->split(1, tidx);
  tv1->split(0, tidy);
  TransformPropagator propagator(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator);

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

  scheduler_utils::parallelizeAllLike(tv1);

  tv2->axis(4)->parallelize(ParallelType::Group);

  // Make sure the reduction expr is converted to GroupedGridReduciton
  // and the non-reduction domains of the output TV are either
  // grouped or parallelized
  GpuLower gpulw(&fusion);
  bool validated = false;
  for (auto expr : KernelExprVisitor::getAllExprs(gpulw.run())) {
    auto grouped_grid_reduction =
        dynamic_cast<kir::GroupedGridReduction*>(expr);
    if (grouped_grid_reduction == nullptr) {
      continue;
    }
    auto out = ir_utils::getTvOutput(grouped_grid_reduction);
    for (auto out_axis : out->getLoopDomain()) {
      auto out_axis_pt = out_axis->getParallelType();
      NVF_CHECK(
          isParallelTypeThread(out_axis_pt) ||
              out_axis_pt == ParallelType::Group,
          "Invalid parallel type of the reduction tensor: ",
          out_axis_pt,
          ". Reduction output tensor: ",
          out->toString());
    }
    validated = true;
  }
  NVF_CHECK(
      validated, "Invalid lowered kernel. No GroupedGridReduction found.");

  std::vector<int64_t> shape({99, 101});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t0_double = t0.to(at::kDouble);
  auto ref = t0_double + t0_double.sum({0}).unsqueeze(0);

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Test grouping of two domains
TEST_F(ReductionTest, CrossIterationGroupedGridAllreduce2) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = set(tv0);
  auto tv2 = sum(tv1, {0});
  auto tv3 = broadcast(tv2, {true, false});
  auto tv4 = add(tv0, tv3);
  fusion.addOutput(tv4);

  const int vec1 = 2;
  const int vec2 = 3;
  const int tidx = 16;
  const int tidy = 8;

  tv1->split(1, vec1);
  tv1->split(1, vec2);
  tv1->split(1, tidx);
  tv1->split(0, tidy);
  TransformPropagator propagator(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator);

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

  scheduler_utils::parallelizeAllLike(tv1);

  tv2->axis(4)->parallelize(ParallelType::Group);
  tv2->axis(5)->parallelize(ParallelType::Group);

  std::vector<int64_t> shape({99, 129});

  // Make sure the reduction expr is converted to GroupedGridReduciton
  // and the non-reduction domains of the output TV are either
  // grouped or parallelized
  GpuLower gpulw(&fusion);
  bool validated = false;
  for (auto expr : KernelExprVisitor::getAllExprs(gpulw.run())) {
    auto grouped_grid_reduction =
        dynamic_cast<kir::GroupedGridReduction*>(expr);
    if (grouped_grid_reduction == nullptr) {
      continue;
    }
    auto out = ir_utils::getTvOutput(grouped_grid_reduction);
    for (auto out_axis : out->getLoopDomain()) {
      auto out_axis_pt = out_axis->getParallelType();
      NVF_CHECK(
          isParallelTypeThread(out_axis_pt) ||
              out_axis_pt == ParallelType::Group,
          "Invalid parallel type of the reduction tensor: ",
          out_axis_pt,
          ". Reduction output tensor: ",
          out->toString());
    }
    validated = true;
  }
  NVF_CHECK(
      validated, "Invalid lowered kernel. No GroupedGridReduction found.");

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t0_double = t0.to(at::kDouble);
  auto ref = t0_double + t0_double.sum({0}).unsqueeze(0);

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Group both expressions and iterations
TEST_F(ReductionTest, CrossIterationGroupedGridAllreduce3) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = add(tv0, IrBuilder::create<Val>(1.0));
  auto tv2 = sum(tv1, {0});
  auto tv3 = broadcast(tv2, {true, false});
  auto tv4 = add(tv1, tv3);

  auto tv5 = add(tv0, IrBuilder::create<Val>(2.0));
  auto tv6 = sum(tv5, {0});
  auto tv7 = broadcast(tv6, {true, false});
  auto tv8 = add(tv5, tv7);

  auto tv9 = add(tv4, tv8);
  fusion.addOutput(tv9);

  groupReductions({tv2, tv6});

  const int vec = 2;
  const int tidx = 32;
  const int tidy = 8;

  tv1->split(1, vec);
  tv1->split(1, tidx);
  tv1->split(0, tidy);
  TransformPropagator propagator(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator);

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

  scheduler_utils::parallelizeAllLike(tv1);

  tv2->axis(4)->parallelize(ParallelType::Group);

  // Make sure the reduction expr is converted to GroupedGridReduciton
  // and the non-reduction domains of the output TV are either
  // grouped or parallelized
  GpuLower gpulw(&fusion);
  bool validated = false;
  for (auto expr : KernelExprVisitor::getAllExprs(gpulw.run())) {
    auto grouped_grid_reduction =
        dynamic_cast<kir::GroupedGridReduction*>(expr);
    if (grouped_grid_reduction == nullptr) {
      continue;
    }
    auto out = ir_utils::getTvOutput(grouped_grid_reduction);
    for (auto out_axis : out->getLoopDomain()) {
      auto out_axis_pt = out_axis->getParallelType();
      NVF_CHECK(
          isParallelTypeThread(out_axis_pt) ||
              out_axis_pt == ParallelType::Group,
          "Invalid parallel type of the reduction tensor: ",
          out_axis_pt,
          ". Reduction output tensor: ",
          out->toString());
    }
    validated = true;
  }
  NVF_CHECK(
      validated, "Invalid lowered kernel. No GroupedGridReduction found.");

  std::vector<int64_t> shape({99, 101});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t0_double = t0.to(at::kDouble);
  auto t4 = t0_double + 1 + (t0_double + 1).sum({0}).unsqueeze(0);
  auto t8 = t0_double + 2 + (t0_double + 2).sum({0}).unsqueeze(0);
  auto ref = t4 + t8;

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// ParallelType::Group with computeAt
TEST_F(ReductionTest, CrossIterationGroupedGridAllreduce4) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = set(tv0);
  auto tv2 = sum(tv1, {0});
  auto tv3 = broadcast(tv2, {true, false});
  auto tv4 = add(tv0, tv3);
  fusion.addOutput(tv4);

  const int vec = 2;
  const int tidx = 32;
  const int tidy = 8;

  tv2->reorder({{0, 1}});
  tv2->split(0, vec);
  tv2->split(0, tidx);
  tv2->split(-1, tidy);

  TransformPropagator propagator(tv2);
  MaxLogicalDomainInfoSpanningTree(tv2).traverse(&propagator);

  tv2->axis(2)->parallelize(ParallelType::Group);

  // This should avoid inlining the grouped domain
  tv0->computeAt(tv4, -1, ComputeAtMode::MostInlined);

  NVF_CHECK(
      tv1->getComputeAtPosition() == 2,
      "Invalid computeAt position: ",
      tv1->toString());
  NVF_CHECK(
      tv2->getComputeAtPosition() == 2,
      "Invalid computeAt position: ",
      tv2->toString());

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

  for (auto tv : fusion.allTvs()) {
    tv->axis(-2)->parallelize(ParallelType::BIDy);
    tv->axis(-1)->parallelize(ParallelType::TIDy);
  }

  // Make sure the reduction expr is converted to GroupedGridReduciton
  // and the non-reduction domains of the output TV are either
  // grouped or parallelized
  GpuLower gpulw(&fusion);
  bool validated = false;
  for (auto expr : KernelExprVisitor::getAllExprs(gpulw.run())) {
    auto grouped_grid_reduction =
        dynamic_cast<kir::GroupedGridReduction*>(expr);
    if (grouped_grid_reduction == nullptr) {
      continue;
    }
    auto out = ir_utils::getTvOutput(grouped_grid_reduction);
    for (auto out_axis : out->getLoopDomain()) {
      auto out_axis_pt = out_axis->getParallelType();
      NVF_CHECK(
          isParallelTypeThread(out_axis_pt) ||
              out_axis_pt == ParallelType::Group,
          "Invalid parallel type of the reduction tensor: ",
          out_axis_pt,
          ". Reduction output tensor: ",
          out->toString());
    }
    validated = true;
  }
  NVF_CHECK(
      validated, "Invalid lowered kernel. No GroupedGridReduction found.");

  std::vector<int64_t> shape({99, 101});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t0_double = t0.to(at::kDouble);
  auto ref = t0_double + t0_double.sum({0}).unsqueeze(0);

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

TEST_F(ReductionTest, CrossIterationGroupedGridAllreduceWelford1) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = set(tv0);
  auto tv2 = Welford(tv1, {0}).avg;
  auto tv3 = broadcast(tv2, {true, false});
  auto tv4 = add(tv0, tv3);
  fusion.addOutput(tv4);

  const int vec = 2;
  const int tidx = 32;
  const int tidy = 8;

  tv1->split(1, vec);
  tv1->split(1, tidx);
  tv1->split(0, tidy);
  TransformPropagator propagator(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator);

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

  scheduler_utils::parallelizeAllLike(tv1);

  tv2->axis(4)->parallelize(ParallelType::Group);

  // Make sure the reduction expr is converted to GroupedGridReduciton
  // and the non-reduction domains of the output TV are either
  // grouped or parallelized
  GpuLower gpulw(&fusion);
  bool validated = false;
  for (auto expr : KernelExprVisitor::getAllExprs(gpulw.run())) {
    auto grouped_grid_reduction = dynamic_cast<kir::GroupedGridWelford*>(expr);
    if (grouped_grid_reduction == nullptr) {
      continue;
    }
    validated = true;
  }
  NVF_CHECK(validated, "Invalid lowered kernel. No GroupedGridWelford found.");

  std::vector<int64_t> shape({99, 101});

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t0_double = t0.to(at::kDouble);
  auto ref = t0_double + t0_double.mean({0}).unsqueeze(0);

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Test grouping of two domains
TEST_F(ReductionTest, CrossIterationGroupedGridAllreduceWelford2) {
  Fusion fusion;
  FusionGuard fg(&fusion);

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

  auto tv1 = set(tv0);
  auto tv2 = Welford(tv1, {0}).avg;
  auto tv3 = broadcast(tv2, {true, false});
  auto tv4 = add(tv0, tv3);
  fusion.addOutput(tv4);

  const int vec1 = 2;
  const int vec2 = 3;
  const int tidx = 16;
  const int tidy = 8;

  tv1->split(1, vec1);
  tv1->split(1, vec2);
  tv1->split(1, tidx);
  tv1->split(0, tidy);
  TransformPropagator propagator(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator);

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

  scheduler_utils::parallelizeAllLike(tv1);

  tv2->axis(4)->parallelize(ParallelType::Group);
  tv2->axis(5)->parallelize(ParallelType::Group);

  std::vector<int64_t> shape({99, 129});

  // Make sure the reduction expr is converted to GroupedGridReduciton
  // and the non-reduction domains of the output TV are either
  // grouped or parallelized
  GpuLower gpulw(&fusion);
  bool validated = false;
  for (auto expr : KernelExprVisitor::getAllExprs(gpulw.run())) {
    auto grouped_grid_reduction = dynamic_cast<kir::GroupedGridWelford*>(expr);
    if (grouped_grid_reduction == nullptr) {
      continue;
    }
    validated = true;
  }
  NVF_CHECK(validated, "Invalid lowered kernel. No GroupedGridWelford found.");

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape, options);

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto outputs = ke.run({t0});

  auto t0_double = t0.to(at::kDouble);
  auto ref = t0_double + t0_double.mean({0}).unsqueeze(0);

  testValidate(
      ke.compiledKernel()->kernel(), outputs, {t0}, {ref}, __LINE__, __FILE__);
}

// Follows the pattern of persistent outer grid welford in batchnorm
TEST_F(ReductionTest, CrossIterationGroupedGridAllreduceWelfordShmoo) {
  struct Params {
    int N;
    int H;
    int W;
    int C;
    int tidx;
    int tidy;
    int vect;
    int persistent_buffer;
    int bidx;
  };

  auto test = [](const Params& params) {
    Fusion fusion;
    FusionGuard fg(&fusion);

    std::vector<bool> bcast_pattern{true, true, true, false};
    std::vector<int64_t> reduction_dims{2, 1, 0};

    auto tv0 = makeContigTensor(4);
    fusion.addInput(tv0);

    auto tv1 = set(tv0);
    auto tvs = Welford(tv1, reduction_dims);
    auto tv2 = tvs.avg;
    auto tv3 = tvs.var_sum;
    auto tv4 = tvs.n;
    auto tv5 = broadcast(tv2, bcast_pattern);
    auto tv6 = broadcast(tv3, bcast_pattern);
    broadcast(tv4, bcast_pattern);
    auto tv8 = sub(tv1, tv5);
    auto tv9 = add(tv8, tv6);
    fusion.addOutput(tv9);

    // Schedule the fusion as it will be done by the persistent
    // scheduler

    tv9->cacheBefore();

    auto transform_ref = tv2;

    transform_ref->merge(0)->merge(0);

    int reduction_pos = 1;

    transform_ref->split(0, params.tidy);
    ++reduction_pos;
    transform_ref->axis(1)->parallelize(ParallelType::TIDy);

    // Persistent buffer
    transform_ref->split(0, params.persistent_buffer);
    ++reduction_pos;

    // Unswitch
    transform_ref->split(0, 1);
    ++reduction_pos;
    transform_ref->axis(1)->parallelize(ParallelType::Unswitch);

    transform_ref->axis(0)->parallelize(ParallelType::BIDy);

    transform_ref->split(reduction_pos, params.vect);
    transform_ref->axis(reduction_pos + 1)
        ->parallelize(ParallelType::Vectorize);

    transform_ref->split(reduction_pos, params.tidx);
    transform_ref->axis(reduction_pos + 1)->parallelize(ParallelType::TIDx);
    transform_ref->split(reduction_pos, params.bidx);
    transform_ref->axis(reduction_pos + 1)->parallelize(ParallelType::BIDx);

    auto transform_ref_rf =
        reduction_scheduler_utils::sortAndRFactor(transform_ref);

    TransformPropagator propagator(transform_ref_rf);
    MaxLogicalDomainInfoSpanningTree(transform_ref_rf).traverse(&propagator);

    int vec_id = std::distance(
        transform_ref_rf->getLoopDomain().begin(),
        std::find_if(
            transform_ref_rf->getLoopDomain().begin(),
            transform_ref_rf->getLoopDomain().end(),
            [](auto id) {
              return id->getParallelType() == ParallelType::Vectorize;
            }));
    transform_ref_rf->axis(vec_id)->parallelize(ParallelType::Serial);

    int unswitch_id = std::distance(
        transform_ref_rf->getLoopDomain().begin(),
        std::find_if(
            transform_ref_rf->getLoopDomain().begin(),
            transform_ref_rf->getLoopDomain().end(),
            [](auto id) {
              return id->getParallelType() == ParallelType::Unswitch;
            }));
    transform_ref_rf->axis(unswitch_id)->parallelize(ParallelType::Serial);

    scheduler_utils::parallelizeAllLike(transform_ref_rf, fusion.allTvs());

    ParallelType vec_pt = ParallelType::Vectorize;
    tv1->axis(vec_id)->parallelize(vec_pt);
    tv9->axis(vec_id)->parallelize(vec_pt);

    transform_ref->axis(vec_id)->parallelize(ParallelType::Group);

    transform_ref_rf->axis(unswitch_id)->parallelize(ParallelType::Unswitch);

    inlineMost();

    // Make sure the reduction expr is converted to GroupedGridReduciton
    // and the non-reduction domains of the output TV are either
    // grouped or parallelized
    GpuLower gpulw(&fusion);
    bool validated = false;
    for (auto expr : KernelExprVisitor::getAllExprs(gpulw.run())) {
      (void)expr; // Suppress unused variable warning
      validated = true;
    }
    NVF_CHECK(
        validated, "Invalid lowered kernel. No GroupedGridWelford found.");

    auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

    const std::vector<int64_t> input_shape{
        params.N, params.H, params.W, params.C};
    auto t0 = at::randn(input_shape, options);

    KernelExecutor ke;
    ke.compile(&fusion, {t0});

    // Skip the rest of this test size if the required number of SMs
    // exceeds the available SM count
    const auto num_required_sms = params.bidx *
        ceilDiv(ceilDiv(params.N * params.H * params.W, params.tidy),
                params.persistent_buffer);
    if (num_required_sms > deviceSMCount()) {
      return;
    }

    auto cg_outputs = ke.run({t0});

    auto t1 = t0.to(at::kDouble);
    auto t2 = t1.mean({0, 1, 2}).unsqueeze(0).unsqueeze(0).unsqueeze(0);
    auto t3 =
        at::var(t1, {0, 1, 2}, false).unsqueeze(0).unsqueeze(0).unsqueeze(0);
    auto t4 = params.N * params.H * params.W;
    auto ref = (t1 - t2 + (t3 * t4));

    testValidate(&fusion, cg_outputs, {t0}, {ref}, __LINE__, __FILE__, "");
  };

  std::vector<Params> base_params;
  base_params.push_back({256, 7, 7, 1, 8, 32, 2, 32, 4});
  base_params.push_back({256, 7, 7, 1, 16, 16, 4, 50, 4});
  base_params.push_back({128, 7, 7, 1, 16, 16, 4, 32, 4});
  base_params.push_back({128, 14, 14, 1, 16, 16, 4, 32, 1});
  base_params.push_back({128, 14, 14, 1, 16, 16, 2, 64, 2});
  base_params.push_back({128, 14, 14, 1, 8, 32, 4, 50, 4});
  base_params.push_back({128, 14, 14, 1, 8, 32, 2, 50, 4});

  std::vector<Params> param_vec;
  for (const auto base_p : base_params) {
    for (const auto c_dim : {512, 1024, 2048}) {
      auto tmp = base_p;
      tmp.C = c_dim;
      param_vec.push_back(tmp);
    }
  }

  for (const auto& params : param_vec) {
    // Shmoo tests can occupy a lot of memory due to allocating many
    // different tensor sizes. So in order to avoid an OOM during this
    // test, we manually clear the allocator after it's reached a certain
    // threshold.
    maybeClearAllocator();
    test(params);
  }
}

TEST_F(ReductionTest, GeluBwdReduction) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  std::vector<int64_t> tensor_shape{4, 128};

  // input
  TensorView* tv0_grad = makeSymbolicTensor(tensor_shape.size());
  fusion.addInput(tv0_grad);
  TensorView* tv1_xvar = makeSymbolicTensor(tensor_shape.size());
  fusion.addInput(tv1_xvar);

  // inter
  const float k_079 = 0.79788456;
  const float k_004 = 0.044715;
  const float k_010 = 0.1070322243;
  auto t8 = mul(tv1_xvar, IrBuilder::create<Val>(k_079));
  auto t9 = mul(tv1_xvar, IrBuilder::create<Val>(k_004));
  auto t10 = mul(t9, tv1_xvar);
  auto t11 = add(t10, IrBuilder::create<Val>(1L));
  auto t12 = mul(t8, t11);
  auto t13 = unaryOp(UnaryOpType::Tanh, t12);
  auto t14 = mul(tv1_xvar, IrBuilder::create<Val>(0.5));
  auto t15 = mul(t13, t13);
  auto t16 = unaryOp(UnaryOpType::Neg, t15);
  auto t17 = add(t16, IrBuilder::create<Val>(1L));
  auto t18 = mul(tv1_xvar, IrBuilder::create<Val>(k_010));
  auto t19 = mul(t18, tv1_xvar);
  auto t20 = add(t19, IrBuilder::create<Val>(k_079));
  auto t21 = mul(t17, t20);
  auto t22 = mul(t14, t21);
  auto t23 = add(t13, IrBuilder::create<Val>(1L));
  auto t24 = mul(t23, IrBuilder::create<Val>(0.5));
  auto t25 = add(t22, t24);
  auto t26 = mul(t25, tv0_grad);
  auto t27 = sum(t26, {0});

  // output
  fusion.addOutput(t26);
  fusion.addOutput(t27);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  // reference values
  at::Tensor at_grad = at::randn(tensor_shape, options);
  at::Tensor at_xvar = at::randn(tensor_shape, options);
  auto at_x = at_xvar.to(c10::ScalarType::Float);

  auto at_output_pointwise = at::gelu_backward(at_grad, at_x, "tanh");
  auto at_output_reduction = at_output_pointwise.sum({0});

  auto cg_results =
      scheduleAndRun(&fusion, SchedulerType::Reduction, {at_grad, at_xvar});
  testValidate(
      &fusion,
      cg_results.outputs,
      {at_grad, at_xvar},
      {at_output_pointwise, at_output_reduction},
      __LINE__,
      __FILE__,
      "",
      cg_results.heuristic_params->lparams);
}

// Test gathering for lookup as is done in the cross_entropy pattern
// See https://github.com/NVIDIA/Fuser/issues/151
TEST_F(ReductionTest, CrossEntropyGatherPattern) {
  const int batch_size = 16384;
  const int num_classes = 60;
  const int tidx = 128;
  const int bidx = 1;

  Fusion fusion;
  FusionGuard fg(&fusion);

  auto log_probs = makeSymbolicTensor(2);
  auto labels = makeSymbolicTensor(1, DataType::Int);
  fusion.addInput(log_probs);
  fusion.addInput(labels);

  auto tv2 = broadcast(labels, {false, true});
  auto tv3 = takeAlongAxis(log_probs, tv2, 1);
  auto tv4 = squeeze(tv3, std::vector<bool>({false, true}));

  fusion.addOutput(tv4);

  tv4->split(0, tidx);
  tv4->split(0, bidx);
  tv4->split(0, 1); // unswitch
  TransformPropagator propagator(tv4);
  MaxLogicalDomainInfoSpanningTree(tv4).traverse(&propagator);

  tv4->axis(0)->parallelize(ParallelType::BIDy);
  tv4->axis(2)->parallelize(ParallelType::BIDx);
  tv4->axis(3)->parallelize(ParallelType::TIDx);
  scheduler_utils::parallelizeAllLike(tv4);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto at_log_probs = at::randn({batch_size, num_classes}, options);
  auto at_labels =
      at::randint(0, num_classes, {batch_size}, options.dtype(at::kLong));

  KernelExecutor ke;
  ke.compile(&fusion, {at_log_probs, at_labels});
  auto cg_outputs = ke.run({at_log_probs, at_labels});

  auto ref = at::gather(at_log_probs, 1, at_labels.unsqueeze(1)).squeeze();

  testValidate(
      &fusion,
      cg_outputs,
      {at_log_probs, at_labels},
      {ref},
      __LINE__,
      __FILE__);
}

TEST_F(ReductionTest, TensorRankLimit) {
  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  std::vector<int64_t> input_shape;
  for (__attribute__((unused)) auto i : arange(12)) {
    input_shape.push_back(3);
  }

  auto tv0 = makeSymbolicTensor(input_shape.size());
  fusion->addInput(tv0);
  auto tv1 = sum(tv0, {3});
  fusion->addOutput(tv1);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn(input_shape, options);

  FusionExecutorCache executor_cache(std::move(fusion));
  auto cg_outputs = executor_cache.runFusionWithInputs({t0});

  testValidate(executor_cache.fusion(), cg_outputs, {t0}, __LINE__, __FILE__);
}

// Test scheduling inner reduction kernels with 1D bulk TMA.
// ndims: Test 2 or 3 dimensions. The first dim is always non-reduce, all the
//        rest are reduce dimensions. Inner dimensions should be merged by the
//        schedule.
// inner_size: The size of the inner-most dimension.
using TmaInnerReductionManualTestParams =
    std::tuple<int64_t, int64_t>; // <ndims, inner_size>

class TmaInnerReductionManualTest
    : public NVFuserFixtureParamTest<TmaInnerReductionManualTestParams> {
 protected:
  void SetUp() override {
    NVFuserFixtureParamTest<TmaInnerReductionManualTestParams>::SetUp();
    NVFUSER_TEST_CUDA_ARCH_GUARD(9, 0);
    enable_options_guard_ = std::make_unique<EnableOptionsGuard>();
    EnableOptionsGuard::getCurOptions().set(EnableOption::TmaReduction);
  }

 private:
  std::unique_ptr<EnableOptionsGuard> enable_options_guard_;
};
TEST_P(TmaInnerReductionManualTest, Basic) {
  [[maybe_unused]] auto dtype = DataType::Float;
  [[maybe_unused]] int64_t dtype_bytes = dataTypeSizeByte(dtype);
  [[maybe_unused]] auto [ndims, inner_size] = GetParam();

  std::vector<int64_t> element_at_each_dim;
  if (ndims == 2) {
    element_at_each_dim = {64, inner_size};
  } else {
    element_at_each_dim = {64, 8, inner_size};
  }

  std::vector<int64_t> reduced_dims;
  int64_t reduced_elem_count = 1;
  for (int i = 1; i < ndims; ++i) {
    reduced_dims.push_back(i);
    reduced_elem_count *= element_at_each_dim[i];
  }

  if (reduced_elem_count * dtype_bytes % 16 != 0) {
    GTEST_SKIP() << "Reduced bytes is not divisible by 16, can't use TMA, "
                    "reduced_elem_count: "
                 << reduced_elem_count << ", dtype_bytes: " << dtype_bytes;
    return;
  }

  const int64_t smem_elems =
      at::cuda::getDeviceProperties(0)->sharedMemPerBlockOptin / dtype_bytes;

  auto fusion_ptr = std::make_unique<Fusion>();
  FusionGuard fg(fusion_ptr.get());
  Fusion& fusion = *fusion_ptr;

  auto tv0 = makeContigTensor(ndims);
  fusion.addInput(tv0);
  auto tv1 = sum(tv0, reduced_dims);
  fusion.addOutput(tv1);
  auto fusion_copy = fusion;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn(element_at_each_dim, options);

  // Phase-1, Set input and output cache and TMA load
  auto tv0smem = tv0->cacheAfter(LoadStoreOpType::CpAsyncBulk);
  tv0smem->setMemoryType(MemoryType::Shared);

  auto redu_tv = tv1->cacheBefore();

  // Phase-2, Merge contiguous reduction dimensions
  // [I, R1, R2] -> [I, R1*R2]
  for (int64_t i = redu_tv->nDims() - 2; i >= 0; --i) {
    if (redu_tv->axis(i)->isReduction() &&
        redu_tv->axis(i + 1)->isReduction()) {
      redu_tv->merge(i, i + 1);
    }
  }

  // Phase-3, Schedule common transformations shared by TMA and non-TMA tensors
  ParallelType outer_type = ParallelType::BIDx;
  const int64_t smem_split_offset = (reduced_elem_count > smem_elems) ? 1 : 0;
  if (smem_split_offset > 0) {
    // Split the reduction axis to fit into shared memory.
    int64_t split_size = std::gcd(smem_elems, reduced_elem_count);

    // [I, R] -> [I, R/split_size, split_size]
    redu_tv->split(1, split_size);
    redu_tv->axis(1)->parallelize(ParallelType::BIDx);

    // Reduction block extent may be larger than 65535, so it needs BIDx.
    outer_type = ParallelType::BIDy;
  }

  redu_tv->axis(0)->parallelize(outer_type);

  // Phase-4, Propagate common transformations to all tensors
  TransformPropagator propagator(redu_tv);
  MaxLogicalDomainInfoSpanningTree(redu_tv).traverse(&propagator);
  scheduler_utils::parallelizeAllLike(redu_tv);

  // Phase-5, Schedule for reduction and rFactor thread local reductions
  const int64_t vect = 4;
  const int64_t tidx = 256;
  const int64_t unroll = 4;

  const int64_t vect_axis = 5 + smem_split_offset;
  const int64_t tid_axis = 4 + smem_split_offset;
  const int64_t unroll_axis = 3 + smem_split_offset;
  const int64_t unswitch_axis = 2 + smem_split_offset;
  const int64_t tma_axis = 1 + smem_split_offset;

  tv0smem->axis(tma_axis)->parallelize(ParallelType::Bulk);

  // [I, R] -> [I, R/vect, vect]
  redu_tv->split(tma_axis, vect);

  // [I, R/vect, vect] -> [I, R/vect/tidx, tidx, vect]
  redu_tv->split(tma_axis, tidx);

  // [I, R/vect/tidx, tidx, vect] -> [I, R/vect/tidx/unroll, unroll, tidx, vect]
  redu_tv->split(tma_axis, unroll);
  redu_tv->split(tma_axis, 1);

  redu_tv->axis(vect_axis)->parallelize(ParallelType::Serial);
  redu_tv->axis(tid_axis)->parallelize(ParallelType::TIDx);
  redu_tv->axis(unroll_axis)->parallelize(ParallelType::Unroll);
  redu_tv->axis(unswitch_axis)->parallelize(ParallelType::Unswitch);
  redu_tv->axis(tma_axis)->parallelize(ParallelType::Serial);

  std::vector<int64_t> rfactor_axes;
  for (int64_t i = 0; i < redu_tv->nDims(); i++) {
    if (redu_tv->axis(i)->isReduction() && !redu_tv->axis(i)->isThread()) {
      rfactor_axes.push_back(i);
    }
  }
  redu_tv->rFactor(rfactor_axes);

  // Step 6: Inline
  inlineMost();

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto cg_outputs = ke.run({t0});
  testValidate(&fusion_copy, cg_outputs, {t0}, __LINE__, __FILE__);
}

INSTANTIATE_TEST_SUITE_P(
    ,
    TmaInnerReductionManualTest,
    testing::Combine(
        testing::Values(2, 3), // ndims
        testing::ValuesIn([] { // inner_size
          std::vector<int64_t> vals(
              Pow2Vals1to1Million.begin(), Pow2Vals1to1Million.end());
          // Add some irregular numbers
          vals.insert(vals.end(), {1024 * 1024 + 8, 1024 * 1024 + 7, 1023});
          return vals;
        }())),
    ([](const testing::TestParamInfo<TmaInnerReductionManualTestParams>& info) {
      auto [ndims, inner_size] = info.param;
      return "ndim_" + std::to_string(ndims) + "_inner_size_" +
          std::to_string(inner_size);
    }));

namespace tma_reduction_check {
bool isTmaParams(const FusionExecutorCache& executor_cache) {
  FusionKernelRuntime* runtime = executor_cache.getMostRecentKernelRuntime();
  const auto& hparams = runtime->schedulerHeuristics()->heuristicsList().at(0);
  return hparams->isA<TmaInnerReductionParams>() ||
      hparams->isA<TmaOuterReductionParams>();
}
bool isOuterTmaParams(const FusionExecutorCache& executor_cache) {
  FusionKernelRuntime* runtime = executor_cache.getMostRecentKernelRuntime();
  const auto& hparams = runtime->schedulerHeuristics()->heuristicsList().at(0);
  return hparams->isA<TmaOuterReductionParams>();
}
} // namespace tma_reduction_check

// <dtype, reduction_size>
using TmaInnerReductionTestParams = std::tuple<DataType, int64_t>;

class TmaInnerReductionTest
    : public NVFuserFixtureParamTest<TmaInnerReductionTestParams> {
 protected:
  void SetUp() override {
    NVFUSER_TEST_CUDA_ARCH_GUARD(9, 0);
    NVFuserFixtureParamTest<TmaInnerReductionTestParams>::SetUp();
    enable_options_guard_ = std::make_unique<EnableOptionsGuard>();
    EnableOptionsGuard::getCurOptions().set(EnableOption::TmaReduction);
  }

  // Check if we expect TMA to be used based on mayUseTma() conditions
  bool expectTmaUsed(DataType dtype, int64_t reduction_size) {
    // TMA requires 16-byte alignment (vectorize_factor > 1)
    int64_t dtype_size_bit = dataTypeSizeBit(dtype);
    int64_t dtype_bytes = dtype_size_bit / 8;
    int64_t min_elems_for_alignment = 16 / dtype_bytes;
    if (reduction_size % min_elems_for_alignment != 0) {
      return false;
    }

    uint64_t total_reduction_bytes = reduction_size * dtype_bytes;

    // Minimum TMA transfer size, below which it seems much slower than non-TMA.
    uint64_t min_tma_bytes = 16384;

    if (total_reduction_bytes < min_tma_bytes) {
      return false;
    }

    // Reduction dim must fit into smem
    auto dev_prop = at::cuda::getCurrentDeviceProperties();
    int64_t smem_elems =
        (dev_prop->sharedMemPerBlockOptin * 8) / dtype_size_bit;
    if (reduction_size > smem_elems) {
      return false;
    }

    return true;
  }

 private:
  std::unique_ptr<EnableOptionsGuard> enable_options_guard_;
};

// Test basic sum reduction with auto-scheduled TMA
TEST_P(TmaInnerReductionTest, Sum) {
  auto dtype = std::get<0>(GetParam());
  auto reduction_size = std::get<1>(GetParam());
  constexpr int64_t iter_size = 1024;

  auto fusion_ptr = std::make_unique<Fusion>();
  FusionGuard fg(fusion_ptr.get());
  Fusion& fusion = *fusion_ptr;

  auto tv0 = makeContigTensor(2, dtype);
  fusion.addInput(tv0);

  auto tv0_cast = maybeCastOp(DataType::Float, tv0);
  auto tv1 = sum(tv0_cast, {1});
  fusion.addOutput(tv1);

  auto unscheduled_fusion_copy = fusion;
  auto options =
      at::TensorOptions().dtype(data_type_to_aten(dtype)).device(at::kCUDA, 0);
  auto t0 = at::randn({iter_size, reduction_size}, options);

  FusionExecutorCache executor_cache(std::move(fusion_ptr));
  auto outputs = executor_cache.runFusionWithInputs({t0});

  if (expectTmaUsed(dtype, reduction_size)) {
    EXPECT_TRUE(tma_reduction_check::isTmaParams(executor_cache))
        << "Expected TMA scheduler for reduction_size=" << reduction_size;
  }

  testValidate(&unscheduled_fusion_copy, outputs, {t0}, __LINE__, __FILE__);
}

INSTANTIATE_TEST_SUITE_P(
    ,
    TmaInnerReductionTest,
    testing::Combine(
        testing::Values(DataType::Float, DataType::BFloat16),
        testing::ValuesIn([] {
          std::vector<int64_t> vals(
              Pow2Vals1to1Million.begin(), Pow2Vals1to1Million.end());
          // Add some irregular numbers
          vals.insert(vals.end(), {1024 * 1024 + 8, 1024 * 1024 + 7, 1023});
          return vals;
        }())),
    ([](const testing::TestParamInfo<TmaInnerReductionTestParams>& info) {
      auto [dtype, reduction_size] = info.param;
      std::ostringstream os;
      os << dtype << "_" << reduction_size;
      return os.str();
    }));

// Test 2D TMA for outer reductions
using TmaOuterReductionManualTestParams =
    std::tuple<int64_t, int64_t>; // <outer_size, iter_size>

class TmaOuterReductionManualTest
    : public NVFuserFixtureParamTest<TmaOuterReductionManualTestParams> {
 protected:
  void SetUp() override {
    NVFuserFixtureParamTest<TmaOuterReductionManualTestParams>::SetUp();
    NVFUSER_TEST_CUDA_ARCH_GUARD(9, 0);
    enable_options_guard_ = std::make_unique<EnableOptionsGuard>();
    EnableOptionsGuard::getCurOptions().set(EnableOption::TmaReduction);
  }

 private:
  std::unique_ptr<EnableOptionsGuard> enable_options_guard_;
};
TEST_P(TmaOuterReductionManualTest, Basic) {
  auto dtype = DataType::Float;
  int64_t dtype_bytes = dataTypeSizeByte(dtype);
  auto [outer_size, iter_size] = GetParam();

  // TMA requires the stride (iter_size * dtype_bytes) to be 16-byte aligned
  if ((iter_size * dtype_bytes) % 16 != 0) {
    GTEST_SKIP() << "Iter dimension bytes not divisible by 16, can't use TMA";
    return;
  }

  std::vector<int64_t> shape = {outer_size, iter_size};

  auto fusion_ptr = std::make_unique<Fusion>();
  FusionGuard fg(fusion_ptr.get());
  Fusion& fusion = *fusion_ptr;

  // Input: [R, I] where R is reduction dim (axis 0), I is iteration dim (axis
  // 1)
  auto tv0 = makeContigTensor(2);
  fusion.addInput(tv0);
  auto tv1 = sum(tv0, {0}); // reduce along axis 0 (outer reduction)
  fusion.addOutput(tv1);
  auto fusion_copy = fusion;

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn(shape, options);

  const int64_t bdimx = 32;
  const int64_t bdimy = 16;
  const int64_t iter_unroll_factor = 4;
  const int64_t redu_unroll_factor = 8;

  int64_t grdim = std::max<int64_t>(
      1, std::min<int64_t>(8, scheduler_utils::lastPow2(outer_size / 256)));

  // TMA tile dimensions = thread dims * unroll factors
  const int64_t tma_tile_i = bdimx * iter_unroll_factor;
  const int64_t tma_tile_r = bdimy * redu_unroll_factor;

  // Skip if outer_size is too small for meaningful 2D TMA test
  if (outer_size < tma_tile_r) {
    GTEST_SKIP() << "outer_size " << outer_size
                 << " is smaller than tma_tile_r " << tma_tile_r;
    return;
  }

  // Skip if iter_size is too small for meaningful 2D TMA test
  if (iter_size < tma_tile_i) {
    GTEST_SKIP() << "iter_size " << iter_size << " is smaller than tma_tile_i "
                 << tma_tile_i;
    return;
  }

  // ========== Phase 1: Create TMA cache in shared memory ==========
  auto tv0smem = tv0->cacheAfter(LoadStoreOpType::CpAsyncBulkTensorTile);
  tv0smem->setMemoryType(MemoryType::Shared);

  // Cache before the output for the grid reduction
  auto redu_tv = tv1->cacheBefore();

  // ========== Phase 2: Schedule TMA tensor with TMA-level tiling ==========
  // [R, I] -> [R/tma_tile_r, tma_tile_r, I]
  tv0smem->split(0, tma_tile_r);

  // -> [R/tma_tile_r, tma_tile_r, I/tma_tile_i, tma_tile_i]
  tv0smem->split(2, tma_tile_i);

  // Split outer reduction for grid parallelization
  // -> [grdim, R', tma_tile_r, I/tma_tile_i, tma_tile_i]
  //       0      1    2           3              4
  tv0smem->split(0, grdim, false);

  // ========== Phase 3: Propagate TMA tiling to all tensors ==========
  TransformPropagator propagator(tv0smem);
  MaxLogicalDomainInfoSpanningTree(tv0smem).traverse(&propagator);

  // ========== Phase 4: Parallelize and finalize TMA tensor ==========
  tv0smem->axis(0)->parallelize(ParallelType::BIDy);
  tv0smem->axis(1)->parallelize(ParallelType::Serial);
  tv0smem->axis(2)->parallelize(ParallelType::Bulk); // reduction tile
  tv0smem->axis(3)->parallelize(ParallelType::BIDx);
  tv0smem->axis(4)->parallelize(ParallelType::Bulk); // iteration tile

  // ========== Phase 5: Sub-split TMA tiles into thread dims ==========
  // Split tma_tile_i into [bdimx, iter_unroll]
  redu_tv->split(4, iter_unroll_factor);

  // Split tma_tile_r into [redu_unroll, bdimy]
  redu_tv->split(2, bdimy);
  // Now: [grdim, R', redu_unroll, bdimy, I/tma_tile_i, bdimx, iter_unroll]
  //       0      1   2            3      4              5      6

  // ========== Phase 6: Parallelize reduction tensor ==========
  redu_tv->axis(0)->parallelize(ParallelType::BIDy);
  redu_tv->axis(1)->parallelize(ParallelType::Serial);
  redu_tv->axis(2)->parallelize(ParallelType::Unroll); // redu_unroll
  redu_tv->axis(3)->parallelize(ParallelType::TIDy); // bdimy
  redu_tv->axis(4)->parallelize(ParallelType::BIDx);
  redu_tv->axis(5)->parallelize(ParallelType::TIDx); // bdimx
  // Use Vectorize so it gets converted to Group for iterGroupedGridReduce
  redu_tv->axis(6)->parallelize(ParallelType::Vectorize); // iter_unroll

  // ========== Phase 7: rFactor for grid reduction ==========
  // The reduction axes that are not thread-parallelized need rFactor
  std::vector<int64_t> rfactor_axes;
  for (int64_t i = 0; i < redu_tv->nDims(); i++) {
    if (redu_tv->axis(i)->isReduction() && !redu_tv->axis(i)->isThread()) {
      rfactor_axes.push_back(i);
    }
  }

  TensorView* ref_tv = redu_tv;
  if (!rfactor_axes.empty()) {
    ref_tv = redu_tv->rFactor(rfactor_axes);
  }

  // ========== Phase 8: Propagate thread-level splits to non-TMA TVs ==========
  std::vector<TensorView*> non_tma_tvs =
      ir_utils::allTvsExcept(&fusion, {tv0smem});
  TransformPropagator non_tma_propagator(ref_tv);
  SetSelector non_tma_selector({non_tma_tvs.begin(), non_tma_tvs.end()});
  MaxLogicalDomainInfoSpanningTree(ref_tv, &non_tma_selector)
      .traverse(&non_tma_propagator);

  // ========== Phase 9: Parallelize with iter-grouped reduction ==========
  // For outer reduction, use iterGroupedGridReduce which is more efficient
  // and has fewer bank conflicts than separate gridReduce calls
  const bool use_iter_grouped_reduction = true; // outer reduction + cross_block
  std::vector<TensorView*> reduction_tvs = {redu_tv};

  if (ref_tv != redu_tv) {
    reduction_scheduler_utils::propagateRFactor(ref_tv, redu_tv, reduction_tvs);
    non_tma_tvs = ir_utils::allTvsExcept(&fusion, {tv0smem});
  }

  // Use propagateParallelization to set up grouped reduction
  reduction_scheduler_utils::propagateParallelization(
      redu_tv,
      ref_tv,
      /*is_unroll_or_vectorization=*/true,
      use_iter_grouped_reduction,
      reduction_tvs,
      /*unroll_vectorizable_cached_tvs=*/{tv1},
      /*selected_tvs=*/non_tma_tvs);

  // ========== Phase 10: Inline ==========
  inlineMost();

  KernelExecutor ke;
  ke.compile(&fusion, {t0});
  auto cg_outputs = ke.run({t0});

  testValidate(&fusion_copy, cg_outputs, {t0}, __LINE__, __FILE__, "");
}

INSTANTIATE_TEST_SUITE_P(
    ,
    TmaOuterReductionManualTest,
    testing::Combine(
        testing::ValuesIn([] { // outer_size
          std::vector<int64_t> vals;
          for (int64_t v = 256; v <= 65536; v *= 4) {
            vals.push_back(v);
          }
          return vals;
        }()),
        testing::ValuesIn([] { // iter_size
          std::vector<int64_t> vals;
          for (int64_t v = 256; v <= 65536; v *= 4) {
            vals.push_back(v);
          }
          return vals;
        }())),
    ([](const testing::TestParamInfo<TmaOuterReductionManualTestParams>& info) {
      auto [outer_size, iter_size] = info.param;
      return "outer_" + std::to_string(outer_size) + "_iter_" +
          std::to_string(iter_size);
    }));

// Test outer reduction with auto-scheduled TMA
using TmaOuterReductionTestParams =
    std::tuple<int64_t, int64_t, int64_t, DataType>;
// <outer_size, iter_size, n_inputs, dtype>

class TmaOuterReductionTest
    : public NVFuserFixtureParamTest<TmaOuterReductionTestParams> {
 protected:
  void SetUp() override {
    NVFUSER_TEST_CUDA_ARCH_GUARD(9, 0);
    NVFuserFixtureParamTest<TmaOuterReductionTestParams>::SetUp();
    enable_options_guard_ = std::make_unique<EnableOptionsGuard>();
    EnableOptionsGuard::getCurOptions().set(EnableOption::TmaReduction);
  }

  // Check if we expect outer TMA to be used based on mayUseTmaOuter conditions
  bool expectOuterTmaUsed(
      int64_t outer_size,
      int64_t iter_size,
      int64_t n_inputs,
      int64_t dtype_bytes) {
    uint64_t total_reduction_bytes = outer_size * dtype_bytes;
    uint64_t min_tma_bytes = 16384;
    if (total_reduction_bytes < min_tma_bytes) {
      return false;
    }
    // The stride (iter_size * dtype_bytes) must be 16-byte aligned
    if ((iter_size * dtype_bytes) % 16 != 0) {
      return false;
    }
    // Check minimum tile smem fits for all inputs, accounting for
    // block reduction workspace + static smem overhead.
    int64_t min_smem_per_input = 16 * 32 * dtype_bytes;
    int64_t smem_overhead = ((512 * dtype_bytes + 127) & ~127) + 128;
    auto* dev_prop = at::cuda::getCurrentDeviceProperties();
    if (min_smem_per_input * n_inputs + smem_overhead >
        (int64_t)dev_prop->sharedMemPerBlockOptin) {
      return false;
    }
    return true;
  }

 private:
  std::unique_ptr<EnableOptionsGuard> enable_options_guard_;
};

TEST_P(TmaOuterReductionTest, Sum) {
  auto [outer_size, iter_size, n_inputs, dtype] = GetParam();
  int64_t dtype_bytes = dataTypeSizeByte(dtype);

  // TMA requires the stride (iter_size * dtype_bytes) to be 16-byte aligned
  if ((iter_size * dtype_bytes) % 16 != 0) {
    GTEST_SKIP() << "Iter dimension bytes not divisible by 16, can't use TMA";
    return;
  }

  auto fusion_ptr = std::make_unique<Fusion>();
  FusionGuard fg(fusion_ptr.get());
  Fusion& fusion = *fusion_ptr;

  // Create n_inputs tensors and sum them element-wise before reducing
  std::vector<TensorView*> inputs;
  for (int64_t i = 0; i < n_inputs; i++) {
    auto tv = makeContigTensor(2, dtype);
    fusion.addInput(tv);
    inputs.push_back(tv);
  }
  auto accum = inputs[0];
  for (int64_t i = 1; i < n_inputs; i++) {
    accum = add(accum, inputs[i]);
  }
  auto reduced = sum(accum, {0});
  fusion.addOutput(reduced);

  auto unscheduled_fusion_copy = fusion;

  auto options =
      at::TensorOptions().dtype(data_type_to_aten(dtype)).device(at::kCUDA, 0);
  std::vector<c10::IValue> aten_inputs;
  for (int64_t i = 0; i < n_inputs; i++) {
    aten_inputs.push_back(at::randn({outer_size, iter_size}, options));
  }

  FusionExecutorCache executor_cache(std::move(fusion_ptr));
  auto outputs = executor_cache.runFusionWithInputs(aten_inputs);

  if (expectOuterTmaUsed(outer_size, iter_size, n_inputs, dtype_bytes)) {
    EXPECT_TRUE(tma_reduction_check::isOuterTmaParams(executor_cache))
        << "Expected outer TMA scheduler for outer_size=" << outer_size
        << " iter_size=" << iter_size << " n_inputs=" << n_inputs;
  }

  testValidate(
      &unscheduled_fusion_copy, outputs, aten_inputs, __LINE__, __FILE__, "");
}

INSTANTIATE_TEST_SUITE_P(
    ,
    TmaOuterReductionTest,
    testing::Values(
        // Size sweep with single input, float
        TmaOuterReductionTestParams{256, 256, 1, DataType::Float},
        TmaOuterReductionTestParams{256, 4096, 1, DataType::Float},
        TmaOuterReductionTestParams{256, 65536, 1, DataType::Float},
        TmaOuterReductionTestParams{4096, 256, 1, DataType::Float},
        TmaOuterReductionTestParams{4096, 4096, 1, DataType::Float},
        TmaOuterReductionTestParams{4096, 65536, 1, DataType::Float},
        TmaOuterReductionTestParams{65536, 256, 1, DataType::Float},
        TmaOuterReductionTestParams{65536, 4096, 1, DataType::Float},
        TmaOuterReductionTestParams{65536, 65536, 1, DataType::Float},
        // Multi-input (exercises smem budget / tile shrinking)
        TmaOuterReductionTestParams{4096, 1024, 2, DataType::Float},
        TmaOuterReductionTestParams{4096, 1024, 3, DataType::Float},
        TmaOuterReductionTestParams{4096, 1024, 5, DataType::Float},
        TmaOuterReductionTestParams{4096, 1024, 2, DataType::Half},
        TmaOuterReductionTestParams{4096, 1024, 3, DataType::Half},
        TmaOuterReductionTestParams{4096, 1024, 5, DataType::Half}),
    ([](const testing::TestParamInfo<TmaOuterReductionTestParams>& info) {
      auto [outer_size, iter_size, n_inputs, dtype] = info.param;
      std::string name = "outer_" + std::to_string(outer_size) + "_iter_" +
          std::to_string(iter_size);
      if (n_inputs > 1) {
        name += "_" + std::to_string(n_inputs) + "inputs";
      }
      if (dtype != DataType::Float) {
        name += "_half";
      }
      return name;
    }));

} // namespace nvfuser