[flang][hlfir] Generate temporary storage in Forall/Where [1/2]

Generate temporary storage inline inside WHERE and FORALL when possible.
A following patch will use the runtime to cover the generic cases.

Reviewed By: vzakhari

Differential Revision: https://reviews.llvm.org/D151247

Author: jeanPerier

flang/include/flang/Optimizer/Builder/TemporaryStorage.h

@@ -93,5 +93,52 @@ class HomogeneousScalarStack {
   /// Temporary storage.
   mlir::Value temp;
 };
+
+/// Structure to hold the value of a single entity.
+class SimpleCopy {
+public:
+  SimpleCopy(mlir::Location loc, fir::FirOpBuilder &builder,
+             hlfir::Entity source, llvm::StringRef tempName);
+
+  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
+                 mlir::Value value) {
+    assert(false && "must not be called: value already set");
+  }
+  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder){};
+  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder) {
+    return copy.getBase();
+  }
+  void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
+
+public:
+  /// Temporary storage for the copy.
+  hlfir::AssociateOp copy;
+};
+
+/// Generic wrapper over the different sorts of temporary storages.
+class TemporaryStorage {
+public:
+  template <typename T>
+  TemporaryStorage(T &&impl) : impl{std::forward<T>(impl)} {}
+
+  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
+                 mlir::Value value) {
+    std::visit([&](auto &temp) { temp.pushValue(loc, builder, value); }, impl);
+  }
+  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder) {
+    std::visit([&](auto &temp) { temp.resetFetchPosition(loc, builder); },
+               impl);
+  }
+  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder) {
+    return std::visit([&](auto &temp) { return temp.fetch(loc, builder); },
+                      impl);
+  }
+  void destroy(mlir::Location loc, fir::FirOpBuilder &builder) {
+    std::visit([&](auto &temp) { temp.destroy(loc, builder); }, impl);
+  }
+
+private:
+  std::variant<HomogeneousScalarStack, SimpleCopy> impl;
+};
 } // namespace fir::factory
 #endif // FORTRAN_OPTIMIZER_BUILDER_TEMPORARYSTORAGE_H

flang/lib/Optimizer/Builder/TemporaryStorage.cpp

@@ -10,8 +10,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "flang/Optimizer/Builder/TemporaryStorage.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/HLFIRTools.h"
-#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"
 #include "flang/Optimizer/Builder/Todo.h"
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 
@@ -133,3 +133,24 @@ hlfir::Entity fir::factory::HomogeneousScalarStack::moveStackAsArrayExpr(
   auto hlfirExpr = builder.create<hlfir::AsExprOp>(loc, temp, mustFree);
   return hlfir::Entity{hlfirExpr};
 }
+
+//===----------------------------------------------------------------------===//
+// fir::factory::SimpleCopy implementation.
+//===----------------------------------------------------------------------===//
+
+fir::factory::SimpleCopy::SimpleCopy(mlir::Location loc,
+                                     fir::FirOpBuilder &builder,
+                                     hlfir::Entity source,
+                                     llvm::StringRef tempName) {
+  // Use hlfir.as_expr and hlfir.associate to create a copy and leave
+  // bufferization deals with how best to make the copy.
+  if (source.isVariable())
+    source = hlfir::Entity{builder.create<hlfir::AsExprOp>(loc, source)};
+  copy = hlfir::genAssociateExpr(loc, builder, source,
+                                 source.getFortranElementType(), tempName);
+}
+
+void fir::factory::SimpleCopy::destroy(mlir::Location loc,
+                                       fir::FirOpBuilder &builder) {
+  builder.create<hlfir::EndAssociateOp>(loc, copy);
+}