Refactoring Fusion Executor, pulling out compiled kernel

CMakeLists.txt

@@ -202,6 +202,7 @@ list(APPEND NVFUSER_SRCS
   ${NVFUSER_SRCS_DIR}/preseg_passes/segment_inplace_update.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_dispatch.cpp
   ${NVFUSER_SRCS_DIR}/runtime/executor_kernel_arg.cpp

benchmarks/cpp/matmul.cpp

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

csrc/cuda_utils.h

@@ -10,6 +10,7 @@
 #include <cuda_runtime.h>
 #include <driver_api.h>
 #include <exceptions.h>
+#include <nvrtc.h>
 
 #define NVFUSER_NVRTC_SAFE_CALL(x)               \
   do {                                           \

csrc/options.h

@@ -39,7 +39,7 @@ enum class DebugDumpOption {
   FusionIrPresched, //!< Dump the segmented Fusion IR before it is scheduled
   // TODO(wujingyue): name the following FusionIrSched
   FusionIr, //!< Dump the Fusion IR before lowering. This is the Fusion IR fed
-            //!< to `KernelExecutor::compileFusion`.
+            //!< to `KernelExecutor::compile`.
   FusionIrGraph, //!< Dump a GraphViz graph of the Fusion IR
   FusionIrMath, //!< Dump just the compute (math) part of the above `FusionIr`
                 //!< for conciseness

csrc/python_frontend/fusion_definition.cpp

@@ -405,7 +405,7 @@ std::vector<at::Tensor> FusionDefinition::execute(
         outputs = user_sched.executor->run(inputs);
       } else {
         // Automatic scheduler was used for UserSchedule.
-        // Pass launch and compile params to compileFusion and runFusion.
+        // Pass launch and compile params to compile and run.
         if (!user_sched.executor->isCompiled()) {
           user_sched.executor->compile(
               user_sched.scheduled_fusion.get(),
@@ -487,10 +487,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);
@@ -516,10 +517,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();
       }
     }
   }