Add support for Viewing custom call text in Graph Viewer.

PiperOrigin-RevId: 831447585

Author: cliveverghese

xprof/convert/BUILD

@@ -149,7 +149,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/utils:custom_call_utils",
+        "@org_xprof//xprof/utils:hlo_module_utils",
+        "@org_xprof//xprof/utils:hlo_proto_to_module",
         "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
+        "@xla//xla/hlo/ir:hlo",
         "@xla//xla/tsl/platform:statusor",
     ],
     alwayslink = 1,

xprof/convert/graph_viewer_processor.cc

@@ -15,12 +15,14 @@ limitations under the License.
 
 #include "xprof/convert/graph_viewer_processor.h"
 
+#include <memory>
 #include <optional>
 #include <string>
 
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/tsl/platform/statusor.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
 #include "xprof/convert/hlo_proto_to_graph_view.h"
@@ -39,6 +41,9 @@ limitations under the License.
 #include "plugin/xprof/protobuf/roofline_model.pb.h"
 #include "plugin/xprof/protobuf/tf_data_stats.pb.h"
 #include "plugin/xprof/protobuf/tf_stats.pb.h"
+#include "xprof/utils/custom_call_utils.h"
+#include "xprof/utils/hlo_module_utils.h"
+#include "xprof/utils/hlo_proto_to_module.h"
 
 namespace xprof {
 
@@ -49,6 +54,7 @@ using ::tensorflow::profiler::GetHloProtoByModuleName;
 using ::tensorflow::profiler::GetParam;
 using ::tensorflow::profiler::GraphViewerParams;
 using ::tensorflow::profiler::kAdjacentNodes;
+using ::tensorflow::profiler::kCustomCallGraphTypeName;
 using ::tensorflow::profiler::kGraphTypeName;
 using ::tensorflow::profiler::ParseGraphViewerParams;
 using ::tensorflow::profiler::SessionSnapshot;
@@ -64,6 +70,16 @@ absl::StatusOr<std::string> ConvertHloProtoToGraphViewer(
     return ConvertHloProtoToGraph(hlo_proto, params.node_name,
                                   params.graph_width, params.render_options,
                                   params.format);
+  } else if (params.type == kCustomCallGraphTypeName) {
+    TF_ASSIGN_OR_RETURN(
+        std::unique_ptr<xla::HloModule> hlo_module,
+        tensorflow::profiler::ConvertHloProtoToModule(hlo_proto));
+    const xla::HloInstruction* hlo_instruction =
+        tensorflow::profiler::FindInstruction(*hlo_module, params.node_name);
+    if (hlo_instruction == nullptr) {
+      return absl::InvalidArgumentError("Hlo Instruction not found.");
+    }
+    return GetCustomCallText(*hlo_instruction);
   } else if (params.type == kAdjacentNodes) {
     return GetAdjacentNodes(hlo_proto, params.node_name);
   } else {

xprof/convert/hlo_proto_to_graph_view.cc

@@ -300,8 +300,10 @@ absl::StatusOr<GraphViewerParams> ParseGraphViewerParams(
     return InvalidArgument("Graph viewer must provide a type option.");
   }
   auto valid_types = {
-      kGraphTypeName,    kJsonTypeName,    kProtoTypeName, kProtoTextTypeName,
-      kShortTxtTypeName, kLongTxtTypeName, kAdjacentNodes,
+      kGraphTypeName,    kJsonTypeName,
+      kProtoTypeName,    kProtoTextTypeName,
+      kShortTxtTypeName, kLongTxtTypeName,
+      kAdjacentNodes,    kCustomCallGraphTypeName,
   };
   if (std::find(valid_types.begin(), valid_types.end(), type.value()) ==
       valid_types.end()) {
@@ -334,6 +336,13 @@ absl::StatusOr<GraphViewerParams> ParseGraphViewerParams(
       params.node_name = node_name.value();
     }
     return params;
+  } else if (type == kCustomCallGraphTypeName) {
+    params.type = type.value();
+    if (std::optional<std::string> node_name =
+            GetParam<std::string>(options, "node_name")) {
+      params.node_name = node_name.value();
+    }
+    return params;
   }
 
   // For txt type.

xprof/convert/tool_options.h

@@ -28,6 +28,7 @@ namespace profiler {
 
 // Default parameter constants for graph viewer.
 static constexpr char kGraphTypeName[] = "graph";
+static constexpr char kCustomCallGraphTypeName[] = "custom_call";
 static constexpr char kAdjacentNodes[] = "adj_nodes";
 static constexpr char kShortTxtTypeName[] = "short_txt";
 static constexpr char kLongTxtTypeName[] = "long_txt";

xprof/utils/BUILD

@@ -650,15 +650,15 @@ cc_library(
     hdrs = ["custom_call_utils.h"],
     deps = [
         ":backend_configs_proto_cc",
-        "//third_party/llvm/llvm-project/mlir:IR",
-        "//third_party/protobuf/json",
-        "//util/task:status",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_protobuf//:json",
+        "@llvm-project//mlir:IR",
+        "@tsl//tsl/platform:protobuf",
         "@xla//xla/hlo/ir:hlo",
         "@xla//xla/pjrt:mlir_to_hlo",
         "@xla//xla/service/llvm_ir:llvm_util",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
     ],
 )
 

xprof/utils/custom_call_utils.cc

@@ -2,18 +2,17 @@
 
 #include <string>
 
-#include "xprof/utils/backend_configs.pb.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/BuiltinOps.h"
-#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
-#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/OwningOpRef.h"
-#include "google/protobuf/json/json.h"
-#include "google/protobuf/json/json.h"
+#include "mlir/include/mlir/IR/BuiltinOps.h"
+#include "mlir/include/mlir/IR/MLIRContext.h"
+#include "mlir/include/mlir/IR/OwningOpRef.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/pjrt/mlir_to_hlo.h"
 #include "xla/service/llvm_ir/llvm_util.h"
-#include "util/task/status_macros.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xprof/utils/backend_configs.pb.h"
 
 namespace xprof {
 
@@ -22,18 +21,18 @@ absl::StatusOr<std::string> GetCustomCallText(
   if (!hlo_instruction.has_backend_config()) {
     return absl::NotFoundError("Backend config not found");
   }
-  google::protobuf::json::ParseOptions options;
+  tsl::protobuf::util::JsonParseOptions options;
   options.ignore_unknown_fields = true;
   BackendConfig config;
-  RETURN_IF_ERROR(google::protobuf::util::JsonStringToMessage(
+  TF_RETURN_IF_ERROR(tsl::protobuf::util::JsonStringToMessage(
       hlo_instruction.raw_backend_config_string(), &config, options));
   if (!config.has_custom_call_config()) {
     return absl::NotFoundError("Custom call config not found");
   }
   CustomCallConfig custom_call_config = config.custom_call_config();
   mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
   context.allowUnregisteredDialects(true);
-  ASSIGN_OR_RETURN(
+  TF_ASSIGN_OR_RETURN(
       mlir::OwningOpRef<mlir::ModuleOp> mlir_op,
       xla::ParseMlirModuleString(
           static_cast<std::string>(custom_call_config.body()), context));