Cleanup disableLaunchParamCache.

NVIDIA · Jan 1, 2025 · b458c7e · b458c7e
1 parent bd1d312
commit b458c7e
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 25 deletions.
diff --git a/csrc/runtime/executor.cpp b/csrc/runtime/executor.cpp
@@ -236,10 +236,7 @@ void KernelExecutor::compile(
   } else if (has_cp_async_bulk) {
     // TMA operations require 32-bit indexing.
     compile_params.index_type = PrimDataType::Int32;
-  }
-
-  // TODO: Is this necessary?
-  if (!compile_params.index_type.has_value()) {
+  } else {
     compile_params.index_type = arg_index_type;
   }
 

diff --git a/csrc/runtime/executor.h b/csrc/runtime/executor.h
@@ -324,12 +324,6 @@ class KernelExecutor : public ExecutorAbstract {
   //! Clear the cached properties of the compiled kernel
   void resetCompiledKernelProperties();
 
-  void disableLaunchParamCache() {
-    if (compiledKernel()) {
-      compiledKernel()->disableLaunchParamCache();
-    }
-  }
-
  private:
   std::unique_ptr<CompiledKernel> compiled_kernel_;
 

diff --git a/csrc/runtime/fusion_executor_cache.cpp b/csrc/runtime/fusion_executor_cache.cpp
@@ -321,7 +321,17 @@ void FusionExecutorCache::profile(bool to_profile) {
 void FusionExecutorCache::disableLaunchParamCache() {
   for (auto& it : kernel_runtimes_) {
     for (auto& kernel_runtime : it.second) {
-      kernel_runtime->disableLaunchParamCache();
+      NVF_CHECK(
+          kernel_runtime->isCompiled(),
+          "Tried to set parameters of executors before they were initialized.");
+      for (auto& executor : kernel_runtime->executors()) {
+        if (auto ke = dynamic_cast<KernelExecutor*>(executor.get())) {
+          NVF_CHECK(
+              ke->compiledKernel(),
+              "Tried to disable parameter cache of uninitialized CompiledKernel.");
+          ke->compiledKernel()->disableLaunchParamCache();
+        }
+      }
     }
   }
 }

diff --git a/csrc/runtime/fusion_kernel_runtime.cpp b/csrc/runtime/fusion_kernel_runtime.cpp
@@ -416,17 +416,6 @@ void FusionKernelRuntime::compileFusionParallel(KernelArgumentHolder args) {
   }
 }
 
-void FusionKernelRuntime::disableLaunchParamCache() {
-  NVF_CHECK(
-      isCompiled(),
-      "Tried to set parameters of executors before they were initialized.");
-  for (auto& executor : executors_) {
-    if (auto ke = dynamic_cast<KernelExecutor*>(executor.get())) {
-      ke->compiledKernel()->disableLaunchParamCache();
-    }
-  }
-}
-
 void FusionKernelRuntime::disableKernelLaunch() {
   NVF_CHECK(
       isCompiled(),

diff --git a/csrc/runtime/fusion_kernel_runtime.h b/csrc/runtime/fusion_kernel_runtime.h
@@ -106,9 +106,6 @@ class FusionKernelRuntime {
     return kernel_time_ms_;
   }
 
-  //! Internal knob for profiling shape inference
-  void disableLaunchParamCache();
-
   //! Internal knob for profiling shape inference
   void disableKernelLaunch();