[flang] Add support of THIS_IMAGE and NUM_IMAGES with PRIF #154081
Merged
+159
−0
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
llvmbot
added
flang
|
@llvm/pr-subscribers-flang-fir-hlfir Author: Jean-Didier PAILLEUX (JDPailleux) ChangesIn relation to the approval and merge of the #76088 specification about multi-image features in Flang. Full diff: https://github.com/llvm/llvm-project/pull/154081.diff 6 Files Affected:
diff --git a/flang/include/flang/Optimizer/Builder/IntrinsicCall.h b/flang/include/flang/Optimizer/Builder/IntrinsicCall.h
index 2afd50410ae82..88c3ada3ff64f 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 f2c76c9e8d978..23bb378c30838 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 319ab1912cd3d..1964ff3a387a5 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},
@@ -7277,6 +7288,20 @@ 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,
@@ -8327,6 +8352,28 @@ 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 eaff6c37ecdbf..fb72fc2089e23 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 0000000000000..ebfce5db0dbfb
--- /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 0000000000000..967b5b6727759
--- /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
+ i = this_image()
+
+ ! CHECK: fir.call @_QMprifPprif_this_image_no_coarray
+ i = this_image(TEAM=team)
+
+end program
|
clementval
reviewed
Aug 18, 2025
clementval
reviewed
Aug 18, 2025
clementval
approved these changes
Aug 18, 2025
JDPailleux
force-pushed
the
jdp/flang/prif_this_image_num_images
branch
from
437018d to
4ceb28b
Compare
|
@clementval Ok for braces and thanks for the review. |
ktras
reviewed
Aug 18, 2025
There was a problem hiding this comment.
Looks great! And I can confirm that flang built from this branch can build Caffeine and link against it when compiling a program with num_images and this_image calls. Here is some output I was able to get:
Compiler version is flang version 22.0.0 ([email protected]:SiPearl/llvm-project.
git 4ceb28b8fea237b281d5f831ec791897d7741c23)
Hello world from image 3 of 4
Hello world from image 1 of 4
Hello world from image 2 of 4
Hello world from image 4 of 4
I had a question on one line of the test, so please check it out.
|
@ktras @JDPailleux We probably talked about this in one of the flang meetings but what is the plan for Caffeine? Is this gonna be proposed to join llvm upstream at some point or is the plan to have a different implementation of PRIF upstream? |
|
@clementval asks:
Thanks for the great question! The short answer is: the question of when/if to upstream a parallel runtime library has not yet been decided. The main idea with PRIF is to have an open runtime interface for the multi-image features. We hope there can eventually be several PRIF library implementations (e.g. potentially optimized in vendor-specific or even proprietary ways) and we want to allow the flang user to mix-and-match parallel runtime libraries to suit their needs. Here's a paper with more details about PRIF: |
|
Ok I see. Thanks for the refresher. The only problem I see without an implementation upstream in LLVM, is that we cannot really test the end to end pipeline. is there any plan to have a CI with caffeine and flang somewhere? |
JDPailleux
force-pushed
the
jdp/flang/prif_this_image_num_images
branch
from
4ceb28b to
14560be
Compare
@clementval Yes, we definitely agree that end to end testing is important. We are discussing how we can best achieve this in a follow-up pull request. |
ktras
approved these changes
Aug 19, 2025
sscalpone
reviewed
Aug 22, 2025
| if (converter && | ||
| !converter->getFoldingContext().languageFeatures().IsEnabled( | ||
| Fortran::common::LanguageFeature::Coarray)) | ||
| fir::emitFatalError(loc, "Coarrays disabled, use '-fcoarray' to enable.", |
There was a problem hiding this comment.
If you could, please change this to something like:
not yet implemented: coarrays are experimental, use '-fcoarray' to enable
to keep the "not yet implemented" pattern in place. Thank you!
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
In relation to the approval and merge of the #76088 specification about multi-image features in Flang.
Here is a PR on adding support for
THIS_IMAGEandNUM_IMAGESin 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.