[go: nahoru, domu]
Skip to content

Commit

Permalink
XProf GPU: Cache CUPTI activity buffer during tracing, process them a…
Browse files Browse the repository at this point in the history 
…fter tracing is stopped.

PiperOrigin-RevId: 631106601
  • Loading branch information
@tensorflower-gardener
tensorflower-gardener committed May 6, 2024
1 parent fedcb89 commit c8aba81
Show file tree
Hide file tree
Showing 7 changed files with 939 additions and 669 deletions.
36 changes: 25 additions & 11 deletions third_party/xla/xla/backends/profiler/gpu/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ tsl_gpu_library(
srcs = tf_additional_device_tracer_srcs(),
copts = tf_profiler_copts() + tsl_copts(),
cuda_deps = [
":cupti_buffer_events",
":cupti_collector",
":cupti_tracer",
":cupti_wrapper",
Expand Down Expand Up @@ -172,15 +173,13 @@ tsl_gpu_library(
copts = tf_profiler_copts() + tsl_copts(),
visibility = ["//visibility:public"],
deps = [
":cupti_buffer_events",
":cupti_collector",
":cupti_interface",
":cupti_utils",
":nvtx_utils",
"@com_google_absl//absl/cleanup",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/container:node_hash_map",
"@com_google_absl//absl/container:node_hash_set",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:optional",
"@local_tsl//tsl/platform:env",
Expand All @@ -190,8 +189,6 @@ tsl_gpu_library(
"@local_tsl//tsl/platform:platform_port",
"@local_tsl//tsl/platform:types",
"@local_tsl//tsl/profiler/backends/cpu:annotation_stack",
"@local_tsl//tsl/profiler/lib:scoped_annotation",
"@local_tsl//tsl/profiler/utils:buffer_pool",
],
)

Expand All @@ -217,7 +214,6 @@ tsl_gpu_library(
"@local_tsl//tsl/platform:types",
"@local_tsl//tsl/profiler/backends/cpu:annotation_stack",
"@local_tsl//tsl/profiler/lib:scoped_annotation",
"@local_tsl//tsl/profiler/utils:buffer_pool",
],
)

Expand Down Expand Up @@ -296,18 +292,15 @@ tsl_gpu_library(
copts = tf_profiler_copts() + tsl_copts(),
visibility = ["//visibility:public"],
deps = [
"@com_google_absl//absl/container:fixed_array",
":cupti_buffer_events",
":cupti_interface",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/container:node_hash_set",
"@com_google_absl//absl/hash",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@local_tsl//tsl/platform:abi",
"@local_tsl//tsl/platform:macros",
"@local_tsl//tsl/platform:mutex",
"@local_tsl//tsl/platform:platform_port",
"@local_tsl//tsl/platform:types",
"@local_tsl//tsl/profiler/protobuf:xplane_proto_cc",
"@local_tsl//tsl/profiler/utils:parse_annotation",
"@local_tsl//tsl/profiler/utils:trace_utils",
Expand All @@ -317,6 +310,27 @@ tsl_gpu_library(
] + if_cuda(["//xla/tsl/cuda:cupti"]),
)

tsl_gpu_library(
name = "cupti_buffer_events",
srcs = if_cuda(["cupti_buffer_events.cc"]),
hdrs = ["cupti_buffer_events.h"],
copts = tf_profiler_copts() + tsl_copts(),
visibility = ["//visibility:public"],
deps = [
":cupti_interface",
"@com_google_absl//absl/container:fixed_array",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:node_hash_set",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@local_tsl//tsl/platform:errors",
"@local_tsl//tsl/platform:mutex",
"@local_tsl//tsl/platform:platform_port",
"@local_tsl//tsl/platform:thread_annotations",
"@local_tsl//tsl/profiler/utils:buffer_pool",
] + if_cuda(["//xla/tsl/cuda:cupti"]),
)

tsl_gpu_library(
name = "cupti_utils",
srcs = if_cuda(["cupti_utils.cc"]),
Expand Down
Loading

0 comments on commit c8aba81

Please sign in to comment.