From ad3aa3f0e50a05f36c5484054860a6a9e13afbf4 Mon Sep 17 00:00:00 2001
From: Jingyue Wu <wujingyue@gmail.com>
Date: Sat, 23 Nov 2024 23:20:59 -0800
Subject: [PATCH] Add a few FUSER_PERF_SCOPEs for isResharding. (#3464)

This is to help track down the host latency regression found in #3421.
---
 csrc/multidevice/utils.cpp             | 4 ++++
 csrc/preseg_passes/optimization_pass.h | 5 +++++
 2 files changed, 9 insertions(+)

diff --git a/csrc/multidevice/utils.cpp b/csrc/multidevice/utils.cpp
index 0c607881e0b..437ea1ba4ae 100644
--- 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
index 78a87a7e3ce..079e5015f44 100644
--- 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(