minor fix

Author: yaochengji

mlir/include/mlir/Dialect/Mesh/Interfaces/ShardingInterface.h

@@ -35,13 +35,13 @@ struct ShardingOption {
       : shardingArray(std::move(shardingArray)), cluster(cluster) {}
 };
 
-// This method aims to retrieve the mesh sharding attribute (MeshShardingAttr)
-// for a given operation result.
-FailureOr<MeshShardingAttr> getMeshShardingAttr(OpResult result,
-                                                bool useOperandSharding);
+// This method retrieves the 'MeshShardingAttr' attribute from a given operation
+// result and includes the 'annotate_for_users' information.
+FailureOr<std::pair<bool, MeshShardingAttr>>
+getMeshShardingAttr(OpResult result);
 
-// This method aims to retrieve the mesh sharding attribute (MeshShardingAttr)
-// for a given operation operand.
+// This method retrieves the 'MeshShardingAttr' attribute from a given operation
+// operand and includes the 'annotate_for_users' information.
 FailureOr<std::pair<bool, MeshShardingAttr>>
 getMeshShardingAttr(OpOperand &opOperand);
 

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

@@ -92,8 +92,8 @@ checkOperandAffineExpr(AffineExpr expr, unsigned numDims) {
 // mesh::getMeshShardingAttr
 //===----------------------------------------------------------------------===//
 
-FailureOr<MeshShardingAttr> mesh::getMeshShardingAttr(OpResult result,
-                                                      bool useOperandSharding) {
+FailureOr<std::pair<bool, MeshShardingAttr>>
+mesh::getMeshShardingAttr(OpResult result) {
   Value val = result.cast<Value>();
   bool anyShardedForDef = llvm::any_of(val.getUsers(), [](Operation *user) {
     auto shardOp = llvm::dyn_cast<mesh::ShardOp>(user);
@@ -108,32 +108,31 @@ FailureOr<MeshShardingAttr> mesh::getMeshShardingAttr(OpResult result,
     if (!val.hasOneUse())
       return failure();
     auto shardOp = llvm::cast<mesh::ShardOp>(*val.getUsers().begin());
-    return shardOp.getShard();
-  } else if (useOperandSharding) {
-    bool anyShardedForUsers = llvm::any_of(val.getUsers(), [](Operation *user) {
-      auto shardOp = llvm::dyn_cast<mesh::ShardOp>(user);
-      if (!shardOp)
-        return false;
-      return shardOp.getAnnotateForUsers();
-    });
-    if (anyShardedForUsers) {
-      SmallVector<ShardOp> shardOps;
-      for (Operation *user : val.getUsers()) {
-        ShardOp shardOp = llvm::dyn_cast<ShardOp>(user);
-        if (shardOp)
-          shardOps.push_back(shardOp);
-      }
-      MeshShardingAttr shardForDef = shardOps[0].getShard();
-      for (size_t i = 1; i < shardOps.size(); ++i) {
-        // TODO: Deduce a reasonable mesh sharding attr for def when they are
-        // different
-        assert(shardOps[i].getShard() == shardForDef &&
-               "only support all shard ops have the same mesh sharding attr");
-      }
-      return shardForDef;
-    }
+    return std::make_pair(false, shardOp.getShard());
   }
 
+  bool anyShardedForUsers = llvm::any_of(val.getUsers(), [](Operation *user) {
+    auto shardOp = llvm::dyn_cast<mesh::ShardOp>(user);
+    if (!shardOp)
+      return false;
+    return shardOp.getAnnotateForUsers();
+  });
+  if (anyShardedForUsers) {
+    SmallVector<ShardOp> shardOps;
+    for (Operation *user : val.getUsers()) {
+      ShardOp shardOp = llvm::dyn_cast<ShardOp>(user);
+      if (shardOp)
+        shardOps.push_back(shardOp);
+    }
+    MeshShardingAttr shardForDef = shardOps[0].getShard();
+    for (size_t i = 1; i < shardOps.size(); ++i) {
+      // TODO: Deduce a reasonable mesh sharding attr for def when they are
+      // different
+      assert(shardOps[i].getShard() == shardForDef &&
+             "only support all shard ops have the same mesh sharding attr");
+    }
+    return std::make_pair(true, shardForDef);
+  }
   return failure();
 }
 
@@ -403,7 +402,9 @@ static LogicalResult addShardOp(OpBuilder &b, OpResult result,
                                 const ShardingOption &shardingOption,
                                 AffineMap map,
                                 ArrayRef<IteratorType> loopTypes) {
-  if (succeeded(getMeshShardingAttr(result, /*useOperandSharding*/ false)))
+  FailureOr<std::pair<bool, MeshShardingAttr>> maybeSharding =
+      getMeshShardingAttr(result);
+  if (succeeded(maybeSharding) && !maybeSharding->first)
     return success();
 
   auto resultType = result.getType().cast<RankedTensorType>();

mlir/lib/Dialect/Mesh/Transforms/ShardingPropagation.cpp

@@ -34,9 +34,10 @@ namespace {
 // Utilities
 //===----------------------------------------------------------------------===//
 
-// This method returns all possible sharding attributes. For example,
-// mustShardings = [shard0, None] and optionalShardings = [None, shard1], the
-// result will be [[shard0, shard1], [shard0, None]]
+// This method retrieves all potential sharding attributes, prioritizing
+// specific shardings. For example, mustShardings = [shard0, None] and
+// optionalShardings = [None, shard1], the result will be [[shard0, shard1],
+// [shard0, None]]
 static SmallVector<SmallVector<MeshShardingAttr>>
 getOrderedPossibleShardingAttrs(ArrayRef<MeshShardingAttr> mustShardings,
                                 ArrayRef<MeshShardingAttr> optionalShardings) {
@@ -92,15 +93,20 @@ LogicalResult visitOp(Operation *op, OpBuilder &builder) {
   }
 
   // collect MeshShardingAttr from results
-  SmallVector<MeshShardingAttr> resultShardings;
-  resultShardings.reserve(op->getNumResults());
+  SmallVector<MeshShardingAttr> allowConflictsResultShardings;
+  allowConflictsResultShardings.resize(op->getNumResults());
+  SmallVector<MeshShardingAttr> resultMustShardings;
+  resultMustShardings.resize(op->getNumResults());
   for (OpResult result : op->getResults()) {
-    FailureOr<MeshShardingAttr> shardAttr =
-        getMeshShardingAttr(result, /*useOperandSharding*/ true);
-    if (succeeded(shardAttr))
-      resultShardings.push_back(*shardAttr);
+    FailureOr<std::pair<bool, MeshShardingAttr>> maybeShardAttr =
+        getMeshShardingAttr(result);
+    if (failed(maybeShardAttr))
+      continue;
+    if (!maybeShardAttr->first)
+      resultMustShardings[result.getResultNumber()] = maybeShardAttr->second;
     else
-      resultShardings.push_back(nullptr);
+      allowConflictsResultShardings[result.getResultNumber()] =
+          maybeShardAttr->second;
   }
 
   // collect MeshShardingAttr from operands
@@ -114,8 +120,7 @@ LogicalResult visitOp(Operation *op, OpBuilder &builder) {
     if (failed(maybeShardAttr))
       continue;
 
-    bool annotateForUsers = maybeShardAttr->first;
-    if (annotateForUsers)
+    if (maybeShardAttr->first)
       operandMustShardings[opOperand.getOperandNumber()] =
           maybeShardAttr->second;
     else
@@ -127,14 +132,22 @@ LogicalResult visitOp(Operation *op, OpBuilder &builder) {
   SmallVector<SmallVector<MeshShardingAttr>> possibleOperandShardingAttrs =
       getOrderedPossibleShardingAttrs(operandMustShardings,
                                       allowConflictsOperandShardings);
+  SmallVector<SmallVector<MeshShardingAttr>> possibleResultShardingAttrs =
+      getOrderedPossibleShardingAttrs(resultMustShardings,
+                                      allowConflictsResultShardings);
   FailureOr<ShardingOption> finalShardingOption = failure();
-  for (ArrayRef<MeshShardingAttr> operandShardings :
-       possibleOperandShardingAttrs) {
-    FailureOr<ShardingOption> shardingOption =
-        shardingOp.getShardingOption(operandShardings, resultShardings);
-    if (succeeded(shardingOption)) {
-      finalShardingOption = shardingOption;
+  for (ArrayRef<MeshShardingAttr> resultShardings :
+       possibleResultShardingAttrs) {
+    if (succeeded(finalShardingOption))
       break;
+    for (ArrayRef<MeshShardingAttr> operandShardings :
+         possibleOperandShardingAttrs) {
+      FailureOr<ShardingOption> shardingOption =
+          shardingOp.getShardingOption(operandShardings, resultShardings);
+      if (succeeded(shardingOption)) {
+        finalShardingOption = shardingOption;
+        break;
+      }
     }
   }