Add methods for checking coexistence, overlap, and prefix without overlap for AxisRef.

This change introduces:
-   `AxisRef::CanCoexist`: Determines if two AxisRefs can be part of the same sharding.
-   `AxisRef::Overlaps`: Checks if two AxisRefs have an overlapping range.
-   `AxisRef::GetPrefixWithoutOverlap`: Returns the portion of an AxisRef that does not overlap with another.
-   `SortAndMergeAxes`: Sorts and merges a vector of AxisRefs.
-   `TruncateAxesByRemovingOverlaps`: Truncates a vector of AxisRefs based on overlaps with another set.

PiperOrigin-RevId: 872818416

Author: Kingston Mandisodza

xla/hlo/ir/mesh_and_axis.cc

@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <algorithm>
 #include <cassert>
+#include <cstddef>
 #include <cstdint>
 #include <iterator>
 #include <memory>
@@ -30,6 +31,7 @@ limitations under the License.
 #include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
+#include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/numbers.h"
 #include "absl/strings/str_cat.h"
@@ -273,6 +275,81 @@ AxisRef::AxisRef(int64_t mesh_axis_index, SubAxis sub_axis_info)
   CHECK_GT(sub_axis_info_->size, 1) << "sub-axis size must be > 1";
 }
 
+namespace {
+
+bool CanSubAxesCoexist(int64_t min_pre_size, int64_t max_pre_size,
+                       int64_t min_next_pre_size, int64_t max_next_pre_size) {
+  if (min_next_pre_size > max_pre_size) {
+    // Sub-axes overlap, check if overlapping and non-overlapping parts are
+    // valid.
+    return min_next_pre_size % max_pre_size == 0 &&
+           max_pre_size % min_pre_size == 0 &&
+           max_next_pre_size % min_next_pre_size == 0;
+  }
+  // Sub-axes don't overlap, check if the gap is valid.
+  return max_pre_size % min_next_pre_size == 0;
+}
+
+}  // namespace
+
+bool AxisRef::CanCoexist(const AxisRef& other) const {
+  if (mesh_axis_index() != other.mesh_axis_index()) {
+    return true;
+  }
+  if (!sub_axis_info_.has_value() || !other.sub_axis_info_.has_value()) {
+    // One of the axes is full
+    return true;
+  }
+  const SubAxis& this_sub = *sub_axis_info_;
+  const SubAxis& other_sub = *other.sub_axis_info_;
+
+  auto [min_pre_size, max_pre_size] =
+      std::minmax(this_sub.pre_size, other_sub.pre_size);
+  int64_t min_next_pre_size =
+      std::min(this_sub.next_pre_size(), other_sub.next_pre_size());
+  int64_t max_next_pre_size =
+      std::max(this_sub.next_pre_size(), other_sub.next_pre_size());
+
+  return CanSubAxesCoexist(min_pre_size, max_pre_size, min_next_pre_size,
+                           max_next_pre_size);
+}
+
+bool AxisRef::Overlaps(const AxisRef& other) const {
+  if (mesh_axis_index() != other.mesh_axis_index()) {
+    return false;
+  }
+  if (!sub_axis_info_.has_value() || !other.sub_axis_info_.has_value()) {
+    // One of the axes is full
+    return true;
+  }
+  const SubAxis& this_sub = *sub_axis_info_;
+  const SubAxis& other_sub = *other.sub_axis_info_;
+
+  return this_sub.pre_size < other_sub.next_pre_size() &&
+         other_sub.pre_size < this_sub.next_pre_size();
+}
+
+std::optional<AxisRef> AxisRef::GetPrefixWithoutOverlap(
+    const AxisRef& other) const {
+  if (!CanCoexist(other)) {
+    return std::nullopt;
+  }
+  if (!Overlaps(other)) {
+    return *this;
+  }
+
+  int64_t this_pre_size =
+      sub_axis_info_.has_value() ? sub_axis_info_->pre_size : 1;
+  int64_t other_pre_size =
+      other.sub_axis_info_.has_value() ? other.sub_axis_info_->pre_size : 1;
+
+  if (this_pre_size >= other_pre_size) {
+    return std::nullopt;
+  }
+  return AxisRef(mesh_axis_index_,
+                 SubAxis{this_pre_size, other_pre_size / this_pre_size});
+}
+
 bool AxisRef::CanCoexistWithoutOverlap(const AxisRef& other) const {
   // Check if the axes are on different mesh dimensions. If so, they can always
   // coexist and never overlap.
@@ -414,4 +491,51 @@ absl::Status ValidateSpanOfAxes(absl::Span<const AxisRef> axes,
   return absl::OkStatus();
 }
 
+void SortAndMergeAxes(std::vector<AxisRef>& axes, const Mesh& mesh) {
+  if (axes.empty()) {
+    return;
+  }
+
+  absl::c_sort(axes);
+
+  auto current = axes.begin();
+  for (auto next = current + 1; next != axes.end(); ++next) {
+    if (current->Overlaps(*next)) {
+      LOG(FATAL) << "Axes should not overlap: " << current->ToString(&mesh)
+                 << " and " << next->ToString(&mesh);
+    }
+    if (current->CanMerge(*next)) {
+      CHECK(current->Merge(*next, mesh));
+    } else {
+      current++;
+      *current = *next;
+    }
+  }
+  axes.erase(current + 1, axes.end());
+}
+
+bool TruncateAxesByRemovingOverlaps(std::vector<AxisRef>& axes,
+                                    absl::Span<const AxisRef> other_axis_refs) {
+  for (size_t i = 0; i < axes.size(); ++i) {
+    std::optional<AxisRef> prefix = axes[i];
+    for (const AxisRef& other : other_axis_refs) {
+      prefix = prefix->GetPrefixWithoutOverlap(other);
+      if (!prefix) {
+        break;
+      }
+    }
+
+    if (!prefix) {
+      axes.erase(axes.begin() + i, axes.end());
+      return true;
+    }
+    if (axes[i] != *prefix) {
+      axes[i] = *prefix;
+      axes.erase(axes.begin() + i + 1, axes.end());
+      return true;
+    }
+  }
+  return false;
+}
+
 }  // namespace xla

