#include <torch/csrc/jit/codegen/cuda/executor.h>
#include <torch/csrc/jit/codegen/cuda/fusion.h>
#include <torch/csrc/jit/codegen/cuda/ir_all_nodes.h>
#include <torch/csrc/jit/codegen/cuda/ir_builder.h>

#include <torch/csrc/jit/codegen/cuda/scheduler/all_schedulers.h>

#include <benchmark/benchmark.h>

#include <benchmarks/cpp/nvfuser/utils.h>

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

static void setup_vit_base_patch16_224_bcast7(Fusion* fusion, void* null) {
  FusionGuard fg(fusion);

  auto t2 = makeContigTensor(3, DataType::Float);
  auto t3 = TensorViewBuilder()
                .shape({-1, -1, 1})
                .dtype(DataType::Float)
                .contiguity({true, true, false})
                .build();
  auto t4 = TensorViewBuilder()
                .shape({-1, -1, 1})
                .dtype(DataType::Float)
                .contiguity({true, true, false})
                .build();
  auto t7 = makeContigTensor(3, DataType::Half);

  fusion->addInput(t2);
  fusion->addInput(t3);
  fusion->addInput(t4);
  fusion->addInput(t7);

  auto t8 = castOp(DataType::Float, t7);
  auto t9 = set(t8);
  auto t10 = sub(t2, t3);
  auto t11 = mul(t10, t4);
  auto t25 = mul(t9, t11);
  auto t26 = sum(t25, {0, 1});
  auto t36 = set(t26);
  auto t27 = sum(t9, {0, 1});
  auto t37 = set(t27);
  auto t39 = castOp(DataType::Half, t11);

  fusion->addOutput(t36);
  fusion->addOutput(t37);
  fusion->addOutput(t39);
}

static void NvFuserScheduler_TIMM_vit_base_patch16_224_bcast7(
    benchmark::State& benchmark_state,
    FusionExecutorCache* fusion_executor_cache,
    void* null) {
  std::vector<int64_t> input_shape{
      benchmark_state.range(0),
      benchmark_state.range(1),
      benchmark_state.range(2)};

  at::manual_seed(0);
  auto fp16_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  auto fp32_options =
      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t2 = at::randn(input_shape, fp32_options);
  auto t3 = at::randn({input_shape[0], input_shape[1], 1}, fp32_options);
  auto t4 = at::randn({input_shape[0], input_shape[1], 1}, fp32_options);
  auto t7 = at::randn(input_shape, fp16_options);

  std::vector<c10::IValue> aten_inputs({t2, t3, t4, t7});
  runBenchmarkIterations(benchmark_state, fusion_executor_cache, aten_inputs);

  // full tensor - float + halfx2 - t2, t7, t39
  // Inner most dimension only - floatx2 - t36, t37
  // Outer two dimensions only - floatx2 - t3, t4
  benchmark_state.SetBytesProcessed(
      int64_t(benchmark_state.iterations()) *
      //           t2 + t7    t3 + t4                 t36 + t37
      t2.numel() * (4 + 2) + t3.numel() * 4 * 2 + input_shape[2] * (4 * 2) +
      // T39
      t2.numel() * 2);
}

NVFUSER_BENCHMARK_DEFINE(
    NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_bcast7,
    setup_vit_base_patch16_224_bcast7,
    NvFuserScheduler_TIMM_vit_base_patch16_224_bcast7,
    nullptr);

// pwise case, broadcasting both sides
NVFUSER_BENCHMARK_RUN(NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_bcast7)
    ->Args({64, 197, 768})
    ->Unit(benchmark::kMicrosecond)
    ->UseManualTime();

