#sdy Move AddAxisOrMergeInserter to Shady OSS utils.

PiperOrigin-RevId: 776780231

Author: Google-ML-Automation

shardy/dialect/sdy/ir/utils.h

@@ -17,6 +17,7 @@ limitations under the License.
 #define SHARDY_DIALECT_SDY_IR_UTILS_H_
 
 #include <cstdint>
+#include <iterator>
 #include <optional>
 #include <string>
 #include <utility>
@@ -498,6 +499,43 @@ bool hasOnlyUsersOfType(Operation* op) {
   });
 }
 
+// Helper class to insert or merge `AxisRefAttr`s into a `SmallVector`.
+// It implements the interface required by `llvm::transform` to be used as an
+// output iterator.
+class AddAxisOrMergeInserter {
+ public:
+  using iterator_category = std::output_iterator_tag;
+  using value_type = void;
+  using difference_type = void;
+  using pointer = void;
+  using reference = void;
+
+  explicit AddAxisOrMergeInserter(SmallVector<AxisRefAttr>* newAxisRefs,
+                                  const MeshAttr* mesh)
+      : axisRefs(newAxisRefs), mesh(mesh) {}
+
+  AddAxisOrMergeInserter& operator=(AxisRefAttr value) {
+    sdy::addAxisOrMerge(*axisRefs, value, *mesh);
+    return *this;
+  }
+
+  AddAxisOrMergeInserter& operator=(ArrayRef<AxisRefAttr> values) {
+    for (AxisRefAttr value : values) {
+      *this = value;
+    }
+    return *this;
+  }
+
+  AddAxisOrMergeInserter& operator*() { return *this; }
+  AddAxisOrMergeInserter& operator++() { return *this; }
+  AddAxisOrMergeInserter& operator++(int) { return *this; }
+
+ private:
+  // Use pointers so that callers like `llvm::transform` can copy the inserter.
+  SmallVector<AxisRefAttr>* axisRefs;
+  const MeshAttr* mesh;
+};
+
 }  // namespace sdy
 }  // namespace mlir
 

shardy/dialect/sdy/ir/utils_test.cc

@@ -39,6 +39,8 @@ namespace {
 using ::testing::ElementsAre;
 using ::testing::IsEmpty;
 
+using AxisRefVector = SmallVector<AxisRefAttr>;
+
 class UtilsTest : public ShardyTestBase {
  protected:
   void SetUp() override {
@@ -430,11 +432,10 @@ TEST_F(UtilsTest, GetAxisSetDiff) {
 TEST_F(UtilsTest, SortAndMergeAxes) {
   MeshAttr mesh = createMesh({{"a", 8}, {"b", 8}, {"c", 8}, {"d", 8}});
 
-  SmallVector<AxisRefAttr> axes = {
-      createSubAxis("a", 1, 2), createSubAxis("a", 4, 2),
-      createSubAxis("b", 1, 2), createSubAxis("b", 2, 2),
-      createSubAxis("b", 4, 2), createAxis("c"),
-      createSubAxis("d", 2, 2), createSubAxis("d", 4, 2)};
+  AxisRefVector axes = {createSubAxis("a", 1, 2), createSubAxis("a", 4, 2),
+                        createSubAxis("b", 1, 2), createSubAxis("b", 2, 2),
+                        createSubAxis("b", 4, 2), createAxis("c"),
+                        createSubAxis("d", 2, 2), createSubAxis("d", 4, 2)};
 
   // Shuffle the vector
   std::random_device rd;
@@ -448,6 +449,68 @@ TEST_F(UtilsTest, SortAndMergeAxes) {
                                 AxisRefIs("c"), SubAxisRefIs("d", 2, 4)));
 }
 
+TEST_F(UtilsTest, AddAxisOrMergeInserterSingleAxisNoMerge) {
+  MeshAttr mesh = createMesh({{"a", 8}, {"b", 8}});
+  AxisRefVector axes = {createSubAxis("a", 1, 2), createAxis("b")};
+  AxisRefVector newAxes;
+  std::transform(axes.begin(), axes.end(),
+                  AddAxisOrMergeInserter(&newAxes, &mesh),
+                  [&](AxisRefAttr axis) {
+                    if (axis.getName() == "b") {
+                      return createSubAxis("c", 4, 2);
+                    }
+                    return axis;
+                  });
+  EXPECT_THAT(newAxes,
+              ElementsAre(SubAxisRefIs("a", 1, 2), SubAxisRefIs("c", 4, 2)));
+}
+
+TEST_F(UtilsTest, AddAxisOrMergeInserterMultipleAxesNoMerge) {
+  MeshAttr mesh = createMesh({{"a", 8}, {"b", 8}, {"c", 8}, {"d", 8}});
+  AxisRefVector axes = {createAxis("a"), createSubAxis("b", 1, 2)};
+  AxisRefVector newAxes;
+  std::transform(axes.begin(), axes.end(),
+                  AddAxisOrMergeInserter(&newAxes, &mesh),
+                  [&](AxisRefAttr axis) -> AxisRefVector {
+                    if (axis.getName() == "b") {
+                      return {createAxis("c"), createAxis("d")};
+                    }
+                    return {};
+                  });
+  EXPECT_THAT(newAxes, ElementsAre(AxisRefIs("c"), AxisRefIs("d")));
+}
+
+TEST_F(UtilsTest, AddAxisOrMergeInserterSingleAxisMerge) {
+  MeshAttr mesh = createMesh({{"a", 8}, {"b", 8}});
+  AxisRefVector axes = {createSubAxis("a", 1, 2), createAxis("b")};
+  AxisRefVector newAxes;
+  std::transform(axes.begin(), axes.end(),
+                  AddAxisOrMergeInserter(&newAxes, &mesh),
+                  [&](AxisRefAttr axis) {
+                    if (axis.getName() == "b") {
+                      return createSubAxis("a", 2, 2);
+                    }
+                    return axis;
+                  });
+  EXPECT_THAT(newAxes, ElementsAre(SubAxisRefIs("a", 1, 4)));
+}
+
+TEST_F(UtilsTest, AddAxisOrMergeInserterMultipleAxesMerge) {
+  MeshAttr mesh = createMesh({{"a", 8}, {"b", 16}});
+  AxisRefVector axes = {createAxis("a"), createSubAxis("b", 8, 2)};
+  AxisRefVector newAxes;
+  std::transform(axes.begin(), axes.end(),
+                  AddAxisOrMergeInserter(&newAxes, &mesh),
+                  [&](AxisRefAttr axis) -> AxisRefVector {
+                    if (axis.getName() == "a") {
+                      return {createSubAxis("b", 2, 2),
+                              createSubAxis("b", 4, 2)};
+                    }
+                    return {axis};
+                  });
+  EXPECT_THAT(newAxes, ElementsAre(SubAxisRefIs("b", 2, 8)));
+}
+
 }  // namespace
 
 }  // namespace sdy