xref: /aosp_15_r20/external/tensorflow/tensorflow/core/profiler/utils/gpu_event_stats.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1 /* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #ifndef TENSORFLOW_CORE_PROFILER_UTILS_GPU_EVENT_STATS_H_
17 #define TENSORFLOW_CORE_PROFILER_UTILS_GPU_EVENT_STATS_H_
18 
19 #include <cstdint>
20 #include <vector>
21 
22 #include "absl/strings/string_view.h"
23 #include "absl/types/optional.h"
24 #include "tensorflow/core/profiler/utils/xplane_visitor.h"
25 
26 namespace tensorflow {
27 namespace profiler {
28 
29 // Stats from a GPU stream XEvent.
30 struct GpuEventStats {
31   explicit GpuEventStats(const XEventVisitor* event);
32 
IsKernelGpuEventStats33   bool IsKernel() const { return !kernel_details.empty(); }
IsMemCpyGpuEventStats34   bool IsMemCpy() const { return !memcpy_details.empty(); }
35 
IsXlaOpGpuEventStats36   bool IsXlaOp() const { return !hlo_op_names.empty(); }
IsTfOpGpuEventStats37   bool IsTfOp() const { return !tf_op_fullname.empty(); }
38 
39   // Stats from TensorFlow.
40   absl::string_view tf_op_fullname;
41   absl::string_view equation;
42   absl::string_view tensor_shapes;
43 
44   // Stats from XLA.
45   std::vector<absl::string_view> hlo_op_names;
46   absl::string_view hlo_module_name;
47   absl::optional<uint64_t> program_id;
48 
49   // Stats from CUPTI.
50   absl::string_view kernel_details;
51   absl::string_view memcpy_details;
52   absl::optional<int64_t> correlation_id;
53 
54   // Stats derived by grouping.
55   absl::optional<int64_t> group_id;
56   bool is_eager = false;
57 };
58 
59 // Stats for a host-side GPU launch XEvent.
60 struct LaunchEventStats {
61   explicit LaunchEventStats(const XEventVisitor* event);
62 
IsLaunchLaunchEventStats63   bool IsLaunch() const {
64     return device_id.has_value() && correlation_id.has_value();
65   }
66 
67   // Stats from CUPTI.
68   absl::optional<int64_t> device_id;
69   absl::optional<int64_t> correlation_id;
70 
71   // Stat derived by grouping.
72   absl::optional<int64_t> group_id;
73 };
74 
75 }  // namespace profiler
76 }  // namespace tensorflow
77 
78 #endif  // TENSORFLOW_CORE_PROFILER_UTILS_GPU_EVENT_STATS_H_
79