[flang-commits] [flang] [flang] Adding lowering of TEAMs features to PRIF in MIF Dialect (PR #165573)

[via flang-commits]

================
@@ -432,6 +464,208 @@ struct MIFCoSumOpConversion : public mlir::OpRewritePattern<mif::CoSumOp> {
   }
 };
 
+/// Convert mif.form_team operation to runtime call of 'prif_form_team'
+struct MIFFormTeamOpConversion
+    : public mlir::OpRewritePattern<mif::FormTeamOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mif::FormTeamOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto mod = op->template getParentOfType<mlir::ModuleOp>();
+    fir::FirOpBuilder builder(rewriter, mod);
+    mlir::Location loc = op.getLoc();
+    mlir::Type errmsgTy = getPRIFErrmsgType(builder);
+    mlir::Type boxTy = fir::BoxType::get(builder.getNoneType());
+    mlir::FunctionType ftype = mlir::FunctionType::get(
+        builder.getContext(),
+        /*inputs*/
+        {builder.getRefType(builder.getI64Type()), boxTy,
+         builder.getRefType(builder.getI32Type()), getPRIFStatType(builder),
+         errmsgTy, errmsgTy},
+        /*results*/ {});
+    mlir::func::FuncOp funcOp =
+        builder.createFunction(loc, getPRIFProcName("form_team"), ftype);
+
+    mlir::Type i64Ty = builder.getI64Type();
+    mlir::Value teamNumber = builder.createTemporary(loc, i64Ty);
+    mlir::Value t =
+        (op.getTeamNumber().getType() == i64Ty)
+            ? op.getTeamNumber()
+            : fir::ConvertOp::create(builder, loc, i64Ty, op.getTeamNumber());
+    fir::StoreOp::create(builder, loc, t, teamNumber);
+
+    mlir::Type i32Ty = builder.getI32Type();
+    mlir::Value newIndex;
+    if (op.getNewIndex()) {
+      newIndex = builder.createTemporary(loc, i32Ty);
+      mlir::Value ni =
+          (op.getNewIndex().getType() == i32Ty)
+              ? op.getNewIndex()
+              : fir::ConvertOp::create(builder, loc, i32Ty, op.getNewIndex());
+      fir::StoreOp::create(builder, loc, ni, newIndex);
+    } else
+      newIndex = fir::AbsentOp::create(builder, loc, builder.getRefType(i32Ty));
+
+    mlir::Value stat = genStatPRIF(builder, loc, op.getStat());
+    auto [errmsgArg, errmsgAllocArg] =
+        genErrmsgPRIF(builder, loc, op.getErrmsg());
+    llvm::SmallVector<mlir::Value> args = fir::runtime::createArguments(
+        builder, loc, ftype, teamNumber, op.getTeamVar(), newIndex, stat,
+        errmsgArg, errmsgAllocArg);
+    fir::CallOp callOp = fir::CallOp::create(builder, loc, funcOp, args);
+    rewriter.replaceOp(op, callOp);
+    return mlir::success();
+  }
+};
+
+/// Convert mif.change_team operation to runtime call of 'prif_change_team'
+struct MIFChangeTeamOpConversion
+    : public mlir::OpRewritePattern<mif::ChangeTeamOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mif::ChangeTeamOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto mod = op->template getParentOfType<mlir::ModuleOp>();
+    fir::FirOpBuilder builder(rewriter, mod);
+
+    mlir::Location loc = op.getLoc();
+    mlir::Type errmsgTy = getPRIFErrmsgType(builder);
+    mlir::Type boxTy = fir::BoxType::get(builder.getNoneType());
+    mlir::FunctionType ftype = mlir::FunctionType::get(
+        builder.getContext(),
+        /*inputs*/ {boxTy, getPRIFStatType(builder), errmsgTy, errmsgTy},
+        /*results*/ {});
+    mlir::func::FuncOp funcOp =
+        builder.createFunction(loc, getPRIFProcName("change_team"), ftype);
+
+    mlir::Value stat = genStatPRIF(builder, loc, op.getStat());
+    auto [errmsgArg, errmsgAllocArg] =
+        genErrmsgPRIF(builder, loc, op.getErrmsg());
+    llvm::SmallVector<mlir::Value> args = fir::runtime::createArguments(
+        builder, loc, ftype, op.getTeam(), stat, errmsgArg, errmsgAllocArg);
+    fir::CallOp callOp = fir::CallOp::create(builder, loc, funcOp, args);
+
+    mlir::Operation *changeOp = op.getOperation();
+    auto &bodyRegion = op.getRegion();
+    mlir::Block &bodyBlock = bodyRegion.front();
+
+    rewriter.moveOpBefore(callOp.getOperation(), changeOp);
----------------
jeanPerier wrote:

I think setting the insertion point before the changeOp would avoid having to swap the ops here and would be easier to read.

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