Merge branch 'main' into zejun/kineto_for_pti_0.11

Author: zejun-chen

libkineto/CMakeLists.txt

@@ -141,14 +141,18 @@ endif()
 if (LIBKINETO_NOCUPTI AND LIBKINETO_NOROCTRACER AND LIBKINETO_NOXPUPTI)
     get_filelist("get_libkineto_cpu_only_srcs(with_api=False)" LIBKINETO_SRCS)
     message(INFO " CUPTI unavailable or disabled - not building GPU profilers")
-elseif(NOT LIBKINETO_NOROCTRACER)
-    get_filelist("get_libkineto_roctracer_srcs(with_api=False)" LIBKINETO_SRCS)
+else()
+  if(NOT LIBKINETO_NOROCTRACER)
+    get_filelist("get_libkineto_roctracer_srcs(with_api=False)" LIBKINETO_roc_SRCS)
     message(INFO " Building with roctracer")
-elseif(DEFINED LIBKINETO_NOXPUPTI AND NOT LIBKINETO_NOXPUPTI)
-    get_filelist("get_libkineto_xpupti_srcs(with_api=False)" LIBKINETO_SRCS)
+  elseif(NOT LIBKINETO_NOCUPTI)
+    get_filelist("get_libkineto_cupti_srcs(with_api=False)" LIBKINETO_cuda_SRCS)
+  endif()
+  if(DEFINED LIBKINETO_NOXPUPTI AND NOT LIBKINETO_NOXPUPTI)
+    get_filelist("get_libkineto_xpupti_srcs(with_api=False)" LIBKINETO_xpu_SRCS)
     message(INFO " Building with xpupti")
-else()
-    get_filelist("get_libkineto_cupti_srcs(with_api=False)" LIBKINETO_SRCS)
+  endif()
+  set(LIBKINETO_SRCS ${LIBKINETO_roc_SRCS} ${LIBKINETO_xpu_SRCS} ${LIBKINETO_cuda_SRCS})
 endif()
 get_filelist("get_libkineto_public_headers()" LIBKINETO_PUBLIC_HEADERS)
 get_filelist("get_libkineto_api_srcs()" LIBKINETO_API_SRCS)

libkineto/include/GenericTraceActivity.h

@@ -47,6 +47,10 @@ class GenericTraceActivity : public ITraceActivity {
     return resource;
   }
 
+  void setDevice(int32_t newDevice) {
+    device = newDevice;
+  }
+
   int32_t getThreadId() const override {
     return threadId;
   }

libkineto/include/IActivityProfiler.h

@@ -130,6 +130,10 @@ class IActivityProfilerSession {
     return "";
   }
 
+  virtual std::unordered_map<std::string, std::string> getMetadata() {
+    return {};
+  }
+
  protected:
   TraceStatus status_ = TraceStatus::READY;
 };

libkineto/include/ThreadUtil.h

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include <sys/stat.h>
 #include <cstdint>
 #include <string>
 #include <utility>
@@ -20,6 +21,7 @@ int32_t threadId();
 bool setThreadName(const std::string& name);
 std::string getThreadName();
 
+int32_t pidNamespace(ino_t& ns);
 int32_t processId(bool cache = true);
 std::string processName(int32_t pid);
 

libkineto/sample_programs/build.sh

@@ -18,5 +18,7 @@ g++ \
   -lpthread \
   -lcuda \
   -lcudart \
+  -lcupti \
+  -lnvperf_host \
   /usr/local/lib/libkineto.a \
   kplay_cu.o

libkineto/sample_programs/kineto_playground.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <cuda_runtime.h>
 #include <iostream>
 #include <string>
 
@@ -14,6 +15,20 @@
 // @lint-ignore-every CLANGTIDY facebook-hte-RelativeInclude
 #include "kineto_playground.cuh"
 
+#define CHECK_CUDA(call)               \
+  do {                                 \
+    cudaError_t status = call;         \
+    if (status != cudaSuccess) {       \
+      fprintf(                         \
+          stderr,                      \
+          "CUDA Error at %s:%d: %s\n", \
+          __FILE__,                    \
+          __LINE__,                    \
+          cudaGetErrorString(status)); \
+      exit(1);                         \
+    }                                  \
+  } while (0)
+
 using namespace kineto;
 
 static const std::string kFileName = "/tmp/kineto_playground_trace.json";
