Re-land with fix -> D55262038: Support CUPTI Range Profiler for H100"…

… for one test failure (#936) Summary: Pull Request resolved: #936 Reviewed By: briancoutinho Differential Revision: D57396669 fbshipit-source-id: 4ab43b6cd2d1c54f80b8f247fd8e55e04a191eec
pytorch · May 23, 2024 · b7b1fbe · b7b1fbe
1 parent f29b01d
commit b7b1fbe
Show file tree

Hide file tree

Showing 10 changed files with 312 additions and 197 deletions.
diff --git a/libkineto/sample_programs/kineto_cupti_profiler.cpp b/libkineto/sample_programs/kineto_cupti_profiler.cpp
@@ -30,7 +30,7 @@ int main() {
     libkineto::ActivityType::CUDA_PROFILER_RANGE,
   };
 
-  libkineto_init(false, true);
+  //libkineto_init(false, true);
   libkineto::api().initProfilerIfRegistered();
 
   // Use a special kineto__cuda_core_flop metric that counts individual

diff --git a/libkineto/sample_programs/kineto_playground.cpp b/libkineto/sample_programs/kineto_playground.cpp
@@ -21,23 +21,27 @@ using namespace kineto;
 static const std::string kFileName = "/tmp/kineto_playground_trace.json";
 
 int main() {
+  std::cout << "Warm up " << std::endl;
   warmup();
 
   // Kineto config
 
   // Empty types set defaults to all types
   std::set<libkineto::ActivityType> types;
-  libkineto_init(false, true);
-  libkineto::api().initProfilerIfRegistered();
 
   auto& profiler = libkineto::api().activityProfiler();
+  libkineto::api().initProfilerIfRegistered();
   profiler.prepareTrace(types);
 
   // Good to warm up after prepareTrace to get cupti initialization to settle
+  std::cout << "Warm up " << std::endl;
   warmup();
+  std::cout << "Start trace" << std::endl;
   profiler.startTrace();
+  std::cout << "Start playground" << std::endl;
   playground();
 
+  std::cout << "Stop Trace" << std::endl;
   auto trace = profiler.stopTrace();
   std::cout << "Stopped and processed trace. Got " << trace->activities()->size() << " activities.";
   trace->save(kFileName);

diff --git a/libkineto/sample_programs/kineto_playground.cu b/libkineto/sample_programs/kineto_playground.cu
@@ -76,6 +76,9 @@ __global__ void square(float* A, int N) {
 
 void playground(void) {
   // Add your experimental CUDA implementation here.
+  basicMemcpyFromDevice();
+  compute();
+  basicMemcpyFromDevice();
 }
 
 void compute(void) {

diff --git a/libkineto/src/CuptiActivityProfiler.cpp b/libkineto/src/CuptiActivityProfiler.cpp
@@ -903,16 +903,21 @@ void CuptiActivityProfiler::configureChildProfilers() {
           derivedConfig_->profileStartTime().time_since_epoch())
           .count();
   for (auto& profiler : profilers_) {
-    LOG(INFO) << "Evaluating whether to run child profiler = " << profiler->name();
+    LOG(INFO) << "[Profiler = " << profiler->name() << "] "
+              << "Evaluating whether to run child profiler." ;
     auto session = profiler->configure(
         start_time_ms,
         derivedConfig_->profileDuration().count(),
         derivedConfig_->profileActivityTypes(),
         *config_);
     if (session) {
-      LOG(INFO) << "Running child profiler " << profiler->name() << " for "
+      LOG(INFO) << "[Profiler = " << profiler->name() << "] "
+                << "Running child profiler " << profiler->name() << " for "
                 << derivedConfig_->profileDuration().count() << " ms";
       sessions_.push_back(std::move(session));
+    } else {
+      LOG(INFO) << "[Profiler = " << profiler->name() << "] "
+                << "Not running child profiler.";
     }
   }
 }

diff --git a/libkineto/src/CuptiCallbackApi.cpp b/libkineto/src/CuptiCallbackApi.cpp
@@ -75,7 +75,7 @@ void CuptiCallbackApi::__callback_switchboard(
    CUpti_CallbackDomain domain,
    CUpti_CallbackId cbid,
    const CUpti_CallbackData* cbInfo) {
-  VLOG(0) << "Callback: domain = " << domain << ", cbid = " << cbid;
+  LOG(INFO) << "Callback: domain = " << domain << ", cbid = " << cbid;
   CallbackList *cblist = nullptr;
 
   switch (domain) {
@@ -87,6 +87,12 @@ void CuptiCallbackApi::__callback_switchboard(
           cblist = &callbacks_.runtime[
             CUDA_LAUNCH_KERNEL - __RUNTIME_CB_DOMAIN_START];
           break;
+#if defined(CUDA_VERSION) && (CUDA_VERSION >= 11080)
+        case CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernelExC_v11060:
+          cblist = &callbacks_.runtime[
+            CUDA_LAUNCH_KERNEL_EXC - __RUNTIME_CB_DOMAIN_START];
+          break;
+#endif
         default:
           break;
       }
@@ -168,9 +174,12 @@ void CuptiCallbackApi::initCallbackApi() {
     cuptiSubscribe(&subscriber_,
       (CUpti_CallbackFunc)callback_switchboard,
       nullptr));
-  if (lastCuptiStatus_ != CUPTI_SUCCESS) {
-    VLOG(1)  << "Failed cuptiSubscribe, status: " << lastCuptiStatus_;
-  }
+
+  // TODO: Remove temporarily to work around static initialization order issue
+  // betweent this and GLOG.
+  // if (lastCuptiStatus_ != CUPTI_SUCCESS) {
+  //   LOG(INFO) << "Failed cuptiSubscribe, status: " << lastCuptiStatus_;
+  // }
 
   initSuccess_ = (lastCuptiStatus_ == CUPTI_SUCCESS);
 #endif

diff --git a/libkineto/src/CuptiCallbackApi.h b/libkineto/src/CuptiCallbackApi.h
@@ -28,6 +28,7 @@ namespace KINETO_NAMESPACE {
 using namespace libkineto;
 
 
+
 /* CuptiCallbackApi : Provides an abstraction over CUPTI callback
  *  interface. This enables various callback functions to be registered
  *  with this class. The class registers a global callback handler that
@@ -54,6 +55,7 @@ class CuptiCallbackApi {
     // can possibly support more callback ids per domain
     //
     __RUNTIME_CB_DOMAIN_START = CUDA_LAUNCH_KERNEL,
+    CUDA_LAUNCH_KERNEL_EXC,  // Used in H100
 
     // Callbacks under Resource CB domain
     RESOURCE_CONTEXT_CREATED,

diff --git a/libkineto/src/CuptiCallbackApiMock.h b/libkineto/src/CuptiCallbackApiMock.h
@@ -17,6 +17,10 @@ enum CUpti_CallbackDomain {
 };
 enum CUpti_CallbackId {
   CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernel_v7000,
+
+#if defined(CUDA_VERSION) && (CUDA_VERSION >= 11080)
+  CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernelExC_v11060,
+#endif
   CUPTI_CBID_RESOURCE_CONTEXT_CREATED,
   CUPTI_CBID_RESOURCE_CONTEXT_DESTROY_STARTING,
 };