Includes instruction source files in the Auto Sharding request proto.

PiperOrigin-RevId: 666104783

Author: tensorflower-gardener

xla/hlo/experimental/auto_sharding/auto_sharding.cc

@@ -2018,6 +2018,7 @@ AutoShardingSolverResult CallSolver(
     request.add_instruction_names(
         absl::StrCat(instruction_name, " (id: ", node_idx, ")"));
     request.add_opcodes(std::string(opcode));
+    request.add_metadata_source_files(instruction->metadata().source_file());
     AutoShardingSolverRequest_Costs ci, di, mi, pi;
     AutoShardingSolverRequest_Names strategy_names;
     std::optional<HloSharding> default_strategy;

xla/hlo/experimental/auto_sharding/auto_sharding.proto

@@ -68,6 +68,7 @@ message AutoShardingSolverRequest {
   repeated Costs value_costs = 15;
   repeated string instruction_names = 16;
   repeated string opcodes = 33;
+  repeated string metadata_source_files = 40;
   repeated Names strategy_names = 32;
   optional SolverTimeout solver_timeout = 17;
   optional Coeff overbudget_coeff = 18;

xla/hlo/experimental/auto_sharding/auto_sharding_solver_test.cc

@@ -141,6 +141,11 @@ AutoShardingSolverRequest DefaultAutoShardingSolverRequest() {
                          1, 0, 1,
                          1, 1, 0}};
   const std::vector<std::string> instruction_names = {"A", "B", "C", "D", "E"};
+  const std::vector<std::string> metadata_source_files = {"attention.py",
+                                                          "convolution.py",
+                                                          "layers.py",
+                                                          "logits.py",
+                                                          "pipeline.py"};
 
   AutoShardingSolverRequest request;
   request.set_num_nodes(5);
@@ -159,6 +164,8 @@ AutoShardingSolverRequest DefaultAutoShardingSolverRequest() {
   AddCosts(request.mutable_value_costs(), v);
   request.mutable_instruction_names()->Add(instruction_names.begin(),
                                            instruction_names.end());
+  request.mutable_metadata_source_files()->Add(metadata_source_files.begin(),
+                                               metadata_source_files.end());
   return request;
 }