Add a few FUSER_PERF_SCOPEs for isResharding. (#3464)

This is to help track down the host latency regression found in #3421.
NVIDIA · Nov 24, 2024 · ad3aa3f · ad3aa3f
1 parent 75bf6af
commit ad3aa3f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/csrc/multidevice/utils.cpp b/csrc/multidevice/utils.cpp
@@ -7,6 +7,8 @@
 // clang-format on
 
 #include <device_lower/utils.h>
+#include <id_model/id_model.h>
+#include <instrumentation.h>
 #include <ir/internal_base_nodes.h>
 #include <ir/iostream.h>
 #include <ir/utils.h>
@@ -176,6 +178,8 @@ bool haveDifferentShardings(
 }
 
 bool isResharding(const Expr* expr) {
+  FUSER_PERF_SCOPE("isResharding");
+
   if (!ir_utils::isTvOp(expr)) {
     return false;
   }

diff --git a/csrc/preseg_passes/optimization_pass.h b/csrc/preseg_passes/optimization_pass.h
@@ -8,6 +8,7 @@
 #pragma once
 #include <debug.h>
 #include <exceptions.h>
+#include <instrumentation.h>
 #include <ir/interface_nodes.h>
 #include <ir/utils.h>
 #include <options.h>
@@ -51,13 +52,17 @@ class OptimizationPass {
     if (!flag_.load()) {
       return;
     }
+
+    FUSER_PERF_SCOPE(DerivedClass::name().c_str());
     DerivedClass::runPass(fusion);
+
     // TODO: skip the logging of the pass where the fusion has not been changed.
     if (isDebugDumpEnabled(DebugDumpOption::PreSegmenterLogging)) {
       debug() << "Fusion after pass: " << DerivedClass::name() << std::endl;
       fusion->printMath();
       debug() << "========================================" << std::endl;
     }
+
 #ifndef NDEBUG
     // cycle detection is only enabled on debug run
     NVF_ERROR(