xla/hlo/ir/mesh_and_axis.h

@@ -166,6 +166,22 @@ class AxisRef {
 
   bool operator!=(const xla::AxisRef& other) const { return !(*this == other); }
 
+  bool operator<(const AxisRef& other) const {
+    if (mesh_axis_index_ != other.mesh_axis_index_) {
+      return mesh_axis_index_ < other.mesh_axis_index_;
+    }
+    if (!sub_axis_info_.has_value()) {
+      return true;
+    }
+    if (!other.sub_axis_info_.has_value()) {
+      return false;
+    }
+    if (sub_axis_info_->pre_size != other.sub_axis_info_->pre_size) {
+      return sub_axis_info_->pre_size < other.sub_axis_info_->pre_size;
+    }
+    return sub_axis_info_->size < other.sub_axis_info_->size;
+  }
+
   template <typename H>
   friend H AbslHashValue(H h, const AxisRef& a) {
     return H::combine(std::move(h), a.mesh_axis_index_, a.sub_axis_info_);
@@ -179,6 +195,29 @@ class AxisRef {
 
   bool CanCoexistWithoutOverlap(const AxisRef& other) const;
 
+  // Returns whether this axis and `other` can coexist in the same mesh:
+  // * If they overlap, then both overlapping and non-overlapping parts must
+  //   be valid axes or sub-axes.
+  // * Otherwise, both axes can be used to shard the same tensor.
+  //
+  // For example:
+  //  "a", "b"           -> true
+  //  "a", "b":(2)2      -> true
+  //  "a", "a"           -> true
+  //  "a", "a":(2)2      -> true
+  //  "a":(1)2, "a":(4)2 -> true
+  //  "a":(1)4, "a":(2)4 -> true
+  //  "a":(1)2, "a":(1)3 -> false
+  //  "a":(1)2, "a":(3)2 -> false
+  //  "a":(1)3, "a":(2)3 -> false
+  bool CanCoexist(const AxisRef& other) const;
+
+  // Returns true if this axis overlaps with `other`.
+  bool Overlaps(const AxisRef& other) const;
+
+  // Returns the largest prefix of this axis that does not overlap with `other`.
+  std::optional<AxisRef> GetPrefixWithoutOverlap(const AxisRef& other) const;
+
   // Returns true if this AxisRef can be merged with the `other`, i.e., they are
   // consecutive sub-axes of same full axis and this sub-axis is major to other.
   bool CanMerge(const AxisRef& other) const;
@@ -209,6 +248,19 @@ absl::Status ValidateSpanOfAxes(absl::Span<const AxisRef> axes,
                                 const Mesh& mesh,
                                 bool allow_mergeable_neighbors = false);
 
+// Sorts and merges the axes in `axes`.
+//
+// The axes are sorted by `operator<` (mesh axis index, then pre-size).
+// Adjacent axes that overlap will cause a fatal error.
+// Adjacent axes that can be merged are merged.
+void SortAndMergeAxes(std::vector<AxisRef>& axes, const Mesh& mesh);
+
+// Removes parts of `axes` that overlap with any axis in `other_axis_refs`.
+//
+// Returns true if `axes` is modified.
+bool TruncateAxesByRemovingOverlaps(std::vector<AxisRef>& axes,
+                                    absl::Span<const AxisRef> other_axis_refs);
+
 }  // namespace xla
 
 #endif  // XLA_HLO_IR_MESH_AND_AXIS_H_

xla/hlo/ir/mesh_and_axis_test.cc

@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/hlo/ir/mesh_and_axis.h"
 
 #include <cstdint>
+#include <optional>
 #include <vector>
 
 #include <gmock/gmock.h>
@@ -432,4 +433,125 @@ TEST(MeshAndAxisTest, AxisRefMerge) {
   EXPECT_EQ(axis_ref4, AxisRef(0, {2, 4}));
 }
 
+TEST(MeshAndAxisTest, CanCoexist_DifferentAxes) {
+  EXPECT_TRUE(AxisRef(0).CanCoexist(AxisRef(1)));
+  EXPECT_TRUE(AxisRef(0, {1, 2}).CanCoexist(AxisRef(1, {1, 2})));
+}
+
+TEST(MeshAndAxisTest, CanCoexist_CompatibleSubAxes) {
+  EXPECT_TRUE(AxisRef(0, {1, 2}).CanCoexist(AxisRef(0, {2, 2})));
+  EXPECT_TRUE(AxisRef(0, {1, 2}).CanCoexist(AxisRef(0, {4, 2})));
+}
+
+TEST(MeshAndAxisTest, CanCoexist_CompatibleOverlappingSubAxes) {
+  EXPECT_TRUE(AxisRef(0, {1, 2}).CanCoexist(AxisRef(0, {1, 2})));
+  EXPECT_TRUE(AxisRef(0, {1, 2}).CanCoexist(AxisRef(0, {1, 4})));
+}
+
+TEST(MeshAndAxisTest, CanCoexist_IncompatibleSubAxes) {
+  EXPECT_FALSE(AxisRef(0, {1, 2}).CanCoexist(AxisRef(0, {1, 3})));
+  EXPECT_FALSE(AxisRef(0, {1, 2}).CanCoexist(AxisRef(0, {3, 2})));
+}
+
+TEST(MeshAndAxisTest, Overlaps_DifferentAxes) {
+  EXPECT_FALSE(AxisRef(0).Overlaps(AxisRef(1)));
+}
+
+TEST(MeshAndAxisTest, Overlaps_FullAxes) {
+  EXPECT_TRUE(AxisRef(0).Overlaps(AxisRef(0)));
+  EXPECT_TRUE(AxisRef(0).Overlaps(AxisRef(0, {1, 2})));
+}
+
+TEST(MeshAndAxisTest, Overlaps_SubAxes_NoOverlap) {
+  EXPECT_FALSE(AxisRef(0, {1, 2}).Overlaps(AxisRef(0, {2, 2})));
+}
+
+TEST(MeshAndAxisTest, Overlaps_SubAxes_Overlap) {
+  EXPECT_TRUE(AxisRef(0, {1, 2}).Overlaps(AxisRef(0, {1, 2})));
+  EXPECT_TRUE(AxisRef(0, {1, 2}).Overlaps(AxisRef(0, {1, 4})));
+  EXPECT_TRUE(AxisRef(0, {1, 4}).Overlaps(AxisRef(0, {2, 2})));
+}
+
+TEST(MeshAndAxisTest, GetPrefixWithoutOverlap_NoOverlap) {
+  EXPECT_EQ(AxisRef(0).GetPrefixWithoutOverlap(AxisRef(1)), AxisRef(0));
+}
+
+TEST(MeshAndAxisTest, GetPrefixWithoutOverlap_FullOverlap) {
+  EXPECT_EQ(AxisRef(0, {1, 4}).GetPrefixWithoutOverlap(AxisRef(0, {1, 2})),
+            std::nullopt);
+}
+
+TEST(MeshAndAxisTest, GetPrefixWithoutOverlap_PartialOverlap) {
+  EXPECT_EQ(AxisRef(0, {1, 4}).GetPrefixWithoutOverlap(AxisRef(0, {2, 2})),
+            AxisRef(0, {1, 2}));
+}
+
+TEST(MeshAndAxisTest, GetPrefixWithoutOverlap_PartialOverlap_NoPrefix) {
+  EXPECT_EQ(AxisRef(0, {1, 4}).GetPrefixWithoutOverlap(AxisRef(0, {1, 2})),
+            std::nullopt);
+}
+
+TEST(MeshAndAxisTest, SortAndMergeAxes) {
+  Mesh mesh({16, 16}, {"x", "y"});
+  std::vector<AxisRef> axes = {AxisRef(0, {2, 2}), AxisRef(0, {4, 2}),
+                               AxisRef(0, {1, 2}), AxisRef(1, {1, 2}),
+                               AxisRef(1, {4, 2}), AxisRef(1, {2, 2})};
+  SortAndMergeAxes(axes, mesh);
+
+  EXPECT_THAT(axes,
+              testing::ElementsAre(AxisRef(0, {1, 8}), AxisRef(1, {1, 8})));
+}
+
+TEST(MeshAndAxisTest, SortAndMergeAxesFull) {
+  Mesh mesh({4}, {"x"});
+  std::vector<AxisRef> axes = {AxisRef(0, {1, 2}), AxisRef(0, {2, 2})};
+  SortAndMergeAxes(axes, mesh);
+
+  EXPECT_THAT(axes, testing::ElementsAre(AxisRef(0)));
+}
+
+TEST(MeshAndAxisTest, TruncateAxesByRemovingOverlaps_PartialOverlap) {
+  std::vector<AxisRef> axes = {AxisRef(0, {1, 4})};
+  std::vector<AxisRef> other = {AxisRef(0, {2, 2})};
+
+  EXPECT_TRUE(TruncateAxesByRemovingOverlaps(axes, other));
+  EXPECT_THAT(axes, testing::ElementsAre(AxisRef(0, {1, 2})));
+}
+
+TEST(MeshAndAxisTest, TruncateAxesByRemovingOverlaps_FullOverlap) {
+  std::vector<AxisRef> axes = {AxisRef(0, {1, 2})};
+  std::vector<AxisRef> other = {AxisRef(0, {1, 4})};
+
+  EXPECT_TRUE(TruncateAxesByRemovingOverlaps(axes, other));
+  EXPECT_THAT(axes, testing::IsEmpty());
+}
+
+TEST(MeshAndAxisTest, TruncateAxesByRemovingOverlaps_MultipleAxes) {
+  std::vector<AxisRef> axes = {AxisRef(0, {1, 4}), AxisRef(1)};
+  std::vector<AxisRef> other = {AxisRef(0, {2, 2})};
+
+  EXPECT_TRUE(TruncateAxesByRemovingOverlaps(axes, other));
+  EXPECT_THAT(axes, testing::ElementsAre(AxisRef(0, {1, 2})));
+}
+
+TEST(MeshAndAxisTest, TruncateAxesByRemovingOverlaps_NoOverlap) {
+  std::vector<AxisRef> axes = {AxisRef(0, {1, 2}), AxisRef(1)};
+  std::vector<AxisRef> other = {AxisRef(2)};
+
+  EXPECT_FALSE(TruncateAxesByRemovingOverlaps(axes, other));
+  EXPECT_THAT(axes, testing::ElementsAre(AxisRef(0, {1, 2}), AxisRef(1)));
+}
+
+TEST(MeshAndAxisTest, OperatorLess) {
+  // "x" < "x":(1)2
+  EXPECT_TRUE(AxisRef(0) < AxisRef(0, {1, 2}));
+  EXPECT_FALSE(AxisRef(0, {1, 2}) < AxisRef(0));
+
+  // "x":(1)2 < "x":(2)2
+  EXPECT_TRUE(AxisRef(0, {1, 2}) < AxisRef(0, {2, 2}));
+
+  // Different axes
+  EXPECT_TRUE(AxisRef(0) < AxisRef(1));
+}
+
 }  // namespace xla