[flang][hlfir] Codegen of hlfir.region_assign where LHS conflicts

jeanPerier · jeanPerier · commit 0446bfcc5ca2 · 2023-06-30T09:20:52.000+02:00
When the analysis of hlfir.region_assign determined that the LHS region evaluation may be impacted by the assignment effects, all LHS must be fully evaluated and saved before any assignment is done. This patch adds TemporaryStorage variants to save address, including vector subscripted entities addresses whose shape must be saved. It uses the DescriptorStack runtime to deal with complex cases inside forall. For the sake of simplicity, this is also used for vector subscripted LHS outside of foralls (each element address is saved as a descriptor on this stack. This is a bit suboptimal, but it is a safe start that will work with all kinds of type (polymorphic, PDTs...) without further work). Another approach would be to saved only the values that are conflicting in the LHS computation, but this would require a much more complex analysis of the LHS region DAG. Differential Revision: https://reviews.llvm.org/D154057
diff --git a/flang/include/flang/Optimizer/Builder/Runtime/TemporaryStack.h b/flang/include/flang/Optimizer/Builder/Runtime/TemporaryStack.h
@@ -30,5 +30,16 @@ void genValueAt(mlir::Location loc, fir::FirOpBuilder &builder,
 void genDestroyValueStack(mlir::Location loc, fir::FirOpBuilder &builder,
                           mlir::Value opaquePtr);
 
+mlir::Value genCreateDescriptorStack(mlir::Location loc,
+                                     fir::FirOpBuilder &builder);
+
+void genPushDescriptor(mlir::Location loc, fir::FirOpBuilder &builder,
+                       mlir::Value opaquePtr, mlir::Value boxValue);
+void genDescriptorAt(mlir::Location loc, fir::FirOpBuilder &builder,
+                     mlir::Value opaquePtr, mlir::Value i,
+                     mlir::Value retValueBox);
+
+void genDestroyDescriptorStack(mlir::Location loc, fir::FirOpBuilder &builder,
+                               mlir::Value opaquePtr);
 } // namespace fir::runtime
 #endif // FORTRAN_OPTIMIZER_BUILDER_RUNTIME_TEMPORARYSTACK_H
diff --git a/flang/include/flang/Optimizer/Builder/TemporaryStorage.h b/flang/include/flang/Optimizer/Builder/TemporaryStorage.h
@@ -120,6 +120,30 @@ class SimpleCopy {
   hlfir::AssociateOp copy;
 };
 
+/// Structure to keep track of a simple mlir::Value. This is useful
+/// when a value does not need an in memory copy because it is
+/// already saved in an SSA value that will be accessible at the fetching
+/// point.
+class SSARegister {
+public:
+  SSARegister(){};
+
+  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
+                 mlir::Value value) {
+    ssaRegister = value;
+  }
+  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder){};
+  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder) {
+    return ssaRegister;
+  }
+  void destroy(mlir::Location loc, fir::FirOpBuilder &builder) {}
+  bool canBeFetchedAfterPush() const { return true; }
+
+public:
+  /// Temporary storage for the copy.
+  mlir::Value ssaRegister;
+};
+
 /// Data structure to stack any kind of values with the same static type and
 /// rank. Each value may have different type parameters, bounds, and dynamic
 /// type. Fetching value N will return a value with the same dynamic type,
@@ -150,6 +174,61 @@ class AnyValueStack {
   mlir::Value retValueBox;
 };
 