static void setup_vit_base_patch16_224_bcast5(Fusion* fusion, void* null) {
  FusionGuard fg(fusion);

  auto t2 = makeContigTensor(3, DataType::Float);
  auto t5 = makeContigTensor(1, DataType::Float);
  auto t3 = makeContigTensor(3, DataType::Half);
  auto t0 = makeContigTensor(1, DataType::Float);
  auto t1 = makeContigTensor(1, DataType::Float);

  fusion->addInput(t2);
  fusion->addInput(t5);
  fusion->addInput(t3);
  fusion->addInput(t0);
  fusion->addInput(t1);

  std::vector<bool> bcast_pattern0({true, true, false});
  std::vector<bool> bcast_pattern1({false, false, true});

  auto t4 = castOp(DataType::Float, t3);
  auto t6 = set(t5);
  auto t7 = broadcast(t6, bcast_pattern0);
  auto t8 = add(t4, t7);
  auto t9 = rand_like(t8);
  auto d34 =
      sub(IrBuilder::create<Double>(1.0), IrBuilder::create<Double>(0.0));
  auto t10 = lt(t9, d34);
  auto t11 = castOp(DataType::Float, t10);
  auto t12 = mul(t8, t11);
  auto b36 = eq(d34, IrBuilder::create<Double>(0.0));
  auto d37 = castOp(DataType::Double, b36);
  auto d38 = add(d37, d34);
  auto d40 = div(IrBuilder::create<Double>(1.0), d38);
  auto t13 = mul(t12, d40);
  auto t14 = set(t13);
  auto t15 = add(t2, t14);
  auto t16 = set(t15);
  auto t36 = sum(t16, {2});
  auto d151 = castOp(DataType::Double, t2->axis(2)->extent());
  auto d152 = mul(IrBuilder::create<Double>(1.0), d151);
  auto t19 = div(t36, d152);
  auto t22 = broadcast(t19, bcast_pattern1);
  auto t23 = sub(t16, t22);
  auto t37 = mul(t23, t23);
  auto t20 = sum(t37, {2});
  auto t24 = broadcast(t20, bcast_pattern1);
  auto d95 = castOp(DataType::Double, t2->axis(2)->extent());
  auto d105 = reciprocal(d95);
  auto t25 = mul(t24, d105);
  auto t26 = add(t25, IrBuilder::create<Double>(1e-6));
  auto t27 = rsqrt(t26);
  auto t28 = mul(t23, t27);
  auto t17 = set(t1);
  auto t29 = broadcast(t17, bcast_pattern0);
  auto t30 = mul(t28, t29);
  auto t18 = set(t0);
  auto t31 = broadcast(t18, bcast_pattern0);
  auto t32 = add(t30, t31);
  auto t33 = set(t32);
  auto t34 = castOp(DataType::Half, t33);

  fusion->addOutput(t16); // full 3d float
  fusion->addOutput(t10); // full 3d bool
  fusion->addOutput(t22); // bcast last dim float
  fusion->addOutput(t27); // bcast last dim float
  fusion->addOutput(t18); // passthrough t0 float
  fusion->addOutput(t17); // passthrough t1 float
  fusion->addOutput(t34); // full 3d half
}

static void NvFuserScheduler_TIMM_vit_base_patch16_224_bcast5(
    benchmark::State& benchmark_state,
    FusionExecutorCache* fusion_executor_cache,
    void* null) {
  std::vector<int64_t> input_shape{
      benchmark_state.range(0),
      benchmark_state.range(1),
      benchmark_state.range(2)};

  at::manual_seed(0);
  auto fp16_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  auto fp32_options =
      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t2 = at::randn(input_shape, fp32_options);
  auto t5 = at::randn({input_shape[2]}, fp32_options);
  auto t3 = at::randn(input_shape, fp16_options);
  auto t0 = at::randn({input_shape[2]}, fp32_options);
  auto t1 = at::randn({input_shape[2]}, fp32_options);

  std::vector<c10::IValue> aten_inputs({t2, t5, t3, t0, t1});
  runBenchmarkIterations(benchmark_state, fusion_executor_cache, aten_inputs);

  // Full tensor - floatx2, halfx2, bool - t2, t16, t3, t34, t16
  // Inner most dim only - floatx5 - t5, t0, t1, t7, t17
  // Outer two dims only - floatx2 - t22, t27

  benchmark_state.SetBytesProcessed(
      int64_t(benchmark_state.iterations()) *
      t2.numel() * (2 * 4 + 2 * 2 + 1) + t5.numel() * 5 * 4 +
      input_shape[0] * input_shape[1] * 2 * 4);
}

NVFUSER_BENCHMARK_DEFINE(
    NvFuserScheduler_TIMM_vit_base_patch16_224_bcast5_NCHW,
    setup_vit_base_patch16_224_bcast5,
    NvFuserScheduler_TIMM_vit_base_patch16_224_bcast5,
    nullptr);

