[flang] Do concurrent locality specifiers

[PeteSteinfeld]

All builds and tests correctly and looks good.

[vdonaldson]

https://reviews.llvm.org/D157360

flang/include/flang/Lower/SymbolMap.h

public:
   std::optional<fir::FortranVariableOpInterface>
   lookupVariableDefinition(semantics::SymbolRef sym) {
     if (auto symBox = lookupSymbol(sym))
       return symBox.getIfFortranVariableOpInterface();
     return std::nullopt;
   }
 
 private:
-  /// Add `symbol` to the current map and bind a `box`.
+  /// Bind `box` to `symRef` in the symbol map.
   void makeSym(semantics::SymbolRef symRef, const SymbolBox &box,
                bool force = false) {
-    const auto *sym = &symRef.get().GetUltimate();
+    auto *sym = symRef->HasLocalLocality() ? &*symRef : &symRef->GetUltimate();
     if (force)
       symbolMapStack.back().erase(sym);
     assert(box && "cannot add an undefined symbol box");
     symbolMapStack.back().try_emplace(sym, box);
   }
 
   llvm::SmallVector<llvm::DenseMap<const semantics::Symbol *, SymbolBox>>
       symbolMapStack;

flang/include/flang/Semantics/symbol.h

return common::visit(
             },
             [](const GenericDetails &x) {
               return x.specific() && x.specific()->HasExplicitInterface();
             },
             [](const auto &) { return false; },
         },
         details_);
   }
