Add a few FUSER_PERF_SCOPEs for isResharding.

This is to help track down the host latency regression found in #3421.
NVIDIA · Nov 23, 2024 · 43630ac · 43630ac
1 parent e57a755
commit 43630ac
Show file tree

Hide file tree

Showing 2 changed files with 6 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,6 +52,7 @@ class OptimizationPass {
     if (!flag_.load()) {
       return;
     }
+    FUSER_PERF_SCOPE(DerivedClass::name());
     DerivedClass::runPass(fusion);
     // TODO: skip the logging of the pass where the fusion has not been changed.
     if (isDebugDumpEnabled(DebugDumpOption::PreSegmenterLogging)) {