// Broadcast on both sides
NVFUSER_BENCHMARK_RUN(NvFuserScheduler_TIMM_vit_base_patch16_224_bcast5_NCHW)
    ->Args({64, 197, 768})
    ->Unit(benchmark::kMicrosecond)
    ->UseManualTime();

static void setup_vit_base_patch16_224_bcast_outer2(
    Fusion* fusion,
    void* null) {
  FusionGuard fg(fusion);

  auto t0 = makeContigTensor(3, DataType::Half);
  auto t2 = makeContigTensor(1, DataType::Float);

  fusion->addInput(t0);
  fusion->addInput(t2);

  auto t1 = castOp(DataType::Float, t0);
  auto t3 = set(t2);
  auto t4 = broadcast(t3, {true, true, false});
  auto t5 = add(t1, t4);
  auto t6 = castOp(DataType::Half, t5);
  auto t7 = castOp(DataType::Half, t3);

  fusion->addOutput(t6);
  fusion->addOutput(t7);
}

static void NvFuserScheduler_TIMM_vit_base_patch16_224_bcast_outer2(
    benchmark::State& benchmark_state,
    FusionExecutorCache* fusion_executor_cache,
    void* null) {
  std::vector<int64_t> input_shape{
      benchmark_state.range(0),
      benchmark_state.range(1),
      benchmark_state.range(2)};

  at::manual_seed(0);
  auto fp16_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  auto fp32_options =
      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(input_shape, fp16_options);
  auto t2 = at::randn({input_shape[2]}, fp32_options);

  std::vector<c10::IValue> aten_inputs({t0, t2});
  runBenchmarkIterations(benchmark_state, fusion_executor_cache, aten_inputs);

  // full tensor - halfx2 - t0, t6
  // inner dimension only - halfx2 - t2, t7
  benchmark_state.SetBytesProcessed(
      int64_t(benchmark_state.iterations()) * t0.numel() * (2 + 2) +
      input_shape[2] * (2 + 4));
}

NVFUSER_BENCHMARK_DEFINE(
    NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_bcast_outer2,
    setup_vit_base_patch16_224_bcast_outer2,
    NvFuserScheduler_TIMM_vit_base_patch16_224_bcast_outer2,
    nullptr);

NVFUSER_BENCHMARK_RUN(
    NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_bcast_outer2)
    ->Args({64, 197, 2304})
    ->Unit(benchmark::kMicrosecond)
    ->UseManualTime();

static void setup_vit_base_patch16_224_norm_inner3(Fusion* fusion, void* null) {
  FusionGuard fg(fusion);

  auto t0 = makeContigTensor(4, DataType::Half);
  fusion->addInput(t0);
  auto d13 = IrBuilder::create<Double>();
  fusion->addInput(d13);

  auto t1 = castOp(DataType::Float, t0);
  auto t2 = set(t1);
  auto t3 = mul(t2, d13);
  auto t4 = set(t3);
  auto t5 = max(t4, {3});
  auto t6 = broadcast(t5, {false, false, false, true});
  auto t7 = sub(t4, t6);
  auto t8 = exp(t7);
  auto t9 = sum(t8, {3});
  auto t10 = broadcast(t9, {false, false, false, true});
  auto t11 = reciprocal(t10);
  auto t12 = mul(t8, t11);
  auto t13 = rand_like(t12);
  auto d79 = sub(IrBuilder::create<Double>(1), IrBuilder::create<Double>(0));
  auto t14 = lt(t13, d79);
  auto t15 = castOp(DataType::Float, t14);
  auto b81 = eq(d79, IrBuilder::create<Double>(0));
  auto d82 = castOp(DataType::Double, b81);
  auto d83 = add(d82, d79);
  auto d85 = div(IrBuilder::create<Double>(1), d83);
  auto t16 = mul(t12, t15);
  auto t17 = mul(t16, d85);
  auto t18 = set(t17);
  auto t19 = castOp(DataType::Half, t18);

  fusion->addOutput(t19);
  fusion->addOutput(t14);
  fusion->addOutput(t12);
  fusion->addOutput(t4);
}

