[Flang][OpenMP] Reapply revert of llvm#155742 with fix for Nekbone

Small addition to the FunctionFiltering pass to handle the new storage argument
for the HLFIR declare operation.

Author: agozillon

flang/include/flang/Lower/AbstractConverter.h

@@ -108,6 +108,21 @@ class AbstractConverter {
   /// added or replaced at the inner-most level of the local symbol map.
   virtual void bindSymbol(SymbolRef sym, const fir::ExtendedValue &exval) = 0;
 
+  /// Binds the symbol's physical storage to a storage descriptor,
+  /// which is the base address of the storage and the offset
+  /// within the storage, where the symbol begins.
+  /// The symbol binding will be added or replaced at the innermost level
+  /// of the local symbol map.
+  virtual void
+  bindSymbolStorage(SymbolRef sym,
+                    Fortran::lower::SymMap::StorageDesc storage) = 0;
+
+  /// Returns the storage descriptor previously bound to this symbol.
+  /// If there is no bound storage, the descriptor will contain
+  /// nullptr base address.
+  virtual Fortran::lower::SymMap::StorageDesc
+  getSymbolStorage(SymbolRef sym) = 0;
+
   /// Override lowering of expression with pre-lowered values.
   /// Associate mlir::Value to evaluate::Expr. All subsequent call to
   /// genExprXXX() will replace any occurrence of an overridden

flang/include/flang/Lower/ConvertVariable.h

@@ -92,10 +92,14 @@ void defineCommonBlocks(
 /// The COMMON block is a global structure. \p commonValue is the base address
 /// of the COMMON block. As the offset from the symbol \p sym, generate the
 /// COMMON block member value (commonValue + offset) for the symbol.
+/// \p commonSize specifies the syze of the COMMON block in bytes.
+/// The size is used to represent a COMMON block reference as
+/// a !fir.ref<!fir.array<SIZExi8>>.
 mlir::Value genCommonBlockMember(AbstractConverter &converter,
                                  mlir::Location loc,
                                  const Fortran::semantics::Symbol &sym,
-                                 mlir::Value commonValue);
+                                 mlir::Value commonValue,
+                                 std::size_t commonSize);
 
 /// Lower a symbol attributes given an optional storage \p and add it to the
 /// provided symbol map. If \preAlloc is not provided, a temporary storage will

flang/include/flang/Lower/SymbolMap.h

@@ -146,14 +146,22 @@ struct SymbolBox : public fir::details::matcher<SymbolBox> {
 class SymMap {
 public:
   using AcDoVar = llvm::StringRef;
+  /// Descriptor of a symbol's storage consists of the base address
+  /// of the storage and the offset within that storage.
+  using StorageDesc = std::pair<mlir::Value, std::uint64_t>;
 
   SymMap() { pushScope(); }
   SymMap(const SymMap &) = delete;
 
-  void pushScope() { symbolMapStack.emplace_back(); }
+  void pushScope() {
+    symbolMapStack.emplace_back();
+    storageMapStack.emplace_back();
+  }
   void popScope() {
     symbolMapStack.pop_back();
     assert(symbolMapStack.size() >= 1);
+    storageMapStack.pop_back();
+    assert(storageMapStack.size() >= 1);
   }
 
   /// Add an extended value to the symbol table.
@@ -287,6 +295,8 @@ class SymMap {
     symbolMapStack.emplace_back();
     assert(symbolMapStack.size() == 1);
     impliedDoStack.clear();
+    storageMapStack.clear();
+    storageMapStack.emplace_back();
   }
 
   friend llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
@@ -315,6 +325,16 @@ class SymMap {
     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.
+  void registerStorage(semantics::SymbolRef sym, StorageDesc storage);
+  /// Lookup the symbol's storage at the innermost level of the symbol table.
+  StorageDesc lookupStorage(semantics::SymbolRef sym);
+  StorageDesc lookupStorage(const semantics::Symbol *sym) {
+    return lookupStorage(*sym);
+  }
+
 private:
   /// Bind `box` to `symRef` in the symbol map.
   void makeSym(semantics::SymbolRef symRef, const SymbolBox &box,
@@ -332,6 +352,10 @@ class SymMap {
   // Implied DO induction variables are not represented as Se::Symbol in
   // Ev::Expr. Keep the variable markers in their own stack.
   llvm::SmallVector<std::pair<AcDoVar, mlir::Value>> impliedDoStack;
+
+  // A stack of maps between the symbols and their storage descriptors.
+  llvm::SmallVector<llvm::DenseMap<const semantics::Symbol *, StorageDesc>>
+      storageMapStack;
 };
 
 /// RAII wrapper for SymMap.

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

@@ -224,7 +224,8 @@ fir::FortranVariableOpInterface
 genDeclare(mlir::Location loc, fir::FirOpBuilder &builder,
            const fir::ExtendedValue &exv, llvm::StringRef name,
            fir::FortranVariableFlagsAttr flags,
-           mlir::Value dummyScope = nullptr,
+           mlir::Value dummyScope = nullptr, mlir::Value storage = nullptr,
+           std::uint64_t storageOffset = 0,
            cuf::DataAttributeAttr dataAttr = {});
 
 /// Generate an hlfir.associate to build a variable from an expression value.

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

@@ -109,10 +109,12 @@ def hlfir_DeclareOp
     attr-dict `:` functional-type(operands, results)
   }];
 
-  let builders = [
-    OpBuilder<(ins "mlir::Value":$memref, "llvm::StringRef":$uniq_name,
-      CArg<"mlir::Value", "{}">:$shape, CArg<"mlir::ValueRange", "{}">:$typeparams,
+  let builders = [OpBuilder<(ins "mlir::Value":$memref,
+      "llvm::StringRef":$uniq_name, CArg<"mlir::Value", "{}">:$shape,
+      CArg<"mlir::ValueRange", "{}">:$typeparams,
       CArg<"mlir::Value", "{}">:$dummy_scope,
+      CArg<"mlir::Value", "{}">:$storage,
+      CArg<"std::uint64_t", "0">:$storage_offset,
       CArg<"fir::FortranVariableFlagsAttr", "{}">:$fortran_attrs,
       CArg<"cuf::DataAttributeAttr", "{}">:$data_attr)>];
 

flang/lib/Lower/Bridge.cpp

@@ -632,6 +632,17 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     addSymbol(sym, exval, /*forced=*/true);
   }
 
+  void bindSymbolStorage(
+      Fortran::lower::SymbolRef sym,
+      Fortran::lower::SymMap::StorageDesc storage) override final {
+    localSymbols.registerStorage(sym, std::move(storage));
+  }
+
+  Fortran::lower::SymMap::StorageDesc
+  getSymbolStorage(Fortran::lower::SymbolRef sym) override final {
+    return localSymbols.lookupStorage(sym);
+  }
+
   void
   overrideExprValues(const Fortran::lower::ExprToValueMap *map) override final {
     exprValueOverrides = map;

flang/lib/Lower/ConvertArrayConstructor.cpp

@@ -315,9 +315,8 @@ class RuntimeTempStrategy : public StrategyBase {
       mlir::Value tempStorage = builder.createHeapTemporary(
           loc, declaredType, tempName, extents, lengths);
       mlir::Value shape = builder.genShape(loc, extents);
-      declare = hlfir::DeclareOp::create(
-          builder, loc, tempStorage, tempName, shape, lengths,
-          /*dummy_scope=*/nullptr, fir::FortranVariableFlagsAttr{});
+      declare = hlfir::DeclareOp::create(builder, loc, tempStorage, tempName,
+                                         shape, lengths);
       initialBoxValue =
           builder.createBox(loc, boxType, declare->getOriginalBase(), shape,
                             /*slice=*/mlir::Value{}, lengths, /*tdesc=*/{});

flang/lib/Lower/ConvertExprToHLFIR.cpp

@@ -1813,10 +1813,8 @@ class HlfirBuilder {
     // Allocate scalar temporary that will be initialized
     // with the values specified by the constructor.
     mlir::Value storagePtr = builder.createTemporary(loc, recTy);
-    auto varOp = hlfir::EntityWithAttributes{hlfir::DeclareOp::create(
-        builder, loc, storagePtr, "ctor.temp", /*shape=*/nullptr,
-        /*typeparams=*/mlir::ValueRange{}, /*dummy_scope=*/nullptr,
-        fir::FortranVariableFlagsAttr{})};
+    auto varOp = hlfir::EntityWithAttributes{
+        hlfir::DeclareOp::create(builder, loc, storagePtr, "ctor.temp")};
 
     // Initialize any components that need initialization.
     mlir::Value box = builder.createBox(loc, fir::ExtendedValue{varOp});

flang/lib/Lower/ConvertVariable.cpp

@@ -1355,6 +1355,7 @@ static void instantiateAlias(Fortran::lower::AbstractConverter &converter,
   mlir::Value bytePtr = fir::CoordinateOp::create(
       builder, loc, i8Ptr, storeAddr, mlir::ValueRange{offset});
   mlir::Value typedPtr = castAliasToPointer(builder, loc, symType, bytePtr);
+  converter.bindSymbolStorage(sym, {storeAddr, off});
   Fortran::lower::StatementContext stmtCtx;
   mapSymbolAttributes(converter, var, symMap, stmtCtx, typedPtr);
   // Default initialization is possible for equivalence members: see
@@ -1594,13 +1595,15 @@ void Fortran::lower::defineCommonBlocks(
 
 mlir::Value Fortran::lower::genCommonBlockMember(
     Fortran::lower::AbstractConverter &converter, mlir::Location loc,
-    const Fortran::semantics::Symbol &sym, mlir::Value commonValue) {
+    const Fortran::semantics::Symbol &sym, mlir::Value commonValue,
+    std::size_t commonSize) {
   fir::FirOpBuilder &builder = converter.getFirOpBuilder();
 
   std::size_t byteOffset = sym.GetUltimate().offset();
   mlir::IntegerType i8Ty = builder.getIntegerType(8);
   mlir::Type i8Ptr = builder.getRefType(i8Ty);
-  mlir::Type seqTy = builder.getRefType(builder.getVarLenSeqTy(i8Ty));
+  fir::SequenceType::Shape shape(1, commonSize);
+  mlir::Type seqTy = builder.getRefType(fir::SequenceType::get(shape, i8Ty));
   mlir::Value base = builder.createConvert(loc, seqTy, commonValue);
 
   mlir::Value offs =
@@ -1609,6 +1612,8 @@ mlir::Value Fortran::lower::genCommonBlockMember(
                                                   mlir::ValueRange{offs});
   mlir::Type symType = converter.genType(sym);
 
+  converter.bindSymbolStorage(sym, {base, byteOffset});
+
   return Fortran::semantics::FindEquivalenceSet(sym) != nullptr
              ? castAliasToPointer(builder, loc, symType, varAddr)
              : builder.createConvert(loc, builder.getRefType(symType), varAddr);
@@ -1637,7 +1642,8 @@ static void instantiateCommon(Fortran::lower::AbstractConverter &converter,
     symMap.addSymbol(common, commonAddr);
   }
 
-  mlir::Value local = genCommonBlockMember(converter, loc, varSym, commonAddr);
+  mlir::Value local =
+      genCommonBlockMember(converter, loc, varSym, commonAddr, common.size());
   Fortran::lower::StatementContext stmtCtx;
   mapSymbolAttributes(converter, var, symMap, stmtCtx, local);
 }
@@ -1909,7 +1915,8 @@ static void genDeclareSymbol(Fortran::lower::AbstractConverter &converter,
       // Declare a local pointer variable.
       auto newBase = hlfir::DeclareOp::create(
           builder, loc, boxAlloc, name, /*shape=*/nullptr, lenParams,
-          /*dummy_scope=*/nullptr, attributes);
+          /*dummy_scope=*/nullptr, /*storage=*/nullptr,
+          /*storage_offset=*/0, attributes);
       mlir::Value nullAddr = builder.createNullConstant(
           loc, llvm::cast<fir::BaseBoxType>(ptrBoxType).getEleTy());
 
@@ -1939,9 +1946,10 @@ static void genDeclareSymbol(Fortran::lower::AbstractConverter &converter,
     mlir::Value dummyScope;
     if (converter.isRegisteredDummySymbol(sym))
       dummyScope = converter.dummyArgsScopeValue();
-    auto newBase =
-        hlfir::DeclareOp::create(builder, loc, base, name, shapeOrShift,
-                                 lenParams, dummyScope, attributes, dataAttr);
+    auto [storage, storageOffset] = converter.getSymbolStorage(sym);
+    auto newBase = hlfir::DeclareOp::create(
+        builder, loc, base, name, shapeOrShift, lenParams, dummyScope, storage,
+        storageOffset, attributes, dataAttr);
     symMap.addVariableDefinition(sym, newBase, force);
     return;
   }
@@ -1999,8 +2007,10 @@ void Fortran::lower::genDeclareSymbol(
       base = genPackArray(converter, sym, exv);
       dummyScope = converter.dummyArgsScopeValue();
     }
-    hlfir::EntityWithAttributes declare = hlfir::genDeclare(
-        loc, builder, base, name, attributes, dummyScope, dataAttr);
+    auto [storage, storageOffset] = converter.getSymbolStorage(sym);
+    hlfir::EntityWithAttributes declare =
+        hlfir::genDeclare(loc, builder, base, name, attributes, dummyScope,
+                          storage, storageOffset, dataAttr);
     symMap.addVariableDefinition(sym, declare.getIfVariableInterface(), force);
     return;
   }

flang/lib/Lower/OpenACC.cpp

@@ -1270,12 +1270,10 @@ mlir::acc::FirstprivateRecipeOp Fortran::lower::createOrGetFirstprivateRecipe(
 
     auto leftDeclOp = hlfir::DeclareOp::create(
         builder, loc, recipe.getCopyRegion().getArgument(0), llvm::StringRef{},
-        shape, llvm::ArrayRef<mlir::Value>{}, /*dummy_scope=*/nullptr,
-        fir::FortranVariableFlagsAttr{});
+        shape);
     auto rightDeclOp = hlfir::DeclareOp::create(
         builder, loc, recipe.getCopyRegion().getArgument(1), llvm::StringRef{},
-        shape, llvm::ArrayRef<mlir::Value>{}, /*dummy_scope=*/nullptr,
-        fir::FortranVariableFlagsAttr{});
+        shape);
 
     hlfir::DesignateOp::Subscripts triplets =
         getSubscriptsFromArgs(recipe.getCopyRegion().getArguments());
@@ -1590,14 +1588,10 @@ static void genCombiner(fir::FirOpBuilder &builder, mlir::Location loc,
       auto shape =
           genShapeFromBoundsOrArgs(loc, builder, seqTy, bounds,
                                    recipe.getCombinerRegion().getArguments());
-      auto v1DeclareOp = hlfir::DeclareOp::create(
-          builder, loc, value1, llvm::StringRef{}, shape,
-          llvm::ArrayRef<mlir::Value>{},
-          /*dummy_scope=*/nullptr, fir::FortranVariableFlagsAttr{});
-      auto v2DeclareOp = hlfir::DeclareOp::create(
-          builder, loc, value2, llvm::StringRef{}, shape,
-          llvm::ArrayRef<mlir::Value>{},
-          /*dummy_scope=*/nullptr, fir::FortranVariableFlagsAttr{});
+      auto v1DeclareOp = hlfir::DeclareOp::create(builder, loc, value1,
+                                                  llvm::StringRef{}, shape);
+      auto v2DeclareOp = hlfir::DeclareOp::create(builder, loc, value2,
+                                                  llvm::StringRef{}, shape);
       hlfir::DesignateOp::Subscripts triplets = getTripletsFromArgs(recipe);
 
       llvm::SmallVector<mlir::Value> lenParamsLeft;