[Mlir-commits] [mlir] [mlir][mesh, mpi] Lower allreduce (PR #144060)

[Christian Ulmann]

================
@@ -41,6 +42,38 @@ struct FoldCast final : public mlir::OpRewritePattern<OpT> {
     return mlir::success();
   }
 };
+
+struct FoldRank final : public mlir::OpRewritePattern<mlir::mpi::CommRankOp> {
+  using mlir::OpRewritePattern<mlir::mpi::CommRankOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::mpi::CommRankOp op,
+                                mlir::PatternRewriter &b) const override {
+    auto comm = op.getComm();
+    if (!comm.getDefiningOp<mlir::mpi::CommWorldOp>()) {
+      return mlir::failure();
+    }
+
+    // Try to get DLTI attribute for MPI:comm_world_rank
+    // If found, set worldRank to the value of the attribute.
+    {
+      auto dltiAttr = dlti::query(op, {"MPI:comm_world_rank"}, false);
+      if (failed(dltiAttr))
+        return mlir::failure();
+      if (!isa<IntegerAttr>(dltiAttr.value())) {
+        return op->emitError()
+               << "Expected an integer attribute for MPI:comm_world_rank";
+      }
----------------
Dinistro wrote:

```suggestion
      if (!isa<IntegerAttr>(dltiAttr.value()))
        return op->emitError()
               << "Expected an integer attribute for MPI:comm_world_rank";
```

https://github.com/llvm/llvm-project/pull/144060