[flang] Adding lowering of TEAMs features to PRIF in MIF Dialect (#165573)

Support for multi-image features has begun to be integrated into LLVM
with the MIF dialect.
In this PR, you will find lowering and operations related to the TEAM
features (`SYNC TEAM`, `GET_TEAM`, `FORM TEAM`, `CHANGE TEAM`,
`TEAM_NUMBER`).

Note regarding the operation for `CHANGE TEAM` : This operation is
partial because it does not support the associated list of coarrays
because the allocation of a coarray and the lowering of PRIF's
`prif_alias_{create|destroy}` procedures are not yet supported in Flang.
This will be integrated later.

Any feedback is welcome.

Author: JDPailleux

flang/include/flang/Lower/AbstractConverter.h

@@ -351,6 +351,11 @@ class AbstractConverter {
 
   virtual Fortran::lower::StatementContext &getFctCtx() = 0;
 
+  /// Generate STAT and ERRMSG from a list of StatOrErrmsg
+  virtual std::pair<mlir::Value, mlir::Value>
+  genStatAndErrmsg(mlir::Location loc,
+                   const std::list<Fortran::parser::StatOrErrmsg> &) = 0;
+
   AbstractConverter(const Fortran::lower::LoweringOptions &loweringOptions)
       : loweringOptions(loweringOptions) {}
   virtual ~AbstractConverter() = default;

flang/include/flang/Lower/Coarray.h renamed to flang/include/flang/Lower/MultiImageFortran.h

@@ -1,13 +1,13 @@
-//===-- Lower/Coarray.h -- image related lowering ---------------*- C++ -*-===//
+//===-- Lower/MultiImageFortran.h -- image related lowering -----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef FORTRAN_LOWER_COARRAY_H
-#define FORTRAN_LOWER_COARRAY_H
+#ifndef FORTRAN_LOWER_MULTIIMAGEFORTRAN_H
+#define FORTRAN_LOWER_MULTIIMAGEFORTRAN_H
 
 #include "flang/Lower/AbstractConverter.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
@@ -33,6 +33,18 @@ namespace pft {
 struct Evaluation;
 } // namespace pft
 
+//===----------------------------------------------------------------------===//
+// Synchronization statements
+//===----------------------------------------------------------------------===//
+
+void genSyncAllStatement(AbstractConverter &, const parser::SyncAllStmt &);
+
+void genSyncImagesStatement(AbstractConverter &,
+                            const parser::SyncImagesStmt &);
+void genSyncMemoryStatement(AbstractConverter &,
+                            const parser::SyncMemoryStmt &);
+void genSyncTeamStatement(AbstractConverter &, const parser::SyncTeamStmt &);
+
 //===----------------------------------------------------------------------===//
 // TEAM constructs
 //===----------------------------------------------------------------------===//
@@ -75,4 +87,4 @@ class CoarrayExprHelper {
 } // namespace lower
 } // namespace Fortran
 
-#endif // FORTRAN_LOWER_COARRAY_H
+#endif // FORTRAN_LOWER_MULTIIMAGEFORTRAN_H

flang/include/flang/Lower/Runtime.h

@@ -57,12 +57,6 @@ void genEventWaitStatement(AbstractConverter &, const parser::EventWaitStmt &);
 void genLockStatement(AbstractConverter &, const parser::LockStmt &);
 void genFailImageStatement(AbstractConverter &);
 void genStopStatement(AbstractConverter &, const parser::StopStmt &);
-void genSyncAllStatement(AbstractConverter &, const parser::SyncAllStmt &);
-void genSyncImagesStatement(AbstractConverter &,
-                            const parser::SyncImagesStmt &);
-void genSyncMemoryStatement(AbstractConverter &,
-                            const parser::SyncMemoryStmt &);
-void genSyncTeamStatement(AbstractConverter &, const parser::SyncTeamStmt &);
 void genUnlockStatement(AbstractConverter &, const parser::UnlockStmt &);
 void genPauseStatement(AbstractConverter &, const parser::PauseStmt &);
 

flang/include/flang/Optimizer/Builder/IntrinsicCall.h

@@ -270,6 +270,7 @@ struct IntrinsicLibrary {
   void genGetEnvironmentVariable(llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genGetGID(mlir::Type resultType,
                         llvm::ArrayRef<mlir::Value> args);
+  mlir::Value genGetTeam(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genGetUID(mlir::Type resultType,
                         llvm::ArrayRef<mlir::Value> args);
   fir::ExtendedValue genHostnm(std::optional<mlir::Type> resultType,
@@ -425,6 +426,8 @@ struct IntrinsicLibrary {
   void genSystemClock(llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genTand(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genTanpi(mlir::Type, llvm::ArrayRef<mlir::Value>);
+  fir::ExtendedValue genTeamNumber(mlir::Type,
+                                   llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genTime(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genTrailz(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genTransfer(mlir::Type,

flang/include/flang/Optimizer/Dialect/MIF/MIFOps.td

@@ -21,6 +21,10 @@ include "flang/Optimizer/Dialect/FIRAttr.td"
 class mif_Op<string mnemonic, list<Trait> traits>
     : Op<MIFDialect, mnemonic, traits>;
 
+class region_Op<string mnemonic, list<Trait> traits = []>
+    : mif_Op<mnemonic, !listconcat(traits, [RecursivelySpeculatable,
+                                            RecursiveMemoryEffects])> {}
+
 //===----------------------------------------------------------------------===//
 // Initialization and Finalization
 //===----------------------------------------------------------------------===//
@@ -174,6 +178,18 @@ def mif_SyncMemoryOp : mif_Op<"sync_memory", [AttrSizedOperandSegments]> {
   }];
 }
 
+def mif_SyncTeamOp : mif_Op<"sync_team", [AttrSizedOperandSegments]> {
+  let summary = "Performs a synchronization of the team, identified by `team`";
+
+  let arguments = (ins AnyRefOrBoxType:$team, Optional<AnyReferenceLike>:$stat,
+      Optional<AnyRefOrBoxType>:$errmsg);
+  let assemblyFormat = [{
+    $team (`stat` $stat^ )?
+    (`errmsg` $errmsg^ )? 
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Collective Operations
 //===----------------------------------------------------------------------===//
@@ -265,4 +281,148 @@ def mif_CoSumOp
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Teams
+//===----------------------------------------------------------------------===//
+
+def mif_FormTeamOp : mif_Op<"form_team", [AttrSizedOperandSegments]> {
+  let summary =
+      "Create a set of sibling teams whose parent team is the current team.";
+  let description = [{
+    Create a new team for each unique `team_number` value specified.
+    Each executing image will belong to the team whose `team_number` is equal 
+    to the value of team-number on that image, and `team_var` becomes defined
+    with a value that identifies that team.
+
+    If `new_index` is specified, the image index of the executing image will take
+    this index in its new team. Otherwise, the new image index is processor 
+    dependent.
+
+    Arguments: 
+      - `team_number`: Shall be a positive integer.
+      - `team_var` : Shall be a variable of type TEAM_TYPE from the intrinsic
+        module ISO_FORTRAN_ENV.
+      - `new_index`(optional): Shall be an integer that correspond to the index that
+        the calling image will have in the new team.
+  }];
+
+  let arguments = (ins AnyIntegerType:$team_number,
+      Arg<fir_BoxType, "", [MemWrite]>:$team_var,
+      Optional<AnyIntegerType>:$new_index,
+      Arg<Optional<AnyReferenceLike>, "", [MemWrite]>:$stat,
+      Arg<Optional<AnyRefOrBoxType>, "", [MemWrite]>:$errmsg);
+
+  let assemblyFormat = [{
+    `team_number` $team_number `team_var` $team_var
+    (`new_index` $new_index^ )?
+    (`stat` $stat^ )?
+    (`errmsg` $errmsg^ )?
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def mif_EndTeamOp : mif_Op<"end_team", [AttrSizedOperandSegments, Terminator,
+                                        ParentOneOf<["ChangeTeamOp"]>]> {
+  let summary = "Changes the current team to the parent team.";
+  let description = [{
+    The END TEAM operation completes the CHANGE TEAM construct and 
+    restores the current team to the team that was current before 
+    the CHANGE TEAM construct. 
+  }];
+
+  let arguments = (ins Arg<Optional<AnyReferenceLike>, "", [MemWrite]>:$stat,
+      Arg<Optional<AnyRefOrBoxType>, "", [MemWrite]>:$errmsg);
+  let builders = [OpBuilder<(ins), [{ /* do nothing */ }]>];
+
+  let assemblyFormat = [{
+    (`stat` $stat^ )? (`errmsg` $errmsg^ )?
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// NOTE: The CHANGE TEAM region will take a coarray association list in
+// argument. However, coarray management and coarray alias creation are not
+// yet supported by the dialect. The argument is therefore not yet supported by
+// this operation and will be added later.
+//===----------------------------------------------------------------------===//
+def mif_ChangeTeamOp
+    : region_Op<"change_team", [AttrSizedOperandSegments,
+                                SingleBlockImplicitTerminator<"EndTeamOp">]> {
+  let summary = "Changes the current team.";
+  let description = [{
+    The CHANGE TEAM construct changes the current team to the specified new
+    team, which must be a child team of the current team.  
+
+    ```
+      mif.change_team %team {
+        %x = fir.convert %i : (index) -> i32
+        ...
+        mif.end_team
+      }
+  }];
+
+  let arguments = (ins AnyRefOrBoxType:$team,
+      Arg<Optional<AnyReferenceLike>, "", [MemWrite]>:$stat,
+      Arg<Optional<AnyRefOrBoxType>, "", [MemWrite]>:$errmsg);
+  let regions = (region SizedRegion<1>:$region);
+
+  let skipDefaultBuilders = 1;
+  let builders =
+      [OpBuilder<(ins "mlir::Value":$team,
+           CArg<"bool", "true">:$ensureTermination,
+           CArg<"llvm::ArrayRef<mlir::NamedAttribute>", "{}">:$attributes)>,
+       OpBuilder<(ins "mlir::Value":$team, "mlir::Value":$stat,
+           "mlir::Value":$errmsg, CArg<"bool", "true">:$ensureTermination,
+           CArg<"llvm::ArrayRef<mlir::NamedAttribute>", "{}">:$attributes)>];
+
+  let extraClassDeclaration = [{
+    /// Get the body of the CHANGE TEAM construct 
+    mlir::Block *getBody() { return &getRegion().front(); }
+  }];
+
+  let assemblyFormat = [{
+    $team (`stat` $stat^)?
+    (`errmsg` $errmsg^)?
+    attr-dict `:` `(` type(operands) `)`
+    custom<ChangeTeamOpBody>($region)
+  }];
+}
+
+def mif_GetTeamOp : mif_Op<"get_team", []> {
+  let summary = "Get the team value for the current or ancestor team.";
+  let description = [{
+    This operation gets the team value for the current or an ancestor team.
+    `level`(optional): If provided, must equal one of the following constants : 
+    `INITIAL_TEAM`, `PARENT_TEAM` or `CURRENT_TEAM` from the module ISO_FORTRAN_ENV.
+    If `level` isn't present or has the value `CURRENT_TEAM` the returned
+    value is the current team.     
+  }];
+
+  let arguments = (ins Optional<AnyIntegerType>:$level);
+  let results = (outs fir_BoxType:$team);
+
+  let assemblyFormat = [{
+    (`level` $level^ )?
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def mif_TeamNumberOp : mif_Op<"team_number", []> {
+  let summary = "Get the team number";
+  let description = [{
+    Argument: `team` is optional and shall be a scalar of type TEAM_TYPE from 
+    module ISO_FORTRAN_ENV and the value identifies the current or an ancestor team. 
+    If `team` is absent, the team specified is the current team.
+  }];
+
+  let arguments = (ins Optional<AnyRefOrBoxType>:$team);
+  let results = (outs I64);
+
+  let assemblyFormat = [{
+    (`team` $team^ )?
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
 #endif // FORTRAN_DIALECT_MIF_MIF_OPS

flang/lib/Lower/Bridge.cpp

@@ -15,7 +15,6 @@
 #include "flang/Lower/Allocatable.h"
 #include "flang/Lower/CUDA.h"
 #include "flang/Lower/CallInterface.h"
-#include "flang/Lower/Coarray.h"
 #include "flang/Lower/ConvertCall.h"
 #include "flang/Lower/ConvertExpr.h"
 #include "flang/Lower/ConvertExprToHLFIR.h"
@@ -26,6 +25,7 @@
 #include "flang/Lower/IO.h"
 #include "flang/Lower/IterationSpace.h"
 #include "flang/Lower/Mangler.h"
+#include "flang/Lower/MultiImageFortran.h"
 #include "flang/Lower/OpenACC.h"
 #include "flang/Lower/OpenMP.h"
 #include "flang/Lower/PFTBuilder.h"
@@ -1111,6 +1111,34 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     return bridge.fctCtx();
   }
 
+  /// Initializes values for STAT and ERRMSG
+  std::pair<mlir::Value, mlir::Value>
+  genStatAndErrmsg(mlir::Location loc,
+                   const std::list<Fortran::parser::StatOrErrmsg>
+                       &statOrErrList) override final {
+    Fortran::lower::StatementContext stmtCtx;
+
+    mlir::Value errMsgExpr, statExpr;
+    for (const Fortran::parser::StatOrErrmsg &statOrErr : statOrErrList) {
+      std::visit(Fortran::common::visitors{
+                     [&](const Fortran::parser::StatVariable &statVar) {
+                       const Fortran::semantics::SomeExpr *expr =
+                           Fortran::semantics::GetExpr(statVar);
+                       statExpr =
+                           fir::getBase(genExprAddr(*expr, stmtCtx, &loc));
+                     },
+                     [&](const Fortran::parser::MsgVariable &errMsgVar) {
+                       const Fortran::semantics::SomeExpr *expr =
+                           Fortran::semantics::GetExpr(errMsgVar);
+                       errMsgExpr =
+                           fir::getBase(genExprBox(loc, *expr, stmtCtx));
+                     }},
+                 statOrErr.u);
+    }
+
+    return {statExpr, errMsgExpr};
+  }
+
   mlir::Value hostAssocTupleValue() override final { return hostAssocTuple; }
 
   /// Record a binding for the ssa-value of the tuple for this function.
@@ -3953,13 +3981,30 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   }
 
   void genFIR(const Fortran::parser::ChangeTeamConstruct &construct) {
-    TODO(toLocation(), "coarray: ChangeTeamConstruct");
+    Fortran::lower::StatementContext stmtCtx;
+    pushActiveConstruct(getEval(), stmtCtx);
+
+    for (Fortran::lower::pft::Evaluation &e :
+         getEval().getNestedEvaluations()) {
+      if (e.getIf<Fortran::parser::ChangeTeamStmt>()) {
+        maybeStartBlock(e.block);
+        setCurrentPosition(e.position);
+        genFIR(e);
+      } else if (e.getIf<Fortran::parser::EndChangeTeamStmt>()) {
+        maybeStartBlock(e.block);
+        setCurrentPosition(e.position);
+        genFIR(e);
+      } else {
+        genFIR(e);
+      }
+    }
+    popActiveConstruct();
   }
   void genFIR(const Fortran::parser::ChangeTeamStmt &stmt) {
-    TODO(toLocation(), "coarray: ChangeTeamStmt");
+    genChangeTeamStmt(*this, getEval(), stmt);
   }
   void genFIR(const Fortran::parser::EndChangeTeamStmt &stmt) {
-    TODO(toLocation(), "coarray: EndChangeTeamStmt");
+    genEndChangeTeamStmt(*this, getEval(), stmt);
   }
 
   void genFIR(const Fortran::parser::CriticalConstruct &criticalConstruct) {

flang/lib/Lower/CMakeLists.txt

@@ -5,7 +5,6 @@ add_flang_library(FortranLower
   Allocatable.cpp
   Bridge.cpp
   CallInterface.cpp
-  Coarray.cpp
   ComponentPath.cpp
   ConvertArrayConstructor.cpp
   ConvertCall.cpp
@@ -23,6 +22,7 @@ add_flang_library(FortranLower
   IterationSpace.cpp
   LoweringOptions.cpp
   Mangler.cpp
+  MultiImageFortran.cpp
   OpenACC.cpp
   OpenMP/Atomic.cpp
   OpenMP/ClauseProcessor.cpp