Move TraceFilterConfig Proto to OSS

plugin/xprof/protobuf/BUILD

@@ -223,6 +223,11 @@ xprof_proto_library(
     srcs = ["smart_suggestion.proto"],
 )
 
+xprof_proto_library(
+    name = "trace_filter_config_proto",
+    srcs = ["trace_filter_config.proto"],
+)
+
 xprof_proto_library(
     name = "protos_all",
     create_go_proto = False,
@@ -260,6 +265,7 @@ xprof_proto_library(
         ":trace_events_old_proto",
         ":trace_events_raw_proto",
         ":smart_suggestion_proto",
+        ":trace_filter_config_proto",
     ],
 )
 
@@ -297,6 +303,7 @@ exports_files(
         "trace_events.proto",
         "trace_events_old.proto",
         "trace_events_raw.proto",
+        "trace_filter_config.proto",
     ],
     visibility = [
         "@org_xprof//plugin/xprof/protobuf:__pkg__",

plugin/xprof/protobuf/trace_filter_config.proto

@@ -0,0 +1,60 @@
+syntax = "proto3";
+
+package tensorflow.profiler;
+
+// TraceEventFilter is a boolean expression on a particular field on the
+// TraceEvent. If all TraceEventFilters (AND condition between each) evaluate
+// as True on a TraceFilter, it means we filter the TraceEvent out.
+message TraceEventFilter {
+  enum Operator {
+    OP_EQ = 0;  // Equal
+
+    OP_LT = 1;  // Less than
+
+    OP_GT = 2;  // Greater than
+
+    OP_LE = 3;  // Less than and equal
+
+    OP_GE = 4;  // Greater than and equal
+
+    OP_REGEX = 5;  // Regex match
+  }
+
+  // The name of a field on the TraceEvent.
+  string field_name = 5;
+  Operator op_id = 6;
+
+  // Whether NOT proposition is set or not.
+  bool negation = 7;
+
+  oneof value {
+    string str_value = 8;
+    string regex_value = 9;  // Only allows using OP_EQ.
+
+    uint64 uint_value = 10;
+    int64 int_value = 12;
+    double double_value = 11;
+  }
+}
+
+message TraceFilterConfig {
+  // List of regexes for device/resource we should match this filter on.
+  // A device/resource is included if it matches any one of the regex.
+  // If empty, assume it means include all.
+  // Device & Resource here matches its definition in the TraceEvent proto
+  // http://google3/third_party/tensorflow/core/profiler/protobuf/trace_events.proto;l=127-137;rcl=590742687
+  repeated string device_regexes = 1;
+  repeated string resource_regexes = 2;
+
+  // Match an event based on the filters below.
+  // Between all the Event Filters (trace_event_filters &
+  // trace_event_arg_filters) it is AND condition. Filters are functionally
+  // complete since they support {AND, NOT} bool operators.
+  repeated TraceEventFilter trace_event_filters = 3;
+  repeated TraceEventFilter trace_event_arg_filters = 4;
+
+  // Whether NOT proposition is set or not on this TraceFilter.
+  // By default False so TraceFilter will filter out events that match it.
+  // When set to True, TraceFilter will only keep events that match it.
+  bool negation = 5;
+}

xprof/convert/trace_viewer/BUILD

@@ -163,6 +163,7 @@ cc_library(
         ":trace_events_to_json",
         ":trace_utils",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@org_xprof//plugin/xprof/protobuf:trace_filter_config_proto_cc",
         "@org_xprof//xprof/convert:tool_options",
     ],
 )
