This CL adds a filter mask to BFCAllocator::AddTraceMe so we can coll…

…ect only memory TraceMe events. PiperOrigin-RevId: 700477323
openxla · Dec 3, 2024 · 698bb37 · 698bb37
1 parent 9c668fd
commit 698bb37
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/xla/tsl/framework/BUILD b/xla/tsl/framework/BUILD
@@ -196,6 +196,7 @@ cc_library(
         ":metrics",
         ":shared_counter",
         "//xla/tsl/lib/core:bits",
+        "//xla/tsl/profiler/utils:trace_filter_utils",
         "//xla/tsl/protobuf:bfc_memory_map_proto_cc",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -205,10 +206,8 @@ cc_library(
         "@com_google_absl//absl/types:optional",
         "@tsl//tsl/platform:env",
         "@tsl//tsl/platform:logging",
-        "@tsl//tsl/platform:macros",
         "@tsl//tsl/platform:numbers",
         "@tsl//tsl/platform:stacktrace",
-        "@tsl//tsl/platform:str_util",
         "@tsl//tsl/platform:strcat",
         "@tsl//tsl/platform:types",
         "@tsl//tsl/profiler/lib:scoped_memory_debug_annotation",

diff --git a/xla/tsl/framework/bfc_allocator.cc b/xla/tsl/framework/bfc_allocator.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "absl/synchronization/mutex.h"
 #include "xla/tsl/framework/allocator.h"
 #include "xla/tsl/framework/allocator_retry.h"
+#include "xla/tsl/profiler/utils/trace_filter_utils.h"
 #include "xla/tsl/protobuf/bfc_memory_map.pb.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/file_system.h"
@@ -47,6 +48,9 @@ limitations under the License.
 
 namespace tsl {
 
+const uint64_t kDefaultMemoryFilterMask = tsl::profiler::TraceMeFiltersToMask(
+    {tsl::profiler::TraceMeFilter::kTraceMemory});
+
 constexpr BFCAllocator::ChunkHandle BFCAllocator::kInvalidChunkHandle;
 
 BFCAllocator::BFCAllocator(std::unique_ptr<SubAllocator> sub_allocator,
@@ -563,7 +567,8 @@ void BFCAllocator::AddTraceMe(absl::string_view traceme_name,
                                {"data_type", annotation.pending_data_type},
                                {"shape", annotation.pending_shape_func()}});
           },
-      /*level=*/tsl::profiler::TraceMeLevel::kInfo);
+      /*level=*/tsl::profiler::TraceMeLevel::kInfo,
+      /*filter_mask=*/kDefaultMemoryFilterMask);
 }
 
 void* BFCAllocator::FindChunkPtr(BinNum bin_num, size_t rounded_bytes,

diff --git a/xla/tsl/profiler/utils/BUILD b/xla/tsl/profiler/utils/BUILD
@@ -510,6 +510,7 @@ filegroup(
     name = "mobile_srcs_no_runtime",
     srcs = [
         "no_init.h",
+        "trace_filter_utils.h",
     ],
     compatible_with = get_compatible_with_portable(),
     visibility = ["//visibility:public"],