Working towards returning KernelArgumentHolder/PolymorphicVal's from fusion execution.

benchmarks/cpp/batch_norm_channels_first.cpp

@@ -99,10 +99,10 @@ static void NvFuserScheduler_BatchNorm(
   at::Tensor at_bias = at::zeros({input_shape[1]}, options);
   at::Tensor at_run_mean = at::zeros({input_shape[1]}, fp32_options);
   at::Tensor at_run_var = at::ones({input_shape[1]}, fp32_options);
-  std::vector<c10::IValue> aten_inputs(
-      {at_x, at_weight, at_bias, at_run_mean, at_run_var});
+  KernelArgumentHolder args = {
+      at_x, at_weight, at_bias, at_run_mean, at_run_var};
 
-  runBenchmarkIterations(benchmark_state, executor_cache, aten_inputs);
+  runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   benchmark_state.SetBytesProcessed(
       int64_t(benchmark_state.iterations()) *

benchmarks/cpp/batch_norm_channels_first_backward.cpp

@@ -112,10 +112,10 @@ static void NvFuserScheduler_BatchNorm_BWD(
   at::Tensor save_mean = at::zeros({input_shape[1]}, fp32_options);
   at::Tensor save_var = at::ones({input_shape[1]}, fp32_options);
 
-  std::vector<c10::IValue> aten_inputs(
-      {input, grad_out, weight, run_mean, run_var, save_mean, save_var});
+  KernelArgumentHolder args = {
+      input, grad_out, weight, run_mean, run_var, save_mean, save_var};
 
-  runBenchmarkIterations(benchmark_state, executor_cache, aten_inputs);
+  runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   benchmark_state.SetBytesProcessed(
       int64_t(benchmark_state.iterations()) *

benchmarks/cpp/batch_norm_channels_last.cpp

@@ -100,10 +100,10 @@ static void NvFuserScheduler_BatchNorm_nhwc(
   at::Tensor at_bias = at::zeros({input_shape[3]}, options);
   at::Tensor at_run_mean = at::zeros({input_shape[3]}, fp32_options);
   at::Tensor at_run_var = at::ones({input_shape[3]}, fp32_options);
-  std::vector<c10::IValue> aten_inputs(
-      {at_x, at_weight, at_bias, at_run_mean, at_run_var});
+  KernelArgumentHolder args = {
+      at_x, at_weight, at_bias, at_run_mean, at_run_var};
 
-  runBenchmarkIterations(benchmark_state, executor_cache, aten_inputs);
+  runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   benchmark_state.SetBytesProcessed(
       int64_t(benchmark_state.iterations()) *

benchmarks/cpp/batch_norm_channels_last_backward.cpp

@@ -113,10 +113,10 @@ static void NvFuserScheduler_BatchNorm_nhwc_BWD(
   at::Tensor save_mean = at::zeros({input_shape[3]}, fp32_options);
   at::Tensor save_var = at::ones({input_shape[3]}, fp32_options);
 
-  std::vector<c10::IValue> aten_inputs(
-      {input, grad_out, weight, run_mean, run_var, save_mean, save_var});
+  KernelArgumentHolder args = {
+      input, grad_out, weight, run_mean, run_var, save_mean, save_var};
 
-  runBenchmarkIterations(benchmark_state, executor_cache, aten_inputs);
+  runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   benchmark_state.SetBytesProcessed(
       int64_t(benchmark_state.iterations()) *

benchmarks/cpp/bert.cpp

@@ -132,10 +132,9 @@ static void NvFuserScheduler_DivMaxSoftDropFwd(
   at::Tensor t0 = at::randn({w, 1, 1, z}, options);
   at::Tensor t1 = at::randn({w, x, y, z}, options);
 
-  std::vector<c10::IValue> at_inputs = {t0, t1};
+  KernelArgumentHolder args = {t0, t1};
 
-  auto bytes =
-      runBenchmarkIterations(benchmark_state, executor_cache, at_inputs);
+  auto bytes = runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   benchmark_state.SetBytesProcessed(
       bytes * int64_t(benchmark_state.iterations()));
@@ -159,10 +158,9 @@ static void NvFuserScheduler_DivMaxSoftDropBwd(
   at::Tensor t2 = at::randn({w, x, y, z}, options);
   at::Tensor t3 = at::randn({w, x, y, z}, options).round().to(at::kBool);
 
-  std::vector<c10::IValue> at_inputs = {t0, t1, t2, t3};
+  KernelArgumentHolder args = {t0, t1, t2, t3};
 
-  auto bytes =
-      runBenchmarkIterations(benchmark_state, executor_cache, at_inputs);
+  auto bytes = runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   // Some reason t1 isn't used, ignore it.
   bytes -=
@@ -244,10 +242,9 @@ static void NvFuserScheduler_BiasDropoutAddLayernormFwd(
   at::Tensor t3 = at::randn({x, y, z}, options);
   at::Tensor t4 = at::randn({z}, options);
 
-  std::vector<c10::IValue> at_inputs = {t0, t1, t2, t3, t4};
+  KernelArgumentHolder args = {t0, t1, t2, t3, t4};
 
-  auto bytes =
-      runBenchmarkIterations(benchmark_state, executor_cache, at_inputs);
+  auto bytes = runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   benchmark_state.SetBytesProcessed(
       bytes * int64_t(benchmark_state.iterations()));
@@ -319,10 +316,9 @@ static void NvFuserScheduler_BiasDropoutAddLayernormBwd1(
   at::Tensor t2 = at::randn({x, y, 1}, options);
   at::Tensor t3 = at::randn({x, y, 1}, options);
 
-  std::vector<c10::IValue> at_inputs = {t0, t1, t2, t3};
+  KernelArgumentHolder args = {t0, t1, t2, t3};
 
-  auto bytes =
-      runBenchmarkIterations(benchmark_state, executor_cache, at_inputs);
+  auto bytes = runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   benchmark_state.SetBytesProcessed(
       bytes * int64_t(benchmark_state.iterations()));
@@ -395,10 +391,9 @@ static void NvFuserScheduler_BiasDropoutAddLayernormBwd2(
   at::Tensor t1 = at::randn({x, y, z}, options);
   at::Tensor t8 = at::randn({x, y, z}, options);
 
-  std::vector<c10::IValue> at_inputs = {t4, t5, t1, t8};
+  KernelArgumentHolder args = {t4, t5, t1, t8};
 
-  auto bytes =
-      runBenchmarkIterations(benchmark_state, executor_cache, at_inputs);
+  auto bytes = runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   benchmark_state.SetBytesProcessed(
       bytes * int64_t(benchmark_state.iterations()));
@@ -451,10 +446,9 @@ static void NvFuserScheduler_BiasDropoutAddLayernormBwd3(
   at::Tensor t0 = at::randn({x, y, z}, options);
   at::Tensor t21 = at::randn({x, y, z}, options);
 
-  std::vector<c10::IValue> at_inputs = {t0, t21};
+  KernelArgumentHolder args = {t0, t21};
 
-  auto bytes =
-      runBenchmarkIterations(benchmark_state, executor_cache, at_inputs);
+  auto bytes = runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   benchmark_state.SetBytesProcessed(
       bytes * int64_t(benchmark_state.iterations()));

benchmarks/cpp/broadcast.cpp

@@ -72,9 +72,9 @@ static void NvFuserScheduler_Broadcast(
 
   at::Tensor t1 = at::randn({iter_size}, options);
 
-  std::vector<c10::IValue> aten_inputs({t0, t1});
+  KernelArgumentHolder args = {t0, t1};
 
-  runBenchmarkIterations(benchmark_state, executor_cache, aten_inputs);
+  runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   benchmark_state.SetBytesProcessed(
       int64_t(benchmark_state.iterations()) *

benchmarks/cpp/gelu_backward.cpp

@@ -77,7 +77,7 @@ static void setupFusion(Fusion* fusion) {
   fusion->addOutput(t27);
 }
 
-static std::vector<c10::IValue> setupInputs() {
+static KernelArgumentHolder setupInputs() {
   at::manual_seed(0);
 
   auto options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
@@ -112,12 +112,12 @@ static void NvFuserScheduler_GeluBackward_AutoSchedule(
     benchmark_state.PauseTiming();
     Fusion fusion;
     setupFusion(&fusion);
-    std::vector<c10::IValue> inputs = setupInputs();
+    KernelArgumentHolder args = setupInputs();
     benchmark_state.ResumeTiming();
 
     // Auto-schedule
     SchedulerEntry::scheduleWith(
-        &fusion, SchedulerType::PointWise, c10::ArrayRef<c10::IValue>(inputs));
+        &fusion, SchedulerType::PointWise, args.toC10Array());
   }
 }
 
@@ -134,10 +134,10 @@ static void NvFuserScheduler_GeluBackward_Lower(
   setupFusion(&fusion);
 
   // inputs
-  std::vector<c10::IValue> inputs = setupInputs();
+  KernelArgumentHolder args = setupInputs();
 
   SchedulerEntry::scheduleWith(
-      &fusion, SchedulerType::PointWise, c10::ArrayRef<c10::IValue>(inputs));
+      &fusion, SchedulerType::PointWise, args.toC10Array());
 
   for (auto _ : benchmark_state) {
     GpuLower(&fusion).run();
@@ -156,14 +156,14 @@ static void NvFuserScheduler_GeluBackward_Compile(
   setupFusion(&fusion);
 
   // inputs
-  std::vector<c10::IValue> inputs = setupInputs();
+  KernelArgumentHolder args = setupInputs();
 
   auto heuristic_params = SchedulerEntry::scheduleWith(
-      &fusion, SchedulerType::PointWise, c10::ArrayRef<c10::IValue>(inputs));
+      &fusion, SchedulerType::PointWise, args.toC10Array());
 
   for (auto _ : benchmark_state) {
     KernelExecutor ke;
-    ke.compile(&fusion, inputs, heuristic_params->lparams);
+    ke.compile(&fusion, args.toC10Array(), heuristic_params->lparams);
   }
 }
 
@@ -179,22 +179,21 @@ static void NvFuserScheduler_GeluBackward_RunFusion(
   setupFusion(&fusion);
 
   // inputs
-  std::vector<c10::IValue> inputs = setupInputs();
+  KernelArgumentHolder args = setupInputs();
 
   // outputs
   std::vector<at::Tensor> outputs;
 
   auto heuristic_params = SchedulerEntry::scheduleWith(
-      &fusion, SchedulerType::PointWise, c10::ArrayRef<c10::IValue>(inputs));
+      &fusion, SchedulerType::PointWise, args.toC10Array());
 
   KernelExecutor ke;
-  ke.compile(&fusion, inputs, heuristic_params->lparams);
+  ke.compile(&fusion, args.toC10Array(), heuristic_params->lparams);
 
   C10_CUDA_CHECK(cudaDeviceSynchronize());
 
   for (auto _ : benchmark_state) {
-    outputs =
-        ke.run(c10::ArrayRef<c10::IValue>(inputs), heuristic_params->lparams);
+    outputs = ke.run(args.toC10Array(), heuristic_params->lparams);
     C10_CUDA_CHECK(cudaDeviceSynchronize());
     clearL2Cache();
   }
@@ -213,16 +212,15 @@ static void NvFuserScheduler_GeluBackward_RunFusion_GpuOnly(
   setupFusion(&fusion);
 
   // inputs
-  std::vector<c10::IValue> inputs = setupInputs();
+  KernelArgumentHolder args = setupInputs();
 
   auto heuristic_params = SchedulerEntry::scheduleWith(
-      &fusion, SchedulerType::PointWise, c10::ArrayRef<c10::IValue>(inputs));
+      &fusion, SchedulerType::PointWise, args.toC10Array());
 
   KernelExecutor ke;
-  ke.compile(&fusion, inputs, heuristic_params->lparams);
+  ke.compile(&fusion, args.toC10Array(), heuristic_params->lparams);
 
-  runBenchmarkIterations(
-      benchmark_state, &ke, inputs, heuristic_params->lparams);
+  runBenchmarkIterations(benchmark_state, &ke, args, heuristic_params->lparams);
 }
 
 BENCHMARK(NvFuserScheduler_GeluBackward_RunFusion_GpuOnly)
@@ -239,21 +237,20 @@ static void NvFuserScheduler_GeluBackward_RunFusion_CpuOnly(
   setupFusion(&fusion);
 
   // inputs
-  std::vector<c10::IValue> inputs = setupInputs();
+  KernelArgumentHolder args = setupInputs();
 
   // outputs
   std::vector<at::Tensor> outputs;
 
   auto heuristic_params = SchedulerEntry::scheduleWith(
-      &fusion, SchedulerType::PointWise, c10::ArrayRef<c10::IValue>(inputs));
+      &fusion, SchedulerType::PointWise, args.toC10Array());
 
   KernelExecutor ke;
   ke.setExecuteKernelFlag(false);
-  ke.compile(&fusion, inputs, heuristic_params->lparams);
+  ke.compile(&fusion, args.toC10Array(), heuristic_params->lparams);
 
   for (auto _ : benchmark_state) {
-    outputs =
-        ke.run(c10::ArrayRef<c10::IValue>(inputs), heuristic_params->lparams);
+    outputs = ke.run(args.toC10Array(), heuristic_params->lparams);
   }
 }
 

benchmarks/cpp/gelu_backward_reduction.cpp

@@ -110,9 +110,9 @@ static void NvFuserScheduler_GeluBackwardReduction(
       (reduction_dim ? at::randn({iter_size, reduction_size}, options)
                      : at::randn({reduction_size, iter_size}, options));
 
-  std::vector<c10::IValue> aten_inputs = {aten_input_grad, aten_input_x};
+  KernelArgumentHolder args = {aten_input_grad, aten_input_x};
 
-  runBenchmarkIterations(benchmark_state, executor_cache, aten_inputs);
+  runBenchmarkIterations(benchmark_state, executor_cache, args);
 
   // inputs: gradient tensor + input tensor
   // outputs: output, output_of_reduction

benchmarks/cpp/heuristic_cache.cpp

@@ -31,15 +31,13 @@ static void NvFuserScheduler_LayerNormBackward_HeuristicCache(
 
   // PreAllocate
   std::unique_ptr<FusionExecutorCache> executor_cache;
-  std::vector<c10::IValue> aten_inputs;
+  KernelArgumentHolder args;
 
   std::vector<int64_t> shape{20, 100, 35, 67};
   std::vector<int64_t> norm_shape{67};
 
   auto runtime = getLayerBackwardNormRuntime(
-      std::move(fusion_ptr), executor_cache, aten_inputs, shape, norm_shape);
-
-  KernelArgumentHolder args(aten_inputs);
+      std::move(fusion_ptr), executor_cache, args, shape, norm_shape);
 
   NVF_ERROR(runtime->getMaybeHeuristicsFor(args).has_value());
 
@@ -56,15 +54,13 @@ static void NvFuserScheduler_LayerNormForward_HeuristicCache(
 
   // PreAllocate
   std::unique_ptr<FusionExecutorCache> executor_cache;
-  std::vector<c10::IValue> aten_inputs;
+  KernelArgumentHolder args;
 
   std::vector<int64_t> shape{20, 100, 35, 67};
   std::vector<int64_t> norm_shape{67};
 
   auto runtime = getLayerForwardNormRuntime(
-      std::move(fusion_ptr), executor_cache, aten_inputs, shape, norm_shape);
-
-  KernelArgumentHolder args(aten_inputs);
+      std::move(fusion_ptr), executor_cache, args, shape, norm_shape);
 
   NVF_ERROR(runtime->getMaybeHeuristicsFor(args).has_value());
 

benchmarks/cpp/heuristic_lookup.cpp

@@ -31,15 +31,13 @@ static void NvFuserScheduler_LayerNormBackward_HeuristicLookup(
 
   // PreAllocate
   std::unique_ptr<FusionExecutorCache> executor_cache;
-  std::vector<c10::IValue> aten_inputs;
+  KernelArgumentHolder args;
 
   std::vector<int64_t> shape{20, 100, 35, 67};
   std::vector<int64_t> norm_shape{67};
 
   auto runtime = getLayerBackwardNormRuntime(
-      std::move(fusion_ptr), executor_cache, aten_inputs, shape, norm_shape);
-
-  KernelArgumentHolder args(aten_inputs);
+      std::move(fusion_ptr), executor_cache, args, shape, norm_shape);
 
   NVF_ERROR(runtime->getMaybeHeuristicsFor(args).has_value());
 
@@ -56,15 +54,13 @@ static void NvFuserScheduler_LayerNormForward_HeuristicLookup(
 
   // PreAllocate
   std::unique_ptr<FusionExecutorCache> executor_cache;
-  std::vector<c10::IValue> aten_inputs;
+  KernelArgumentHolder args;
 
   std::vector<int64_t> shape{20, 100, 35, 67};
   std::vector<int64_t> norm_shape{67};
 
   auto runtime = getLayerForwardNormRuntime(
-      std::move(fusion_ptr), executor_cache, aten_inputs, shape, norm_shape);
-
-  KernelArgumentHolder args(aten_inputs);
+      std::move(fusion_ptr), executor_cache, args, shape, norm_shape);
 
   NVF_ERROR(runtime->getMaybeHeuristicsFor(args).has_value());