+/// Data structure to stack any kind of variables with the same static type and
+/// rank. Each variable may have different type parameters, bounds, and dynamic
+/// type. Fetching variable N will return a variable with the same address,
+/// dynamic type, bounds, and type parameters as the Nth variable that was
+/// pushed. It is implemented using runtime.
+class AnyVariableStack {
+public:
+  AnyVariableStack(mlir::Location loc, fir::FirOpBuilder &builder,
+                   mlir::Type valueStaticType);
+
+  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
+                 mlir::Value value);
+  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder);
+  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder);
+  void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
+  bool canBeFetchedAfterPush() const { return true; }
+
+private:
+  /// Keep the original variable type.
+  mlir::Type variableStaticType;
+  /// Runtime cookie created by the runtime. It is a pointer to an opaque
+  /// runtime data structure that manages the stack.
+  mlir::Value opaquePtr;
+  /// Counter to keep track of the fetching position.
+  Counter counter;
+  /// Pointer box passed to the runtime when fetching the values.
+  mlir::Value retValueBox;
+};
+
+class TemporaryStorage;
+
+/// Data structure to stack vector subscripted entity shape and
+/// element addresses. AnyVariableStack allows saving vector subscripted
+/// entities element addresses, but when saving several vector subscripted
+/// entities on a stack, and if the context does not allow retrieving the
+/// vector subscript entities shapes, these shapes must be saved too.
+class AnyVectorSubscriptStack : public AnyVariableStack {
+public:
+  AnyVectorSubscriptStack(mlir::Location loc, fir::FirOpBuilder &builder,
+                          mlir::Type valueStaticType,
+                          bool shapeCanBeSavedAsRegister, int rank);
+  void pushShape(mlir::Location loc, fir::FirOpBuilder &builder,
+                 mlir::Value shape);
+  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder);
+  mlir::Value fetchShape(mlir::Location loc, fir::FirOpBuilder &builder);
+  void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
+  bool canBeFetchedAfterPush() const { return true; }
+
+private:
+  std::unique_ptr<TemporaryStorage> shapeTemp;
+  // If the shape is saved inside a descriptor (as extents),
+  // keep track of the descriptor type.
+  std::optional<mlir::Type> boxType;
+};
+
 /// Generic wrapper over the different sorts of temporary storages.
 class TemporaryStorage {
 public:
@@ -178,8 +257,15 @@ class TemporaryStorage {
                       impl);
   }
 
+  template <typename T>
+  T &cast() {
+    return std::get<T>(impl);
+  }
+
 private:
-  std::variant<HomogeneousScalarStack, SimpleCopy, AnyValueStack> impl;
+  std::variant<HomogeneousScalarStack, SimpleCopy, SSARegister, AnyValueStack,
+               AnyVariableStack, AnyVectorSubscriptStack>
+      impl;
 };
 } // namespace fir::factory
 #endif // FORTRAN_OPTIMIZER_BUILDER_TEMPORARYSTORAGE_H
diff --git a/flang/lib/Optimizer/Builder/Runtime/TemporaryStack.cpp b/flang/lib/Optimizer/Builder/Runtime/TemporaryStack.cpp
@@ -56,3 +56,52 @@ void fir::runtime::genDestroyValueStack(mlir::Location loc,
   auto args = fir::runtime::createArguments(builder, loc, funcType, opaquePtr);
   builder.create<fir::CallOp>(loc, func, args);
 }
