Associates names with individual input sharding combinations (rather …

…than strategies).

PiperOrigin-RevId: 681639794

xla/hlo/experimental/auto_sharding/auto_sharding.h

@@ -207,6 +207,7 @@ bool HasReduceScatterOpportunity(const HloInstruction* inst,
                                  const ConstInstructionSet& modified);
 
 HloSharding GetReduceScatterOutput(const HloInstruction* ins,
+                                   const InputShardings& input_shardings,
                                    const ShardingStrategy& strategy,
                                    const ClusterEnvironment& cluster_env);
 

xla/hlo/experimental/auto_sharding/auto_sharding_cost_graph.cc

@@ -123,16 +123,26 @@ CostGraph::CostGraph(const StrategyGroups& strategy_groups,
                                        node_lens_[dst_idx]);
     absl::flat_hash_map<std::string, NodeStrategyIdx>
         src_strategy_name_to_idx_map;
-    for (NodeStrategyIdx i = 0; i < node_lens_[src_idx]; ++i) {
+    const auto& src_strategy_input_shardings =
+        src_strategy_group.GetStrategyInputShardings();
+    for (size_t iid = 0; iid < src_strategy_input_shardings.size(); ++iid) {
+      const InputShardings& input_shardings = src_strategy_input_shardings[iid];
+      NodeStrategyIdx i =
+          src_strategy_group.GetStrategyIdxForInputShardings(iid);
       const ShardingStrategy& strategy = src_strategy_group.GetStrategy(i);
       if (strategy.communication_cost > 0) {
-        src_strategy_name_to_idx_map[strategy.name] = i;
+        src_strategy_name_to_idx_map[input_shardings.name] = i;
       }
     }
-    for (NodeStrategyIdx i = 0; i < node_lens_[dst_idx]; ++i) {
+    const auto& dst_strategy_input_shardings =
+        dst_strategy_group.GetStrategyInputShardings();
+    for (size_t iid = 0; iid < dst_strategy_input_shardings.size(); ++iid) {
+      const InputShardings& input_shardings = dst_strategy_input_shardings[iid];
+      NodeStrategyIdx i =
+          dst_strategy_group.GetStrategyIdxForInputShardings(iid);
       const ShardingStrategy& dst_strategy = dst_strategy_group.GetStrategy(i);
       if (dst_strategy.communication_cost > 0) {
-        auto it = src_strategy_name_to_idx_map.find(dst_strategy.name);
+        auto it = src_strategy_name_to_idx_map.find(input_shardings.name);
         if (it != src_strategy_name_to_idx_map.end()) {
           const auto& src_strategy = src_strategy_group.GetStrategy(it->second);
           CHECK_LE(std::abs(src_strategy.communication_cost -

xla/hlo/experimental/auto_sharding/auto_sharding_dot_handler.cc

@@ -335,12 +335,12 @@ void HandlerBase::AppendNewStrategy(const std::string& name,
   }
 
   strategy_group_->AddStrategy(
-      ShardingStrategy({name, output_spec, compute_cost, communication_cost,
+      ShardingStrategy({output_spec, compute_cost, communication_cost,
                         static_cast<double>(ByteSizeOfShapeWithSharding(
                             ins_->shape(), output_spec)),
                         communication_resharding_costs,
                         memory_resharding_costs}),
-      {input_specs.begin(), input_specs.end()});
+      {name, {input_specs.begin(), input_specs.end()}});
 }
 
 // Given lhs and rhs dim maps, infers a sharding for the output by relying
@@ -467,7 +467,7 @@ void HandlerBase::SortStrategies() {
       [](const std::pair<ShardingStrategy, InputShardings>& s1,
          const std::pair<ShardingStrategy, InputShardings>& s2) {
         if (s1.first.memory_cost == s2.first.memory_cost) {
-          return s1.first.name < s2.first.name;
+          return s1.second.name < s2.second.name;
         } else {
           return s1.first.memory_cost < s2.first.memory_cost;
         }

xla/hlo/experimental/auto_sharding/auto_sharding_strategy.cc

@@ -346,14 +346,14 @@ BuildStrategyAndCost(
               ByteSizeOfShapeWithSharding(ins->shape(), scatter_sharding);
 
           InputShardings input_shardings_optional(
-              {data_sharding, indices_sharding, update_sharding});
+              {name, {data_sharding, indices_sharding, update_sharding}});
           std::pair<ReshardingCosts, ReshardingCosts> resharding_costs =
               GenerateReshardingCostsAndMissingShardingsForAllOperands(
                   ins, scatter_sharding, strategy_map, cluster_env, call_graph,
                   input_shardings_optional);
 
           strategy_group->AddStrategy(
-              ShardingStrategy({name, scatter_sharding, compute_cost,
+              ShardingStrategy({scatter_sharding, compute_cost,
                                 communication_cost, memory_cost,
                                 std::move(resharding_costs.first),
                                 std::move(resharding_costs.second)}),
@@ -397,15 +397,14 @@ BuildStrategyAndCost(
           double memory_cost =
               ByteSizeOfShapeWithSharding(gather_shape, output_sharding);
           InputShardings input_shardings_optional(
-              {data_sharding, indices_sharding});
+              {output_sharding.ToString(), {data_sharding, indices_sharding}});
           std::pair<ReshardingCosts, ReshardingCosts> resharding_costs =
               GenerateReshardingCostsAndMissingShardingsForAllOperands(
                   ins, output_sharding, strategy_map, cluster_env, call_graph,
                   input_shardings_optional);
 
           strategy_group->AddStrategy(
-              ShardingStrategy({std::string(output_sharding.ToString()),
-                                output_sharding, compute_cost,
+              ShardingStrategy({output_sharding, compute_cost,
                                 communication_cost, memory_cost,
                                 std::move(resharding_costs.first),
                                 std::move(resharding_costs.second)}),
@@ -565,14 +564,13 @@ BuildStrategyAndCost(
               MemoryReshardingCostVector(src_strategy_group, operand->shape(),
                                          input_spec, cluster_env);
           strategy_group->AddStrategy(
-              ShardingStrategy({name,
-                                output_spec,
+              ShardingStrategy({output_spec,
                                 compute_cost,
                                 communication_cost,
                                 memory_cost,
                                 {communication_resharding_costs},
                                 {memory_resharding_costs}}),
-              {input_spec});
+              {name, {input_spec}});
         }
         break;
       }
@@ -685,9 +683,9 @@ BuildStrategyAndCost(
             if (k == follow_idx ||
                 ToString(ins->operand(k)->shape().dimensions()) ==
                     ToString(operand->shape().dimensions())) {
-              input_shardings.push_back(input_spec);
+              input_shardings.shardings.push_back(input_spec);
             } else {
-              input_shardings.push_back(std::nullopt);
+              input_shardings.shardings.push_back(std::nullopt);
             }
           }
           if (!output_spec.has_value()) {
@@ -703,11 +701,10 @@ BuildStrategyAndCost(
                   input_shardings);
 
           strategy_group->AddStrategy(
-              ShardingStrategy({name, *output_spec, compute_cost,
-                                communication_cost, memory_cost,
-                                std::move(resharding_costs.first),
+              ShardingStrategy({*output_spec, compute_cost, communication_cost,
+                                memory_cost, std::move(resharding_costs.first),
                                 std::move(resharding_costs.second)}),
-              {input_spec});
+              {name, {input_spec}});
         }
 
         if (strategy_group->GetStrategies().empty()) {
@@ -1094,7 +1091,7 @@ BuildStrategyAndCost(
           const InputShardings& input_shardings = strategy_input_shardings[iid];
           const ShardingStrategy& strategy =
               strategy_group->GetStrategyForInputShardings(iid);
-          if (strategy.name == stra_names[idx]) {
+          if (input_shardings.name == stra_names[idx]) {
             new_strategies.push_back({strategy, input_shardings});
           }
         }

xla/hlo/experimental/auto_sharding/auto_sharding_strategy.h

@@ -77,12 +77,37 @@ using ReshardingCache =
     ConstInstructionMap<std::vector<std::pair<HloSharding, HloInstruction*>>>;
 // Resharding costs for each operand
 using ReshardingCosts = std::vector<std::vector<double>>;
-// Optional shardings for each operand
-using InputShardings = std::vector<std::optional<HloSharding>>;
+
+// A named vector of optional shardings for each operand.
+struct InputShardings {
+  std::string name;
+  std::vector<std::optional<HloSharding>> shardings;
+
+  std::string ToString() const {
+    std::string str = absl::StrCat(name, " ");
+    for (const auto& s : shardings) {
+      if (!s.has_value()) {
+        absl::StrAppend(&str, "[*],");
+      } else if (s->IsReplicated()) {
+        absl::StrAppend(&str, "[R],");
+      } else {
+        if (s->ReplicateOnLastTileDim()) {
+          absl::StrAppend(
+              &str, "[", absl::StrJoin(s->tile_assignment().dimensions(), ", "),
+              "]last_tile_dim_replicate,");
+        } else {
+          absl::StrAppend(
+              &str, "[", absl::StrJoin(s->tile_assignment().dimensions(), ", "),
+              "],");
+        }
+      }
+    }
+    return str;
+  }
+};
 
 // One sharding strategy
 struct ShardingStrategy {
-  std::string name;
   HloSharding output_sharding;
   double compute_cost;
   double communication_cost;
@@ -94,9 +119,7 @@ struct ShardingStrategy {
   ReshardingCosts communication_resharding_costs;
   ReshardingCosts memory_resharding_costs;
 
-  std::string ToString() const {
-    return absl::StrCat(name, ", ", output_sharding.ToString());
-  }
+  std::string ToString() const { return output_sharding.ToString(); }
 
   std::string ToStringLong() const {
     std::vector<std::string> communication_resharding_vector_strings;
@@ -119,15 +142,15 @@ struct ShardingStrategy {
         "{", absl::StrJoin(memory_resharding_vector_strings, ", "), "}");
 
     return absl::StrCat(
-        name, ", ", output_sharding.ToString(), ", compute_cost=", compute_cost,
+        output_sharding.ToString(), ", compute_cost=", compute_cost,
         ", communication_cost=", communication_cost,
         ", memory_cost=", memory_cost,
         ", communication_resharding_costs=", communication_resharding_cost_str,
         ", memory_resharding_costs=", memory_resharding_cost_str);
   }
 
   bool operator==(const ShardingStrategy& other) const {
-    return name == other.name && output_sharding == other.output_sharding &&
+    return output_sharding == other.output_sharding &&
            compute_cost == other.compute_cost &&
            communication_cost == other.communication_cost &&
            memory_cost == other.memory_cost &&
@@ -221,25 +244,8 @@ struct StrategyGroup {
     }
     if (!is_tuple) {
       for (const auto& input_shardings : strategy_input_shardings) {
-        std::string input_sharding_str = "{";
-        for (const auto& s : input_shardings) {
-          if (!s.has_value()) {
-            input_sharding_str += "[*],";
-          } else if (s->IsReplicated()) {
-            input_sharding_str += "[R],";
-          } else {
-            if (s->ReplicateOnLastTileDim()) {
-              input_sharding_str +=
-                  "[" + absl::StrJoin(s->tile_assignment().dimensions(), ", ") +
-                  "]last_tile_dim_replicate,";
-            } else {
-              input_sharding_str +=
-                  "[" + absl::StrJoin(s->tile_assignment().dimensions(), ", ") +
-                  "],";
-            }
-          }
-        }
-        input_sharding_str += "}\n";
+        const std::string input_sharding_str =
+            absl::StrCat("{", input_shardings.ToString(), "}\n");
         absl::StrAppend(&str, indent, "Input Sharding ", input_sharding_str);
       }
     }
@@ -313,6 +319,10 @@ struct StrategyGroup {
     return strategies[strategy_idx];
   }
 
+  size_t GetStrategyIdxForInputShardings(size_t input_sharding_idx) const {
+    return input_sharding_idx_to_strategy_idx[input_sharding_idx];
+  }
+
   const InputShardings& GetInputShardings(size_t input_sharding_idx) const {
     return strategy_input_shardings[input_sharding_idx];
   }

xla/hlo/experimental/auto_sharding/auto_sharding_util.cc

@@ -851,8 +851,8 @@ void RemoveDuplicatedStrategy(StrategyGroup& strategy_group) {
       continue;
     }
     std::string key = strategy.output_sharding.ToString();
-    if (!input_shardings.empty()) {
-      for (const auto& sharding : input_shardings) {
+    if (!input_shardings.shardings.empty()) {
+      for (const auto& sharding : input_shardings.shardings) {
         key += "/" + (sharding.has_value() ? sharding->ToString() : "none");
       }
     }