[XLA] Re-populate source_file, source_line, source_end_line, source_column, source_end_column to enable correct source code tracing in xprof.

PiperOrigin-RevId: 853006348

Author: Google-ML-Automation

xla/hlo/translate/mhlo_to_hlo/location_exporter.cc

@@ -143,6 +143,11 @@ xla::OpMetadata CreateOpMetadataFromLocation(
     auto result = frame_index_builder->AddCallStackAndGetFirstFrameId(loc);
     if (result.last_frame_id != mlir::StackFrameIndexBuilder::kInvalidIndex) {
       metadata.set_stack_frame_id(result.last_frame_id);
+      metadata.set_source_file(result.last_frame_file);
+      metadata.set_source_line(result.last_frame_line);
+      metadata.set_source_end_line(result.last_frame_end_line);
+      metadata.set_source_column(result.last_frame_column);
+      metadata.set_source_end_column(result.last_frame_end_column);
       return metadata;
     }
   }

xla/hlo/translate/mhlo_to_hlo/tests/location_to_op_metadata.mlir

@@ -62,7 +62,7 @@ func.func @main(%arg0: !mhlo.token) -> !mhlo.token {
 }
 
 // CHECK: after-all
-// CHECK-SAME: metadata={op_name="name(anothername)" stack_frame_id=1}
+// CHECK-SAME: metadata={op_name="name(anothername)" source_file="file_name" source_line=2 source_end_line=2 source_column=8 source_end_column=8 stack_frame_id=1}
 
 // -----
 

xla/hlo/translate/mhlo_to_hlo/tests/location_to_stacktrace.mlir

@@ -33,7 +33,12 @@ module @main attributes {mhlo.cross_program_prefetches = [], mhlo.is_dynamic = f
 // CHECK-NEXT: }
 // CHECK-NEXT: metadata {
 // CHECK-NEXT:   op_name: "name(anothername)"
+// CHECK-NEXT:   source_file: "file_name"
+// CHECK-NEXT:   source_line: 2
 // CHECK-NEXT:   stack_frame_id: 1
+// CHECK-NEXT:   source_end_line: 2
+// CHECK-NEXT:   source_column: 8
+// CHECK-NEXT:   source_end_column: 8
 // CHECK-NEXT: }
 
 // CHECK: stack_frame_index {
@@ -74,7 +79,12 @@ module @main attributes {mhlo.cross_program_prefetches = [], mhlo.is_dynamic = f
 // CHECK-NEXT: metadata {
 // CHECK-NEXT:   op_type: "atype"
 // CHECK-NEXT:   op_name: "name(anothername)"
+// CHECK-NEXT:   source_file: "file_name_2"
+// CHECK-NEXT:   source_line: 3
 // CHECK-NEXT:   stack_frame_id: 2
+// CHECK-NEXT:   source_end_line: 3
+// CHECK-NEXT:   source_column: 4
+// CHECK-NEXT:   source_end_column: 4
 // CHECK-NEXT: }
 
 // CHECK: stack_frame_index {