[Mlir-commits] [mlir] [mlir][mesh, MPI] Mesh2mpi (PR #104566)

[Chengji Yao]

================
@@ -0,0 +1,240 @@
+//===- MeshToMPI.cpp - Mesh to MPI  dialect conversion -----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a translation of Mesh communication ops tp MPI ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/MeshToMPI/MeshToMPI.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/MPI/IR/MPI.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Mesh/IR/MeshOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#define DEBUG_TYPE "mesh-to-mpi"
+#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
+
+namespace mlir {
+#define GEN_PASS_DEF_CONVERTMESHTOMPIPASS
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::mesh;
+
+namespace {
+
+// This pattern converts the mesh.update_halo operation to MPI calls
+struct ConvertUpdateHaloOp
+    : public mlir::OpRewritePattern<mlir::mesh::UpdateHaloOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::mesh::UpdateHaloOp op,
----------------
yaochengji wrote:

> I don't know enough about bufferization and tensor optimization to tell if converting to MPI right after spmdization (using to_memref/tensor) would disallow any optimization possible otherwise. Any insights are welcome.

In my understanding, usually we prefer to perform optimization on tensor type than memref type because RAW dependency is difficult to detect on memref type.


> I made this work on memref because its semantics are memref-like, e.g. the halos are updated in-place
Even if most of case it is updated in place, we could still let it support tensor type and make it "in-place" after bufferization.

Combine the two points together, I would suggest that make `updateHalo` support tensor type will make it easier to optimize the IR containing both `updateHalo` and other ops in mesh dialect. Because we only need to handle pure tensor types in that case.

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