MLIR  22.0.0git
MPIToLLVM.cpp
Go to the documentation of this file.
1 //===- MPIToLLVM.cpp - MPI to LLVM dialect conversion ---------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 //
10 // Copyright (C) by Argonne National Laboratory
11 // See COPYRIGHT in top-level directory
12 // of MPICH source repository.
13 //
14 
19 #include "mlir/Dialect/DLTI/DLTI.h"
24 #include <memory>
25 
26 using namespace mlir;
27 
28 namespace {
29 
30 template <typename Op, typename... Args>
31 static Op getOrDefineGlobal(ModuleOp &moduleOp, const Location loc,
32  ConversionPatternRewriter &rewriter, StringRef name,
33  Args &&...args) {
34  Op ret;
35  if (!(ret = moduleOp.lookupSymbol<Op>(name))) {
36  ConversionPatternRewriter::InsertionGuard guard(rewriter);
37  rewriter.setInsertionPointToStart(moduleOp.getBody());
38  ret = Op::create(rewriter, loc, std::forward<Args>(args)...);
39  }
40  return ret;
41 }
42 
43 static LLVM::LLVMFuncOp getOrDefineFunction(ModuleOp &moduleOp,
44  const Location loc,
45  ConversionPatternRewriter &rewriter,
46  StringRef name,
47  LLVM::LLVMFunctionType type) {
48  return getOrDefineGlobal<LLVM::LLVMFuncOp>(
49  moduleOp, loc, rewriter, name, name, type, LLVM::Linkage::External);
50 }
51 
52 std::pair<Value, Value> getRawPtrAndSize(const Location loc,
53  ConversionPatternRewriter &rewriter,
54  Value memRef, Type elType) {
55  Type ptrType = LLVM::LLVMPointerType::get(rewriter.getContext());
56  Value dataPtr =
57  LLVM::ExtractValueOp::create(rewriter, loc, ptrType, memRef, 1);
58  Value offset = LLVM::ExtractValueOp::create(rewriter, loc,
59  rewriter.getI64Type(), memRef, 2);
60  Value resPtr =
61  LLVM::GEPOp::create(rewriter, loc, ptrType, elType, dataPtr, offset);
62  Value size;
63  if (cast<LLVM::LLVMStructType>(memRef.getType()).getBody().size() > 3) {
64  size = LLVM::ExtractValueOp::create(rewriter, loc, memRef,
65  ArrayRef<int64_t>{3, 0});
66  size = LLVM::TruncOp::create(rewriter, loc, rewriter.getI32Type(), size);
67  } else {
68  size = arith::ConstantIntOp::create(rewriter, loc, 1, 32);
69  }
70  return {resPtr, size};
71 }
72 
73 /// When lowering the mpi dialect to functions calls certain details
74 /// differ between various MPI implementations. This class will provide
75 /// these in a generic way, depending on the MPI implementation that got
76 /// selected by the DLTI attribute on the module.
77 class MPIImplTraits {
78  ModuleOp &moduleOp;
79 
80 public:
81  /// Instantiate a new MPIImplTraits object according to the DLTI attribute
82  /// on the given module. Default to MPICH if no attribute is present or
83  /// the value is unknown.
84  static std::unique_ptr<MPIImplTraits> get(ModuleOp &moduleOp);
85 
86  explicit MPIImplTraits(ModuleOp &moduleOp) : moduleOp(moduleOp) {}
87 
88  virtual ~MPIImplTraits() = default;
89 
90  ModuleOp &getModuleOp() { return moduleOp; }
91 
92  /// Gets or creates MPI_COMM_WORLD as a Value.
93  /// Different MPI implementations have different communicator types.
94  /// Using i64 as a portable, intermediate type.
95  /// Appropriate cast needs to take place before calling MPI functions.
96  virtual Value getCommWorld(const Location loc,
97  ConversionPatternRewriter &rewriter) = 0;
98 
99  /// Type converter provides i64 type for communicator type.
100  /// Converts to native type, which might be ptr or int or whatever.
101  virtual Value castComm(const Location loc,
102  ConversionPatternRewriter &rewriter, Value comm) = 0;
103 
104  /// Get the MPI_STATUS_IGNORE value (typically a pointer type).
105  virtual intptr_t getStatusIgnore() = 0;
106 
107  /// Get the MPI_IN_PLACE value (void *).
108  virtual void *getInPlace() = 0;
109 
110  /// Gets or creates an MPI datatype as a value which corresponds to the given
111  /// type.
112  virtual Value getDataType(const Location loc,
113  ConversionPatternRewriter &rewriter, Type type) = 0;
114 
115  /// Gets or creates an MPI_Op value which corresponds to the given
116  /// enum value.
117  virtual Value getMPIOp(const Location loc,
118  ConversionPatternRewriter &rewriter,
119  mpi::MPI_ReductionOpEnum opAttr) = 0;
120 };
121 
122 //===----------------------------------------------------------------------===//
123 // Implementation details for MPICH ABI compatible MPI implementations
124 //===----------------------------------------------------------------------===//
125 
126 class MPICHImplTraits : public MPIImplTraits {
127  static constexpr int MPI_FLOAT = 0x4c00040a;
128  static constexpr int MPI_DOUBLE = 0x4c00080b;
129  static constexpr int MPI_INT8_T = 0x4c000137;
130  static constexpr int MPI_INT16_T = 0x4c000238;
131  static constexpr int MPI_INT32_T = 0x4c000439;
132  static constexpr int MPI_INT64_T = 0x4c00083a;
133  static constexpr int MPI_UINT8_T = 0x4c00013b;
134  static constexpr int MPI_UINT16_T = 0x4c00023c;
135  static constexpr int MPI_UINT32_T = 0x4c00043d;
136  static constexpr int MPI_UINT64_T = 0x4c00083e;
137  static constexpr int MPI_MAX = 0x58000001;
138  static constexpr int MPI_MIN = 0x58000002;
139  static constexpr int MPI_SUM = 0x58000003;
140  static constexpr int MPI_PROD = 0x58000004;
141  static constexpr int MPI_LAND = 0x58000005;
142  static constexpr int MPI_BAND = 0x58000006;
143  static constexpr int MPI_LOR = 0x58000007;
144  static constexpr int MPI_BOR = 0x58000008;
145  static constexpr int MPI_LXOR = 0x58000009;
146  static constexpr int MPI_BXOR = 0x5800000a;
147  static constexpr int MPI_MINLOC = 0x5800000b;
148  static constexpr int MPI_MAXLOC = 0x5800000c;
149  static constexpr int MPI_REPLACE = 0x5800000d;
150  static constexpr int MPI_NO_OP = 0x5800000e;
151 
152 public:
153  using MPIImplTraits::MPIImplTraits;
154 
155  ~MPICHImplTraits() override = default;
156 
157  Value getCommWorld(const Location loc,
158  ConversionPatternRewriter &rewriter) override {
159  static constexpr int MPI_COMM_WORLD = 0x44000000;
160  return LLVM::ConstantOp::create(rewriter, loc, rewriter.getI64Type(),
161  MPI_COMM_WORLD);
162  }
163 
164  Value castComm(const Location loc, ConversionPatternRewriter &rewriter,
165  Value comm) override {
166  return LLVM::TruncOp::create(rewriter, loc, rewriter.getI32Type(), comm);
167  }
168 
169  intptr_t getStatusIgnore() override { return 1; }
170 
171  void *getInPlace() override { return reinterpret_cast<void *>(-1); }
172 
173  Value getDataType(const Location loc, ConversionPatternRewriter &rewriter,
174  Type type) override {
175  int32_t mtype = 0;
176  if (type.isF32())
177  mtype = MPI_FLOAT;
178  else if (type.isF64())
179  mtype = MPI_DOUBLE;
180  else if (type.isInteger(64) && !type.isUnsignedInteger())
181  mtype = MPI_INT64_T;
182  else if (type.isInteger(64))
183  mtype = MPI_UINT64_T;
184  else if (type.isInteger(32) && !type.isUnsignedInteger())
185  mtype = MPI_INT32_T;
186  else if (type.isInteger(32))
187  mtype = MPI_UINT32_T;
188  else if (type.isInteger(16) && !type.isUnsignedInteger())
189  mtype = MPI_INT16_T;
190  else if (type.isInteger(16))
191  mtype = MPI_UINT16_T;
192  else if (type.isInteger(8) && !type.isUnsignedInteger())
193  mtype = MPI_INT8_T;
194  else if (type.isInteger(8))
195  mtype = MPI_UINT8_T;
196  else
197  assert(false && "unsupported type");
198  return LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
199  mtype);
200  }
201 
202  Value getMPIOp(const Location loc, ConversionPatternRewriter &rewriter,
203  mpi::MPI_ReductionOpEnum opAttr) override {
204  int32_t op = MPI_NO_OP;
205  switch (opAttr) {
206  case mpi::MPI_ReductionOpEnum::MPI_OP_NULL:
207  op = MPI_NO_OP;
208  break;
209  case mpi::MPI_ReductionOpEnum::MPI_MAX:
210  op = MPI_MAX;
211  break;
212  case mpi::MPI_ReductionOpEnum::MPI_MIN:
213  op = MPI_MIN;
214  break;
215  case mpi::MPI_ReductionOpEnum::MPI_SUM:
216  op = MPI_SUM;
217  break;
218  case mpi::MPI_ReductionOpEnum::MPI_PROD:
219  op = MPI_PROD;
220  break;
221  case mpi::MPI_ReductionOpEnum::MPI_LAND:
222  op = MPI_LAND;
223  break;
224  case mpi::MPI_ReductionOpEnum::MPI_BAND:
225  op = MPI_BAND;
226  break;
227  case mpi::MPI_ReductionOpEnum::MPI_LOR:
228  op = MPI_LOR;
229  break;
230  case mpi::MPI_ReductionOpEnum::MPI_BOR:
231  op = MPI_BOR;
232  break;
233  case mpi::MPI_ReductionOpEnum::MPI_LXOR:
234  op = MPI_LXOR;
235  break;
236  case mpi::MPI_ReductionOpEnum::MPI_BXOR:
237  op = MPI_BXOR;
238  break;
239  case mpi::MPI_ReductionOpEnum::MPI_MINLOC:
240  op = MPI_MINLOC;
241  break;
242  case mpi::MPI_ReductionOpEnum::MPI_MAXLOC:
243  op = MPI_MAXLOC;
244  break;
245  case mpi::MPI_ReductionOpEnum::MPI_REPLACE:
246  op = MPI_REPLACE;
247  break;
248  }
249  return LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(), op);
250  }
251 };
252 
253 //===----------------------------------------------------------------------===//
254 // Implementation details for OpenMPI
255 //===----------------------------------------------------------------------===//
256 class OMPIImplTraits : public MPIImplTraits {
257  LLVM::GlobalOp getOrDefineExternalStruct(const Location loc,
258  ConversionPatternRewriter &rewriter,
259  StringRef name,
260  LLVM::LLVMStructType type) {
261 
262  return getOrDefineGlobal<LLVM::GlobalOp>(
263  getModuleOp(), loc, rewriter, name, type, /*isConstant=*/false,
264  LLVM::Linkage::External, name,
265  /*value=*/Attribute(), /*alignment=*/0, 0);
266  }
267 
268 public:
269  using MPIImplTraits::MPIImplTraits;
270 
271  ~OMPIImplTraits() override = default;
272 
273  Value getCommWorld(const Location loc,
274  ConversionPatternRewriter &rewriter) override {
275  auto context = rewriter.getContext();
276  // get external opaque struct pointer type
277  auto commStructT =
278  LLVM::LLVMStructType::getOpaque("ompi_communicator_t", context);
279  StringRef name = "ompi_mpi_comm_world";
280 
281  // make sure global op definition exists
282  getOrDefineExternalStruct(loc, rewriter, name, commStructT);
283 
284  // get address of symbol
285  auto comm = LLVM::AddressOfOp::create(rewriter, loc,
287  SymbolRefAttr::get(context, name));
288  return LLVM::PtrToIntOp::create(rewriter, loc, rewriter.getI64Type(), comm);
289  }
290 
291  Value castComm(const Location loc, ConversionPatternRewriter &rewriter,
292  Value comm) override {
293  return LLVM::IntToPtrOp::create(
294  rewriter, loc, LLVM::LLVMPointerType::get(rewriter.getContext()), comm);
295  }
296 
297  intptr_t getStatusIgnore() override { return 0; }
298 
299  void *getInPlace() override { return reinterpret_cast<void *>(1); }
300 
301  Value getDataType(const Location loc, ConversionPatternRewriter &rewriter,
302  Type type) override {
303  StringRef mtype;
304  if (type.isF32())
305  mtype = "ompi_mpi_float";
306  else if (type.isF64())
307  mtype = "ompi_mpi_double";
308  else if (type.isInteger(64) && !type.isUnsignedInteger())
309  mtype = "ompi_mpi_int64_t";
310  else if (type.isInteger(64))
311  mtype = "ompi_mpi_uint64_t";
312  else if (type.isInteger(32) && !type.isUnsignedInteger())
313  mtype = "ompi_mpi_int32_t";
314  else if (type.isInteger(32))
315  mtype = "ompi_mpi_uint32_t";
316  else if (type.isInteger(16) && !type.isUnsignedInteger())
317  mtype = "ompi_mpi_int16_t";
318  else if (type.isInteger(16))
319  mtype = "ompi_mpi_uint16_t";
320  else if (type.isInteger(8) && !type.isUnsignedInteger())
321  mtype = "ompi_mpi_int8_t";
322  else if (type.isInteger(8))
323  mtype = "ompi_mpi_uint8_t";
324  else
325  assert(false && "unsupported type");
326 
327  auto context = rewriter.getContext();
328  // get external opaque struct pointer type
329  auto typeStructT =
330  LLVM::LLVMStructType::getOpaque("ompi_predefined_datatype_t", context);
331  // make sure global op definition exists
332  getOrDefineExternalStruct(loc, rewriter, mtype, typeStructT);
333  // get address of symbol
334  return LLVM::AddressOfOp::create(rewriter, loc,
336  SymbolRefAttr::get(context, mtype));
337  }
338 
339  Value getMPIOp(const Location loc, ConversionPatternRewriter &rewriter,
340  mpi::MPI_ReductionOpEnum opAttr) override {
341  StringRef op;
342  switch (opAttr) {
343  case mpi::MPI_ReductionOpEnum::MPI_OP_NULL:
344  op = "ompi_mpi_no_op";
345  break;
346  case mpi::MPI_ReductionOpEnum::MPI_MAX:
347  op = "ompi_mpi_max";
348  break;
349  case mpi::MPI_ReductionOpEnum::MPI_MIN:
350  op = "ompi_mpi_min";
351  break;
352  case mpi::MPI_ReductionOpEnum::MPI_SUM:
353  op = "ompi_mpi_sum";
354  break;
355  case mpi::MPI_ReductionOpEnum::MPI_PROD:
356  op = "ompi_mpi_prod";
357  break;
358  case mpi::MPI_ReductionOpEnum::MPI_LAND:
359  op = "ompi_mpi_land";
360  break;
361  case mpi::MPI_ReductionOpEnum::MPI_BAND:
362  op = "ompi_mpi_band";
363  break;
364  case mpi::MPI_ReductionOpEnum::MPI_LOR:
365  op = "ompi_mpi_lor";
366  break;
367  case mpi::MPI_ReductionOpEnum::MPI_BOR:
368  op = "ompi_mpi_bor";
369  break;
370  case mpi::MPI_ReductionOpEnum::MPI_LXOR:
371  op = "ompi_mpi_lxor";
372  break;
373  case mpi::MPI_ReductionOpEnum::MPI_BXOR:
374  op = "ompi_mpi_bxor";
375  break;
376  case mpi::MPI_ReductionOpEnum::MPI_MINLOC:
377  op = "ompi_mpi_minloc";
378  break;
379  case mpi::MPI_ReductionOpEnum::MPI_MAXLOC:
380  op = "ompi_mpi_maxloc";
381  break;
382  case mpi::MPI_ReductionOpEnum::MPI_REPLACE:
383  op = "ompi_mpi_replace";
384  break;
385  }
386  auto context = rewriter.getContext();
387  // get external opaque struct pointer type
388  auto opStructT =
389  LLVM::LLVMStructType::getOpaque("ompi_predefined_op_t", context);
390  // make sure global op definition exists
391  getOrDefineExternalStruct(loc, rewriter, op, opStructT);
392  // get address of symbol
393  return LLVM::AddressOfOp::create(rewriter, loc,
395  SymbolRefAttr::get(context, op));
396  }
397 };
398 
399 std::unique_ptr<MPIImplTraits> MPIImplTraits::get(ModuleOp &moduleOp) {
400  auto attr = dlti::query(*&moduleOp, {"MPI:Implementation"}, true);
401  if (failed(attr))
402  return std::make_unique<MPICHImplTraits>(moduleOp);
403  auto strAttr = dyn_cast<StringAttr>(attr.value());
404  if (strAttr && strAttr.getValue() == "OpenMPI")
405  return std::make_unique<OMPIImplTraits>(moduleOp);
406  if (!strAttr || strAttr.getValue() != "MPICH")
407  moduleOp.emitWarning() << "Unknown \"MPI:Implementation\" value in DLTI ("
408  << strAttr.getValue() << "), defaulting to MPICH";
409  return std::make_unique<MPICHImplTraits>(moduleOp);
410 }
411 
412 //===----------------------------------------------------------------------===//
413 // InitOpLowering
414 //===----------------------------------------------------------------------===//
415 
416 struct InitOpLowering : public ConvertOpToLLVMPattern<mpi::InitOp> {
418 
419  LogicalResult
420  matchAndRewrite(mpi::InitOp op, OpAdaptor adaptor,
421  ConversionPatternRewriter &rewriter) const override {
422  Location loc = op.getLoc();
423 
424  // ptrType `!llvm.ptr`
425  Type ptrType = LLVM::LLVMPointerType::get(rewriter.getContext());
426 
427  // instantiate nullptr `%nullptr = llvm.mlir.zero : !llvm.ptr`
428  auto nullPtrOp = LLVM::ZeroOp::create(rewriter, loc, ptrType);
429  Value llvmnull = nullPtrOp.getRes();
430 
431  // grab a reference to the global module op:
432  auto moduleOp = op->getParentOfType<ModuleOp>();
433 
434  // LLVM Function type representing `i32 MPI_Init(ptr, ptr)`
435  auto initFuncType =
436  LLVM::LLVMFunctionType::get(rewriter.getI32Type(), {ptrType, ptrType});
437  // get or create function declaration:
438  LLVM::LLVMFuncOp initDecl =
439  getOrDefineFunction(moduleOp, loc, rewriter, "MPI_Init", initFuncType);
440 
441  // replace init with function call
442  rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, initDecl,
443  ValueRange{llvmnull, llvmnull});
444 
445  return success();
446  }
447 };
448 
449 //===----------------------------------------------------------------------===//
450 // FinalizeOpLowering
451 //===----------------------------------------------------------------------===//
452 
453 struct FinalizeOpLowering : public ConvertOpToLLVMPattern<mpi::FinalizeOp> {
455 
456  LogicalResult
457  matchAndRewrite(mpi::FinalizeOp op, OpAdaptor adaptor,
458  ConversionPatternRewriter &rewriter) const override {
459  // get loc
460  Location loc = op.getLoc();
461 
462  // grab a reference to the global module op:
463  auto moduleOp = op->getParentOfType<ModuleOp>();
464 
465  // LLVM Function type representing `i32 MPI_Finalize()`
466  auto initFuncType = LLVM::LLVMFunctionType::get(rewriter.getI32Type(), {});
467  // get or create function declaration:
468  LLVM::LLVMFuncOp initDecl = getOrDefineFunction(
469  moduleOp, loc, rewriter, "MPI_Finalize", initFuncType);
470 
471  // replace init with function call
472  rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, initDecl, ValueRange{});
473 
474  return success();
475  }
476 };
477 
478 //===----------------------------------------------------------------------===//
479 // CommWorldOpLowering
480 //===----------------------------------------------------------------------===//
481 
482 struct CommWorldOpLowering : public ConvertOpToLLVMPattern<mpi::CommWorldOp> {
484 
485  LogicalResult
486  matchAndRewrite(mpi::CommWorldOp op, OpAdaptor adaptor,
487  ConversionPatternRewriter &rewriter) const override {
488  // grab a reference to the global module op:
489  auto moduleOp = op->getParentOfType<ModuleOp>();
490  auto mpiTraits = MPIImplTraits::get(moduleOp);
491  // get MPI_COMM_WORLD
492  rewriter.replaceOp(op, mpiTraits->getCommWorld(op.getLoc(), rewriter));
493 
494  return success();
495  }
496 };
497 
498 //===----------------------------------------------------------------------===//
499 // CommSplitOpLowering
500 //===----------------------------------------------------------------------===//
501 
502 struct CommSplitOpLowering : public ConvertOpToLLVMPattern<mpi::CommSplitOp> {
504 
505  LogicalResult
506  matchAndRewrite(mpi::CommSplitOp op, OpAdaptor adaptor,
507  ConversionPatternRewriter &rewriter) const override {
508  // grab a reference to the global module op:
509  auto moduleOp = op->getParentOfType<ModuleOp>();
510  auto mpiTraits = MPIImplTraits::get(moduleOp);
511  Type i32 = rewriter.getI32Type();
512  Type ptrType = LLVM::LLVMPointerType::get(op->getContext());
513  Location loc = op.getLoc();
514 
515  // get communicator
516  Value comm = mpiTraits->castComm(loc, rewriter, adaptor.getComm());
517  auto one = LLVM::ConstantOp::create(rewriter, loc, i32, 1);
518  auto outPtr =
519  LLVM::AllocaOp::create(rewriter, loc, ptrType, comm.getType(), one);
520 
521  // int MPI_Comm_split(MPI_Comm comm, int color, int key, MPI_Comm * newcomm)
522  auto funcType =
523  LLVM::LLVMFunctionType::get(i32, {comm.getType(), i32, i32, ptrType});
524  // get or create function declaration:
525  LLVM::LLVMFuncOp funcDecl = getOrDefineFunction(moduleOp, loc, rewriter,
526  "MPI_Comm_split", funcType);
527 
528  auto callOp =
529  LLVM::CallOp::create(rewriter, loc, funcDecl,
530  ValueRange{comm, adaptor.getColor(),
531  adaptor.getKey(), outPtr.getRes()});
532 
533  // load the communicator into a register
534  Value res = LLVM::LoadOp::create(rewriter, loc, i32, outPtr.getResult());
535  res = LLVM::SExtOp::create(rewriter, loc, rewriter.getI64Type(), res);
536 
537  // if retval is checked, replace uses of retval with the results from the
538  // call op
539  SmallVector<Value> replacements;
540  if (op.getRetval())
541  replacements.push_back(callOp.getResult());
542 
543  // replace op
544  replacements.push_back(res);
545  rewriter.replaceOp(op, replacements);
546 
547  return success();
548  }
549 };
550 
551 //===----------------------------------------------------------------------===//
552 // CommRankOpLowering
553 //===----------------------------------------------------------------------===//
554 
555 struct CommRankOpLowering : public ConvertOpToLLVMPattern<mpi::CommRankOp> {
557 
558  LogicalResult
559  matchAndRewrite(mpi::CommRankOp op, OpAdaptor adaptor,
560  ConversionPatternRewriter &rewriter) const override {
561  // get some helper vars
562  Location loc = op.getLoc();
563  MLIRContext *context = rewriter.getContext();
564  Type i32 = rewriter.getI32Type();
565 
566  // ptrType `!llvm.ptr`
567  Type ptrType = LLVM::LLVMPointerType::get(context);
568 
569  // grab a reference to the global module op:
570  auto moduleOp = op->getParentOfType<ModuleOp>();
571 
572  auto mpiTraits = MPIImplTraits::get(moduleOp);
573  // get communicator
574  Value comm = mpiTraits->castComm(loc, rewriter, adaptor.getComm());
575 
576  // LLVM Function type representing `i32 MPI_Comm_rank(ptr, ptr)`
577  auto rankFuncType =
578  LLVM::LLVMFunctionType::get(i32, {comm.getType(), ptrType});
579  // get or create function declaration:
580  LLVM::LLVMFuncOp initDecl = getOrDefineFunction(
581  moduleOp, loc, rewriter, "MPI_Comm_rank", rankFuncType);
582 
583  // replace with function call
584  auto one = LLVM::ConstantOp::create(rewriter, loc, i32, 1);
585  auto rankptr = LLVM::AllocaOp::create(rewriter, loc, ptrType, i32, one);
586  auto callOp = LLVM::CallOp::create(rewriter, loc, initDecl,
587  ValueRange{comm, rankptr.getRes()});
588 
589  // load the rank into a register
590  auto loadedRank =
591  LLVM::LoadOp::create(rewriter, loc, i32, rankptr.getResult());
592 
593  // if retval is checked, replace uses of retval with the results from the
594  // call op
595  SmallVector<Value> replacements;
596  if (op.getRetval())
597  replacements.push_back(callOp.getResult());
598 
599  // replace all uses, then erase op
600  replacements.push_back(loadedRank.getRes());
601  rewriter.replaceOp(op, replacements);
602 
603  return success();
604  }
605 };
606 
607 //===----------------------------------------------------------------------===//
608 // SendOpLowering
609 //===----------------------------------------------------------------------===//
610 
611 struct SendOpLowering : public ConvertOpToLLVMPattern<mpi::SendOp> {
613 
614  LogicalResult
615  matchAndRewrite(mpi::SendOp op, OpAdaptor adaptor,
616  ConversionPatternRewriter &rewriter) const override {
617  // get some helper vars
618  Location loc = op.getLoc();
619  MLIRContext *context = rewriter.getContext();
620  Type i32 = rewriter.getI32Type();
621  Type elemType = op.getRef().getType().getElementType();
622 
623  // ptrType `!llvm.ptr`
624  Type ptrType = LLVM::LLVMPointerType::get(context);
625 
626  // grab a reference to the global module op:
627  auto moduleOp = op->getParentOfType<ModuleOp>();
628 
629  // get MPI_COMM_WORLD, dataType and pointer
630  auto [dataPtr, size] =
631  getRawPtrAndSize(loc, rewriter, adaptor.getRef(), elemType);
632  auto mpiTraits = MPIImplTraits::get(moduleOp);
633  Value dataType = mpiTraits->getDataType(loc, rewriter, elemType);
634  Value comm = mpiTraits->castComm(loc, rewriter, adaptor.getComm());
635 
636  // LLVM Function type representing `i32 MPI_send(data, count, datatype, dst,
637  // tag, comm)`
638  auto funcType = LLVM::LLVMFunctionType::get(
639  i32, {ptrType, i32, dataType.getType(), i32, i32, comm.getType()});
640  // get or create function declaration:
641  LLVM::LLVMFuncOp funcDecl =
642  getOrDefineFunction(moduleOp, loc, rewriter, "MPI_Send", funcType);
643 
644  // replace op with function call
645  auto funcCall = LLVM::CallOp::create(rewriter, loc, funcDecl,
646  ValueRange{dataPtr, size, dataType,
647  adaptor.getDest(),
648  adaptor.getTag(), comm});
649  if (op.getRetval())
650  rewriter.replaceOp(op, funcCall.getResult());
651  else
652  rewriter.eraseOp(op);
653 
654  return success();
655  }
656 };
657 
658 //===----------------------------------------------------------------------===//
659 // RecvOpLowering
660 //===----------------------------------------------------------------------===//
661 
662 struct RecvOpLowering : public ConvertOpToLLVMPattern<mpi::RecvOp> {
664 
665  LogicalResult
666  matchAndRewrite(mpi::RecvOp op, OpAdaptor adaptor,
667  ConversionPatternRewriter &rewriter) const override {
668  // get some helper vars
669  Location loc = op.getLoc();
670  MLIRContext *context = rewriter.getContext();
671  Type i32 = rewriter.getI32Type();
672  Type i64 = rewriter.getI64Type();
673  Type elemType = op.getRef().getType().getElementType();
674 
675  // ptrType `!llvm.ptr`
676  Type ptrType = LLVM::LLVMPointerType::get(context);
677 
678  // grab a reference to the global module op:
679  auto moduleOp = op->getParentOfType<ModuleOp>();
680 
681  // get MPI_COMM_WORLD, dataType, status_ignore and pointer
682  auto [dataPtr, size] =
683  getRawPtrAndSize(loc, rewriter, adaptor.getRef(), elemType);
684  auto mpiTraits = MPIImplTraits::get(moduleOp);
685  Value dataType = mpiTraits->getDataType(loc, rewriter, elemType);
686  Value comm = mpiTraits->castComm(loc, rewriter, adaptor.getComm());
687  Value statusIgnore = LLVM::ConstantOp::create(rewriter, loc, i64,
688  mpiTraits->getStatusIgnore());
689  statusIgnore =
690  LLVM::IntToPtrOp::create(rewriter, loc, ptrType, statusIgnore);
691 
692  // LLVM Function type representing `i32 MPI_Recv(data, count, datatype, dst,
693  // tag, comm)`
694  auto funcType =
695  LLVM::LLVMFunctionType::get(i32, {ptrType, i32, dataType.getType(), i32,
696  i32, comm.getType(), ptrType});
697  // get or create function declaration:
698  LLVM::LLVMFuncOp funcDecl =
699  getOrDefineFunction(moduleOp, loc, rewriter, "MPI_Recv", funcType);
700 
701  // replace op with function call
702  auto funcCall = LLVM::CallOp::create(
703  rewriter, loc, funcDecl,
704  ValueRange{dataPtr, size, dataType, adaptor.getSource(),
705  adaptor.getTag(), comm, statusIgnore});
706  if (op.getRetval())
707  rewriter.replaceOp(op, funcCall.getResult());
708  else
709  rewriter.eraseOp(op);
710 
711  return success();
712  }
713 };
714 
715 //===----------------------------------------------------------------------===//
716 // AllReduceOpLowering
717 //===----------------------------------------------------------------------===//
718 
719 struct AllReduceOpLowering : public ConvertOpToLLVMPattern<mpi::AllReduceOp> {
721 
722  LogicalResult
723  matchAndRewrite(mpi::AllReduceOp op, OpAdaptor adaptor,
724  ConversionPatternRewriter &rewriter) const override {
725  Location loc = op.getLoc();
726  MLIRContext *context = rewriter.getContext();
727  Type i32 = rewriter.getI32Type();
728  Type i64 = rewriter.getI64Type();
729  Type elemType = op.getSendbuf().getType().getElementType();
730 
731  // ptrType `!llvm.ptr`
732  Type ptrType = LLVM::LLVMPointerType::get(context);
733  auto moduleOp = op->getParentOfType<ModuleOp>();
734  auto mpiTraits = MPIImplTraits::get(moduleOp);
735  auto [sendPtr, sendSize] =
736  getRawPtrAndSize(loc, rewriter, adaptor.getSendbuf(), elemType);
737  auto [recvPtr, recvSize] =
738  getRawPtrAndSize(loc, rewriter, adaptor.getRecvbuf(), elemType);
739 
740  // If input and output are the same, request in-place operation.
741  if (adaptor.getSendbuf() == adaptor.getRecvbuf()) {
742  sendPtr = LLVM::ConstantOp::create(
743  rewriter, loc, i64,
744  reinterpret_cast<int64_t>(mpiTraits->getInPlace()));
745  sendPtr = LLVM::IntToPtrOp::create(rewriter, loc, ptrType, sendPtr);
746  }
747 
748  Value dataType = mpiTraits->getDataType(loc, rewriter, elemType);
749  Value mpiOp = mpiTraits->getMPIOp(loc, rewriter, op.getOp());
750  Value commWorld = mpiTraits->castComm(loc, rewriter, adaptor.getComm());
751 
752  // 'int MPI_Allreduce(const void *sendbuf, void *recvbuf, int count,
753  // MPI_Datatype datatype, MPI_Op op, MPI_Comm comm)'
754  auto funcType = LLVM::LLVMFunctionType::get(
755  i32, {ptrType, ptrType, i32, dataType.getType(), mpiOp.getType(),
756  commWorld.getType()});
757  // get or create function declaration:
758  LLVM::LLVMFuncOp funcDecl =
759  getOrDefineFunction(moduleOp, loc, rewriter, "MPI_Allreduce", funcType);
760 
761  // replace op with function call
762  auto funcCall = LLVM::CallOp::create(
763  rewriter, loc, funcDecl,
764  ValueRange{sendPtr, recvPtr, sendSize, dataType, mpiOp, commWorld});
765 
766  if (op.getRetval())
767  rewriter.replaceOp(op, funcCall.getResult());
768  else
769  rewriter.eraseOp(op);
770 
771  return success();
772  }
773 };
774 
775 //===----------------------------------------------------------------------===//
776 // ConvertToLLVMPatternInterface implementation
777 //===----------------------------------------------------------------------===//
778 
779 /// Implement the interface to convert Func to LLVM.
780 struct FuncToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
782  /// Hook for derived dialect interface to provide conversion patterns
783  /// and mark dialect legal for the conversion target.
784  void populateConvertToLLVMConversionPatterns(
785  ConversionTarget &target, LLVMTypeConverter &typeConverter,
786  RewritePatternSet &patterns) const final {
788  }
789 };
790 } // namespace
791 
792 //===----------------------------------------------------------------------===//
793 // Pattern Population
794 //===----------------------------------------------------------------------===//
795 
798  // Using i64 as a portable, intermediate type for !mpi.comm.
799  // It would be nicer to somehow get the right type directly, but TLDI is not
800  // available here.
801  converter.addConversion([](mpi::CommType type) {
802  return IntegerType::get(type.getContext(), 64);
803  });
804  patterns.add<CommRankOpLowering, CommSplitOpLowering, CommWorldOpLowering,
805  FinalizeOpLowering, InitOpLowering, SendOpLowering,
806  RecvOpLowering, AllReduceOpLowering>(converter);
807 }
808 
810  registry.addExtension(+[](MLIRContext *ctx, mpi::MPIDialect *dialect) {
811  dialect->addInterfaces<FuncToLLVMDialectInterface>();
812  });
813 }
Attributes are known-constant values of operations.
Definition: Attributes.h:25
IntegerType getI64Type()
Definition: Builders.cpp:64
IntegerType getI32Type()
Definition: Builders.cpp:62
MLIRContext * getContext() const
Definition: Builders.h:55
This class implements a pattern rewriter for use with ConversionPatterns.
void replaceOp(Operation *op, ValueRange newValues) override
Replace the given operation with the new values.
void eraseOp(Operation *op) override
PatternRewriter hook for erasing a dead operation.
This class describes a specific conversion target.
Utility class for operation conversions targeting the LLVM dialect that match exactly one source oper...
Definition: Pattern.h:209
ConvertOpToLLVMPattern(const LLVMTypeConverter &typeConverter, PatternBenefit benefit=1)
Definition: Pattern.h:215
Base class for dialect interfaces providing translation to LLVM IR.
ConvertToLLVMPatternInterface(Dialect *dialect)
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
bool addExtension(TypeID extensionID, std::unique_ptr< DialectExtensionBase > extension)
Add the given extension to the registry.
Conversion from types to the LLVM IR dialect.
Definition: TypeConverter.h:35
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:63
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:429
This provides public APIs that all operations should have.
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:519
void addConversion(FnT &&callback)
Register a conversion function.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
bool isF64() const
Definition: Types.cpp:41
bool isF32() const
Definition: Types.cpp:40
bool isUnsignedInteger() const
Return true if this is an unsigned integer type (with the specified width).
Definition: Types.cpp:88
bool isInteger() const
Return true if this is an integer type (with the specified width).
Definition: Types.cpp:56
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:105
static ConstantIntOp create(OpBuilder &builder, Location location, int64_t value, unsigned width)
Definition: ArithOps.cpp:258
NestedPattern Op(FilterFunctionType filter=defaultFilterFunction)
FailureOr< Attribute > query(Operation *op, ArrayRef< DataLayoutEntryKey > keys, bool emitError=false)
Perform a DLTI-query at op, recursively querying each key of keys on query interface-implementing att...
Definition: DLTI.cpp:537
void populateMPIToLLVMConversionPatterns(LLVMTypeConverter &converter, RewritePatternSet &patterns)
Definition: MPIToLLVM.cpp:796
void registerConvertMPIToLLVMInterface(DialectRegistry &registry)
Definition: MPIToLLVM.cpp:809
detail::InFlightRemark failed(Location loc, RemarkOpts opts)
Report an optimization remark that failed.
Definition: Remarks.h:491
Include the generated interface declarations.
const FrozenRewritePatternSet & patterns
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
LLVM::LLVMFuncOp getOrDefineFunction(gpu::GPUModuleOp moduleOp, Location loc, OpBuilder &b, StringRef name, LLVM::LLVMFunctionType type)
Find or create an external function declaration in the given module.