Initial Cupti PM sampling support for gpu profiler #24406

cashton-nvidia · 2025-03-31T16:35:47Z

Adds perf metric sampling to the Cupti gpu tracer.

Adds a unit test which validates expected number of flops are sampled.

google-cla · 2025-03-31T16:35:52Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

rengoog

Thanks @cashton-nvidia ! 2 more nits: other than that LGTM.

.clang-format

…-sampling' into add-cupti-pm-sampling

xla/backends/profiler/gpu/cupti_tracer.cc

allanrenucci · 2025-04-03T13:48:43Z

xla/backends/profiler/gpu/cupti_interface.h

+// CUPTI PM sampling headers added in 12.6
+// Note - bug fix for PM sampling added after 12.6.2
+#if CUDA_VERSION >= 12060
+#define CUPTI_PM_SAMPLING 1
+#else
+#define CUPTI_PM_SAMPLING 0
+#endif


#if, #ifdef, or #ifndef macros are banned in XLA unless they cannot be avoided. Do you think you can refactor the code to avoid preprocessor conditionals (or at least significantly reduce the number of usages)? E.g.

Return a absl::Unimplemented errors (or equivalent)

Skip tests. I.e. if (!CUPTI_PM_SAMPLING) { GTEST_SKIP() }

Are #define macros banned as well? Or just the #if test?

Do you have an example of how XLA currently handles versioning in external dependencies where APIs / external headers & symbols are possibly not present, and does this without macros?

If not, first thoughts -
I'd imagine most / all of the #if cases in the tracer can be replaced with unimplemented errors.

The interface though seems more problematic - the new headers need to be avoided. Since the preprocessor runs first, I'm not aware of any tricks except #if to avoid those includes.

Beyond including the headers, there's also the issue of the wrapper itself. I could use #if #else to have alternate wrappers for these new possibly missing functions, or I could probably write some weak symbol replacements for the missing symbols. I might be able to get the compiler to completely eliminate the references to the missing symbols and not need this, but I think that would rely on the compiler optimizations more than we can guarantee. All of those options are more complicated to maintain than what's there currently.

I've done what I can - couldn't completely remove #if in a couple of cases - I needed it to include the headers, and to define a couple of types before use, but everywhere else I used compiler elimination of dead code to avoid errors due to missing functions / types. Can you take a look?

Are #define macros banned as well? Or just the #if test?

Just #if tests.

I needed it to include the headers, and to define a couple of types before use, but everywhere else I used compiler elimination of dead code to avoid errors due to missing functions / types. Can you take a look?

Right, guarding new headers can't be avoided. Overall LGTM.

allanrenucci · 2025-04-03T14:18:02Z

xla/backends/profiler/gpu/profile_with_cuda_kernels.h