static void NvFuserScheduler_TIMM_vit_base_patch16_224_norm_inner3(
    benchmark::State& benchmark_state,
    FusionExecutorCache* fusion_executor_cache,
    void* null) {
  std::vector<int64_t> input_shape{
      benchmark_state.range(0),
      benchmark_state.range(1),
      benchmark_state.range(2),
      benchmark_state.range(2)};

  at::manual_seed(0);
  auto fp16_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);

  auto t0 = at::randn(input_shape, fp16_options);

  std::vector<c10::IValue> aten_inputs({t0, 0.125});
  runBenchmarkIterations(benchmark_state, fusion_executor_cache, aten_inputs);

  // Full tensors - floatx2, half x2, bool - t12, t4, t0, t19, t14
  benchmark_state.SetBytesProcessed(
      int64_t(benchmark_state.iterations()) * t0.numel() * 13);
}

NVFUSER_BENCHMARK_DEFINE(
    NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_norm_inner3,
    setup_vit_base_patch16_224_norm_inner3,
    NvFuserScheduler_TIMM_vit_base_patch16_224_norm_inner3,
    nullptr);

// Norm inner dim
NVFUSER_BENCHMARK_RUN(
    NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_norm_inner3)
    ->Args({64, 12, 197})
    ->Unit(benchmark::kMicrosecond)
    ->UseManualTime();

static void setup_vit_base_patch16_224_bcast_outer6(
    Fusion* fusion,
    void* null) {
  FusionGuard fg(fusion);

  auto t0 = makeContigTensor(3, DataType::Half);
  auto t2 = makeContigTensor(1, DataType::Float);

  fusion->addInput(t0);
  fusion->addInput(t2);

  auto t1 = castOp(DataType::Float, t0);
  auto t3 = set(t2);
  auto t4 = broadcast(t3, {true, true, false});
  auto t5 = add(t1, t4);
  auto t6 = set(t5);
  auto t7 = mul(t6, IrBuilder::create<Double>(0.707106));
  auto t8 = erf(t7);
  auto t9 = add(IrBuilder::create<Double>(1), t8);
  auto t10 = mul(IrBuilder::create<Double>(0.5), t9);
  auto t11 = mul(t6, t10);
  auto t12 = rand_like(t11);
  auto d66 = sub(IrBuilder::create<Double>(1), IrBuilder::create<Double>(0));
  auto t13 = lt(t12, d66);
  auto t14 = castOp(DataType::Float, t13);
  auto t15 = mul(t11, t14);
  auto b68 = eq(d66, IrBuilder::create<Double>(0));
  auto d69 = castOp(DataType::Double, b68);
  auto d70 = add(d69, d66);
  auto d72 = div(IrBuilder::create<Double>(1), d70);
  auto t16 = mul(t15, d72);
  auto t17 = set(t16);
  auto t18 = castOp(DataType::Half, t17);
  auto t19 = castOp(DataType::Half, t3);

  fusion->addOutput(t18);
  fusion->addOutput(t13);
  fusion->addOutput(t6);
  fusion->addOutput(t19);
}

static void NvFuserScheduler_TIMM_vit_base_patch16_224_bcast_outer6(
    benchmark::State& benchmark_state,
    FusionExecutorCache* fusion_executor_cache,
    void* null) {
  std::vector<int64_t> input_shape{
      benchmark_state.range(0),
      benchmark_state.range(1),
      benchmark_state.range(2)};

  at::manual_seed(0);
  auto fp16_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  auto fp32_options =
      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(input_shape, fp16_options);
  auto t2 = at::randn({input_shape[2]}, fp32_options);

  std::vector<c10::IValue> aten_inputs({t0, t2});
  runBenchmarkIterations(benchmark_state, fusion_executor_cache, aten_inputs);
  // full tensors - float, halfx2, bool - t6, t0, t18, t13
  // inner dimension only - float, half - t2, t19
  benchmark_state.SetBytesProcessed(
      int64_t(benchmark_state.iterations()) * t0.numel() * (2 + 2 + 1 + 4) +
      input_shape[2] * (4 + 2));
}

NVFUSER_BENCHMARK_DEFINE(
    NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_bcast_outer6,
    setup_vit_base_patch16_224_bcast_outer6,
    NvFuserScheduler_TIMM_vit_base_patch16_224_bcast_outer6,
    nullptr);

NVFUSER_BENCHMARK_RUN(
    NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_bcast_outer6)
    // First size is original, the rest are variations to check perf
    // reliability.
    ->Args({64, 197, 3 * 1024})
    ->Args({64, 197, 2 * 1024})
    ->Args({64, 197, 1024})
    ->Args({64, 197, 512})
    ->Args({3, 1024, 64 * 197})
    ->Args({2, 1024, 64 * 197})
    ->Args({1, 1024, 64 * 197})
    ->Args({2, 256, 64 * 197})
    ->Unit(benchmark::kMicrosecond)
    ->UseManualTime();

// Reverse the broadcast dimensions to check for consistency in scheduling.
static void setup_vit_base_patch16_224_bcast_inner6(
    Fusion* fusion,
    void* null) {
  FusionGuard fg(fusion);

  auto t0 = makeContigTensor(3, DataType::Half);
  auto t2 = makeContigTensor(2, DataType::Float);

  fusion->addInput(t0);
  fusion->addInput(t2);

  auto t1 = castOp(DataType::Float, t0);
  auto t3 = set(t2);
  auto t4 = broadcast(t3, {false, false, true});
  auto t5 = add(t1, t4);
  auto t6 = set(t5);
  auto t7 = mul(t6, IrBuilder::create<Double>(0.707106));
  auto t8 = erf(t7);
  auto t9 = add(IrBuilder::create<Double>(1), t8);
  auto t10 = mul(IrBuilder::create<Double>(0.5), t9);
  auto t11 = mul(t6, t10);
  auto t12 = rand_like(t11);
  auto d66 = sub(IrBuilder::create<Double>(1), IrBuilder::create<Double>(0));
  auto t13 = lt(t12, d66);
  auto t14 = castOp(DataType::Float, t13);
  auto t15 = mul(t11, t14);
  auto b68 = eq(d66, IrBuilder::create<Double>(0));
  auto d69 = castOp(DataType::Double, b68);
  auto d70 = add(d69, d66);
  auto d72 = div(IrBuilder::create<Double>(1), d70);
  auto t16 = mul(t15, d72);
  auto t17 = set(t16);
  auto t18 = castOp(DataType::Half, t17);
  auto t19 = castOp(DataType::Half, t3);

  fusion->addOutput(t18);
  fusion->addOutput(t13);
  fusion->addOutput(t6);
  fusion->addOutput(t19);
}

static void NvFuserScheduler_TIMM_vit_base_patch16_224_bcast_inner6(
    benchmark::State& benchmark_state,
    FusionExecutorCache* fusion_executor_cache,
    void* null) {
  std::vector<int64_t> input_shape{
      benchmark_state.range(0),
      benchmark_state.range(1),
      benchmark_state.range(2)};

  at::manual_seed(0);
  auto fp16_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  auto fp32_options =
      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(input_shape, fp16_options);
  auto t2 = at::randn({input_shape[0], input_shape[1]}, fp32_options);

  std::vector<c10::IValue> aten_inputs({t0, t2});
  runBenchmarkIterations(benchmark_state, fusion_executor_cache, aten_inputs);

  // full tensors - float, halfx2, bool - t6, t0, t18, t13
  // outer two dimensions only - float, half - t2, t19
  benchmark_state.SetBytesProcessed(
      int64_t(benchmark_state.iterations()) * t0.numel() * (2 + 2 + 1 + 4) +
      input_shape[0] * input_shape[1] * (4 + 2));
}

NVFUSER_BENCHMARK_DEFINE(
    NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_bcast_inner6,
    setup_vit_base_patch16_224_bcast_inner6,
    NvFuserScheduler_TIMM_vit_base_patch16_224_bcast_inner6,
    nullptr);

NVFUSER_BENCHMARK_RUN(
    NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_bcast_inner6)
    ->Args({64, 197, 3 * 1024})
    ->Args({64, 197, 2 * 1024})
    ->Args({64, 197, 1024})
    ->Args({64, 197, 512})
    ->Args({3, 1024, 64 * 197})
    ->Args({2, 1024, 64 * 197})
    ->Args({1, 1024, 64 * 197})
    ->Args({2, 256, 64 * 197})
    ->Unit(benchmark::kMicrosecond)
    ->UseManualTime();

