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

[via flang-commits]

================
@@ -1114,6 +1165,159 @@ 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(builder, loc, 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();
+  if (result.getType() != op.getType())
+    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()) {
+      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();
+      std::string imageIndexName =
+          fir::unwrapPassByRefType(team.getType()).isInteger()
----------------
jeanPerier wrote:

For team_number, this does not look consistent with the operation interface.
The operation defines `team_number` as a different operation argument as `team`, so I think you should use `getTeamNumber` somewhere to deal with it instead of detecting the type of `team`.

Note that maybe you could also change the operation design to remove the team_number operand since lowering is also forwarding both TEAM and TEAM_NUMBER argument to the mif.image_index operand (hence this probably works end to end). I am not necessarly advising it.

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