Integrate Smart Suggestion into Xprof

xprof/convert/BUILD

@@ -857,12 +857,12 @@ cc_library(
     name = "xplane_to_tools_data",
     srcs = ["xplane_to_tools_data.cc"],
     hdrs = ["xplane_to_tools_data.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
     deps = [
         ":compute_inference_latency",
         ":hlo_to_tools_data",
         ":multi_xplanes_to_op_stats",
         ":multi_xspace_to_inference_stats",
-        ":op_stats_processor",
         ":op_stats_to_hlo_stats",
         ":op_stats_to_input_pipeline_analysis",
         ":op_stats_to_op_profile",
@@ -899,6 +899,11 @@ cc_library(
         "@org_xprof//plugin/xprof/protobuf:roofline_model_proto_cc",
         "@org_xprof//plugin/xprof/protobuf:tf_data_stats_proto_cc",
         "@org_xprof//plugin/xprof/protobuf:tf_stats_proto_cc",
+        "@org_xprof//xprof/convert/smart_suggestion:all_rules",
+        "@org_xprof//xprof/convert/smart_suggestion:signal_provider",
+        "@org_xprof//xprof/convert/smart_suggestion:smart_suggestion_engine",
+        "@org_xprof//xprof/convert/smart_suggestion:smart_suggestion_rule_factory",
+        "@org_xprof//xprof/convert/smart_suggestion:tool_data_provider_impl",
         "@org_xprof//xprof/convert/trace_viewer:trace_events",
         "@org_xprof//xprof/convert/trace_viewer:trace_events_to_json",
         "@org_xprof//xprof/convert/trace_viewer:trace_options",

xprof/convert/smart_suggestion/BUILD

@@ -9,6 +9,7 @@ package(
 cc_library(
     name = "signal_provider",
     hdrs = ["signal_provider.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
     deps = [
         ":tool_data_provider",
         "//util/task:status",
@@ -22,6 +23,7 @@ cc_library(
 cc_library(
     name = "smart_suggestion_rule",
     hdrs = ["smart_suggestion_rule.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
     deps = [
         ":signal_provider",
         "@com_google_absl//absl/status:statusor",
@@ -32,6 +34,7 @@ cc_library(
 cc_library(
     name = "tool_data_provider",
     hdrs = ["tool_data_provider.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
     deps = [
         "@com_google_absl//absl/status:statusor",
         "@org_xprof//plugin/xprof/protobuf:input_pipeline_proto_cc",
@@ -42,6 +45,7 @@ cc_library(
 cc_library(
     name = "tool_data_provider_impl",
     hdrs = ["tool_data_provider_impl.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
     deps = [
         ":tool_data_provider",
         "//util/task:status",
@@ -59,6 +63,7 @@ cc_library(
 cc_library(
     name = "input_bound_rule",
     hdrs = ["input_bound_rule.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
     deps = [
         ":signal_provider",
         ":smart_suggestion_rule",
@@ -70,6 +75,7 @@ cc_library(
 cc_library(
     name = "host_processing_bound_rule",
     hdrs = ["host_processing_bound_rule.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
     deps = [
         ":input_bound_rule",
         ":signal_provider",
@@ -85,6 +91,7 @@ cc_library(
 cc_library(
     name = "data_transfer_bound_rule",
     hdrs = ["data_transfer_bound_rule.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
     deps = [
         ":input_bound_rule",
         ":signal_provider",
@@ -97,6 +104,42 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "smart_suggestion_rule_factory",
+    hdrs = ["smart_suggestion_rule_factory.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
+    deps = [
+        ":smart_suggestion_rule",
+    ],
+)
+
+cc_library(
+    name = "all_rules",
+    hdrs = ["all_rules.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
+    deps = [
+        ":data_transfer_bound_rule",
+        ":host_processing_bound_rule",
+        ":input_bound_rule",
+        ":smart_suggestion_rule_factory",
+    ],
+)
+
+cc_library(
+    name = "smart_suggestion_engine",
+    srcs = ["smart_suggestion_engine.cc"],
+    hdrs = ["smart_suggestion_engine.h"],
+    compatible_with = ["//buildenv/target:non_prod"],
+    deps = [
+        ":signal_provider",
+        ":smart_suggestion_rule",
+        ":smart_suggestion_rule_factory",
+        "//util/task:status",
+        "@com_google_absl//absl/status:statusor",
+        "@org_xprof//plugin/xprof/protobuf:smart_suggestion_proto_cc",
+    ],
+)
+
 cc_test(
     name = "host_processing_bound_rule_test",
     srcs = ["host_processing_bound_rule_test.cc"],

xprof/convert/smart_suggestion/all_rules.h

@@ -0,0 +1,39 @@
+/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_ALL_RULES_H_
+#define THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_ALL_RULES_H_
+
+#include "xprof/convert/smart_suggestion/data_transfer_bound_rule.h"
+#include "xprof/convert/smart_suggestion/host_processing_bound_rule.h"
+#include "xprof/convert/smart_suggestion/input_bound_rule.h"
+#include "xprof/convert/smart_suggestion/smart_suggestion_rule_factory.h"
+
+namespace tensorflow {
+namespace profiler {
+
+// Registers all smart suggestion rules.
+inline void RegisterAllRules(SmartSuggestionRuleFactory* f) {
+  // go/keep-sorted start
+  f->Register<DataTransferBoundRule>();
+  f->Register<HostProcessingBoundRule>();
+  f->Register<InputBoundRule>();
+  // go/keep-sorted end
+}
+
+}  // namespace profiler
+}  // namespace tensorflow
+
+#endif  // THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_ALL_RULES_H_

xprof/convert/smart_suggestion/smart_suggestion_engine.cc

@@ -0,0 +1,46 @@
+/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "xprof/convert/smart_suggestion/smart_suggestion_engine.h"
+
+#include <optional>
+
+#include "absl/status/statusor.h"
+#include "xprof/convert/smart_suggestion/signal_provider.h"
+#include "xprof/convert/smart_suggestion/smart_suggestion_rule.h"
+#include "xprof/convert/smart_suggestion/smart_suggestion_rule_factory.h"
+#include "plugin/xprof/protobuf/smart_suggestion.pb.h"
+#include "util/task/status_macros.h"
+
+namespace tensorflow {
+namespace profiler {
+
+absl::StatusOr<SmartSuggestionReport> SmartSuggestionEngine::Run(
+    const SignalProvider& signal_provider,
+    const SmartSuggestionRuleFactory& rule_factory) const {
+  SmartSuggestionReport report;
+
+  const auto& rules = rule_factory.GetAllRules();
+  for (const auto& rule : rules) {
+    ASSIGN_OR_RETURN(std::optional<SmartSuggestion> suggestion,
+                     rule->Apply(signal_provider));
+    if (suggestion.has_value()) {
+      *report.add_suggestions() = *suggestion;
+    }
+  }
+  return report;
+}
+
+}  // namespace profiler
+}  // namespace tensorflow

xprof/convert/smart_suggestion/smart_suggestion_engine.h

@@ -0,0 +1,42 @@
+/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_SMART_SUGGESTION_ENGINE_H_
+#define THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_SMART_SUGGESTION_ENGINE_H_
+
+#include "absl/status/statusor.h"
+#include "xprof/convert/smart_suggestion/signal_provider.h"
+#include "xprof/convert/smart_suggestion/smart_suggestion_rule_factory.h"
+#include "plugin/xprof/protobuf/smart_suggestion.pb.h"
+
+namespace tensorflow {
+namespace profiler {
+
+// Engine to generate smart suggestions.
+class SmartSuggestionEngine {
+ public:
+  explicit SmartSuggestionEngine() = default;
+
+  // Generates smart suggestions based on the provided signal provider and rule
+  // factory.
+  absl::StatusOr<SmartSuggestionReport> Run(
+      const SignalProvider& signal_provider,
+      const SmartSuggestionRuleFactory& rule_factory) const;
+};
+
+}  // namespace profiler
+}  // namespace tensorflow
+
+#endif  // THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_SMART_SUGGESTION_ENGINE_H_

xprof/convert/smart_suggestion/smart_suggestion_rule_factory.h

@@ -0,0 +1,47 @@
+/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_RULE_FACTORY_H_
+#define THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_RULE_FACTORY_H_
+
+#include <memory>
+#include <vector>
+
+#include "xprof/convert/smart_suggestion/smart_suggestion_rule.h"
+
+namespace tensorflow {
+namespace profiler {
+
+// Factory class to manage smart suggestion rules.
+class SmartSuggestionRuleFactory {
+ public:
+  // Registers a rule with the factory.
+  template <typename RuleType>
+  void Register() {
+    rules_.push_back(std::make_unique<RuleType>());
+  }
+
+  const std::vector<std::unique_ptr<SmartSuggestionRule>>& GetAllRules() const {
+    return rules_;
+  }
+
+ private:
+  std::vector<std::unique_ptr<SmartSuggestionRule>> rules_;
+};
+
+}  // namespace profiler
+}  // namespace tensorflow
+
+#endif  // THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_RULE_FACTORY_H_

xprof/convert/xplane_to_tools_data.cc

@@ -53,6 +53,11 @@ limitations under the License.
 #include "xprof/convert/profile_processor.h"
 #include "xprof/convert/profile_processor_factory.h"
 #include "xprof/convert/repository.h"
+#include "xprof/convert/smart_suggestion/all_rules.h"
+#include "xprof/convert/smart_suggestion/signal_provider.h"
+#include "xprof/convert/smart_suggestion/smart_suggestion_engine.h"
+#include "xprof/convert/smart_suggestion/smart_suggestion_rule_factory.h"
+#include "xprof/convert/smart_suggestion/tool_data_provider_impl.h"
 #include "xprof/convert/tool_options.h"
 #include "xprof/convert/trace_viewer/trace_events.h"
 #include "xprof/convert/trace_viewer/trace_events_to_json.h"
@@ -391,6 +396,34 @@ absl::Status RunMapReduce(xprof::ProfileProcessor* processor,
   return processor->Reduce(session_snapshot, map_output_files);
 }
 
+absl::StatusOr<std::string> ConvertMultiXSpacesToSmartSuggestion(
+    const SessionSnapshot& session_snapshot) {
+  SmartSuggestionEngine engine;
+  SmartSuggestionRuleFactory rule_factory;
+  RegisterAllRules(&rule_factory);
+
+  auto tool_data_provider =
+      std::make_unique<ToolDataProviderImpl>(session_snapshot);
+  SignalProvider signal_provider(std::move(tool_data_provider));
+
+  TF_ASSIGN_OR_RETURN(SmartSuggestionReport report,
+                      engine.Run(signal_provider, rule_factory));
+  std::string json_output;
+  tsl::protobuf::util::JsonPrintOptions opts;
+  opts.always_print_primitive_fields = true;
+  // Perform the Proto to JSON conversion.
+  auto encode_status =
+      tsl::protobuf::util::MessageToJsonString(report, &json_output, opts);
+  if (!encode_status.ok()) {
+    const auto& error_message = encode_status.message();
+    return tsl::errors::Internal(
+        "Could not convert smart suggestion report to json. Error: ",
+        absl::string_view(error_message.data(), error_message.length()));
+  }
+  // Return the generated JSON string.
+  return json_output;
+}
+
 }  // namespace
 
 absl::StatusOr<std::string> ConvertMultiXSpacesToToolData(
@@ -435,6 +468,8 @@ absl::StatusOr<std::string> ConvertMultiXSpacesToToolData(
     return PreprocessXSpace(session_snapshot);
   } else if (tool_name == "inference_profile") {
     return ConvertMultiXSpacesToInferenceStats(session_snapshot, options);
+  } else if (tool_name == "smart_suggestion") {
+    return ConvertMultiXSpacesToSmartSuggestion(session_snapshot);
   } else {
     return tsl::errors::InvalidArgument(
         "Can not find tool: ", tool_name,