static void setup_vit_base_patch16_224_LN_BWD(Fusion* fusion, void* null) {
  FusionGuard fg(fusion);

  auto t0 = makeContigTensor(3, DataType::Bool);
  fusion->addInput(t0);

  auto t1 = makeContigTensor(3, DataType::Half);
  fusion->addInput(t1);

  auto t2 = castOp(DataType::Float, t1);

  auto t3 = makeContigTensor(3, DataType::Half);
  fusion->addInput(t3);

  auto t4 = castOp(DataType::Float, t3);

  auto d35 = t3->axis(2)->extent();

  auto t5 = TensorViewBuilder()
                .shape({-1, -1, 1})
                .dtype(DataType::Float)
                .contiguity({true, true, false})
                .build();
  fusion->addInput(t5);

  auto t6 = TensorViewBuilder()
                .shape({-1, -1, 1})
                .dtype(DataType::Float)
                .contiguity({true, true, false})
                .build();
  fusion->addInput(t6);

  auto t7 = makeContigTensor(1, DataType::Half);
  fusion->addInput(t7);

  auto t8 = castOp(DataType::Float, t7);

  auto t9 = makeContigTensor(1, DataType::Half);
  fusion->addInput(t9);

  auto t11 = sub(t4, t5);
  auto t12 = mul(t11, t6);

  auto t13 = broadcast(t8, {true, true, false});
  auto t14 = mul(t2, t13);
  auto t15 = mul(d35, t14);
  auto t16 = sum(t14, {2});
  auto t17 = broadcast(t16, {false, false, true});
  auto t18 = mul(t14, t12);
  auto t19 = sum(t18, {2});
  auto t20 = broadcast(t19, {false, false, true});

  auto t40 = castOp(DataType::Half, t12);
  auto t41 = castOp(DataType::Float, t40);
  auto t42 = castOp(DataType::Half, t20);
  auto t43 = castOp(DataType::Float, t42);
  auto t21 = mul(t42, t43);

  auto t38 = castOp(DataType::Half, t15);
  auto t39 = castOp(DataType::Float, t38);
  auto t44 = castOp(DataType::Half, t17);
  auto t45 = castOp(DataType::Float, t44);
  auto t22 = sub(t39, t45);

  auto t23 = sub(t22, t21);

  auto d87 = reciprocal(d35);
  auto t24 = mul(d87, t6);

  auto t25 = mul(t24, t23);
  auto t26 = mul(t2, t41);
  auto t27 = sum(t26, {0, 1});
  auto t28 = sum(t2, {0, 1});

  auto t29 = castOp(DataType::Float, t0);
  auto t30 = mul(t25, t29);

  auto d33 = IrBuilder::create<Double>();
  fusion->addInput(d33);
  auto t31 = mul(t30, d33);
  auto t32 = sum(t31, {0, 1});
  auto t33 = castOp(DataType::Half, t32);
  auto t34 = castOp(DataType::Half, t31);
  auto t35 = castOp(DataType::Half, t25);
  auto t36 = castOp(DataType::Half, t27);
  auto t37 = castOp(DataType::Half, t28);

  fusion->addOutput(t33);
  fusion->addOutput(t34);
  fusion->addOutput(t35);
  fusion->addOutput(t36);
  fusion->addOutput(t37);
}

static void NvFuserScheduler_TIMM_vit_base_patch16_224_LN_BWD(
    benchmark::State& benchmark_state,
    FusionExecutorCache* fusion_executor_cache,
    void* null) {
  std::vector<int64_t> input_shape{
      benchmark_state.range(0),
      benchmark_state.range(1),
      benchmark_state.range(2)};

  at::manual_seed(0);
  // auto bool_options = at::TensorOptions().dtype(at::kBool).device(at::kCUDA,
  // 0);
  auto fp16_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  auto fp32_options =
      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn(input_shape, fp16_options).to(at::kBool);
  auto t1 = at::randn(input_shape, fp16_options);
  auto t3 = at::randn(input_shape, fp16_options);
  auto t5 = at::randn({input_shape[0], input_shape[1], 1}, fp32_options);
  auto t6 = at::randn({input_shape[0], input_shape[1], 1}, fp32_options);
  auto t7 = at::randn({input_shape[2]}, fp16_options);
  auto t9 = at::randn({input_shape[2]}, fp16_options);

  std::vector<c10::IValue> aten_inputs({t0, t1, t3, t5, t6, t7, t9, 1.0});
  runBenchmarkIterations(benchmark_state, fusion_executor_cache, aten_inputs);

  // Full tensors - bool, halfx4 - t0, t1, t3, t34, t35
  // Outer two dimensions - floatx2 - t5, t6
  // Inner dimension - halfx5 - t7, t9, t33, t36, t37
  benchmark_state.SetBytesProcessed(
      int64_t(benchmark_state.iterations()) * ((t0.numel() * (4 * 2 + 1))) +
      (t5.numel() * 4 * 2) + (t7.numel() * 5 * 2));
}

