Refactoring Fusion Executor, pulling out compiled kernel

CMakeLists.txt

@@ -198,6 +198,7 @@ list(APPEND NVFUSER_SRCS
   ${NVFUSER_SRCS_DIR}/preseg_passes/reorder_sharded_axis.cpp
   ${NVFUSER_SRCS_DIR}/rng.cpp
   ${NVFUSER_SRCS_DIR}/runtime/allocations.cpp
+  ${NVFUSER_SRCS_DIR}/runtime/compiled_kernel.cpp
   ${NVFUSER_SRCS_DIR}/runtime/executor.cpp
   ${NVFUSER_SRCS_DIR}/runtime/executor_kernel_arg.cpp
   ${NVFUSER_SRCS_DIR}/runtime/executor_params.cpp

benchmarks/cpp/matmul.cpp

@@ -178,7 +178,8 @@ static void SingleMatmulBase(
   FusionExecutor fe;
   fe.compileFusion(fusion, args, launch_constraints, cparams);
   NVF_CHECK(
-      getBankConflictInfo(fe.kernel(), launch_constraints).empty(),
+      getBankConflictInfo(fe.compiledKernel()->kernel(), launch_constraints)
+          .empty(),
       "Shared memory bank conflict not removed.");
 
   std::vector<c10::IValue> aten_inputs({inputs.first, inputs.second});
@@ -359,7 +360,7 @@ static void SingleMatmulPartitionedK(
   auto lparams = LaunchParams();
   fe.compileFusion(fusion, args, lparams, cparams);
   NVF_CHECK(
-      getBankConflictInfo(fe.kernel(), lparams).empty(),
+      getBankConflictInfo(fe.compiledKernel()->kernel(), lparams).empty(),
       "Shared memory bank conflict not removed.");
 
   // Warm up run
@@ -466,7 +467,9 @@ static void NvFuserScheduler_MatmulSplitKReduction(
       fusion, args, heuristic_params->lparams, heuristic_params->cparams);
 
   NVF_CHECK(
-      getBankConflictInfo(fe.kernel(), heuristic_params->lparams).empty(),
+      getBankConflictInfo(
+          fe.compiledKernel()->kernel(), heuristic_params->lparams)
+          .empty(),
       "Shared memory bank conflict not removed.");
 
   // Warm up run

csrc/evaluator_common.cpp

@@ -184,7 +184,11 @@ void PrecomputedValues::bindValues(
     const std::vector<Val*>& inputs,
     const KernelArgumentHolder& args) {
   NVF_ERROR(
-      args.size() == inputs.size(), "kernel inputs size does not match args");
+      args.size() == inputs.size(),
+      "kernel inputs size does not match args: ",
+      args.size(),
+      " vs ",
+      inputs.size());
 
   for (const auto i : c10::irange((int64_t)inputs.size())) {
     const auto input = inputs[i];

csrc/python_frontend/fusion_definition.cpp

@@ -416,10 +416,11 @@ std::string FusionDefinition::lastCudaCode(
 
   if (!override_user_schedule && (user_exec != nullptr)) {
     if (intrinsic_code) {
-      result = user_exec->getStructuredCode(
-          user_exec->kernelString(), user_exec->kernel()->indexType());
+      result = user_exec->compiledKernel()->getStructuredCode(
+          user_exec->compiledKernel()->kernelString(),
+          user_exec->compiledKernel()->kernel()->indexType());
     } else {
-      result = user_exec->kernelString();
+      result = user_exec->compiledKernel()->kernelString();
     }
   } else {
     result = scheds->auto_gen_schedules->getMostRecentCode(intrinsic_code);
@@ -445,10 +446,11 @@ std::string FusionDefinition::cudaCodeFor(
           scheds, user_sched_id.value(), device);
       auto user_exec = user_sched.executor.get();
       if (intrinsic_code) {
-        return user_exec->getStructuredCode(
-            user_exec->kernelString(), user_exec->kernel()->indexType());
+        return user_exec->compiledKernel()->getStructuredCode(
+            user_exec->compiledKernel()->kernelString(),
+            user_exec->compiledKernel()->kernel()->indexType());
       } else {
-        return user_exec->kernelString();
+        return user_exec->compiledKernel()->kernelString();
       }
     }
   }