+
+mlir::Value fir::runtime::genCreateDescriptorStack(mlir::Location loc,
+                                                   fir::FirOpBuilder &builder) {
+  mlir::func::FuncOp func =
+      fir::runtime::getRuntimeFunc<mkRTKey(CreateDescriptorStack)>(loc,
+                                                                   builder);
+  mlir::FunctionType funcType = func.getFunctionType();
+  mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);
+  mlir::Value sourceLine =
+      fir::factory::locationToLineNo(builder, loc, funcType.getInput(1));
+  auto args = fir::runtime::createArguments(builder, loc, funcType, sourceFile,
+                                            sourceLine);
+  return builder.create<fir::CallOp>(loc, func, args).getResult(0);
+}
+
+void fir::runtime::genPushDescriptor(mlir::Location loc,
+                                     fir::FirOpBuilder &builder,
+                                     mlir::Value opaquePtr,
+                                     mlir::Value boxDescriptor) {
+  mlir::func::FuncOp func =
+      fir::runtime::getRuntimeFunc<mkRTKey(PushDescriptor)>(loc, builder);
+  mlir::FunctionType funcType = func.getFunctionType();
+  auto args = fir::runtime::createArguments(builder, loc, funcType, opaquePtr,
+                                            boxDescriptor);
+  builder.create<fir::CallOp>(loc, func, args);
+}
+
+void fir::runtime::genDescriptorAt(mlir::Location loc,
+                                   fir::FirOpBuilder &builder,
+                                   mlir::Value opaquePtr, mlir::Value i,
+                                   mlir::Value retDescriptorBox) {
+  mlir::func::FuncOp func =
+      fir::runtime::getRuntimeFunc<mkRTKey(DescriptorAt)>(loc, builder);
+  mlir::FunctionType funcType = func.getFunctionType();
+  auto args = fir::runtime::createArguments(builder, loc, funcType, opaquePtr,
+                                            i, retDescriptorBox);
+  builder.create<fir::CallOp>(loc, func, args);
+}
+
+void fir::runtime::genDestroyDescriptorStack(mlir::Location loc,
+                                             fir::FirOpBuilder &builder,
+                                             mlir::Value opaquePtr) {
+  mlir::func::FuncOp func =
+      fir::runtime::getRuntimeFunc<mkRTKey(DestroyDescriptorStack)>(loc,
+                                                                    builder);
+  mlir::FunctionType funcType = func.getFunctionType();
+  auto args = fir::runtime::createArguments(builder, loc, funcType, opaquePtr);
+  builder.create<fir::CallOp>(loc, func, args);
+}
diff --git a/flang/lib/Optimizer/Builder/TemporaryStorage.cpp b/flang/lib/Optimizer/Builder/TemporaryStorage.cpp
@@ -231,3 +231,127 @@ void fir::factory::AnyValueStack::destroy(mlir::Location loc,
                                           fir::FirOpBuilder &builder) {
   fir::runtime::genDestroyValueStack(loc, builder, opaquePtr);
 }
