From b458c7e4b8dcbed35bf61d86b347a0107f4e8a0a Mon Sep 17 00:00:00 2001 From: Christian Sarofeen Date: Wed, 1 Jan 2025 09:20:25 -0800 Subject: [PATCH] Cleanup disableLaunchParamCache. --- csrc/runtime/executor.cpp | 5 +---- csrc/runtime/executor.h | 6 ------ csrc/runtime/fusion_executor_cache.cpp | 12 +++++++++++- csrc/runtime/fusion_kernel_runtime.cpp | 11 ----------- csrc/runtime/fusion_kernel_runtime.h | 3 --- 5 files changed, 12 insertions(+), 25 deletions(-) diff --git a/csrc/runtime/executor.cpp b/csrc/runtime/executor.cpp index 745a5f9550b..3f72f43e419 100644 --- 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 index 290fed42a45..57d99e6b862 100644 --- 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 compiled_kernel_; diff --git a/csrc/runtime/fusion_executor_cache.cpp b/csrc/runtime/fusion_executor_cache.cpp index e4dc470665c..04ed4396e70 100644 --- 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(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 index 8031b951f37..74d0b08abe3 100644 --- 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(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 index 327c7bff8fc..4b898b454f0 100644 --- 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();