NVFUSER_BENCHMARK_DEFINE(
    NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_LN_BWD,
    setup_vit_base_patch16_224_LN_BWD,
    NvFuserScheduler_TIMM_vit_base_patch16_224_LN_BWD,
    nullptr);

NVFUSER_BENCHMARK_RUN(NvFuserScheduler_TIMM_NCHW_vit_base_patch16_224_LN_BWD)
    ->Args({128, 197, 768})
    ->Unit(benchmark::kMicrosecond)
    ->UseManualTime();

static void nhwc_seresnet152d_transpose65(Fusion* fusion, void* null) {
  FusionGuard fg(fusion);

  auto t2 = makeContigTensor(4, DataType::Half);
  auto t5 = makeContigTensor(4, DataType::Half);
  auto t7 = makeContigTensor(4, DataType::Half);
  auto t9 = makeContigTensor(4, DataType::Half);
  auto t4 = makeConcreteTensor({}, DataType::Half);

  fusion->addInput(t2);
  fusion->addInput(t5);
  fusion->addInput(t7);
  fusion->addInput(t9);
  fusion->addInput(t4);

  auto d86 = IrBuilder::create<Double>(0);

  auto t3 = castOp(DataType::Float, t2);
  auto t6 = castOp(DataType::Float, t5);
  auto t8 = castOp(DataType::Float, t7);
  auto t10 = castOp(DataType::Float, t9);
  auto t11 = add(t8, t10);
  auto t12 = set(t11);
  auto t13 = set(t6);
  auto t14 = lt(t13, d86);
  auto t15 = broadcast(t4, {true, true, true, true});
  auto t16 = where(t14, t15, t12);
  auto t17 = set(t16);
  auto t29 = castOp(DataType::Half, t17);
  auto t18 = mul(t17, t3);
  auto t19 = permute(t18, {0, 2, 3, 1});
  auto t30 = castOp(DataType::Half, t19);

  fusion->addOutput(t29);
  fusion->addOutput(t30);
}

static void NvFuserScheduler_nhwc_seresnet152d_transpose65(
    benchmark::State& benchmark_state,
    FusionExecutorCache* fusion_executor_cache,
    void* null) {
  std::vector<int64_t> input_shape{
      benchmark_state.range(0),
      benchmark_state.range(2),
      benchmark_state.range(2),
      benchmark_state.range(1)};

  at::manual_seed(0);
  auto fp16_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);

  auto t2 = at::randn(input_shape, fp16_options);
  auto t5 = at::randn(input_shape, fp16_options);
  auto t7 = at::randn(input_shape, fp16_options);
  auto t9 = at::randn(input_shape, fp16_options);
  // Need zero dim tensor don't know how to do that, so just going to reduce a
  // 1D tensor
  auto t4 = at::randn({2}, fp16_options).sum();

  std::vector<c10::IValue> aten_inputs({t2, t5, t7, t9, t4});
  runBenchmarkIterations(benchmark_state, fusion_executor_cache, aten_inputs);

  // Full tensors - halfx6 - t2, t5, t7, t9, t29, t30
  benchmark_state.SetBytesProcessed(
      int64_t(benchmark_state.iterations()) * t2.numel() * 6 * 2);
}

NVFUSER_BENCHMARK_DEFINE(
    NvFuserScheduler_TIMM_nhwc_seresnet152d_transpose65,
    nhwc_seresnet152d_transpose65,
    NvFuserScheduler_nhwc_seresnet152d_transpose65,
    nullptr);

// Norm inner dim Half version of vit_base_patch16_224_norm_inner3
NVFUSER_BENCHMARK_RUN(NvFuserScheduler_TIMM_nhwc_seresnet152d_transpose65)
    ->Args({128, 12, 197})
    ->Unit(benchmark::kMicrosecond)
    ->UseManualTime();