Create dispatch system for executors

CMakeLists.txt

@@ -203,6 +203,7 @@ list(APPEND NVFUSER_SRCS
   ${NVFUSER_SRCS_DIR}/rng.cpp
   ${NVFUSER_SRCS_DIR}/runtime/allocations.cpp
   ${NVFUSER_SRCS_DIR}/runtime/executor.cpp
+  ${NVFUSER_SRCS_DIR}/runtime/executor_dispatch.cpp
   ${NVFUSER_SRCS_DIR}/runtime/executor_kernel_arg.cpp
   ${NVFUSER_SRCS_DIR}/runtime/executor_params.cpp
   ${NVFUSER_SRCS_DIR}/runtime/executor_utils.cpp

benchmarks/cpp/utils.cpp

@@ -192,7 +192,8 @@ int64_t runBenchmarkIterations(
 
   const auto& compile_log = executor_cache->getMostRecentExecutorInfo();
   auto params = toString(compile_log.params);
-  auto lparams = toString(compile_log.fusion_executor->lastLaunchParams());
+  auto lparams = toString(
+      compile_log.fusion_executor->as<KernelExecutor>()->lastLaunchParams());
   // Only set if not segmented. In the case of segmented fusions,
   // this could be confusing as the log would refect only the last
   // segment. Revisit if necessary.
@@ -223,19 +224,18 @@ int64_t runBenchmarkIterations(
 
 int64_t runBenchmarkIterations(
     benchmark::State& benchmark_state,
-    KernelExecutor* fusion_executor,
+    KernelExecutor* ke,
     std::vector<c10::IValue>& aten_inputs,
     const LaunchParams& launch_constraints,
     CompileParams compile_params) {
   int64_t io_bytes = getSizeOfInputs(aten_inputs);
   {
     // Warm-up run
-    auto cg_outputs =
-        fusion_executor->run(aten_inputs, launch_constraints, compile_params);
+    auto cg_outputs = ke->run(aten_inputs, launch_constraints, compile_params);
     io_bytes += getSizeOfOutputs(cg_outputs);
   }
 
-  auto lparams = toString(fusion_executor->lastLaunchParams());
+  auto lparams = toString(ke->lastLaunchParams());
   benchmark_state.SetLabel(lparams);
 
   // Sync everything up before we start
@@ -246,8 +246,7 @@ int64_t runBenchmarkIterations(
     clearL2Cache();
     FusionProfiler::start();
     FusionProfiler::createSegments(1);
-    auto cg_outputs =
-        fusion_executor->run(aten_inputs, launch_constraints, compile_params);
+    auto cg_outputs = ke->run(aten_inputs, launch_constraints, compile_params);
     FusionProfiler::stop();
     benchmark_state.SetIterationTime(
         FusionProfiler::profile().kernel_time_ms / 1000.0);

benchmarks/cpp/utils.h

@@ -48,7 +48,7 @@ int64_t runBenchmarkIterations(
 //! kernel time is added to benchmark_state.
 int64_t runBenchmarkIterations(
     benchmark::State& benchmark_state,
-    KernelExecutor* fusion_executor,
+    KernelExecutor* ke,
     std::vector<c10::IValue>& aten_inputs,
     const LaunchParams& launch_constraints = LaunchParams(),
     CompileParams compile_params = CompileParams());

csrc/fusion.cpp

@@ -807,13 +807,6 @@ bool Fusion::hasDynamicTransform() {
   return !ir_utils::getTVsWithDynamicTransform(this).empty();
 }
 
-bool isExpressionEvaluated(Fusion* fusion) {
-  return std::all_of(
-      fusion->outputs().begin(), fusion->outputs().end(), [&fusion](Val* out) {
-        return fusion->getOutputAlias(out).type == AllocationType::Evaluate;
-      });
-}
-
 namespace {
 std::vector<TensorView*> findAllTvs(Fusion* fusion) {
   auto used_vals = fusion->usedMathVals();

csrc/fusion.h

@@ -506,7 +506,4 @@ void Fusion::manage(std::string key, T data) {
   });
 }
 
-// Returns true if all fusion outputs are expression evaluated.
-bool isExpressionEvaluated(Fusion* fusion);
-
 } // namespace nvfuser

csrc/fusion_profiler.cpp

@@ -326,17 +326,15 @@ SegmentProfiler::SegmentProfiler(uint32_t id, bool cupti_disabled)
       output_bytes_(0),
       kernel_profile_state_(ProfilerState::Ready) {}
 
-void SegmentProfiler::startCompile(int device) {
-  device_ = device;
+void SegmentProfiler::startCompile() {
   compile_timer_.start();
 }
 
 void SegmentProfiler::stopCompile() {
   compile_timer_.stop();
 }
 
-void SegmentProfiler::startKernel(int device) {
-  device_ = device;
+void SegmentProfiler::startKernel() {
   NVF_CHECK(
       kernel_profile_state_ == ProfilerState::Ready,
       "ProfilerState is not Ready!",

csrc/fusion_profiler.h

@@ -186,12 +186,15 @@ class SegmentProfiler {
  public:
   SegmentProfiler(uint32_t id, bool cupti_disabled);
 
-  void startCompile(int device);
+  void startCompile();
   void stopCompile();
 
-  void startKernel(int device);
+  void startKernel();
   void stopKernel();
 
+  void setDevice(int64_t device) {
+    device_ = (int)device;
+  }
   void inputBytesAccessed(int64_t bytes);
   void outputBytesAccessed(int64_t bytes);
 
@@ -219,13 +222,13 @@ class SegmentProfiler {
  private:
   bool cupti_disabled_;
 
-  int device_;
-  uint32_t segment_id_;
+  int device_ = -1;
+  uint32_t segment_id_ = -1;
 
   HostTimer compile_timer_;
-  int64_t input_bytes_;
-  int64_t output_bytes_;
-  std::string scheduler_;
+  int64_t input_bytes_ = -1;
+  int64_t output_bytes_ = -1;
+  std::string scheduler_ = "None";
   ProfilerState kernel_profile_state_;
 };