+
+//===----------------------------------------------------------------------===//
+// fir::factory::AnyVariableStack implementation.
+//===----------------------------------------------------------------------===//
+
+fir::factory::AnyVariableStack::AnyVariableStack(mlir::Location loc,
+                                                 fir::FirOpBuilder &builder,
+                                                 mlir::Type variableStaticType)
+    : variableStaticType{variableStaticType},
+      counter{loc, builder,
+              builder.createIntegerConstant(loc, builder.getI64Type(), 0),
+              /*stackThroughLoops=*/true} {
+  opaquePtr = fir::runtime::genCreateDescriptorStack(loc, builder);
+  mlir::Type storageType =
+      hlfir::getFortranElementOrSequenceType(variableStaticType);
+  mlir::Type ptrType = fir::PointerType::get(storageType);
+  mlir::Type boxType;
+  if (hlfir::isPolymorphicType(variableStaticType))
+    boxType = fir::ClassType::get(ptrType);
+  else
+    boxType = fir::BoxType::get(ptrType);
+  retValueBox = builder.createTemporary(loc, boxType);
+}
+
+void fir::factory::AnyVariableStack::pushValue(mlir::Location loc,
+                                               fir::FirOpBuilder &builder,
+                                               mlir::Value variable) {
+  hlfir::Entity entity{variable};
+  mlir::Type storageElementType =
+      hlfir::getFortranElementType(retValueBox.getType());
+  auto [box, maybeCleanUp] =
+      hlfir::convertToBox(loc, builder, entity, storageElementType);
+  fir::runtime::genPushDescriptor(loc, builder, opaquePtr, fir::getBase(box));
+  if (maybeCleanUp)
+    (*maybeCleanUp)();
+}
+
+void fir::factory::AnyVariableStack::resetFetchPosition(
+    mlir::Location loc, fir::FirOpBuilder &builder) {
+  counter.reset(loc, builder);
+}
+
+mlir::Value fir::factory::AnyVariableStack::fetch(mlir::Location loc,
+                                                  fir::FirOpBuilder &builder) {
+  mlir::Value indexValue = counter.getAndIncrementIndex(loc, builder);
+  fir::runtime::genDescriptorAt(loc, builder, opaquePtr, indexValue,
+                                retValueBox);
+  hlfir::Entity retBox{builder.create<fir::LoadOp>(loc, retValueBox)};
+  // The runtime always tracks variable as address, but the form of the variable
+  // that was saved may be different (raw address, fir.boxchar), ensure
+  // the returned variable has the same form of the one that was saved.
+  if (mlir::isa<fir::BaseBoxType>(variableStaticType))
+    return builder.createConvert(loc, variableStaticType, retBox);
+  if (mlir::isa<fir::BoxCharType>(variableStaticType))
+    return hlfir::genVariableBoxChar(loc, builder, retBox);
+  mlir::Value rawAddr = genVariableRawAddress(loc, builder, retBox);
+  return builder.createConvert(loc, variableStaticType, rawAddr);
+}
+
+void fir::factory::AnyVariableStack::destroy(mlir::Location loc,
+                                             fir::FirOpBuilder &builder) {
+  fir::runtime::genDestroyDescriptorStack(loc, builder, opaquePtr);
+}
+
+//===----------------------------------------------------------------------===//
+// fir::factory::AnyVectorSubscriptStack implementation.
+//===----------------------------------------------------------------------===//
+
+fir::factory::AnyVectorSubscriptStack::AnyVectorSubscriptStack(
+    mlir::Location loc, fir::FirOpBuilder &builder,
+    mlir::Type variableStaticType, bool shapeCanBeSavedAsRegister, int rank)
+    : AnyVariableStack{loc, builder, variableStaticType} {
+  if (shapeCanBeSavedAsRegister) {
+    shapeTemp =
+        std::unique_ptr<TemporaryStorage>(new TemporaryStorage{SSARegister{}});
+    return;
+  }
+  // The shape will be tracked as the dimension inside a descriptor because
+  // that is the easiest from a lowering point of view, and this is an
+  // edge case situation that will probably not very well be exercised.
+  mlir::Type type =
+      fir::BoxType::get(builder.getVarLenSeqTy(builder.getI32Type(), rank));
+  boxType = type;
+  shapeTemp = std::unique_ptr<TemporaryStorage>(
+      new TemporaryStorage{AnyVariableStack{loc, builder, type}});
+}
+
+void fir::factory::AnyVectorSubscriptStack::pushShape(
+    mlir::Location loc, fir::FirOpBuilder &builder, mlir::Value shape) {
+  if (boxType) {
+    // The shape is saved as a dimensions inside a descriptors.
+    mlir::Type refType = fir::ReferenceType::get(
+        hlfir::getFortranElementOrSequenceType(*boxType));
+    mlir::Value null = builder.createNullConstant(loc, refType);
+    mlir::Value descriptor =
+        builder.create<fir::EmboxOp>(loc, *boxType, null, shape);
+    shapeTemp->pushValue(loc, builder, descriptor);
+    return;
+  }
+  // Otherwise, simply keep track of the fir.shape itself, it is invariant.
+  shapeTemp->cast<SSARegister>().pushValue(loc, builder, shape);
+}
+
+void fir::factory::AnyVectorSubscriptStack::resetFetchPosition(
+    mlir::Location loc, fir::FirOpBuilder &builder) {
+  static_cast<AnyVariableStack *>(this)->resetFetchPosition(loc, builder);
+  shapeTemp->resetFetchPosition(loc, builder);
+}
+
+mlir::Value
+fir::factory::AnyVectorSubscriptStack::fetchShape(mlir::Location loc,
+                                                  fir::FirOpBuilder &builder) {
+  if (boxType) {
+    hlfir::Entity descriptor{shapeTemp->fetch(loc, builder)};
+    return hlfir::genShape(loc, builder, descriptor);
+  }
+  return shapeTemp->cast<SSARegister>().fetch(loc, builder);
+}
+
+void fir::factory::AnyVectorSubscriptStack::destroy(
+    mlir::Location loc, fir::FirOpBuilder &builder) {
+  static_cast<AnyVariableStack *>(this)->destroy(loc, builder);
+  shapeTemp->destroy(loc, builder);
+}
diff --git a/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp b/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp
diff --git a/flang/test/HLFIR/order_assignments/lhs-conflicts-codegen.fir b/flang/test/HLFIR/order_assignments/lhs-conflicts-codegen.fir
