[flang][MLIR][OpenMP] Extend delayed privatization for allocatables

flang/include/flang/Lower/SymbolMap.h

@@ -15,6 +15,7 @@
 
 #include "flang/Common/reference.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
+#include "flang/Optimizer/Builder/VariableShadow.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/FortranVariableInterface.h"
 #include "flang/Optimizer/Support/Matcher.h"
@@ -77,7 +78,8 @@ struct SymbolBox : public fir::details::matcher<SymbolBox> {
 
   using VT =
       std::variant<Intrinsic, FullDim, Char, CharFullDim, PointerOrAllocatable,
-                   Box, fir::FortranVariableOpInterface, None>;
+                   Box, fir::FortranVariableOpInterface,
+                   hlfir::FortranVariableShadow, None>;
 
   //===--------------------------------------------------------------------===//
   // Constructors
@@ -97,11 +99,13 @@ struct SymbolBox : public fir::details::matcher<SymbolBox> {
   /// scalar. For an array, this is the address of the first element in the
   /// array, etc.
   mlir::Value getAddr() const {
-    return match([](const None &) { return mlir::Value{}; },
-                 [](const fir::FortranVariableOpInterface &x) {
-                   return fir::FortranVariableOpInterface(x).getBase();
-                 },
-                 [](const auto &x) { return x.getAddr(); });
+    return match(
+        [](const None &) { return mlir::Value{}; },
+        [](const fir::FortranVariableOpInterface &x) {
+          return fir::FortranVariableOpInterface(x).getBase();
+        },
+        [](const hlfir::FortranVariableShadow &x) { return x.getBase(); },
+        [](const auto &x) { return x.getAddr(); });
   }
 
   std::optional<fir::FortranVariableOpInterface>

flang/include/flang/Optimizer/Builder/BoxValue.h

@@ -13,6 +13,8 @@
 #ifndef FORTRAN_OPTIMIZER_BUILDER_BOXVALUE_H
 #define FORTRAN_OPTIMIZER_BUILDER_BOXVALUE_H
 
+#include "flang/Optimizer/Builder/Todo.h"
+#include "flang/Optimizer/Builder/VariableShadow.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Support/FatalError.h"
 #include "flang/Optimizer/Support/Matcher.h"
@@ -479,7 +481,8 @@ class ExtendedValue : public details::matcher<ExtendedValue> {
 public:
   using VT =
       std::variant<UnboxedValue, CharBoxValue, ArrayBoxValue, CharArrayBoxValue,
-                   ProcBoxValue, BoxValue, MutableBoxValue, PolymorphicValue>;
+                   ProcBoxValue, BoxValue, MutableBoxValue, PolymorphicValue,
+                   hlfir::FortranVariableShadow>;
 
   ExtendedValue() : box{UnboxedValue{}} {}
   template <typename A, typename = std::enable_if_t<
@@ -516,6 +519,11 @@ class ExtendedValue : public details::matcher<ExtendedValue> {
                  [](const fir::CharBoxValue &box) -> unsigned { return 0; },
                  [](const fir::ProcBoxValue &box) -> unsigned { return 0; },
                  [](const fir::PolymorphicValue &box) -> unsigned { return 0; },
+                 [](const hlfir::FortranVariableShadow &box) -> unsigned {
+                   TODO(box.getValue().getLoc(),
+                        "rank(): FortranVariableShadow");
+                   return 0;
+                 },
                  [](const auto &box) -> unsigned { return box.rank(); });
   }
 

flang/include/flang/Optimizer/Builder/HLFIRTools.h

@@ -14,6 +14,7 @@
 #define FORTRAN_OPTIMIZER_BUILDER_HLFIRTOOLS_H
 
 #include "flang/Optimizer/Builder/BoxValue.h"
+#include "flang/Optimizer/Builder/VariableShadow.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FortranVariableInterface.h"
 #include "flang/Optimizer/HLFIR/HLFIRDialect.h"
@@ -55,8 +56,12 @@ class Entity : public mlir::Value {
   }
   Entity(fir::FortranVariableOpInterface variable)
       : mlir::Value(variable.getBase()) {}
+  Entity(FortranVariableShadow variableShadow)
+      : mlir::Value(variableShadow.getValue()) {}
+
   bool isValue() const { return isFortranValue(*this); }
   bool isVariable() const { return !isValue(); }
+  bool isMutableShadow() const;
   bool isMutableBox() const { return hlfir::isBoxAddressType(getType()); }
   bool isProcedurePointer() const {
     return fir::isBoxProcAddressType(getType());
@@ -74,6 +79,10 @@ class Entity : public mlir::Value {
   /// Is this an assumed ranked entity?
   bool isAssumedRank() const { return getRank() == -1; }
 
+  bool isFortranVariableShadow() const {
+    return this->getType().isa<fir::ShadowType>();
+  }
+
   /// Return the rank of this entity or -1 if it is an assumed rank.
   int getRank() const {
     mlir::Type type = fir::unwrapPassByRefType(fir::unwrapRefType(getType()));
@@ -152,6 +161,14 @@ class Entity : public mlir::Value {
     return this->getDefiningOp<fir::FortranVariableOpInterface>();
   }
 
+  fir::ShadowType getIfVariableShadow() const {
+    fir::ExtractValueOp op = this->getDefiningOp<fir::ExtractValueOp>();
+    if (!op)
+      return fir::ShadowType();
+
+    return op.getAdt().getType().dyn_cast<fir::ShadowType>();
+  }
+
   // Return a "declaration" operation for this variable if visible,
   // or the "declaration" operation of the allocatable/pointer this
   // variable was dereferenced from (if it is visible).
@@ -174,8 +191,13 @@ class Entity : public mlir::Value {
   }
 
   bool isAllocatable() const {
-    auto varIface = getIfVariableInterface();
-    return varIface ? varIface.isAllocatable() : false;
+    if (auto varIface = getIfVariableInterface())
+      return varIface.isAllocatable();
+
+    if (auto shadowType = getIfVariableShadow())
+      return shadowType.getAllocatable();
+
+    return false;
   }
 
   bool isPointer() const {
@@ -185,7 +207,7 @@ class Entity : public mlir::Value {
 
   // Get the entity as an mlir SSA value containing all the shape, type
   // parameters and dynamic shape information.
-  mlir::Value getBase() const { return *this; }
+  mlir::Value getBase() const;
 
   // Get the entity as a FIR base. This may not carry the shape and type
   // parameters information, and even when it is a box with shape information.
@@ -231,6 +253,9 @@ translateToExtendedValue(mlir::Location loc, fir::FirOpBuilder &builder,
 fir::ExtendedValue
 translateToExtendedValue(mlir::Location loc, fir::FirOpBuilder &builder,
                          fir::FortranVariableOpInterface fortranVariable);
+fir::ExtendedValue
+translateToExtendedValue(mlir::Location loc, fir::FirOpBuilder &builder,
+                         hlfir::FortranVariableShadow fortranVariableShadow);
 
 /// Generate declaration for a fir::ExtendedValue in memory.
 fir::FortranVariableOpInterface

flang/include/flang/Optimizer/Builder/VariableShadow.h

@@ -0,0 +1,47 @@
+//===-- VariableShadow.h ----------------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Dialect/FortranVariableInterface.h"
+#include "mlir/IR/Value.h"
+
+#ifndef FORTRAN_OPTIMIZER_BUILDER_VARIABLESHADOW_H
+#define FORTRAN_OPTIMIZER_BUILDER_VARIABLESHADOW_H
+
+namespace hlfir {
+
+class FortranVariableShadow {
+public:
+  FortranVariableShadow(fir::FirOpBuilder &builder,
+                        mlir::BlockArgument shadowingVal,
+                        fir::FortranVariableOpInterface shadowedVariable);
+
+  mlir::BlockArgument getValue() const { return shadowingVal; }
+
+  mlir::Value getBase() const { return base; }
+
+  mlir::Value getFirBase() const { return firBase; }
+
+  fir::FortranVariableOpInterface getShadowedVariable() const {
+    return shadowedVariable;
+  }
+
+private:
+  mlir::BlockArgument shadowingVal;
+  mlir::Value base;
+  mlir::Value firBase;
+  fir::FortranVariableOpInterface shadowedVariable;
+};
+
+} // namespace hlfir
+
+#endif // FORTRAN_OPTIMIZER_BUILDER_VARIABLESHADOW_H

flang/include/flang/Optimizer/Dialect/FIRTypes.td

@@ -576,6 +576,65 @@ def fir_VoidType : FIR_Type<"Void", "void"> {
   let genStorageClass = 0;
 }
 
+def fir_VariableShadowType : FIR_Type<"Shadow", "shadow"> {
+  let summary = "Type for Fortran-variables shadow values";
+
+  let description = [{
+    A type to encapulate required data for a Fortran-variable shadow value. A
+    shadow type encapsulates the base type as well as the original base type. It
+    also encapsulates information other information about the shadowed value,
+    such as whether it is allocatable or not.
+
+    For example, if we want to shadow a value produced a `hlfir.declare`
+    instruction similar to the following:
+    ```
+    %1:2 = hlfir.declare %0
+      {fortran_attrs = #fir.var_attrs<allocatable>}
+      : (!fir.ref<!fir.box<!fir.heap<i32>>>)
+        -> (!fir.ref<!fir.box<!fir.heap<i32>>>, !fir.ref<!fir.box<!fir.heap<i32>>>)
+    ```
+
+    We would need to somehow create a value of the type:
+    ```
+    !fir.shadow<!fir.ref<!fir.box<!fir.heap<i32>>>,
+                !fir.ref<!fir.box<!fir.heap<i32>>>,
+                allocatable : true>
+    ```
+
+    The mechanism to create this value can be, for example, through a
+    combincation `fir.undefined`, `fir.insert_value` instructions similar to what
+    would be done for other composite types like `tuple`s:
+    ```
+    %2 = fir.undefined !fir.shadow<!fir.ref<!fir.box<!fir.heap<i32>>>,
+                                   !fir.ref<!fir.box<!fir.heap<i32>>>,
+                                   allocatable : true>
+    %3 = fir.insert_value %2, %1#0, [0 : index] : <... type info omitted ...>
+    %4 = fir.insert_value %3, %3#1, [1 : index] : <... type info omitted ...>
+    ```
+
+    From that point on, you can use `%4` as the shadow of the originally declared
+    variable being shadowed `%1`.
+
+    For more info about the concept of variable shadows, check `VariableShadow.h`.
+  }];
+
+  let parameters = (ins "mlir::Type":$baseTy,
+                        "mlir::Type":$firBaseTy,
+                        "bool":$allocatable);
+
+  let assemblyFormat = [{
+    `<` $baseTy `,` $firBaseTy `,` `allocatable` `:` $allocatable `>`
+  }];
+
+  let extraClassDeclaration = [{
+    mlir::Type getType(unsigned idx) const {
+      if (idx == 0) return getBaseTy();
+      if (idx == 1) return getFirBaseTy();
+      return mlir::Type();
+    }
+  }];
+}
+
 // Whether a type is a BaseBoxType
 def IsBaseBoxTypePred
         : CPred<"$_self.isa<::fir::BaseBoxType>()">;
@@ -590,13 +649,17 @@ def AnyRealLike : TypeConstraint<Or<[FloatLike.predicate,
     fir_RealType.predicate]>, "any real">;
 def AnyIntegerType : Type<AnyIntegerLike.predicate, "any integer">;
 
+def IsShadowTypePred
+        : CPred<"$_self.isa<::fir::ShadowType>()">;
+
 // Composable types
 // Note that we include both fir_ComplexType and AnyComplex, so we can use both
 // the FIR ComplexType and the MLIR ComplexType (the former is used to represent
 // Fortran complex and the latter for C++ std::complex).
 def AnyCompositeLike : TypeConstraint<Or<[fir_RecordType.predicate,
     fir_SequenceType.predicate, fir_ComplexType.predicate, AnyComplex.predicate,
-    fir_VectorType.predicate, IsTupleTypePred, fir_CharacterType.predicate]>,
+    fir_VectorType.predicate, IsTupleTypePred, fir_CharacterType.predicate,
+    IsShadowTypePred]>,
     "any composite">;
 
 // Reference types

flang/lib/Lower/Bridge.cpp

@@ -468,6 +468,10 @@ class FirConverter : public Fortran::lower::AbstractConverter {
           return hlfir::translateToExtendedValue(getCurrentLocation(),
                                                  getFirOpBuilder(), x);
         },
+        [&](const hlfir::FortranVariableShadow &x) -> fir::ExtendedValue {
+          return hlfir::translateToExtendedValue(getCurrentLocation(),
+                                                 getFirOpBuilder(), x);
+        },
         [](const auto &box) -> fir::ExtendedValue { return box; });
   }
 
@@ -1009,6 +1013,16 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         return v.match(
             [&](const Fortran::lower::SymbolBox::Intrinsic &)
                 -> Fortran::lower::SymbolBox { return v; },
+            [&](const hlfir::FortranVariableShadow &box)
+                -> Fortran::lower::SymbolBox {
+              auto exv =
+                  hlfir::translateToExtendedValue(toLocation(), *builder, box);
+              return exv.match(
+                  [](mlir::Value x) -> Fortran::lower::SymbolBox {
+                    return Fortran::lower::SymbolBox::Intrinsic{x};
+                  },
+                  [](auto x) -> Fortran::lower::SymbolBox { return x; });
+            },
             [](const auto &) -> Fortran::lower::SymbolBox { return {}; });
 
       return {};
@@ -1066,6 +1080,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     assert(lowerToHighLevelFIR());
     hlfir::Entity lhs{dst};
     hlfir::Entity rhs{src};
+
     // Temporary_lhs is set to true in hlfir.assign below to avoid user
     // assignment to be used and finalization to be called on the LHS.
     // This may or may not be correct but mimics the current behaviour

flang/lib/Lower/ConvertExpr.cpp

@@ -981,6 +981,9 @@ class ScalarExprLowering {
           },
           [&](const fir::ProcBoxValue &toBox) {
             TODO(loc, "procedure pointer component in derived type assignment");
+          },
+          [&](const hlfir::FortranVariableShadow &toBox) {
+            TODO(loc, "FortranVariableShadow in derived type assignment");
           });
     }
     return res;