diff --git a/xprof/convert/BUILD b/xprof/convert/BUILD index 6ce85250..25dfc846 100644 --- a/xprof/convert/BUILD +++ b/xprof/convert/BUILD @@ -1379,11 +1379,38 @@ cc_test( ], ) +cc_library( + name = "smart_suggestion_processor", + srcs = ["smart_suggestion_processor.cc"], + hdrs = ["smart_suggestion_processor.h"], + deps = [ + ":profile_processor", + ":profile_processor_factory", + ":repository", + ":tool_options", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:string_view", + "@com_google_protobuf//:json", + "@org_xprof//plugin/xprof/protobuf:smart_suggestion_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", + "@tsl//tsl/profiler/protobuf:xplane_proto_cc", + "@xla//xla/tsl/platform:errors", + "@xla//xla/tsl/platform:statusor", + ], + alwayslink = 1, +) + cc_library( name = "xplane_to_tools_data", srcs = ["xplane_to_tools_data.cc"], hdrs = ["xplane_to_tools_data.h"], deps = [ + ":event_time_fraction_analyzer_processor", ":framework_op_stats_processor", ":graph_viewer_processor", ":hlo_stats_processor", @@ -1398,16 +1425,13 @@ cc_library( ":multi_xspace_to_inference_stats", ":op_profile_processor", ":op_stats_to_op_profile", - ":op_stats_to_overview_page", - ":op_stats_to_pod_viewer", - ":op_stats_to_roofline_model", - ":op_stats_to_tf_stats", ":overview_page_processor", ":pod_viewer_processor", ":preprocess_single_host_xplane", ":process_megascale_dcn", ":repository", ":roofline_model_processor", + ":smart_suggestion_processor", ":streaming_trace_viewer_processor", ":tool_options", ":trace_viewer_processor", @@ -2306,6 +2330,7 @@ cc_library( ":xspace_to_event_time_fraction_analyzer", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/time", "@com_google_protobuf//:protobuf", "@org_xprof//plugin/xprof/protobuf:event_time_fraction_analyzer_proto_cc", "@tsl//tsl/profiler/protobuf:xplane_proto_cc", diff --git a/xprof/convert/smart_suggestion_processor.cc b/xprof/convert/smart_suggestion_processor.cc new file mode 100644 index 00000000..5ebfca3a --- /dev/null +++ b/xprof/convert/smart_suggestion_processor.cc @@ -0,0 +1,107 @@ +/* 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_processor.h" + +#include +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "google/protobuf/json/json.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/statusor.h" +#include "tsl/profiler/protobuf/xplane.pb.h" +#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 "plugin/xprof/protobuf/smart_suggestion.pb.h" + +namespace xprof { +namespace { + +using ::tensorflow::profiler::RegisterAllRules; +using ::tensorflow::profiler::SessionSnapshot; +using ::tensorflow::profiler::SignalProvider; +using ::tensorflow::profiler::SmartSuggestionEngine; +using ::tensorflow::profiler::SmartSuggestionReport; +using ::tensorflow::profiler::SmartSuggestionRuleFactory; +using ::tensorflow::profiler::ToolDataProviderImpl; +using ::tensorflow::profiler::ToolOptions; +using ::tensorflow::profiler::XSpace; + +std::unique_ptr CreateSmartSuggestionProcessor( + const ToolOptions& options) { + return std::make_unique(options); +} + +RegisterProfileProcessor smart_suggestion_processor_registration( + "smart_suggestion", CreateSmartSuggestionProcessor); + +} // namespace + +absl::StatusOr SmartSuggestionProcessor::Map( + const SessionSnapshot& session_snapshot, const std::string& hostname, + const XSpace& xspace) { + return absl::UnimplementedError( + "SmartSuggestionProcessor::Map is not intended to be used."); +} + +absl::Status SmartSuggestionProcessor::Reduce( + const SessionSnapshot& session_snapshot, + const std::vector& map_output_files) { + return absl::UnimplementedError( + "SmartSuggestionProcessor::Reduce is not intended to be used."); +} + +absl::Status SmartSuggestionProcessor::ProcessSession( + const SessionSnapshot& session_snapshot, const ToolOptions& options) { + SmartSuggestionEngine engine; + SmartSuggestionRuleFactory rule_factory; + RegisterAllRules(&rule_factory); + + auto tool_data_provider = + std::make_unique(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_fields_with_no_presence = true; + 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())); + } + + data_ = json_output; + return absl::OkStatus(); +} + +} // namespace xprof diff --git a/xprof/convert/smart_suggestion_processor.h b/xprof/convert/smart_suggestion_processor.h new file mode 100644 index 00000000..6748a913 --- /dev/null +++ b/xprof/convert/smart_suggestion_processor.h @@ -0,0 +1,52 @@ +/* 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_PROCESSOR_H_ +#define THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_PROCESSOR_H_ + +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "tsl/profiler/protobuf/xplane.pb.h" +#include "xprof/convert/profile_processor.h" +#include "xprof/convert/repository.h" +#include "xprof/convert/tool_options.h" + +namespace xprof { + +class SmartSuggestionProcessor : public ProfileProcessor { + public: + explicit SmartSuggestionProcessor( + const tensorflow::profiler::ToolOptions& options) {} + + absl::StatusOr Map( + const tensorflow::profiler::SessionSnapshot& session_snapshot, + const std::string& hostname, + const tensorflow::profiler::XSpace& xspace) override; + + absl::Status Reduce( + const tensorflow::profiler::SessionSnapshot& session_snapshot, + const std::vector& map_output_files) override; + + absl::Status ProcessSession( + const tensorflow::profiler::SessionSnapshot& session_snapshot, + const tensorflow::profiler::ToolOptions& options) override; +}; + +} // namespace xprof + +#endif // THIRD_PARTY_XPROF_CONVERT_SMART_SUGGESTION_PROCESSOR_H_