Address antiagainst's PR comment

Author: sogartar

mlir/include/mlir/IR/Dialect.h

@@ -216,6 +216,9 @@ class Dialect {
         {TypeID::get<ConcreteT>(), InterfaceT::getInterfaceID()});
   }
 
+  // Declare the same interface for multiple types.
+  // Example:
+  // declarePromisedInterfaces<FunctionOpInterface, MyFuncType1, MyFuncType2>()
   template <typename InterfaceT, typename... ConcreteT>
   void declarePromisedInterfaces() {
     (declarePromisedInterface<ConcreteT, InterfaceT>(), ...);

mlir/lib/Dialect/Linalg/Transforms/AllInterfaces.cpp

@@ -1,4 +1,4 @@
-//===- AllInterfaces.cpp - --------------------------------------*- C++ -*-===//
+//===- AllInterfaces.cpp - ------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

mlir/lib/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.cpp

@@ -59,6 +59,15 @@ static ReductionKind getReductionKind(Operation *op) {
       .Case([](arith::OrIOp op) { return ReductionKind::BitwiseOr; })
       .Case([](arith::XOrIOp op) { return ReductionKind::BitwiseXor; })
       .Case([](arith::AndIOp op) { return ReductionKind::Sum; })
+      // TODO: handle signless, signed and unsigned types properly.
+      // It is assumed that the element type of the collective operands and
+      // result drive the meaning of the reduction kind, whether it is signed
+      // or unsigned.
+      // The reduction op inside the linalg op may have different result type
+      // from the element type of the linalg op's result.
+      // Also signed and unsigned Arith dialect ops may accept signed, unsigned
+      // or signless operands.
+      // Maybe expand the reduction kinds.
       .Case([](arith::MaxUIOp op) { return ReductionKind::Max; })
       .Case([](arith::MinUIOp op) { return ReductionKind::Min; })
       .Case([](arith::MaxSIOp op) { return ReductionKind::Max; })
@@ -67,7 +76,7 @@ static ReductionKind getReductionKind(Operation *op) {
       .Default([](Operation *op) { return ReductionKind::Generic; });
 }
 
-static std::optional<Operation *> getReductionOp(LinalgOp op) {
+static std::optional<Operation *> getCombinerOp(LinalgOp op) {
   SmallVector<Operation *> combinerOps;
   Value reducedValue = matchReduction(op.getRegionOutputArgs(), 0, combinerOps);
   if (!reducedValue || combinerOps.size() != 1) {
@@ -78,10 +87,16 @@ static std::optional<Operation *> getReductionOp(LinalgOp op) {
 }
 
 static ReductionKind getReductionKindOfLinalgOp(LinalgOp op) {
-  std::optional<Operation *> reductionOp = getReductionOp(op);
+  std::optional<Operation *> reductionOp = getCombinerOp(op);
   if (!reductionOp) {
     return ReductionKind::Generic;
   }
+  Type resultElementType =
+      llvm::cast<RankedTensorType>(op->getResult(0).getType()).getElementType();
+  // TODO: handle case when result type of the reduction op does not match the
+  // element type of the result tensor.
+  // Would it makes sense at all?
+  assert(resultElementType == reductionOp.value()->getResult(0).getType());
   return getReductionKind(reductionOp.value());
 }
 
@@ -276,9 +291,8 @@ struct StructuredOpShardingInterface
         });
     if (!allIndexingMapsAreProjectedPermutation) {
       // TODO: handle non-projected permutations.
-      op->emitOpError()
-          << "Only projected permutation indexing maps are supported.";
-      return failure();
+      return op->emitOpError()
+             << "supports indexing maps that are only projected permutation.";
     }
 
     SmallVector<utils::IteratorType> loopIteratorTypes =

mlir/lib/Dialect/Mesh/Interfaces/ShardingInterface.cpp

@@ -539,8 +539,9 @@ static bool areValuesCompatibleWithFullReplicationShardings(
   if (std::size(values) != std::size(shardings)) {
     return false;
   }
-  return llvm::all_of(llvm::zip(std::forward<ValueRange>(values),
-                                std::forward<MeshShardingAttrRage>(shardings)),
+  return llvm::all_of(llvm::zip_equal(
+                          std::forward<ValueRange>(values),
+                          std::forward<MeshShardingAttrRage>(shardings)),
                       [](auto valueAndSharding) {
                         return isValueCompatibleWithFullReplicationSharding(
                             std::get<0>(valueAndSharding),
@@ -588,11 +589,9 @@ ShardingArray mesh::getMeshAxisAssignmentForLoopIterators(
   SmallVector<MeshShardingAttr> operatorAndResultShardings;
   operatorAndResultShardings.reserve(operandShardings.size() +
                                      resultShardings.size());
-  operatorAndResultShardings.insert(operatorAndResultShardings.end(),
-                                    operandShardings.begin(),
-                                    operandShardings.end());
+  llvm::append_range(operatorAndResultShardings, operandShardings);
   for (auto [sharding, affineMap] :
-       llvm::zip(operatorAndResultShardings, indexingMaps)) {
+       llvm::zip_equal(operatorAndResultShardings, indexingMaps)) {
     if (!sharding) {
       continue;
     }
@@ -602,6 +601,12 @@ ShardingArray mesh::getMeshAxisAssignmentForLoopIterators(
           meshAxesAssignmentForTensorAxis.asArrayRef(), indexingExpr,
           meshAxisAssignmentForLoopIterators);
     }
+    // Missing trailing split axes means replication on those tensor dimensions.
+    for (unsigned i = sharding.getSplitAxes().size();
+         i < affineMap.getNumResults(); ++i) {
+      updateMeshAxisAssignmentForLoopIterators(
+          {}, affineMap.getResults()[i], meshAxisAssignmentForLoopIterators);
+    }
   }
 
   ShardingArray res;
@@ -619,7 +624,7 @@ bool mesh::isAtLeastOneReductionIteratorSharded(
     ArrayRef<utils::IteratorType> loopIteratorTypes,
     ArrayRef<SmallVector<MeshAxis>> meshAxisAssignmentForLoopIterators) {
   for (auto [loopIteratorType, meshAxisAssignment] :
-       llvm::zip(loopIteratorTypes, meshAxisAssignmentForLoopIterators)) {
+       llvm::zip_equal(loopIteratorTypes, meshAxisAssignmentForLoopIterators)) {
     if (loopIteratorType == utils::IteratorType::reduction &&
         !meshAxisAssignment.empty()) {
       return true;
@@ -633,10 +638,9 @@ SmallVector<MeshAxis> mesh::getReductionMeshAxes(
     ArrayRef<SmallVector<MeshAxis>> meshAxisAssignmentForLoopIterators) {
   SmallVector<MeshAxis> meshAxes;
   for (auto [loopIteratorType, meshAxisAssignment] :
-       llvm::zip(loopIteratorTypes, meshAxisAssignmentForLoopIterators)) {
+       llvm::zip_equal(loopIteratorTypes, meshAxisAssignmentForLoopIterators)) {
     if (loopIteratorType == utils::IteratorType::reduction) {
-      meshAxes.insert(meshAxes.end(), meshAxisAssignment.begin(),
-                      meshAxisAssignment.end());
+      llvm::append_range(meshAxes, meshAxisAssignment);
     }
   }
   return meshAxes;
@@ -651,7 +655,7 @@ void mesh::spmdizeTriviallyShardableOperation(
   Operation *newOp = builder.clone(op, spmdizationMap);
   // Set the result types to the sharded counterparts.
   for (auto [oldResult, newResult, sharding] :
-       llvm::zip(op.getResults(), newOp->getResults(), resultShardings)) {
+       llvm::zip_equal(op.getResults(), newOp->getResults(), resultShardings)) {
     newResult.setType(shardType(newResult.getType(),
                                 getMesh(&op, sharding.getMesh(), symbolTable),
                                 sharding));