[Flang] Coarray allocation/deallocation + basic intrinsics for coarrays

flang-rt/lib/runtime/CMakeLists.txt

@@ -23,6 +23,7 @@ set(supported_sources
   assign.cpp
   buffer.cpp
   character.cpp
+  coarray.cpp
   connection.cpp
   copy.cpp
   derived-api.cpp

flang-rt/lib/runtime/coarray.cpp

@@ -0,0 +1,36 @@
+//===-- runtime/coarray.cpp -----------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Runtime/coarray.h"
+#include "flang-rt/runtime/descriptor.h"
+#include "flang-rt/runtime/type-info.h"
+
+namespace Fortran::runtime {
+
+extern "C" {
+RT_EXT_API_GROUP_BEGIN
+
+void RTDEF(ComputeLastUcobound)(
+    int num_images, const Descriptor &lcobounds, const Descriptor &ucobounds) {
+  int corank = ucobounds.GetDimension(0).Extent();
+  int64_t *lcobounds_ptr = (int64_t *)lcobounds.raw().base_addr;
+  int64_t *ucobounds_ptr = (int64_t *)ucobounds.raw().base_addr;
+  int64_t index = 1;
+  for (int i = 0; i < corank - 1; i++) {
+    index *= ucobounds_ptr[i] - lcobounds_ptr[i] + 1;
+  }
+  if (corank == 1)
+    ucobounds_ptr[0] = num_images - lcobounds_ptr[0] + 1;
+  else if (index < num_images)
+    ucobounds_ptr[corank - 1] =
+        (num_images / index) + (num_images % index != 0);
+}
+
+RT_EXT_API_GROUP_END
+}
+} // namespace Fortran::runtime

flang/include/flang/Lower/MultiImageFortran.h

@@ -59,6 +59,25 @@ void genEndChangeTeamStmt(AbstractConverter &, pft::Evaluation &eval,
 void genFormTeamStatement(AbstractConverter &, pft::Evaluation &eval,
                           const parser::FormTeamStmt &);
 
+//===----------------------------------------------------------------------===//
+// COARRAY utils
+//===----------------------------------------------------------------------===//
+
+mlir::DenseI64ArrayAttr genLowerCoBounds(AbstractConverter &converter,
+                                         mlir::Location loc,
+                                         const semantics::Symbol &sym);
+
+mlir::DenseI64ArrayAttr genUpperCoBounds(AbstractConverter &converter,
+                                         mlir::Location loc,
+                                         const semantics::Symbol &sym);
+
+mlir::Value genAllocateCoarray(
+    AbstractConverter &converter, mlir::Location loc,
+    const semantics::Symbol &sym, mlir::Value addr,
+    const std::optional<Fortran::parser::AllocateCoarraySpec> &allocSpec =
+        std::nullopt,
+    mlir::Value errMsg = {}, bool hasStat = false);
+
 //===----------------------------------------------------------------------===//
 // COARRAY expressions
 //===----------------------------------------------------------------------===//

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

@@ -219,6 +219,7 @@ struct IntrinsicLibrary {
                                          llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genCmplx(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genConjg(mlir::Type, llvm::ArrayRef<mlir::Value>);
+  fir::ExtendedValue genCoshape(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genCount(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   void genCpuTime(llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genCshift(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
@@ -335,6 +336,8 @@ struct IntrinsicLibrary {
   mlir::Value genIeeeUnordered(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genIeeeValue(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genIeor(mlir::Type, llvm::ArrayRef<mlir::Value>);
+  fir::ExtendedValue genImageIndex(mlir::Type,
+                                   llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genIndex(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genIor(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genIparity(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
@@ -349,6 +352,8 @@ struct IntrinsicLibrary {
   mlir::Value genIshft(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genIshftc(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genLbound(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
+  fir::ExtendedValue genLcobound(mlir::Type,
+                                 llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genLeadz(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genLen(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genLenTrim(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
@@ -452,6 +457,8 @@ struct IntrinsicLibrary {
                                   llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genTrim(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genUbound(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
+  fir::ExtendedValue genUcobound(mlir::Type,
+                                 llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genUnlink(std::optional<mlir::Type> resultType,
                                llvm::ArrayRef<fir::ExtendedValue> args);
   fir::ExtendedValue genUnpack(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);

flang/include/flang/Optimizer/Builder/MIFCommon.h

@@ -0,0 +1,26 @@
+//===-- MIFCommon.h -------------------------------------------------------===//
+//
+// 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_OPTIMIZER_TRANSFORMS_MIFCOMMON_H_
+#define FORTRAN_OPTIMIZER_TRANSFORMS_MIFCOMMON_H_
+
+#include "flang/Lower/AbstractConverter.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/MIF/MIFOps.h"
+#include "flang/Runtime/coarray.h"
+#include "mlir/IR/BuiltinOps.h"
+
+static constexpr llvm::StringRef coarrayHandleSuffix = "_coarray_handle";
+
+namespace mif {
+
+std::string getFullUniqName(mlir::Value addr);
+
+} // namespace mif
+
+#endif // FORTRAN_OPTIMIZER_TRANSFORMS_MIFCOMMON_H_

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

@@ -17,6 +17,7 @@
 include "flang/Optimizer/Dialect/MIF/MIFDialect.td"
 include "flang/Optimizer/Dialect/FIRTypes.td"
 include "flang/Optimizer/Dialect/FIRAttr.td"
+include "mlir/IR/BuiltinAttributes.td"
 
 class mif_Op<string mnemonic, list<Trait> traits>
     : Op<MIFDialect, mnemonic, traits>;
@@ -101,11 +102,13 @@ def mif_ThisImageOp : mif_Op<"this_image", [AttrSizedOperandSegments]> {
     value 5, `this_image(A)` has the value [5, 0, 0].
   }];
 
-  let arguments = (ins Optional<fir_BoxType>:$coarray,
-      Optional<AnyInteger>:$dim, Optional<AnyRefOrBoxType>:$team);
-  let results = (outs I32:$res);
+  let arguments = (ins Optional<AnyType>:$coarray, Optional<AnyInteger>:$dim,
+      Optional<AnyRefOrBoxType>:$team);
+  let results = (outs AnyType:$res);
 
-  let builders = [OpBuilder<(ins "mlir::Value":$coarray, "mlir::Value":$team)>,
+  let builders = [OpBuilder<(ins "mlir::Value":$coarray, "mlir::Value":$dim,
+                      "mlir::Value":$team)>,
+                  OpBuilder<(ins "mlir::Value":$coarray, "mlir::Value":$team)>,
                   OpBuilder<(ins "mlir::Value":$team)>];
 
   let hasVerifier = 1;
@@ -117,6 +120,120 @@ def mif_ThisImageOp : mif_Op<"this_image", [AttrSizedOperandSegments]> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Coarray Queries
+//===----------------------------------------------------------------------===//
+
+def mif_ImageIndexOp : mif_Op<"image_index", [AttrSizedOperandSegments]> {
+  let summary = "Image index from cosubscripts.";
+  let description = [{
+    Arguments:
+      - `coarray`: Shall be a coarray of any type.
+      - `sub`: rank-one integer array of size equal to the corank of `coarray`.
+      - `team`: Shall be a scalar of type `team_type` from ISO_FORTRAN_ENV.
+      - `team_number`: It shall identify the initial team or a sibling team 
+        of the current team.
+
+    Usage:
+    - Case(1) : `call image_index(coarray, sub)`
+    - Case(2) : `call image_index(coarray, sub, team)`
+    - Case(3) : `call image_index(coarray, sub, team_number)`
+
+    Result: If the value of `sub` is a valid sequence of cosubscripts for `coarray` in the
+      team specified by `team` or `team_number`, or the current team if neither `team` nor
+      `team_number` appears, the result is the index of the corresponding image in that team.
+      Otherwise, the result is zero.
+  }];
+
+  let arguments = (ins AnyType:$coarray, fir_BoxType:$sub,
+      Optional<AnyRefOrBoxType>:$team, Optional<AnyInteger>:$team_number);
+  let builders = [OpBuilder<(ins "mlir::Value":$coarray, "mlir::Value":$sub,
+      "mlir::Value":$team)>];
+
+  let results = (outs I32);
+
+  let assemblyFormat = [{
+    `coarray` $coarray `sub` $sub 
+    ( `team` $team^ )? 
+    ( `team_number` $team_number^ )? 
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def mif_LcoboundOp : mif_Op<"lcobound", [NoMemoryEffect]> {
+  let summary = "Returns the lower cobound(s) associated with a coarray.";
+  let description = [{
+    This operation returns the lower cobound(s) associated with a coarray.
+    Arguments: 
+      - `coarray`: Shall be a coarray of any type.
+      - `dim`(optional) : Shall be an integer scalar. Its value shall be in the range of
+          `1 <= DIM <= N`, where `N` is the corank of the coarray.
+    Results:
+      - Case(1): If `dim` is present, the result is an integer scalar equal to
+        the lower cobound for codimension `dim`.
+      - Case(2): `dim` is absent, so the result is an array whose size matches
+        the corank of the indicated coarray.
+  }];
+
+  let arguments = (ins AnyType:$coarray, Optional<AnyInteger>:$dim);
+  let results = (outs AnyType);
+
+  let builders = [OpBuilder<(ins "mlir::Value":$coarray, "mlir::Value":$dim)>,
+                  OpBuilder<(ins "mlir::Value":$coarray)>];
+
+  let assemblyFormat = [{
+    `coarray` $coarray 
+    ( `dim` $dim^ )? 
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def mif_UcoboundOp : mif_Op<"ucobound", [NoMemoryEffect]> {
+  let summary = "Returns the upper cobound(s) associated with a coarray.";
+  let description = [{
+    This operation returns the upper cobound(s) associated with a coarray.
+    Arguments: 
+      - `coarray`: Shall be a coarray of any type.
+      - `dim`(optional) : Shall be an integer scalar. Its value shall be in the range of
+          `1 <= DIM <= N`, where `N` is the corank of the coarray.
+    Results:
+      - Case(1): If `dim` is present, the result is an integer scalar equal to
+        the upper cobound for codimension `dim`.
+      - Case(2): `dim` is absent, so the result is an array whose size matches
+        the corank of the indicated coarray.
+  }];
+
+  let arguments = (ins AnyType:$coarray, Optional<AnyInteger>:$dim);
+  let results = (outs AnyType);
+
+  let builders = [OpBuilder<(ins "mlir::Value":$coarray, "mlir::Value":$dim)>,
+                  OpBuilder<(ins "mlir::Value":$coarray)>];
+
+  let assemblyFormat = [{
+    `coarray` $coarray 
+    ( `dim` $dim^ )? 
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def mif_CoshapeOp : mif_Op<"coshape", [NoMemoryEffect]> {
+  let summary = "Return the sizes of codimensions of a coarray.";
+  let description = [{
+    Argument: `coarray`: Shall be a coarray of any type.
+    Result : Is an array whose size matches the corank of the indicated coarray and
+       returns `UCOBOUND - LCOBOUND + 1`.
+  }];
+
+  let arguments = (ins AnyType:$coarray);
+  let builders = [OpBuilder<(ins "mlir::Value":$coarray)>];
+  let results = (outs AnyBoxedArray);
+
+  let assemblyFormat = [{
+    `coarray` $coarray 
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Synchronization
 //===----------------------------------------------------------------------===//
@@ -425,4 +542,98 @@ def mif_TeamNumberOp : mif_Op<"team_number", []> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Allocation and Deallocation
+//===----------------------------------------------------------------------===//
+
+def mif_AllocaCoarrayOp
+    : mif_Op<"alloca_coarray", [AttrSizedOperandSegments,
+                                MemoryEffects<[MemAlloc<DefaultResource>]>]> {
+  let summary = "Perform the allocation of a coarray and provide a "
+                "corresponding coarray descriptor";
+
+  let description = [{
+    This operation allocates a coarray and provides the corresponding 
+    coarray descriptor. This call is collective over the current team.
+  }];
+
+  let arguments = (ins StrAttr:$uniq_name,
+      Arg<fir_ReferenceType, "", [MemRead, MemWrite]>:$box,
+      DenseI64ArrayAttr:$lcobounds, DenseI64ArrayAttr:$ucobounds,
+      Arg<Optional<AnyReferenceLike>, "", [MemWrite]>:$stat,
+      Arg<Optional<AnyRefOrBoxType>, "", [MemWrite]>:$errmsg);
+
+  let builders = [OpBuilder<(ins "mlir::Value":$box, "llvm::StringRef":$symName,
+                      "mlir::DenseI64ArrayAttr":$lcobounds,
+                      "mlir::DenseI64ArrayAttr":$ucobounds, "mlir::Value":$stat,
+                      "mlir::Value":$errmsg)>,
+                  OpBuilder<(ins "mlir::Value":$box, "llvm::StringRef":$symName,
+                      "mlir::DenseI64ArrayAttr":$lcobounds,
+                      "mlir::DenseI64ArrayAttr":$ucobounds)>];
+
+  let assemblyFormat = [{
+    $box 
+    (`stat` $stat^ )?
+    (`errmsg` $errmsg^ )? 
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def mif_AllocaOp
+    : mif_Op<"alloca", [AttrSizedOperandSegments,
+                        MemoryEffects<[MemAlloc<DefaultResource>]>]> {
+  let summary =
+      "Perform a non-collective allocation of remotely accessible storage";
+
+  let description = [{
+    A non-collective allocation of remotely accessible storage (must be done 
+    for an allocatable component of a coarray).
+  }];
+
+  let arguments = (ins Arg<fir_ReferenceType, "", [MemRead, MemWrite]>:$box,
+      Arg<Optional<AnyReferenceLike>, "", [MemWrite]>:$stat,
+      Arg<Optional<AnyRefOrBoxType>, "", [MemWrite]>:$errmsg);
+
+  let assemblyFormat = [{
+    $box 
+    (`stat` $stat^ )?
+    (`errmsg` $errmsg^ )? 
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def mif_DeallocaCoarrayOp
+    : mif_Op<"dealloca_coarray", [AttrSizedOperandSegments,
+                                  MemoryEffects<[MemFree<DefaultResource>]>]> {
+  let summary = "Perform the deallocation of a coarray";
+  let description = [{
+    This call releases memory allocated by `mif_AllocaCoarrayOp` for a coarray.
+  }];
+  let arguments = (ins Arg<fir_ReferenceType, "", [MemFree]>:$coarray,
+      Arg<Optional<AnyReferenceLike>, "", [MemWrite]>:$stat,
+      Arg<Optional<AnyRefOrBoxType>, "", [MemWrite]>:$errmsg);
+
+  let assemblyFormat = [{
+    $coarray (`stat` $stat^ )?  (`errmsg` $errmsg^ )?
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def mif_DeallocaOp
+    : mif_Op<"dealloca", [AttrSizedOperandSegments,
+                          MemoryEffects<[MemFree<DefaultResource>]>]> {
+  let summary = "Perform the deallocation of a coarray";
+  let description = [{
+    This call releases memory allocated by `mif_AllocaOp`.
+  }];
+  let arguments = (ins Arg<fir_ReferenceType, "", [MemFree]>:$addr,
+      Arg<Optional<AnyReferenceLike>, "", [MemWrite]>:$stat,
+      Arg<Optional<AnyRefOrBoxType>, "", [MemWrite]>:$errmsg);
+
+  let assemblyFormat = [{
+    $addr (`stat` $stat^ )?  (`errmsg` $errmsg^ )?
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
 #endif // FORTRAN_DIALECT_MIF_MIF_OPS

flang/include/flang/Optimizer/Transforms/MIFOpConversion.h

@@ -20,7 +20,9 @@ class LLVMTypeConverter;
 namespace mif {
 
 /// Patterns that convert MIF operations to runtime calls.
-void populateMIFOpConversionPatterns(mlir::RewritePatternSet &patterns);
+void populateMIFOpConversionPatterns(const fir::LLVMTypeConverter &converter,
+                                     mlir::DataLayout &dl,
+                                     mlir::RewritePatternSet &patterns);
 
 } // namespace mif