Add a dump option to dump a fusion IR graph (#3603)

`NVFUSER_DUMP=fusion_ir_graph` saves the dot representation of a fusion before lowering to a file named like `__tmp_fusion_ir_graph_inner_persistent_f0_c1_r0_g0.dot`. Example visualization: #3498 (comment)
NVIDIA · Dec 17, 2024 · ffd186e · ffd186e
1 parent d623221
commit ffd186e
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 8 deletions.
diff --git a/csrc/options.cpp b/csrc/options.cpp
@@ -112,12 +112,13 @@ std::unordered_map<DebugDumpOption, std::vector<std::string>> Options<
       {"expr_sort_verbose", DebugDumpOption::ExprSortVerbose},
       {"ftrace", DebugDumpOption::FunctionTrace},
       {"fusion_args", DebugDumpOption::FusionArgs},
-      {"fusion_ir_original", DebugDumpOption::FusionIrOriginal},
-      {"fusion_ir_concretized", DebugDumpOption::FusionIrConcretized},
-      {"fusion_ir_preseg", DebugDumpOption::FusionIrPreseg},
-      {"fusion_ir_presched", DebugDumpOption::FusionIrPresched},
       {"fusion_ir", DebugDumpOption::FusionIr},
+      {"fusion_ir_concretized", DebugDumpOption::FusionIrConcretized},
+      {"fusion_ir_graph", DebugDumpOption::FusionIrGraph},
       {"fusion_ir_math", DebugDumpOption::FusionIrMath},
+      {"fusion_ir_original", DebugDumpOption::FusionIrOriginal},
+      {"fusion_ir_presched", DebugDumpOption::FusionIrPresched},
+      {"fusion_ir_preseg", DebugDumpOption::FusionIrPreseg},
       {"global_zeroed_memory", DebugDumpOption::GlobalZeroedMemory},
       {"host_ir", DebugDumpOption::HostIr},
       {"index_type", DebugDumpOption::IndexType},

diff --git a/csrc/options.h b/csrc/options.h
@@ -40,6 +40,7 @@ enum class DebugDumpOption {
   // TODO(wujingyue): name the following FusionIrSched
   FusionIr, //!< Dump the Fusion IR before lowering. This is the Fusion IR fed
             //!< to `KernelExecutor::compileFusion`.
+  FusionIrGraph, //!< Dump a GraphViz graph of the Fusion IR
   FusionIrMath, //!< Dump just the compute (math) part of the above `FusionIr`
                 //!< for conciseness
   KernelIr, //!< Dump the compiler Kernel IR

diff --git a/csrc/runtime/executor.cpp b/csrc/runtime/executor.cpp
@@ -16,6 +16,7 @@
 #include <global_allocator.h>
 #include <instrumentation.h>
 #include <ir/all_nodes.h>
+#include <ir/graphviz.h>
 #include <ir/utils.h>
 #include <iter_visitor.h>
 #include <kernel_ir.h>
@@ -304,6 +305,9 @@ void KernelExecutor::compile(
   NVF_ERROR(
       !fusion->outputs().empty(), "No output found for this kernel, aborting.");
 
+  createKernelId(
+      scheduler_type, fusion_id_, concrete_id_, runtime_id_, group_id_);
+
   // TODO: refactor the options_ passed through
   options_.device = c10::Device(c10::DeviceType::CUDA, args.getDeviceIndex());
 
@@ -346,10 +350,21 @@ void KernelExecutor::compile(
     }
   }
 
+  if (isDebugDumpEnabled(DebugDumpOption::FusionIrMath)) {
+    fusion->printMath();
+  }
+
   if (isDebugDumpEnabled(DebugDumpOption::FusionIr)) {
     fusion->print();
-  } else if (isDebugDumpEnabled(DebugDumpOption::FusionIrMath)) {
-    fusion->printMath();
+  }
+
+  if (isDebugDumpEnabled(DebugDumpOption::FusionIrGraph)) {
+    std::stringstream file_name;
+    file_name << "__tmp_fusion_ir_graph_" << kernel_id_ << ".dot";
+    IrGraphGenerator::print(
+        fusion,
+        file_name.str().c_str(),
+        IrGraphGenerator::DetailLevel::ComputeOnly);
   }
 
   //! Force index_type to int and disable magic zero if we detect that the
@@ -418,8 +433,7 @@ void KernelExecutor::compile(
   for (const auto& hook : post_lowering_hooks_) {
     hook(kernel);
   }
-  createKernelId(
-      scheduler_type, fusion_id_, concrete_id_, runtime_id_, group_id_);
+
   setUsedTVs();
 
   if (isDebugDumpEnabled(DebugDumpOption::KernelIr)) {