[flang] Add support of THIS_IMAGE and NUM_IMAGES with PRIF (#154081)

In relation to the approval and merge of the
https://github.com/llvm/llvm-project/pull/76088 specification about
multi-image features in Flang.
Here is a PR on adding support for `THIS_IMAGE` and `NUM_IMAGES` in
conformance with the PRIF specification.
For `THIS_IMAGE`, the lowering to the subroutine containing the coarray
argument is not present in this PR, and will be in a future one.

6 files changed, 159 insertions, 0 deletions

diff --git a/flang/include/flang/Optimizer/Builder/IntrinsicCall.h b/flang/include/flang/Optimizer/Builder/IntrinsicCall.h
index 2afd504..88c3ada 100644
--- a/flang/include/flang/Optimizer/Builder/IntrinsicCall.h
+++ b/flang/include/flang/Optimizer/Builder/IntrinsicCall.h
@@ -378,6 +378,8 @@ struct IntrinsicLibrary {
   fir::ExtendedValue genNorm2(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genNot(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genNull(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
+  fir::ExtendedValue genNumImages(mlir::Type,
+                                  llvm::ArrayRef<fir::ExtendedValue>);
   template <typename OpTy>
   mlir::Value genNVVMTime(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genPack(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
@@ -449,6 +451,8 @@ struct IntrinsicLibrary {
   fir::ExtendedValue genTranspose(mlir::Type,
                                   llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genThisGrid(mlir::Type, llvm::ArrayRef<mlir::Value>);
+  fir::ExtendedValue genThisImage(mlir::Type,
+                                  llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genThisThreadBlock(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genThisWarp(mlir::Type, llvm::ArrayRef<mlir::Value>);
   void genThreadFence(llvm::ArrayRef<fir::ExtendedValue>);
@@ -563,6 +567,15 @@ struct IntrinsicLibrary {
 
   void setResultMustBeFreed() { resultMustBeFreed = true; }
 
+  // Check support of coarray features
+  void checkCoarrayEnabled() {
+    if (converter &&
+        !converter->getFoldingContext().languageFeatures().IsEnabled(
+            Fortran::common::LanguageFeature::Coarray))
+      fir::emitFatalError(loc, "Coarrays disabled, use '-fcoarray' to enable.",
+                          false);
+  }
+
   fir::FirOpBuilder &builder;
   mlir::Location loc;
   bool resultMustBeFreed = false;
diff --git a/flang/include/flang/Optimizer/Builder/Runtime/Coarray.h b/flang/include/flang/Optimizer/Builder/Runtime/Coarray.h
index f2c76c9..23bb378 100644
--- a/flang/include/flang/Optimizer/Builder/Runtime/Coarray.h
+++ b/flang/include/flang/Optimizer/Builder/Runtime/Coarray.h
@@ -37,5 +37,17 @@ namespace fir::runtime {
 /// Generate Call to runtime prif_init
 mlir::Value genInitCoarray(fir::FirOpBuilder &builder, mlir::Location loc);
 
+/// Generate Call to runtime prif_num_images
+mlir::Value getNumImages(fir::FirOpBuilder &builder, mlir::Location loc);
+
+/// Generate Call to runtime prif_num_images_with_team or
+/// prif_num_images_with_team_number
+mlir::Value getNumImagesWithTeam(fir::FirOpBuilder &builder, mlir::Location loc,
+                                 mlir::Value team);
+
+/// Generate Call to runtime prif_this_image_no_coarray
+mlir::Value getThisImage(fir::FirOpBuilder &builder, mlir::Location loc,
+                         mlir::Value team = {});
+
 } // namespace fir::runtime
 #endif // FORTRAN_OPTIMIZER_BUILDER_RUNTIME_COARRAY_H
diff --git a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
index bc3086e..8aacdb1 100644
--- a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
+++ b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
@@ -25,6 +25,7 @@
 #include "flang/Optimizer/Builder/Runtime/Allocatable.h"
 #include "flang/Optimizer/Builder/Runtime/CUDA/Descriptor.h"
 #include "flang/Optimizer/Builder/Runtime/Character.h"
+#include "flang/Optimizer/Builder/Runtime/Coarray.h"
 #include "flang/Optimizer/Builder/Runtime/Command.h"
 #include "flang/Optimizer/Builder/Runtime/Derived.h"
 #include "flang/Optimizer/Builder/Runtime/Exceptions.h"
@@ -778,6 +779,10 @@ static constexpr IntrinsicHandler handlers[]{
      /*isElemental=*/false},
     {"not", &I::genNot},
     {"null", &I::genNull, {{{"mold", asInquired}}}, /*isElemental=*/false},
+    {"num_images",
+     &I::genNumImages,
+     {{{"team", asAddr}, {"team_number", asAddr}}},
+     /*isElemental*/ false},
     {"pack",
      &I::genPack,
      {{{"array", asBox},
@@ -947,6 +952,12 @@ static constexpr IntrinsicHandler handlers[]{
     {"tand", &I::genTand},
     {"tanpi", &I::genTanpi},
     {"this_grid", &I::genThisGrid, {}, /*isElemental=*/false},
+    {"this_image",
+     &I::genThisImage,
+     {{{"coarray", asBox},
+       {"dim", asAddr},
+       {"team", asBox, handleDynamicOptional}}},
+     /*isElemental=*/false},
     {"this_thread_block", &I::genThisThreadBlock, {}, /*isElemental=*/false},
     {"this_warp", &I::genThisWarp, {}, /*isElemental=*/false},
     {"threadfence", &I::genThreadFence, {}, /*isElemental=*/false},
@@ -7279,6 +7290,19 @@ IntrinsicLibrary::genNull(mlir::Type, llvm::ArrayRef<fir::ExtendedValue> args) {
   return fir::MutableBoxValue(boxStorage, mold->nonDeferredLenParams(), {});
 }
 
+// NUM_IMAGES
+fir::ExtendedValue
+IntrinsicLibrary::genNumImages(mlir::Type resultType,
+                               llvm::ArrayRef<fir::ExtendedValue> args) {
+  checkCoarrayEnabled();
+  assert(args.size() == 0 || args.size() == 1);
+
+  if (args.size())
+    return fir::runtime::getNumImagesWithTeam(builder, loc,
+                                              fir::getBase(args[0]));
+  return fir::runtime::getNumImages(builder, loc);
+}
+
 // CLOCK, CLOCK64, GLOBALTIMER
 template <typename OpTy>
 mlir::Value IntrinsicLibrary::genNVVMTime(mlir::Type resultType,
@@ -8329,6 +8353,27 @@ mlir::Value IntrinsicLibrary::genThisGrid(mlir::Type resultType,
   return res;
 }
 
+// THIS_IMAGE
+fir::ExtendedValue
+IntrinsicLibrary::genThisImage(mlir::Type resultType,
+                               llvm::ArrayRef<fir::ExtendedValue> args) {
+  checkCoarrayEnabled();
+  assert(args.size() >= 1 && args.size() <= 3);
+  const bool coarrayIsAbsent = args.size() == 1;
+  mlir::Value team =
+      !isStaticallyAbsent(args, args.size() - 1)
+          ? fir::getBase(args[args.size() - 1])
+          : builder
+                .create<fir::AbsentOp>(loc,
+                                       fir::BoxType::get(builder.getNoneType()))
+                .getResult();
+
+  if (!coarrayIsAbsent)
+    TODO(loc, "this_image with coarray argument.");
+  mlir::Value res = fir::runtime::getThisImage(builder, loc, team);
+  return builder.createConvert(loc, resultType, res);
+}
+
 // THIS_THREAD_BLOCK
 mlir::Value
 IntrinsicLibrary::genThisThreadBlock(mlir::Type resultType,
diff --git a/flang/lib/Optimizer/Builder/Runtime/Coarray.cpp b/flang/lib/Optimizer/Builder/Runtime/Coarray.cpp
index eaff6c3..fb72fc2 100644
--- a/flang/lib/Optimizer/Builder/Runtime/Coarray.cpp
+++ b/flang/lib/Optimizer/Builder/Runtime/Coarray.cpp
@@ -27,3 +27,60 @@ mlir::Value fir::runtime::genInitCoarray(fir::FirOpBuilder &builder,
   builder.create<fir::CallOp>(loc, funcOp, args);
   return builder.create<fir::LoadOp>(loc, result);
 }
+
+/// Generate Call to runtime prif_num_images
+mlir::Value fir::runtime::getNumImages(fir::FirOpBuilder &builder,
+                                       mlir::Location loc) {
+  mlir::Value result = builder.createTemporary(loc, builder.getI32Type());
+  mlir::FunctionType ftype =
+      PRIF_FUNCTYPE(builder.getRefType(builder.getI32Type()));
+  mlir::func::FuncOp funcOp =
+      builder.createFunction(loc, PRIFNAME_SUB("num_images"), ftype);
+  llvm::SmallVector<mlir::Value> args =
+      fir::runtime::createArguments(builder, loc, ftype, result);
+  builder.create<fir::CallOp>(loc, funcOp, args);
+  return builder.create<fir::LoadOp>(loc, result);
+}
+
+/// Generate Call to runtime prif_num_images_with_{team|team_number}
+mlir::Value fir::runtime::getNumImagesWithTeam(fir::FirOpBuilder &builder,
+                                               mlir::Location loc,
+                                               mlir::Value team) {
+  bool isTeamNumber = fir::unwrapPassByRefType(team.getType()).isInteger();
+  std::string numImagesName = isTeamNumber
+                                  ? PRIFNAME_SUB("num_images_with_team_number")
+                                  : PRIFNAME_SUB("num_images_with_team");
+
+  mlir::Value result = builder.createTemporary(loc, builder.getI32Type());
+  mlir::Type refTy = builder.getRefType(builder.getI32Type());
+  mlir::FunctionType ftype =
+      isTeamNumber
+          ? PRIF_FUNCTYPE(builder.getRefType(builder.getI64Type()), refTy)
+          : PRIF_FUNCTYPE(fir::BoxType::get(builder.getNoneType()), refTy);
+  mlir::func::FuncOp funcOp = builder.createFunction(loc, numImagesName, ftype);
+
+  if (!isTeamNumber)
+    team = builder.createBox(loc, team);
+  llvm::SmallVector<mlir::Value> args =
+      fir::runtime::createArguments(builder, loc, ftype, team, result);
+  builder.create<fir::CallOp>(loc, funcOp, args);
+  return builder.create<fir::LoadOp>(loc, result);
+}
+
+/// Generate Call to runtime prif_this_image_no_coarray
+mlir::Value fir::runtime::getThisImage(fir::FirOpBuilder &builder,
+                                       mlir::Location loc, mlir::Value team) {
+  mlir::Type refTy = builder.getRefType(builder.getI32Type());
+  mlir::Type boxTy = fir::BoxType::get(builder.getNoneType());
+  mlir::FunctionType ftype = PRIF_FUNCTYPE(boxTy, refTy);
+  mlir::func::FuncOp funcOp =
+      builder.createFunction(loc, PRIFNAME_SUB("this_image_no_coarray"), ftype);
+
+  mlir::Value result = builder.createTemporary(loc, builder.getI32Type());
+  mlir::Value teamArg =
+      !team ? builder.create<fir::AbsentOp>(loc, boxTy) : team;
+  llvm::SmallVector<mlir::Value> args =
+      fir::runtime::createArguments(builder, loc, ftype, teamArg, result);
+  builder.create<fir::CallOp>(loc, funcOp, args);
+  return builder.create<fir::LoadOp>(loc, result);
+}
diff --git a/flang/test/Lower/Coarray/num_images.f90 b/flang/test/Lower/Coarray/num_images.f90
new file mode 100644
index 0000000..ebfce5d
--- /dev/null
+++ b/flang/test/Lower/Coarray/num_images.f90
@@ -0,0 +1,18 @@
+! RUN: %flang_fc1 -emit-hlfir -fcoarray %s -o - | FileCheck %s
+
+program test
+  use iso_fortran_env
+  integer :: i
+  integer :: team_number 
+  type(team_type) :: team
+
+  ! CHECK: fir.call @_QMprifPprif_num_images
+  i = num_images()
+
+  ! CHECK: fir.call @_QMprifPprif_num_images_with_team_number
+  i = num_images(TEAM_NUMBER=team_number)
+
+  ! CHECK: fir.call @_QMprifPprif_num_images_with_team
+  i = num_images(TEAM=team)
+
+end program
diff --git a/flang/test/Lower/Coarray/this_image.f90 b/flang/test/Lower/Coarray/this_image.f90
new file mode 100644
index 0000000..143504b
--- /dev/null
+++ b/flang/test/Lower/Coarray/this_image.f90
@@ -0,0 +1,14 @@
+! RUN: %flang_fc1 -emit-hlfir -fcoarray %s -o - | FileCheck %s
+
+program test
+  use iso_fortran_env
+  integer :: i
+  type(team_type) :: team
+
+  ! CHECK: fir.call @_QMprifPprif_this_image_no_coarray
+  i = this_image()
+
+  ! CHECK: fir.call @_QMprifPprif_this_image_no_coarray
+  i = this_image(TEAM=team)
+
+end program