@@ -23,6 +38,7 @@ int main() {
   warmup();
 
   // Kineto config
+  libkineto_init(false, true);
 
   // Empty types set defaults to all types
   std::set<libkineto::ActivityType> types;
@@ -38,6 +54,7 @@ int main() {
   profiler.startTrace();
   std::cout << "Start playground" << std::endl;
   playground();
+  CHECK_CUDA(cudaDeviceSynchronize());
 
   std::cout << "Stop Trace" << std::endl;
   auto trace = profiler.stopTrace();

libkineto/sample_programs/kineto_playground.cu

@@ -76,7 +76,7 @@ __global__ void square(float* A, int N) {
 
 void playground(void) {
   // Add your experimental CUDA implementation here.
-  basicMemcpyFromDevice();
+  basicMemcpyToDevice();
   compute();
   basicMemcpyFromDevice();
 }

libkineto/src/AbstractConfig.cpp

@@ -116,8 +116,8 @@ int64_t AbstractConfig::toInt64(const string& val) const {
 }
 
 bool AbstractConfig::toBool(string& val) const {
-  const std::array<string, 8> bool_vals{
-      "n", "y", "no", "yes", "f", "t", "false", "true"};
+  const std::array<string, 10> bool_vals{
+      "n", "y", "no", "yes", "f", "t", "false", "true", "0", "1"};
   const string lower_val = toLower(val);
   for (int i = 0; i < bool_vals.size(); i++) {
     if (lower_val == bool_vals[i]) {

libkineto/src/ApproximateClock.h

@@ -18,8 +18,6 @@
 #include <functional>
 #include <type_traits>
 
-namespace libkineto {
-
 #if defined(__i386__) || defined(__x86_64__) || defined(__amd64__)
 #define KINETO_RDTSC
 #if defined(_MSC_VER)
@@ -42,6 +40,8 @@ namespace libkineto {
 #define KINETO_UNUSED __attribute__((__unused__))
 #endif //_MSC_VER
 
+namespace libkineto {
+
 using time_t = int64_t;
 using steady_clock_t = std::conditional_t<
     std::chrono::high_resolution_clock::is_steady,

libkineto/src/CuptiActivityProfiler.cpp

@@ -300,10 +300,27 @@ void CuptiActivityProfiler::logGpuVersions() {
 #endif
 }
 
+namespace {
+
+const std::unordered_set<std::string>& getLoggerMedataAllowList() {
+  static const std::unordered_set<std::string> kLoggerMedataAllowList{
+      "with_stack", "with_modules", "record_shapes", "profile_memory"};
+  return kLoggerMedataAllowList;
+}
+
+} // namespace
+
 void CuptiActivityProfiler::processTraceInternal(ActivityLogger& logger) {
   LOG(INFO) << "Processing " << traceBuffers_->cpu.size() << " CPU buffers";
   VLOG(0) << "Profile time range: " << captureWindowStartTime_ << " - "
           << captureWindowEndTime_;
+
+  // Pass metadata within the trace to the logger observer.
+  for (const auto& pair : metadata_) {
+    if (getLoggerMedataAllowList().count(pair.first) > 0) {
+      LOGGER_OBSERVER_ADD_METADATA(pair.first, pair.second);
+    }
+  }
   for (auto& pair : versionMetadata_) {
     addMetadata(pair.first, pair.second);
   }
@@ -319,6 +336,9 @@ void CuptiActivityProfiler::processTraceInternal(ActivityLogger& logger) {
         device_properties.push_back(props);
       }
     }
+    for (const auto& [key, value] : session->getMetadata()) {
+      addMetadata(key, value);
+    }
   }
   logger.handleTraceStart(
       metadata_, fmt::format("{}", fmt::join(device_properties, ",")));
@@ -429,7 +449,7 @@ void CuptiActivityProfiler::processCpuTrace(
     return;
   }
   setCpuActivityPresent(true);
-
+  bool warn_once = false;
   CpuGpuSpanPair& span_pair =
       recordTraceSpan(cpuTrace.span, cpuTrace.gpuOpCount);
   TraceSpan& cpu_span = span_pair.first;
@@ -442,18 +462,23 @@ void CuptiActivityProfiler::processCpuTrace(
               const std::unique_ptr<GenericTraceActivity>>::value,
           "handleActivity is unsafe and relies on the caller to maintain not "
           "only lifetime but also address stability.");
-      if (act->type() == ActivityType::USER_ANNOTATION &&
-          act->duration() <= 0) {
+      if (act->duration() < 0) {
         act->endTime = captureWindowEndTime_;
         act->addMetadata("finished", "false");
-      } else {
-        act->addMetadata("finished", "true");
       }
       logger.handleActivity(*act);
     }
     clientActivityTraceMap_[act->correlationId()] = &span_pair;
     activityMap_[act->correlationId()] = act.get();
-
+    if (act->deviceId() == 0) {
+      if (!warn_once) {
+        LOG(WARNING)
+            << "CPU activity with pid 0 detected. This is likely due to the python stack"
+               " tracer not being able to determine the pid for an event. Overriding pid to main thread pid";
+      }
+      act->setDevice(processId());
+      warn_once = true;
+    }
     recordThreadInfo(act->resourceId(), act->getThreadId(), act->deviceId());
   }
   logger.handleTraceSpan(cpu_span);
@@ -564,8 +589,8 @@ inline static bool isBlockListedRuntimeCbid(CUpti_CallbackId cbid) {
   if (cbid == CUPTI_RUNTIME_TRACE_CBID_cudaGetDevice_v3020 ||
       cbid == CUPTI_RUNTIME_TRACE_CBID_cudaSetDevice_v3020 ||
       cbid == CUPTI_RUNTIME_TRACE_CBID_cudaGetLastError_v3020 ||
-      // Support cudaEventRecord and cudaEventSynchronize, revisit if others are
-      // needed
+      // Support cudaEventRecord and cudaEventSynchronize, revisit if others
+      // are needed
       cbid == CUPTI_RUNTIME_TRACE_CBID_cudaEventCreate_v3020 ||
       cbid == CUPTI_RUNTIME_TRACE_CBID_cudaEventCreateWithFlags_v3020 ||
       cbid == CUPTI_RUNTIME_TRACE_CBID_cudaEventDestroy_v3020) {
@@ -1322,20 +1347,21 @@ const time_point<system_clock> CuptiActivityProfiler::performRunLoopStep(
           || cupti_.stopCollection
 #endif // HAS_CUPTI || HAS_ROCTRACER
       ) {
-        // Update runloop state first to prevent further updates to shared state
+        // Update runloop state first to prevent further updates to shared
+        // state
         LOG(INFO) << "Tracing complete.";
         VLOG_IF(1, currentIter >= 0)
             << "This state change was invoked by application's step() call";
 
         // currentIter >= 0 means this is called from the step() api of
-        // the profile in pytorch main thread, it should be executed in another
-        // thread in case pytorch main thread is blocked
+        // the profile in pytorch main thread, it should be executed in
+        // another thread in case pytorch main thread is blocked
         if (currentIter >= 0) {
           // if collectTraceThread_ is already running, there's no need to
           // execute collectTrace twice.
-          // Do not call collectTrace when profilerThread_ is collecting Trace.
-          // Otherwise, libkineto::api().client()->stop will be called twice,
-          // which leads to an unrecoverable ::c10:Error at
+          // Do not call collectTrace when profilerThread_ is collecting
+          // Trace. Otherwise, libkineto::api().client()->stop will be called
+          // twice, which leads to an unrecoverable ::c10:Error at
           // disableProfiler
           if (!collectTraceThread_ && !getCollectTraceState()) {
             std::lock_guard<std::recursive_mutex> guard(mutex_);
@@ -1404,13 +1430,33 @@ void CuptiActivityProfiler::finalizeTrace(
     iterationCountMap_.clear();
   }
 
+  // Thread & stream info
+  for (const auto& pair : resourceInfo_) {
+    const auto& resource = pair.second;
+    logger.handleResourceInfo(resource, captureWindowStartTime_);
+  }
+
+  bool use_default_device_info = true;
+  for (auto& session : sessions_) {
+    auto device_info = session->getDeviceInfo();
+    if (device_info != nullptr) {
+      use_default_device_info = false;
+      logger.handleDeviceInfo(*device_info, captureWindowStartTime_);
+    }
+
+    auto resource_infos = session->getResourceInfos();
+    for (const auto& resource_info : resource_infos) {
+      logger.handleResourceInfo(resource_info, captureWindowStartTime_);
+    }
+  }
+
   // Process names
   int32_t pid = processId();
   string process_name = processName(pid);
   if (!process_name.empty()) {
     logger.handleDeviceInfo(
         {pid, pid, process_name, "CPU"}, captureWindowStartTime_);
-    if (!cpuOnly_) {
+    if (!cpuOnly_ && use_default_device_info) {
       // Usually, GPU events use device id as pid (0-7).
       // In some cases, CPU sockets are numbered starting from 0.
       // In the worst case, 8 CPU sockets + 8 GPUs, so the max GPU ID is 15.
@@ -1428,24 +1474,6 @@ void CuptiActivityProfiler::finalizeTrace(
     }
   }
 
-  // Thread & stream info
-  for (auto pair : resourceInfo_) {
-    const auto& resource = pair.second;
-    logger.handleResourceInfo(resource, captureWindowStartTime_);
-  }
-
-  for (auto& session : sessions_) {
-    auto device_info = session->getDeviceInfo();
-    if (device_info != nullptr) {
-      logger.handleDeviceInfo(*device_info, captureWindowStartTime_);
-    }
-
-    auto resource_infos = session->getResourceInfos();
-    for (auto resource_info : resource_infos) {
-      logger.handleResourceInfo(resource_info, captureWindowStartTime_);
-    }
-  }
-
   for (const auto& iterations : traceSpans_) {
     for (const auto& span_pair : iterations.second) {
       const TraceSpan& gpu_span = span_pair.second;