[XLA:GPU]: Add some debugging information for triton collective kernels

PiperOrigin-RevId: 900676743

Author: sohaibiftikhar

xla/service/BUILD

@@ -2318,6 +2318,7 @@ xla_cc_test(
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",

xla/service/hlo_creation_utils.cc

@@ -998,9 +998,11 @@ std::unique_ptr<HloModule> NewModuleWithFusion(
   HloComputation::Builder entry_builder("entry");
   std::vector<HloInstruction*> entry_parameters =
       build_parameter_instructions(entry_builder);
-  HloInstruction* fusion_instruction = entry_builder.AddInstruction(
-      HloInstruction::CreateFusion(instruction->shape(), fusion_kind,
-                                   entry_parameters, fused_computation));
+  HloInstruction* fusion_instruction =
+      entry_builder.AddInstruction(HloInstruction::CreateFusion(
+          instruction->shape(), fusion_kind, entry_parameters,
+          fused_computation,
+          /*prefix=*/absl::StrCat(instruction->name(), "-")));
 
   hlo_module->AddEntryComputation(entry_builder.Build(fusion_instruction));
 

xla/service/hlo_creation_utils_test.cc

@@ -20,6 +20,7 @@ limitations under the License.
 
 #include <gtest/gtest.h>
 #include "absl/log/check.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "xla/array2d.h"
@@ -602,6 +603,8 @@ TEST_F(HloCreationUtilsTest, NewModuleWithFusion) {
   EXPECT_EQ(to_apply->name(), "apply_op");
   EXPECT_EQ(to_apply->num_parameters(), 2);
   EXPECT_EQ(to_apply->root_instruction()->opcode(), HloOpcode::kAdd);
+  EXPECT_EQ(fusion_module->entry_computation()->root_instruction()->name(),
+            absl::StrCat(all_reduce_start->name(), "-", "fusion"));
 }
 }  // namespace
 }  // namespace xla

xla/stream_executor/cuda/cuda_executor.cc

@@ -252,7 +252,7 @@ absl::StatusOr<CUfunction> GetModuleFunction(Context* context, CUmodule module,
   TF_RETURN_IF_ERROR(cuda::ToStatus(
       cuModuleGetFunction(&function, module, kernel_name),
       absl::StrCat(xla::XlaFormatDevice(context->device_ordinal()),
-                   "Failed to get module function")));
+                   "Failed to get module function ", kernel_name)));
   return function;
 }