@@ -173,6 +174,7 @@ cc_test(
     deps = [
         ":trace_events_filter_interface",
         ":trace_options",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
         "@org_xprof//xprof/convert:tool_options",

xprof/convert/trace_viewer/trace_options.cc

@@ -1,16 +1,34 @@
 #include "xprof/convert/trace_viewer/trace_options.h"
 
+#include <cstdint>
 #include <memory>
 
 #include "absl/container/flat_hash_set.h"
 #include "xprof/convert/tool_options.h"
 #include "xprof/convert/trace_viewer/trace_events_filter_interface.h"
 #include "xprof/convert/trace_viewer/trace_events_to_json.h"
 #include "xprof/convert/trace_viewer/trace_utils.h"
+#include "plugin/xprof/protobuf/trace_filter_config.pb.h"
 
 namespace tensorflow {
 namespace profiler {
+namespace filter_internal {
 
+class TraceEventsFilter : public TraceEventsFilterInterface {
+ public:
+  explicit TraceEventsFilter(const TraceOptions& options) : options_(options) {}
+
+  void SetUp(const Trace& trace) override;
+
+  bool Filter(const TraceEvent& event) override;
+
+ private:
+  const TraceOptions options_;
+
+  TraceDeviceType device_type_ = TraceDeviceType::kUnknownDevice;
+  absl::flat_hash_set<uint32_t /*device_id*/> tpu_noncore_devices_;
+  absl::flat_hash_set<uint32_t /*device_id*/> tpu_core_devices_;
+};
 void TraceEventsFilter::SetUp(const Trace& trace) {
   for (const auto& [device_id, device] : trace.devices()) {
     if (IsTpuCoreDeviceName(device.name())) {
@@ -39,6 +57,8 @@ bool TraceEventsFilter::Filter(const TraceEvent& event) {
   return false;
 }
 
+}  // namespace filter_internal
+
 TraceOptions TraceOptionsFromToolOptions(const ToolOptions& tool_options) {
   TraceOptions options;
   options.full_dma =
@@ -61,8 +81,8 @@ JsonTraceOptions::Details TraceOptionsToDetails(TraceDeviceType device_type,
 }
 
 std::unique_ptr<tensorflow ::profiler::TraceEventsFilterInterface>
-TraceEventsFilterFromTraceOptions(const TraceOptions& options) {
-  return std::make_unique<TraceEventsFilter>(options);
+CreateTraceEventsFilterFromTraceOptions(const TraceOptions& options) {
+  return std::make_unique<filter_internal::TraceEventsFilter>(options);
 }
 
 }  // namespace profiler

xprof/convert/trace_viewer/trace_options.h

@@ -3,12 +3,14 @@
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 
 #include "absl/container/flat_hash_set.h"
 #include "xprof/convert/tool_options.h"
 #include "xprof/convert/trace_viewer/trace_events_filter_interface.h"
 #include "xprof/convert/trace_viewer/trace_events_to_json.h"
 #include "xprof/convert/trace_viewer/trace_utils.h"
+#include "plugin/xprof/protobuf/trace_filter_config.pb.h"
 
 namespace tensorflow {
 namespace profiler {
@@ -29,6 +31,9 @@ inline constexpr const char* kShowHloCostModel = "show_hlo_cost_model";
 // Options used to select TraceEvents (e.g., for visualization or further
 // processing).
 struct TraceOptions {
+  // TraceFilterConfig is a filtering config written by the user for filtering
+  // trace events.
+  std::optional<tensorflow::profiler::TraceFilterConfig> trace_filter_config;
   // Visualize tensor core overlays. This is low-level debug information for
   // XLA and xprof devs. so don't show by default.
   bool overlay = false;
@@ -58,22 +63,6 @@ struct TraceOptions {
   bool show_hlo_cost_model = false;
 };
 
-class TraceEventsFilter : public TraceEventsFilterInterface {
- public:
-  explicit TraceEventsFilter(const TraceOptions& options) : options_(options) {}
-
-  void SetUp(const Trace& trace) override;
-
-  bool Filter(const TraceEvent& event) override;
-
- private:
-  const TraceOptions options_;
-
-  TraceDeviceType device_type_ = TraceDeviceType::kUnknownDevice;
-  absl::flat_hash_set<uint32_t /*device_id*/> tpu_noncore_devices_;
-  absl::flat_hash_set<uint32_t /*device_id*/> tpu_core_devices_;
-};
-
 // Returns TraceOptions for use when the generated events will be visualized.
 // Parses the tool options from a request.
 TraceOptions TraceOptionsFromToolOptions(const ToolOptions& options);
@@ -85,7 +74,7 @@ JsonTraceOptions::Details TraceOptionsToDetails(TraceDeviceType device_type,
 // Returns a trace events filter configured to filter events according to the
 // given options.
 std::unique_ptr<tensorflow::profiler::TraceEventsFilterInterface>
-TraceEventsFilterFromTraceOptions(const TraceOptions& options);
+CreateTraceEventsFilterFromTraceOptions(const TraceOptions& options);
 
 inline bool IsTpuTrace(const Trace& trace) {
   for (const auto& [device_id, device] : trace.devices()) {

xprof/convert/trace_viewer/trace_options_test.cc

@@ -1,14 +1,35 @@
 #include "xprof/convert/trace_viewer/trace_options.h"
 
+#include <cstdint>
 #include <memory>
 
 #include "testing/base/public/gmock.h"
 #include "<gtest/gtest.h>"
+#include "absl/container/flat_hash_set.h"
 #include "xprof/convert/tool_options.h"
 #include "xprof/convert/trace_viewer/trace_events_filter_interface.h"
 
 namespace tensorflow {
 namespace profiler {
+namespace filter_internal {
+// Test Peer Class for TraceEventsFilter.
+class TraceEventsFilter : public TraceEventsFilterInterface {
+ public:
+  explicit TraceEventsFilter(const TraceOptions& options) : options_(options) {}
+
+  void SetUp(const Trace& trace) override;
+
+  bool Filter(const TraceEvent& event) override;
+
+ private:
+  const TraceOptions options_;
+
+  TraceDeviceType device_type_ = TraceDeviceType::kUnknownDevice;
+  absl::flat_hash_set<uint32_t /*device_id*/> tpu_noncore_devices_;
+  absl::flat_hash_set<uint32_t /*device_id*/> tpu_core_devices_;
+};
+}  // namespace filter_internal
+
 namespace {
 
 using ::testing::IsEmpty;
@@ -62,14 +83,14 @@ TEST(TraceOptionsTest, IsTpuTraceTest) {
 TEST(TraceOptionsTest, TraceEventsFilterFromTraceOptionsTest) {
   TraceOptions options;
   std::unique_ptr<TraceEventsFilterInterface> filter =
-      TraceEventsFilterFromTraceOptions(options);
+      CreateTraceEventsFilterFromTraceOptions(options);
   EXPECT_NE(filter, nullptr);
 }
 
 TEST(TraceEventsFilterTest, FilterTest) {
   TraceOptions options;
   options.full_dma = false;
-  TraceEventsFilter filter(options);
+  filter_internal::TraceEventsFilter filter(options);
 
   Trace trace;
   Device& tpu_device = (*trace.mutable_devices())[0];
@@ -95,15 +116,15 @@ TEST(TraceEventsFilterTest, FilterTest) {
 
   // With full_dma=true, no events should be filtered.
   options.full_dma = true;
-  TraceEventsFilter full_dma_filter(options);
+  filter_internal::TraceEventsFilter full_dma_filter(options);
   full_dma_filter.SetUp(trace);
   EXPECT_FALSE(full_dma_filter.Filter(flow_event));
   EXPECT_FALSE(full_dma_filter.Filter(non_flow_event));
 }
 
 TEST(TraceEventsFilterTest, NonTpuTraceTest) {
   TraceOptions options;
-  TraceEventsFilter filter(options);
+  filter_internal::TraceEventsFilter filter(options);
 
   Trace trace;
   Device& gpu_device = (*trace.mutable_devices())[0];

xprof/convert/xplane_to_tools_data.cc

@@ -147,7 +147,7 @@ absl::StatusOr<std::string> ConvertXSpaceToTraceEvents(
     // Trace smaller than threshold will be disabled from streaming.
     constexpr int64_t kDisableStreamingThreshold = 500000;
     auto trace_events_filter =
-        TraceEventsFilterFromTraceOptions(profiler_trace_options);
+        CreateTraceEventsFilterFromTraceOptions(profiler_trace_options);
     TF_RETURN_IF_ERROR(trace_container.LoadFromLevelDbTable(
         *sstable_path, std::move(trace_events_filter),
         std::move(visibility_filter), kDisableStreamingThreshold));