@@ -0,0 +1,29 @@
+/* Copyright 2024 The OpenXLA Authors.


Please add missing #ifndef header guard.

Ahh is pragma once not sufficient?

This is part of Google C++ style guide: https://google.github.io/styleguide/cppguide.html#The__define_Guard

It triggers a blocking LINT error internally. Please remove #pragma once in favor of #ifndef header guard

…to add-cupti-pm-sampling

allanrenucci · 2025-04-04T09:19:36Z

xla/backends/profiler/gpu/profile_with_cuda_kernels.h

@@ -0,0 +1,29 @@
+/* Copyright 2024 The OpenXLA Authors.


This is part of Google C++ style guide: https://google.github.io/styleguide/cppguide.html#The__define_Guard

It triggers a blocking LINT error internally. Please remove #pragma once in favor of #ifndef header guard

allanrenucci · 2025-04-04T09:23:59Z

xla/backends/profiler/gpu/profile_with_cuda_kernels_test.cc

+
+using xla::profiler::CuptiInterface;
+using xla::profiler::CuptiTracer;
+using xla::profiler::CuptiTracerCollectorOptions;
+using xla::profiler::CuptiTracerOptions;
+using xla::profiler::CuptiWrapper;
+
+// Needed to create different cupti tracer for each test cases.
+class TestableCuptiTracer : public CuptiTracer {
+ public:
+  explicit TestableCuptiTracer()
+      : CuptiTracer(new CuptiErrorManager(std::make_unique<CuptiWrapper>())) {}
+};
+
+
+namespace {


Can this be moved into the anonymous namespace?

Suggested change

using xla::profiler::CuptiInterface;

using xla::profiler::CuptiTracer;

using xla::profiler::CuptiTracerCollectorOptions;

using xla::profiler::CuptiTracerOptions;

using xla::profiler::CuptiWrapper;

// Needed to create different cupti tracer for each test cases.

class TestableCuptiTracer : public CuptiTracer {

public:

explicit TestableCuptiTracer()

: CuptiTracer(new CuptiErrorManager(std::make_unique<CuptiWrapper>())) {}

};

namespace {

namespace {

using xla::profiler::CuptiInterface;

using xla::profiler::CuptiTracer;

using xla::profiler::CuptiTracerCollectorOptions;

using xla::profiler::CuptiTracerOptions;

using xla::profiler::CuptiWrapper;

// Needed to create different cupti tracer for each test cases.

class TestableCuptiTracer : public CuptiTracer {

public:

explicit TestableCuptiTracer()

: CuptiTracer(new CuptiErrorManager(std::make_unique<CuptiWrapper>())) {}

};

Thanks for the style guide link. I'd looked through xla source for a coding conventions reference and didn't find anything, though some sub parts of it referenced other conventions. (I'm surrpised pragma once is forbidden since every compiler I'm aware of supports it, but I know it's technically non-standard.) Will make these changes.

allanrenucci · 2025-04-04T09:27:07Z

xla/backends/profiler/gpu/mock_cupti.h

@@ -91,8 +91,113 @@ class MockCupti : public xla::profiler::CuptiInterface {
  MOCK_METHOD(CUptiResult, GetGraphExecId,
              (CUgraphExec graph_exec, uint32_t* graph_id), (override));

+#if CUPTI_PM_SAMPLING


can we remove this one?

I missed this - I'm surprised it worked. Tested by building with CUPTI_PM_SAMPLING forced to both 0 and 1, and it worked. Must not be any live dependencies on the mock code? Fixed.

allanrenucci · 2025-04-04T09:29:29Z

xla/backends/profiler/gpu/cupti_tracer.cc

+    if (VLOG_IS_ON(2)) {
+      LOG(INFO) << "(Profiling::PM Sampling) Top of decode loop";
+    }


Suggested change

if (VLOG_IS_ON(2)) {

LOG(INFO) << "(Profiling::PM Sampling) Top of decode loop";

}

VLOG(2) << "(Profiling::PM Sampling) Top of decode loop";

I wasn't clear on what vlog resolved to - if it's equivalent to log(info), I'll switch all these to the simpler code where possible

Yes, it is equivalent to LOG(INFO). I'd say you only want to use if (VLOG_IS_ON(...)) to guard against evaluating an expensive statement. E.g.

if (VLOG_IS_ON(1)) { Foo foo = Expensive(); LOG(INFO) << foo; }

But you can often rewrite as:

VLOG(1) << Expensive();

which is equivalent.

allanrenucci · 2025-04-04T09:29:54Z

xla/backends/profiler/gpu/cupti_tracer.cc

+      if (VLOG_IS_ON(2)) {
+        LOG(INFO) << "(Profiling::PM Sampling)  Beginning decode for device " <<
+            dev->device_id_;
+      }


Suggested change

if (VLOG_IS_ON(2)) {

LOG(INFO) << "(Profiling::PM Sampling) Beginning decode for device " <<

dev->device_id_;

}

VLOG(2) << "(Profiling::PM Sampling) Beginning decode for device " <<

dev->device_id_;

allanrenucci · 2025-04-04T10:18:40Z

xla/backends/profiler/gpu/cupti_tracer.cc

+        "restart session";
+  }
+  if (ret != CUPTI_SUCCESS) {
+    return tsl::errors::Internal("CUPTI error during "


Please replace all tsl::errors:: usages with absl:: equivalents. E.g. absl::InternalError.

allanrenucci · 2025-04-04T10:22:34Z

xla/backends/profiler/gpu/cupti_tracer.h

+class PmSamplingDevice {
+  // Internal state
+  // Declared roughly in order of initialization
+  char const* chipName_ = nullptr;


std::string?

The CUPTI APIs use char *, so this avoids converting to / from strings. This chip name is produced by CUPTI and then later passed to other APIs. It's not used by other XLA / xprof code. Do we still need to convert it?

I am asking because it looks there is a memory leak.

We have chipName_ = strdup(p.pChipName) which IIUC creates a copy of p.pChipName. Should the copy be be freed when PmSamplingDevice is destroyed? If chipName_ is stored as a std::string it will be freed automatically when PmSamplingDevice is destroyed.

This chip name is produced by CUPTI and then later passed to other APIs.

Note you can just call std::string::data to get back a const char* when calling other APIs.

Ahh, true. The string size would be small, ~6 bytes, but would be lost in this case. Switching seemed more complex but w/ it being so simple I hadn't considered the leak.

allanrenucci · 2025-04-04T10:25:57Z

xla/backends/profiler/gpu/profile_with_cuda_kernels_test.cc

+    auto start_walltime_ns = absl::GetCurrentTimeNanos();
+    auto start_gputime_ns = CuptiTracer::GetTimestamp();


Please don't use auto here. It only saves a few characters.

https://google.github.io/styleguide/cppguide.html#Type_deduction

allanrenucci · 2025-04-04T10:27:56Z

xla/backends/profiler/gpu/profile_with_cuda_kernels_test.cc

+std::atomic_uint64_t atomic_total_fp64 = 0;
+std::atomic_uint64_t atomic_total_read = 0;
+std::atomic_uint64_t atomic_total_write = 0;
+bool skip_first = true;


Should this be std::atomic_bool?

allanrenucci · 2025-04-04T10:33:34Z

xla/backends/profiler/gpu/cupti_interface.h

+// CUPTI PM sampling headers added in 12.6
+// Note - bug fix for PM sampling added after 12.6.2
+#if CUDA_VERSION >= 12060
+#define CUPTI_PM_SAMPLING 1
+#else
+#define CUPTI_PM_SAMPLING 0
+#endif


Are #define macros banned as well? Or just the #if test?

Just #if tests.

I needed it to include the headers, and to define a couple of types before use, but everywhere else I used compiler elimination of dead code to avoid errors due to missing functions / types. Can you take a look?

Right, guarding new headers can't be avoided. Overall LGTM.

rengoog approved these changes Mar 31, 2025

View reviewed changes

.clang-format Outdated Show resolved Hide resolved

cashton-nvidia marked this pull request as draft March 31, 2025 21:05

Initial Cupti PM sampling support for gpu profiler

f2e570b

cashton-nvidia force-pushed the add-cupti-pm-sampling branch from 19630ba to f2e570b Compare April 1, 2025 15:56

Merge branch 'openxla:main' into add-cupti-pm-sampling

90d2785

cashton-nvidia marked this pull request as ready for review April 1, 2025 15:59

cashton-nvidia added 2 commits April 1, 2025 15:00

Add test for driver version 12.6 for PM sampling support

4cf990a

Merge remote-tracking branch 'refs/remotes/upstream-fork/add-cupti-pm…

4e3e795

…-sampling' into add-cupti-pm-sampling

cashton-nvidia marked this pull request as draft April 2, 2025 17:40

cashton-nvidia marked this pull request as ready for review April 2, 2025 17:40

cashton-nvidia requested a review from rengoog April 2, 2025 17:40

Merge branch 'openxla:main' into add-cupti-pm-sampling

528ecf4

rengoog approved these changes Apr 2, 2025

View reviewed changes

jiayongle approved these changes Apr 2, 2025

View reviewed changes

allanrenucci suggested changes Apr 3, 2025

View reviewed changes

allanrenucci reviewed Apr 3, 2025

View reviewed changes

cashton-nvidia added 2 commits April 3, 2025 16:15

Switch how CUPTI versioning is handled, remove #if whereever possible

77b55bf

Merge remote-tracking branch 'upstream-fork/add-cupti-pm-sampling' in…

8bb420f

…to add-cupti-pm-sampling

allanrenucci suggested changes Apr 4, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Initial Cupti PM sampling support for gpu profiler #24406

Initial Cupti PM sampling support for gpu profiler #24406

cashton-nvidia commented Mar 31, 2025

google-cla bot commented Mar 31, 2025

rengoog left a comment

allanrenucci Apr 3, 2025

cashton-nvidia Apr 3, 2025

cashton-nvidia Apr 3, 2025

allanrenucci Apr 4, 2025

allanrenucci Apr 3, 2025

cashton-nvidia Apr 3, 2025

allanrenucci Apr 4, 2025

allanrenucci Apr 4, 2025

allanrenucci Apr 4, 2025

cashton-nvidia Apr 4, 2025

allanrenucci Apr 4, 2025

cashton-nvidia Apr 4, 2025

allanrenucci Apr 4, 2025

cashton-nvidia Apr 4, 2025

allanrenucci Apr 4, 2025

allanrenucci Apr 4, 2025

allanrenucci Apr 4, 2025

allanrenucci Apr 4, 2025

cashton-nvidia Apr 4, 2025

allanrenucci Apr 4, 2025

cashton-nvidia Apr 4, 2025

allanrenucci Apr 4, 2025

allanrenucci Apr 4, 2025

allanrenucci Apr 4, 2025

		auto start_walltime_ns = absl::GetCurrentTimeNanos();
		auto start_gputime_ns = CuptiTracer::GetTimestamp();

Initial Cupti PM sampling support for gpu profiler #24406

Are you sure you want to change the base?

Initial Cupti PM sampling support for gpu profiler #24406

Conversation

cashton-nvidia commented Mar 31, 2025

google-cla bot commented Mar 31, 2025

rengoog left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment