[Mlir-commits] [mlir] [mlir][mpi] Lowering Mpi To LLVM (PR #127053)

[Rolf Morel]

================
@@ -184,29 +185,53 @@ struct OMPIImplTraits {
 //===----------------------------------------------------------------------===//
 // When lowering the mpi dialect to functions calls certain details
 // differ between various MPI implementations. This class will provide
-// these in a gnereic way, depending on the MPI implementation that got
-// included.
+// these in a generic way, depending on the MPI implementation that got
+// selected by the DLTI attribute on the module.
 //===----------------------------------------------------------------------===//
 struct MPIImplTraits {
+  enum MPIImpl { MPICH, OMPI };
+
+  // Get the MPI implementation from a DLTI attribute on the module.
+  // Default to MPICH (and ABI compatible).
+  static MPIImpl getMPIImpl(mlir::ModuleOp &moduleOp) {
+    auto attr = dlti::query(*&moduleOp, {"MPI:Implementation"}, true);
+    if (failed(attr)) {
+      return MPICH;
+    }
+    auto strAttr = dyn_cast<StringAttr>(attr.value());
+    if (strAttr && strAttr.getValue() == "OpenMPI") {
+      return OMPI;
+    }
+    return MPICH;
+  }
+
   // get/create MPI_COMM_WORLD as a mlir::Value
   static mlir::Value getCommWorld(mlir::ModuleOp &moduleOp,
                                   const mlir::Location loc,
                                   mlir::ConversionPatternRewriter &rewriter) {
-    // TODO: dispatch based on the MPI implementation
+    if (MPIImplTraits::getMPIImpl(moduleOp) == OMPI) {
----------------
rolfmorel wrote:

Depending on how often `getCommWorld` and `getStatusIgnore` get invoked, I would recommend "pulling up" the call to `getMPIImpl` as far as possible, e.g. if feasible to `runOnOperation` / outermost function of your code which is first handed the ModuleOp. This is because there are instances where calling `dlti::query` could be expensive.

Having said that, these particular calls to `dlti::query` are really the cheapest possible invocation of it (no traversal of ancestors - no nested attributes), so I very much doubt it would be problematic unless `getCommWorld` and `getStatusIgnore` get invoked often.

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