Define AxisRefsList as SmallVector<SmallVector<AxisRefAttr> and use it.

PiperOrigin-RevId: 689864870
openxla · Oct 25, 2024 · c048282 · c048282
1 parent e5d73bd
commit c048282
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/shardy/dialect/sdy/transforms/propagation/aggressive_factor_propagation.cc b/shardy/dialect/sdy/transforms/propagation/aggressive_factor_propagation.cc
@@ -52,7 +52,7 @@ UpdateTensorShardings AggressiveFactorPropagation::propagateFactorShardings(
   }
 
   // Find the compatible major axes ignoring conflicts.
-  SmallVector<SmallVector<AxisRefAttr>> axesPerFactor;
+  AxesPerFactor axesPerFactor;
   axesPerFactor.reserve(factorSizes.size());
   bool allElementsAreEmpty = true;
   for (int64_t i = 0; i < factorSizes.size(); ++i) {

diff --git a/shardy/dialect/sdy/transforms/propagation/sharding_projection.h b/shardy/dialect/sdy/transforms/propagation/sharding_projection.h
@@ -27,6 +27,8 @@ limitations under the License.
 namespace mlir {
 namespace sdy {
 
+using AxesPerFactor = SmallVector<SmallVector<AxisRefAttr>>;
+
 // Returns true if the `oldAxes` is a strict prefix of `newAxes`,
 bool shouldUpdate(ArrayRef<AxisRefAttr> oldAxes, ArrayRef<AxisRefAttr> newAxes);