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

Adds delayed privatization support for allocatables. In order to explain
the problem this diff is trying to solve, it would be useful to see an
example of `firstprivate` for an alloctable and how it is currently
emitted by flang **without** delayed privatization.

Consider the following Fortran code:
```fortran
subroutine delayed_privatization_allocatable
  implicit none
  integer, allocatable :: var1

!$omp parallel firstprivate(var1)
  var1 = 10
!$omp end parallel
end subroutine
```

You would get something like this (again no delayed privatization yet):
```mlir
    ...
    %3:2 = hlfir.declare %0 {fortran_attrs = #fir.var_attrs<allocatable>, ...} : ...
    omp.parallel {
      // Allocation logic
      ...
      %9 = fir.load %3#1 : !fir.ref<!fir.box<!fir.heap<i32>>>
      ...
      fir.if %12 {
        ...
      } else {
        ...
      }
      %13:2 = hlfir.declare %8 {fortran_attrs = #fir.var_attrs<allocatable>, ...} : ...

      // Copy logic
      %14 = fir.load %13#0 : !fir.ref<!fir.box<!fir.heap<i32>>>
      ...
      fir.if %17 {
        %24 = fir.load %3#0 : !fir.ref<!fir.box<!fir.heap<i32>>>
        ...
      }
      ...
      omp.terminator
    }
```

Note that for both the original and the private declaration, the
allocation and copy logic use both elements (e.g. `%3#0` and `%3#1`).
This poses the following problem for delayed privatization: how can
capture both values in order to pass them to the outlined privatizer?
The main issue is that `hlfir.declare` returns 2 SSA values that not
encapsulated together under a single composite value.

Some possible solutions are:
1. Change the return type of `hlfir.declare`. However, this would be
   very invasive and therefore does not sound like a good idea.
2. Use the built-in `tuple` type to "pack" both values together. This is
   reasonable but not very flexible since we won't be able to express
   other information about the encapsulated values such as whether it is
   alloctable.
3. Introduce a new concept/type that allows us to do the encapsulation
   in a more flexible way. This is the "variable shadow" concept
   introduced by this PR.

A "variable shadow" is a composite value of a special type:
`!fir.shadow` that has some special properties:
* It can be constructed only from `BlockArgument`s.
* It always has the type `!fir.shadow`.
* It packs the required information needed by the delayed privatizer
  that correspond to the result of `hlfir.declare` operation.

We don't introduce any special opertions to deal with shadow values.
Rather we treat them similar to other composites. If you want to get a
certain component, use `fir.extract_value`. If you want to populate a
certain component, use `fir.insert_value`.

Author: ergawy

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;