diff --git a/flang/include/flang/Optimizer/Builder/IntrinsicCall.h b/flang/include/flang/Optimizer/Builder/IntrinsicCall.h index 2afd50410ae8..88c3ada3ff64 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); mlir::Value genNot(mlir::Type, llvm::ArrayRef); fir::ExtendedValue genNull(mlir::Type, llvm::ArrayRef); + fir::ExtendedValue genNumImages(mlir::Type, + llvm::ArrayRef); template mlir::Value genNVVMTime(mlir::Type, llvm::ArrayRef); fir::ExtendedValue genPack(mlir::Type, llvm::ArrayRef); @@ -449,6 +451,8 @@ struct IntrinsicLibrary { fir::ExtendedValue genTranspose(mlir::Type, llvm::ArrayRef); mlir::Value genThisGrid(mlir::Type, llvm::ArrayRef); + fir::ExtendedValue genThisImage(mlir::Type, + llvm::ArrayRef); mlir::Value genThisThreadBlock(mlir::Type, llvm::ArrayRef); mlir::Value genThisWarp(mlir::Type, llvm::ArrayRef); void genThreadFence(llvm::ArrayRef); @@ -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 f2c76c9e8d97..23bb378c3083 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 bc3086e28f75..8aacdb1815e3 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 args) { return fir::MutableBoxValue(boxStorage, mold->nonDeferredLenParams(), {}); } +// NUM_IMAGES +fir::ExtendedValue +IntrinsicLibrary::genNumImages(mlir::Type resultType, + llvm::ArrayRef 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 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 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(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 eaff6c37ecdb..fb72fc2089e2 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(loc, funcOp, args); return builder.create(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 args = + fir::runtime::createArguments(builder, loc, ftype, result); + builder.create(loc, funcOp, args); + return builder.create(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 args = + fir::runtime::createArguments(builder, loc, ftype, team, result); + builder.create(loc, funcOp, args); + return builder.create(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(loc, boxTy) : team; + llvm::SmallVector args = + fir::runtime::createArguments(builder, loc, ftype, teamArg, result); + builder.create(loc, funcOp, args); + return builder.create(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 000000000000..ebfce5db0dbf --- /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 000000000000..143504b5f922 --- /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