[flang-commits] [flang] [Flang][MIF] Adding support of intrinsics with coarray argument (PR #192944)

[via flang-commits]

================
@@ -1113,6 +1161,156 @@ struct MIFDeallocCoarrayOpConversion
   }
 };
 
+/// Convert mif.coshape operation to runtime call of 'prif_coshape'
+struct MIFCoshapeOpConversion : public mlir::OpRewritePattern<mif::CoshapeOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mif::CoshapeOp 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 i64Ty = builder.getI64Type();
+    mlir::Type boxTy = fir::BoxType::get(builder.getNoneType());
+
+    mlir::FunctionType ftype =
+        mlir::FunctionType::get(builder.getContext(),
+                                /*inputs*/ {boxTy, genBoxedSequenceType(i64Ty)},
+                                /*results*/ {});
+    mlir::func::FuncOp funcOp =
+        builder.createFunction(loc, getPRIFProcName("coshape"), ftype);
+
+    mlir::Value coarrayHandle = getCoarrayHandle(builder, loc, op.getCoarray());
+    std::int64_t corank = getCorank(op.getCoarray());
+    mlir::Type resultType = fir::SequenceType::get(
+        static_cast<fir::SequenceType::Extent>(corank), i64Ty);
+    mlir::Value result =
+        builder.createBox(loc, builder.createTemporary(loc, resultType));
+
+    llvm::SmallVector<mlir::Value> args = fir::runtime::createArguments(
+        builder, loc, ftype, coarrayHandle, result);
+    fir::CallOp::create(builder, loc, funcOp, args);
+    result = fir::ConvertOp::create(builder, loc, genBoxedSequenceType(i64Ty),
+                                    result);
+    rewriter.replaceOp(op, result);
+    return mlir::success();
+  }
+};
+
+template <class T>
+mlir::LogicalResult CoboundOpConversion(T op, mlir::PatternRewriter &rewriter,
+                                        const std::string &prefix) {
+  auto mod = op->template getParentOfType<mlir::ModuleOp>();
+  fir::FirOpBuilder builder(rewriter, mod);
+  mlir::Location loc = op.getLoc();
+  mlir::Type i64Ty = builder.getI64Type();
+  mlir::Type boxTy = fir::BoxType::get(builder.getNoneType());
+
+  mlir::Value coarrayHandle = getCoarrayHandle(builder, loc, op.getCoarray());
+  mlir::Type i32Ty = builder.getI32Type();
+  mlir::FunctionType ftype = mlir::FunctionType::get(
+      builder.getContext(),
+      /*inputs*/
+      {boxTy, builder.getRefType(i32Ty), builder.getRefType(i64Ty)},
+      /*results*/ {});
+  mlir::func::FuncOp funcOp =
+      builder.createFunction(loc, getPRIFProcName(prefix + "_with_dim"), ftype);
+
+  mlir::Value result = builder.createTemporary(loc, i64Ty);
+  mlir::Value dim = builder.createTemporary(loc, i32Ty);
+  mlir::Value d = fir::ConvertOp::create(builder, loc, i32Ty, op.getDim());
+  fir::StoreOp::create(builder, loc, d, dim);
+
+  llvm::SmallVector<mlir::Value> args = fir::runtime::createArguments(
+      builder, loc, ftype, coarrayHandle, dim, result);
+  fir::CallOp::create(builder, loc, funcOp, args);
+  result = fir::LoadOp::create(builder, loc, result).getResult();
+  result = builder.createConvert(loc, op.getType(), result);
+  rewriter.replaceOp(op, result);
+  return mlir::success();
+}
+
+/// Convert mif.lcobound operation to runtime call of
+/// 'prif_lcobound_{with|no}_dim'
+struct MIFLcoboundOpConversion
+    : public mlir::OpRewritePattern<mif::LcoboundOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mif::LcoboundOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    return CoboundOpConversion(op, rewriter, "lcobound");
+  }
+};
+
+/// Convert mif.ucobound operation to runtime call of
+/// 'prif_ucobound_{with|no}_dim'
+struct MIFUcoboundOpConversion
+    : public mlir::OpRewritePattern<mif::UcoboundOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mif::UcoboundOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    return CoboundOpConversion(op, rewriter, "ucobound");
+  }
+};
+
+/// Convert mif.image_index operation to runtime call of
+/// 'prif_image_index[_with_team[_number]]'
+struct MIFImageIndexOpConversion
+    : public mlir::OpRewritePattern<mif::ImageIndexOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mif::ImageIndexOp 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 i64Ty = builder.getI64Type();
+    mlir::Type i32Ty = builder.getI32Type();
+    mlir::Type resTy = builder.getRefType(i32Ty);
+    mlir::Type boxTy = fir::BoxType::get(builder.getNoneType());
+    mlir::Value result = builder.createTemporary(loc, i32Ty);
+
+    mlir::func::FuncOp funcOp;
+    llvm::SmallVector<mlir::Value> args;
+    mlir::Value coarrayHandle = getCoarrayHandle(builder, loc, op.getCoarray());
+    if (!op.getTeam() && !op.getTeamNumber()) {
+      mlir::FunctionType ftype = mlir::FunctionType::get(
+          builder.getContext(),
+          /*inputs*/ {boxTy, genBoxedSequenceType(i64Ty), resTy},
+          /*results*/ {});
+      funcOp =
+          builder.createFunction(loc, getPRIFProcName("image_index"), ftype);
+      args = fir::runtime::createArguments(builder, loc, ftype, coarrayHandle,
+                                           op.getSub(), result);
+    } else {
+      mlir::Value team = op.getTeam() ? op.getTeam() : op.getTeamNumber();
+      std::string imageIndexName =
+          op.getTeamNumber() ? getPRIFProcName("image_index_with_team_number")
+                             : getPRIFProcName("image_index_with_team");
+      mlir::Type teamTy =
+          op.getTeamNumber() ? builder.getRefType(i64Ty) : boxTy;
+      mlir::FunctionType ftype = mlir::FunctionType::get(
+          builder.getContext(),
+          /*inputs*/ {boxTy, genBoxedSequenceType(i64Ty), teamTy, resTy},
+          /*results*/ {});
+      funcOp = builder.createFunction(loc, imageIndexName, ftype);
+
+      if (fir::isa_ref_type(team.getType()))
+        team = builder.createBox(loc, team);
+      args = fir::runtime::createArguments(builder, loc, ftype, coarrayHandle,
+                                           op.getSub(), team, result);
+    }
----------------
jeanPerier wrote:

@JDPailleux, I think that could be because the lowering or operation builder was not changed to dispatch the integer arguments to team_number operand instead of the team operand of the mif::ImageIndexOp.

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