[flang][OpenACC] remap component references in structured constructs (#171501)

OpenACC data clauses of structured constructs may contain component
references (`obj%comp`, or `obj%array(i:j:k)`, ...).

This changes allows using the ACC dialect data operation result for such
clauses every time the component is referred to inside the scope of the
construct.

The bulk of the change is to add the ability to map
`evaluate::Component` to mlir values in the symbol map used in lowering.
This is done by adding the `ComponentMap` helper class to the lowering
symbol map, and using it to override `evaluate::Component` reference
lowering in expression lowering (ConvertExprToHLFIR.cpp).

Some changes are made in Lower/Support/Utils.h in order to set-up/expose
the hashing/equality helpers needed to use `evaluate::Component` in
llvm::DenseMap.

In OpenACC.cpp, `genPrivatizationRecipes` and `genDataOperandOperations`
are merged to unify the processing of Designator in data clauses.

New code is added to unwrap the rightmost `evaluate::Component`, if any,
and remap it.

Note that when the right most part is an array reference on a component
(`obj%array(i:j:k)`), the whole component `obj%array(` is remapped, and
the array reference is dealt with a bound operand like for whole object
array.

After this patch, all designators in data clauses on structured
constructs will be remapped, except for array reference in private/first
private/reduction (the result type of the related operation needs to be
changed and the recipe adapted to "offset back"), component reference in
reduction (this patch is adding a TODO for it), and device_ptr
(previously bypassed remapping because of issues with descriptors,
should be OK to lift it in a further patch).

Note that this patch assumes that it is illegal to have code with
intermediate variable indexing in the component reference (e.g.
`array(i)%comp`) where the value of the index would be changed inside
the region (e.g., `i` is assigned a new value), or where the component
would be used with different indices meant to have the same value as the
one used in the clause (e.g. `array(i)%comp` where `j` is meant to be
the same as `i`). I will try to add a warning in semantics for such
questionable/risky usages.

Author: jeanPerier

flang/include/flang/Lower/AbstractConverter.h

@@ -16,6 +16,7 @@
 #include "flang/Lower/LoweringOptions.h"
 #include "flang/Lower/PFTDefs.h"
 #include "flang/Lower/Support/Utils.h"
+#include "flang/Lower/SymbolMap.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Dialect/FIRAttr.h"
 #include "flang/Semantics/symbol.h"

flang/include/flang/Lower/Support/Utils.h

@@ -14,7 +14,6 @@
 #define FORTRAN_LOWER_SUPPORT_UTILS_H
 
 #include "flang/Common/indirection.h"
-#include "flang/Lower/IterationSpace.h"
 #include "flang/Parser/char-block.h"
 #include "flang/Semantics/tools.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -23,10 +22,25 @@
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringRef.h"
 
+namespace Fortran::evaluate {
+class Component;
+class ArrayRef;
+} // namespace Fortran::evaluate
+
 namespace Fortran::lower {
 using SomeExpr = Fortran::evaluate::Expr<Fortran::evaluate::SomeType>;
+using ExplicitSpaceArrayBases =
+    std::variant<const semantics::Symbol *, const evaluate::Component *,
+                 const evaluate::ArrayRef *>;
+// FIXME: needed for privatizeSymbol that does not belong to this header.
+class AbstractConverter;
+class SymMap;
 } // end namespace Fortran::lower
 
+namespace fir {
+class FirOpBuilder;
+}
+
 //===----------------------------------------------------------------------===//
 // Small inline helper functions to deal with repetitive, clumsy conversions.
 //===----------------------------------------------------------------------===//
@@ -89,12 +103,15 @@ A flatZip(const A &container1, const A &container2) {
 
 namespace Fortran::lower {
 unsigned getHashValue(const Fortran::lower::SomeExpr *x);
-unsigned getHashValue(const Fortran::lower::ExplicitIterSpace::ArrayBases &x);
+unsigned getHashValue(const Fortran::lower::ExplicitSpaceArrayBases &x);
+unsigned getHashValue(const Fortran::evaluate::Component *x);
 
 bool isEqual(const Fortran::lower::SomeExpr *x,
              const Fortran::lower::SomeExpr *y);
-bool isEqual(const Fortran::lower::ExplicitIterSpace::ArrayBases &x,
-             const Fortran::lower::ExplicitIterSpace::ArrayBases &y);
+bool isEqual(const Fortran::lower::ExplicitSpaceArrayBases &x,
+             const Fortran::lower::ExplicitSpaceArrayBases &y);
+bool isEqual(const Fortran::evaluate::Component *x,
+             const Fortran::evaluate::Component *y);
 
 template <typename OpType, typename OperandsStructType>
 void privatizeSymbol(
@@ -125,6 +142,24 @@ struct DenseMapInfo<const Fortran::lower::SomeExpr *> {
     return Fortran::lower::isEqual(lhs, rhs);
   }
 };
+
+// DenseMapInfo for pointers to Fortran::evaluate::Component.
+template <>
+struct DenseMapInfo<const Fortran::evaluate::Component *> {
+  static inline const Fortran::evaluate::Component *getEmptyKey() {
+    return reinterpret_cast<Fortran::evaluate::Component *>(~0);
+  }
+  static inline const Fortran::evaluate::Component *getTombstoneKey() {
+    return reinterpret_cast<Fortran::evaluate::Component *>(~0 - 1);
+  }
+  static unsigned getHashValue(const Fortran::evaluate::Component *v) {
+    return Fortran::lower::getHashValue(v);
+  }
+  static bool isEqual(const Fortran::evaluate::Component *lhs,
+                      const Fortran::evaluate::Component *rhs) {
+    return Fortran::lower::isEqual(lhs, rhs);
+  }
+};
 } // namespace llvm
 
 #endif // FORTRAN_LOWER_SUPPORT_UTILS_H

flang/include/flang/Lower/SymbolMap.h

@@ -14,6 +14,7 @@
 #define FORTRAN_LOWER_SYMBOLMAP_H
 
 #include "flang/Common/reference.h"
+#include "flang/Lower/Support/Utils.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/FortranVariableInterface.h"
@@ -134,6 +135,41 @@ struct SymbolBox : public fir::details::matcher<SymbolBox> {
   VT box;
 };
 
+/// Helper class to map `Fortran::evaluate::Component` references to IR values.
+/// This is used when the evaluation of a component reference must be
+/// overridden with a pre-computed address.
+class ComponentMap {
+public:
+  void insert(const Fortran::evaluate::Component &component,
+              fir::FortranVariableOpInterface definingOp) {
+    auto iter = componentMap.find(&component);
+    if (iter != componentMap.end()) {
+      iter->second = definingOp;
+      return;
+    }
+    componentStorage.push_back(
+        std::make_unique<Fortran::evaluate::Component>(component));
+    componentMap.insert({componentStorage.back().get(), definingOp});
+  }
+
+  std::optional<fir::FortranVariableOpInterface>
+  lookup(const Fortran::evaluate::Component *component) const {
+    auto iter = componentMap.find(component);
+    if (iter != componentMap.end())
+      return iter->second;
+    return std::nullopt;
+  }
+
+  LLVM_DUMP_METHOD void dump() const;
+
+private:
+  llvm::DenseMap<const Fortran::evaluate::Component *,
+                 fir::FortranVariableOpInterface>
+      componentMap;
+  llvm::SmallVector<std::unique_ptr<Fortran::evaluate::Component>>
+      componentStorage;
+};
+
 //===----------------------------------------------------------------------===//
 // Map of symbol information
 //===----------------------------------------------------------------------===//
@@ -156,12 +192,15 @@ class SymMap {
   void pushScope() {
     symbolMapStack.emplace_back();
     storageMapStack.emplace_back();
+    componentMapStack.emplace_back();
   }
   void popScope() {
     symbolMapStack.pop_back();
     assert(symbolMapStack.size() >= 1);
     storageMapStack.pop_back();
     assert(storageMapStack.size() >= 1);
+    componentMapStack.pop_back();
+    assert(componentMapStack.size() >= 1);
   }
 
   /// Add an extended value to the symbol table.
@@ -301,6 +340,8 @@ class SymMap {
     impliedDoStack.clear();
     storageMapStack.clear();
     storageMapStack.emplace_back();
+    componentMapStack.clear();
+    componentMapStack.emplace_back();
   }
 
   friend llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
@@ -329,6 +370,33 @@ class SymMap {
     return std::nullopt;
   }
 
+  /// Register a mapping from a front-end component reference to the FIR
+  /// variable that should be used to implement it. This is used to override
+  /// the default lowering of component references in specific contexts.
+  void addComponentOverride(const Fortran::evaluate::Component &component,
+                            fir::FortranVariableOpInterface definingOp) {
+    assert(!componentMapStack.empty() && "component map stack is empty");
+    if (!componentMapStack.back())
+      componentMapStack.back() = std::make_unique<ComponentMap>();
+    componentMapStack.back().value()->insert(component, definingOp);
+  }
+
+  /// Lookup an overridden FIR variable definition for a given component
+  /// reference, if any.
+  std::optional<fir::FortranVariableOpInterface>
+  lookupComponentOverride(const Fortran::evaluate::Component &component) const {
+    for (auto jmap = componentMapStack.rbegin(),
+              jend = componentMapStack.rend();
+         jmap != jend; ++jmap) {
+      if (*jmap) {
+        auto iter = (**jmap)->lookup(&component);
+        if (iter != std::nullopt)
+          return iter;
+      }
+    }
+    return std::nullopt;
+  }
+
   /// Register the symbol's storage at the innermost level
   /// of the symbol table. If the storage is already registered,
   /// it will be replaced.
@@ -360,6 +428,12 @@ class SymMap {
   // A stack of maps between the symbols and their storage descriptors.
   llvm::SmallVector<llvm::DenseMap<const semantics::Symbol *, StorageDesc>>
       storageMapStack;
+
+  // A stack of maps from front-end component references to the FIR variables
+  // that should be used to implement them. This allows overriding component
+  // references in specific lowering contexts.
+  llvm::SmallVector<std::optional<std::unique_ptr<ComponentMap>>>
+      componentMapStack;
 };
 
 /// RAII wrapper for SymMap.

flang/lib/Lower/ConvertExprToHLFIR.cpp

@@ -368,6 +368,8 @@ class HlfirDesignatorBuilder {
 
   fir::FortranVariableOpInterface
   gen(const Fortran::evaluate::Component &component) {
+    if (auto remapped = symMap.lookupComponentOverride(component))
+      return *remapped;
     if (Fortran::semantics::IsAllocatableOrPointer(component.GetLastSymbol()))
       return genWholeAllocatableOrPointerComponent(component);
     PartInfo partInfo;
@@ -477,6 +479,8 @@ class HlfirDesignatorBuilder {
 
   fir::FortranVariableOpInterface genWholeAllocatableOrPointerComponent(
       const Fortran::evaluate::Component &component) {
+    if (auto remapped = symMap.lookupComponentOverride(component))
+      return *remapped;
     // Generate whole allocatable or pointer component reference. The
     // hlfir.designate result will be a pointer/allocatable.
     PartInfo partInfo;
@@ -539,7 +543,8 @@ class HlfirDesignatorBuilder {
       // components need special care to deal with the array%array_comp(indices)
       // case.
       if (Fortran::semantics::IsAllocatableOrObjectPointer(
-              &component->GetLastSymbol()))
+              &component->GetLastSymbol()) ||
+          symMap.lookupComponentOverride(*component))
         baseType = visit(*component, partInfo);
       else
         baseType = hlfir::getFortranElementOrSequenceType(
@@ -686,6 +691,15 @@ class HlfirDesignatorBuilder {
                                  partInfo.typeParams);
       return partInfo.base->getElementOrSequenceType();
     }
+    if (auto remapped = symMap.lookupComponentOverride(component)) {
+      // Do not generate field for the designate if the component
+      // is overridden, the override value is already addressing
+      // the component.
+      partInfo.base = *remapped;
+      hlfir::genLengthParameters(loc, getBuilder(), *partInfo.base,
+                                 partInfo.typeParams);
+      return partInfo.base->getElementOrSequenceType();
+    }
     // This function must be called from contexts where the component is not the
     // base of an ArrayRef. In these cases, the component cannot be an array
     // if the base is an array. The code below determines the shape of the