From 412c37d1427c3f77ae5c7feb19c91d6c797f31ed Mon Sep 17 00:00:00 2001 From: Jacob Hinkle Date: Tue, 28 Jan 2025 08:31:57 -0500 Subject: [PATCH] Print scheduled fusion instead of kernel for NVFUSER_DUMP=fusion_ir PR #3468 changed to using `CompiledKernel` and in the shuffle, we switched from using the incoming fusion before compilation to the lowered `kir::Kernel`. This PR just moves the printing to occur just before lowering, inside the constructor for `CompiledKernel`. I believe this is enough to restore the previous behavior. Fixes #3765 --- csrc/runtime/compiled_kernel.cpp | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/csrc/runtime/compiled_kernel.cpp b/csrc/runtime/compiled_kernel.cpp index 46dd9e4eda1..b41103c3957 100644 --- a/csrc/runtime/compiled_kernel.cpp +++ b/csrc/runtime/compiled_kernel.cpp @@ -1158,6 +1158,13 @@ NVF_API CompiledKernel::CompiledKernel( lowered_(std::make_unique(fusion, compile_params)), device_(device) { FUSER_PERF_SCOPE("CompiledKernel::CompiledKernel"); + + if (isDebugDumpEnabled(DebugDumpOption::FusionIr)) { + fusion->print(); + } else if (isDebugDumpEnabled(DebugDumpOption::FusionIrMath)) { + fusion->printMath(); + } + // TODO: No hooks can be sent because this is in the constructor for (const auto& hook : pre_lowering_hooks) { hook(lowered_.get()); @@ -1202,12 +1209,6 @@ void CompiledKernel::compile(int64_t block_size) { // it under these circumstances. disable_parameter_cache_ = requiresDisabledParamCache(kernel()); - if (isDebugDumpEnabled(DebugDumpOption::FusionIr)) { - kernel()->print(); - } else if (isDebugDumpEnabled(DebugDumpOption::FusionIrMath)) { - kernel()->printMath(); - } - if (isDebugDumpEnabled(DebugDumpOption::FusionIrGraph)) { std::stringstream file_name; file_name << "__tmp_fusion_ir_graph_" << kernel_id_ << ".dot";