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

[Frank Schlimbach]

================
@@ -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,
----------------
fschlimb wrote:

I made this work on memref because its semantics are memref-like, e.g. the halos are updated in-place, the input memref gets mutated, it does not return a new tensor or memref. In general we can of course think about adding a updateHalo variant wich operates on tensors and simply returns a new tensor.

Generally we could do the same for other ops, but I thought we should do that once we see need for it. `updateHalo` is a very special operation which mostly applies to array computations, and is probably less relevant in the tensor/AI world. Currently I do not see that any of the other operations have memref semantics.

Notice: within spmdization and for `updateHalo` specifically, an `spmdize` for relevant ops (like an inplace  `array.insert_slice`) would insert `bufferization.to_memref` and `bufferization.to_tensor` ops appropriately around `updateHalo`. In our (@tkarna) experience this approach works fine (even with one-shot.bufferize) when using `restrict=true` in `bufferization.to_tensor`.

Wrt to when to apply this pass: generally it is probably a good idea to do the MPI conversion after bufferization. This will require changes to the op spec so that they are allowed to accept memrefs. 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.

I am planning to add `send` and `recv` in a follow-up PR. 





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