+  bool HasLocalLocality() const {
+    return test(Flag::LocalityLocal) || test(Flag::LocalityLocalInit);
+  }
 
   bool operator==(const Symbol &that) const { return this == &that; }
   bool operator!=(const Symbol &that) const { return !(*this == that); }
 
   int Rank() const { return RankImpl(); }
 
   int Corank() const {
     return common::visit(

flang/lib/Lower/Bridge.cpp

struct IncrementLoopInfo {
 
   bool isStructured() const { return !headerBlock; }
 
   mlir::Type getLoopVariableType() const {
     assert(loopVariable && "must be set");
     return fir::unwrapRefType(loopVariable.getType());
   }
 
+  bool hasLocalitySpecs() const {
+    return !localSymList.empty() || !localInitSymList.empty() ||
+           !sharedSymList.empty();
+  }
+
   // Data members common to both structured and unstructured loops.
   const Fortran::semantics::Symbol &loopVariableSym;
   const Fortran::lower::SomeExpr *lowerExpr;
   const Fortran::lower::SomeExpr *upperExpr;
   const Fortran::lower::SomeExpr *stepExpr;
   const Fortran::lower::SomeExpr *maskExpr = nullptr;
   bool isUnordered; // do concurrent, forall
+  llvm::SmallVector<const Fortran::semantics::Symbol *> localSymList;
   llvm::SmallVector<const Fortran::semantics::Symbol *> localInitSymList;
   llvm::SmallVector<const Fortran::semantics::Symbol *> sharedSymList;
   mlir::Value loopVariable = nullptr;
   mlir::Value stepValue = nullptr; // possible uses in multiple blocks
 
   // Data members for structured loops.
   fir::DoLoopOp doLoop = nullptr;
 
for (const Fortran::parser::ConcurrentControl &control :
          std::get<std::list<Fortran::parser::ConcurrentControl>>(header.t))
       incrementLoopNestInfo.emplace_back(
           *std::get<0>(control.t).symbol, std::get<1>(control.t),
           std::get<2>(control.t), std::get<3>(control.t), /*isUnordered=*/true);
     IncrementLoopInfo &info = incrementLoopNestInfo.back();
     info.maskExpr = Fortran::semantics::GetExpr(
         std::get<std::optional<Fortran::parser::ScalarLogicalExpr>>(header.t));
     for (const Fortran::parser::LocalitySpec &x : localityList) {
+      if (const auto *localList =
+              std::get_if<Fortran::parser::LocalitySpec::Local>(&x.u))
+        for (const Fortran::parser::Name &x : localList->v)
+          info.localSymList.push_back(x.symbol);
       if (const auto *localInitList =
               std::get_if<Fortran::parser::LocalitySpec::LocalInit>(&x.u))
         for (const Fortran::parser::Name &x : localInitList->v)
           info.localInitSymList.push_back(x.symbol);
       if (const auto *sharedList =
               std::get_if<Fortran::parser::LocalitySpec::Shared>(&x.u))
         for (const Fortran::parser::Name &x : sharedList->v)
           info.sharedSymList.push_back(x.symbol);
-      if (std::get_if<Fortran::parser::LocalitySpec::Local>(&x.u))
-        TODO(toLocation(), "do concurrent locality specs not implemented");
     }
     return incrementLoopNestInfo;
   }
 
+  /// Create DO CONCURRENT construct symbol bindings and generate LOCAL_INIT
+  /// assignments.
+  void handleLocalitySpecs(const IncrementLoopInfo &info) {
+    Fortran::semantics::SemanticsContext &semanticsContext =
+        bridge.getSemanticsContext();
+    for (const Fortran::semantics::Symbol *sym : info.localSymList)
+      createHostAssociateVarClone(*sym);
+    for (const Fortran::semantics::Symbol *sym : info.localInitSymList) {
+      createHostAssociateVarClone(*sym);
+      const auto *hostDetails =
+          sym->detailsIf<Fortran::semantics::HostAssocDetails>();
+      assert(hostDetails && "missing locality spec host symbol");
+      const Fortran::semantics::Symbol *hostSym = &hostDetails->symbol();
+      Fortran::evaluate::ExpressionAnalyzer ea{semanticsContext};
+      Fortran::evaluate::Assignment assign{
+          ea.Designate(Fortran::evaluate::DataRef{*sym}).value(),
+          ea.Designate(Fortran::evaluate::DataRef{*hostSym}).value()};
+      if (Fortran::semantics::IsPointer(*sym))
+        assign.u = Fortran::evaluate::Assignment::BoundsSpec{};
+      genAssignment(assign);
+    }
+    for (const Fortran::semantics::Symbol *sym : info.sharedSymList) {
+      const auto *hostDetails =
+          sym->detailsIf<Fortran::semantics::HostAssocDetails>();
+      copySymbolBinding(hostDetails->symbol(), *sym);
+    }
+  }
+
   /// Generate FIR for a DO construct. There are six variants:
   ///  - unstructured infinite and while loops
   ///  - structured and unstructured increment loops
   ///  - structured and unstructured concurrent loops
   void genFIR(const Fortran::parser::DoConstruct &doConstruct) {
     setCurrentPositionAt(doConstruct);
     // Collect loop nest information.
     // Generate begin loop code directly for infinite and while loops.
auto genControlValue = [&](const Fortran::lower::SomeExpr *expr,
       if (expr)
         return builder->createConvert(loc, controlType,
                                       createFIRExpr(loc, expr, stmtCtx));
 
       if (info.hasRealControl)
         return builder->createRealConstant(loc, controlType, 1u);
       return builder->createIntegerConstant(loc, controlType, 1); // step
     };
-    auto handleLocalitySpec = [&](IncrementLoopInfo &info) {
-      // Generate Local Init Assignments
-      for (const Fortran::semantics::Symbol *sym : info.localInitSymList) {
-        const auto *hostDetails =
-            sym->detailsIf<Fortran::semantics::HostAssocDetails>();
-        assert(hostDetails && "missing local_init variable host variable");
-        const Fortran::semantics::Symbol &hostSym = hostDetails->symbol();
-        (void)hostSym;
-        TODO(loc, "do concurrent locality specs not implemented");
-      }
-      // Handle shared locality spec
-      for (const Fortran::semantics::Symbol *sym : info.sharedSymList) {
-        const auto *hostDetails =
-            sym->detailsIf<Fortran::semantics::HostAssocDetails>();
-        assert(hostDetails && "missing shared variable host variable");
-        const Fortran::semantics::Symbol &hostSym = hostDetails->symbol();
-        copySymbolBinding(hostSym, *sym);
-      }
-    };
     for (IncrementLoopInfo &info : incrementLoopNestInfo) {
       info.loopVariable =
           genLoopVariableAddress(loc, info.loopVariableSym, info.isUnordered);
       mlir::Value lowerValue = genControlValue(info.lowerExpr, info);
       mlir::Value upperValue = genControlValue(info.upperExpr, info);
       info.stepValue = genControlValue(info.stepExpr, info);
 
       // Structured loop - generate fir.do_loop.
for (IncrementLoopInfo &info : incrementLoopNestInfo) {
           mlir::Value maskCond = createFIRExpr(loc, info.maskExpr, stmtCtx);
           stmtCtx.finalizeAndReset();
           mlir::Value maskCondCast =
               builder->createConvert(loc, builder->getI1Type(), maskCond);
           auto ifOp = builder->create<fir::IfOp>(loc, maskCondCast,
                                                  /*withElseRegion=*/false);
           builder->setInsertionPointToStart(&ifOp.getThenRegion().front());
         }
-        handleLocalitySpec(info);
+        if (info.hasLocalitySpecs())
+          handleLocalitySpecs(info);
         continue;
       }
 
       // Unstructured loop preheader - initialize tripVariable and loopVariable.
       mlir::Value tripCount;
       if (info.hasRealControl) {
         auto diff1 =
             builder->create<mlir::arith::SubFOp>(loc, upperValue, lowerValue);
for (IncrementLoopInfo &info : incrementLoopNestInfo) {
         mlir::Value maskCond = createFIRExpr(loc, info.maskExpr, stmtCtx);
         stmtCtx.finalizeAndReset();
         genConditionalBranch(maskCond, info.bodyBlock, latchBlock);
       } else {
         genConditionalBranch(cond, info.bodyBlock, info.exitBlock);
         if (&info != &incrementLoopNestInfo.back()) // not innermost
           startBlock(info.bodyBlock); // preheader block of enclosed dimension
       }
-      if (!info.localInitSymList.empty()) {
+      if (info.hasLocalitySpecs()) {
         mlir::OpBuilder::InsertPoint insertPt = builder->saveInsertionPoint();
         builder->setInsertionPointToStart(info.bodyBlock);
-        handleLocalitySpec(info);
+        handleLocalitySpecs(info);
         builder->restoreInsertionPoint(insertPt);
       }
     }
   }
 
   /// Generate FIR to end a structured or unstructured increment loop nest.
   void genFIRIncrementLoopEnd(IncrementLoopNestInfo &incrementLoopNestInfo) {
     assert(!incrementLoopNestInfo.empty() && "empty loop nest");

flang/lib/Lower/ConvertExpr.cpp

 }
 
 // Helper to get the ultimate first symbol. This works around the fact that
 // symbol resolution in the front end doesn't always resolve a symbol to its
 // ultimate symbol but may leave placeholder indirections for use and host
 // associations.
 template <typename A>
 const Fortran::semantics::Symbol &getFirstSym(const A &obj) {
-  return obj.GetFirstSymbol().GetUltimate();
+  const Fortran::semantics::Symbol &sym = obj.GetFirstSymbol();
+  return sym.HasLocalLocality() ? sym : sym.GetUltimate();
 }
 
 // Helper to get the ultimate last symbol.
 template <typename A>
 const Fortran::semantics::Symbol &getLastSym(const A &obj) {
-  return obj.GetLastSymbol().GetUltimate();
+  const Fortran::semantics::Symbol &sym = obj.GetLastSymbol();
+  return sym.HasLocalLocality() ? sym : sym.GetUltimate();
 }
 
 // Return true if TRANSPOSE should be lowered without a runtime call.
 static bool
 isTransposeOptEnabled(const Fortran::lower::AbstractConverter &converter) {
   return optimizeTranspose &&
          converter.getLoweringOptions().getOptimizeTranspose();
 }

flang/lib/Lower/SymbolMap.cpp

exv.match([&](const fir::UnboxedValue &v) { addSymbol(sym, v, force); },
             [&](const fir::PolymorphicValue &v) { makeSym(sym, v, force); },
             [](auto) {
               llvm::report_fatal_error("value not added to symbol table");
             });
 }
 
 Fortran::lower::SymbolBox
 Fortran::lower::SymMap::lookupSymbol(Fortran::semantics::SymbolRef symRef) {
-  Fortran::semantics::SymbolRef sym = symRef.get().GetUltimate();
+  auto *sym = symRef->HasLocalLocality() ? &*symRef : &symRef->GetUltimate();
   for (auto jmap = symbolMapStack.rbegin(), jend = symbolMapStack.rend();
        jmap != jend; ++jmap) {
-    auto iter = jmap->find(&*sym);
+    auto iter = jmap->find(sym);
     if (iter != jmap->end())
       return iter->second;
   }
   return SymbolBox::None{};
 }
 
 Fortran::lower::SymbolBox Fortran::lower::SymMap::shallowLookupSymbol(
     Fortran::semantics::SymbolRef symRef) {
+  auto *sym = symRef->HasLocalLocality() ? &*symRef : &symRef->GetUltimate();
   auto &map = symbolMapStack.back();
-  auto iter = map.find(&symRef.get().GetUltimate());
+  auto iter = map.find(sym);
   if (iter != map.end())
     return iter->second;
   return SymbolBox::None{};
 }
 
 /// Skip one level when looking up the symbol. The use case is such as looking
 /// up the host variable symbol box by skipping the associated level in
 /// host-association in OpenMP code.
 Fortran::lower::SymbolBox Fortran::lower::SymMap::lookupOneLevelUpSymbol(
     Fortran::semantics::SymbolRef symRef) {
-  Fortran::semantics::SymbolRef sym = symRef.get().GetUltimate();
+  auto *sym = symRef->HasLocalLocality() ? &*symRef : &symRef->GetUltimate();
   auto jmap = symbolMapStack.rbegin();
   auto jend = symbolMapStack.rend();
   if (jmap == jend)
     return SymbolBox::None{};
   // Skip one level in symbol map stack.
   for (++jmap; jmap != jend; ++jmap) {
-    auto iter = jmap->find(&*sym);
+    auto iter = jmap->find(sym);
     if (iter != jmap->end())
       return iter->second;
   }
   return SymbolBox::None{};
 }
 
 mlir::Value
 Fortran::lower::SymMap::lookupImpliedDo(Fortran::lower::SymMap::AcDoVar var) {

flang/test/Lower/loops.f90

subroutine loop_test
   ! CHECK-NOT: fir.do_loop
   do x = 1.5, 3.5, 0.3
     xsum = xsum + 1
   enddo
   ! CHECK: fir.call @_FortranAioBeginExternalFormattedOutput
   print '(" F:",X,F3.1,A,I2)', x, ' -', xsum
 end subroutine loop_test
 
+! CHECK-LABEL: c.func @_QPlis
+subroutine lis(n)
+  ! CHECK-DAG: fir.alloca i32 {bindc_name = "m"}
+  ! CHECK-DAG: fir.alloca i32 {bindc_name = "j"}
+  ! CHECK-DAG: fir.alloca i32 {bindc_name = "i"}
+  ! CHECK-DAG: fir.alloca i8 {bindc_name = "i"}
+  ! CHECK-DAG: fir.alloca i32 {bindc_name = "j", uniq_name = "_QFlisEj"}
+  ! CHECK-DAG: fir.alloca i32 {bindc_name = "k", uniq_name = "_QFlisEk"}
+  ! CHECK-DAG: fir.alloca !fir.box<!fir.ptr<!fir.array<?x?x?xi32>>> {bindc_name = "p", uniq_name = "_QFlisEp"}
+  ! CHECK-DAG: fir.alloca !fir.array<?x?x?xi32>, %{{.*}}, %{{.*}}, %{{.*}} {bindc_name = "a", fir.target, uniq_name = "_QFlisEa"}
+  ! CHECK-DAG: fir.alloca !fir.array<?x?xi32>, %{{.*}}, %{{.*}} {bindc_name = "r", uniq_name = "_QFlisEr"}
+  ! CHECK-DAG: fir.alloca !fir.array<?x?xi32>, %{{.*}}, %{{.*}} {bindc_name = "s", uniq_name = "_QFlisEs"}
+  ! CHECK-DAG: fir.alloca !fir.array<?x?xi32>, %{{.*}}, %{{.*}} {bindc_name = "t", uniq_name = "_QFlisEt"}
+  integer, target    :: a(n,n,n) ! operand via p
+  integer            :: r(n,n)   ! result, unspecified locality
+  integer            :: s(n,n)   ! shared locality
+  integer            :: t(n,n)   ! local locality
+  integer, pointer   :: p(:,:,:) ! local_init locality
+
+  p => a
+  ! CHECK:     fir.do_loop %arg1 = %c0{{.*}} to %{{.*}} step %c1{{.*}} unordered iter_args(%arg2 = %{{.*}}) -> (!fir.array<?x?xi32>) {
+  ! CHECK:       fir.do_loop %arg3 = %c0{{.*}} to %{{.*}} step %c1{{.*}} unordered iter_args(%arg4 = %arg2) -> (!fir.array<?x?xi32>) {
+  ! CHECK:       }
+  ! CHECK:     }
+  r = 0
+
+  ! CHECK:     fir.do_loop %arg1 = %{{.*}} to %{{.*}} step %{{.*}} unordered {
+  ! CHECK:       fir.do_loop %arg2 = %{{.*}} to %{{.*}} step %c1{{.*}} iter_args(%arg3 = %{{.*}}) -> (index, i32) {
+  ! CHECK:       }
+  ! CHECK:     }
+  do concurrent (integer(kind=1)::i=n:1:-1)
+    do j = 1,n
+      a(i,j,:) = 2*(i+j)
+      s(i,j)   = -i-j
+    enddo
+  enddo
+
+  ! CHECK:     fir.do_loop %arg1 = %{{.*}} to %{{.*}} step %c1{{.*}} unordered {
+  ! CHECK:       fir.do_loop %arg2 = %{{.*}} to %{{.*}} step %c1{{.*}} unordered {
+  ! CHECK:         fir.if %{{.*}} {
+  ! CHECK:           %[[V_95:[0-9]+]] = fir.alloca !fir.array<?x?xi32>, %{{.*}}, %{{.*}} {bindc_name = "t", pinned, uniq_name = "_QFlisEt"}
+  ! CHECK:           %[[V_96:[0-9]+]] = fir.alloca !fir.box<!fir.ptr<!fir.array<?x?x?xi32>>> {bindc_name = "p", pinned, uniq_name = "_QFlisEp"}
+  ! CHECK:           fir.store %{{.*}} to %[[V_96]] : !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?x?xi32>>>>
+  ! CHECK:           fir.do_loop %arg3 = %{{.*}} to %{{.*}} step %c1{{.*}} iter_args(%arg4 = %{{.*}}) -> (index, i32) {
+  ! CHECK:             fir.do_loop %arg5 = %{{.*}} to %{{.*}} step %c1{{.*}} unordered {
+  ! CHECK:               fir.load %[[V_96]] : !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?x?xi32>>>>
+  ! CHECK:               fir.convert %[[V_95]] : (!fir.ref<!fir.array<?x?xi32>>) -> !fir.ref<!fir.array<?xi32>>
+  ! CHECK:             }
+  ! CHECK:           }
+  ! CHECK:           fir.convert %[[V_95]] : (!fir.ref<!fir.array<?x?xi32>>) -> !fir.ref<!fir.array<?xi32>>
+  ! CHECK:         }
+  ! CHECK:       }
+  ! CHECK:     }
+  do concurrent (i=1:n,j=1:n,i.ne.j) local(t) local_init(p) shared(s)
+    do k=1,n
+      do concurrent (m=1:n)
+        t(k,m) = p(k,m,k)
+      enddo
+    enddo
+    r(i,j) = t(i,j) + s(i,j)
+  enddo
+
+  print*, sum(r) ! n=6 -> 210
+end
+
 ! CHECK-LABEL: print_nothing
 subroutine print_nothing(k1, k2)
   if (k1 > 0) then
     ! CHECK: br [[header:\^bb[0-9]+]]
     ! CHECK: [[header]]
     do while (k1 > k2)
       print*, k1, k2 ! no output
       k2 = k2 + 1
       ! CHECK: br [[header]]
     end do
   end if
 end
 
   call loop_test
+  call lis(6)
   call print_nothing(2, 2)
 end