[flang][NFC] use llvm.intr.stacksave/restore instead of opaque calls (llvm#108562)

The new LLVM stack save/restore intrinsic operations are more convenient
than function calls because they do not add function declarations to the
module and therefore do not block the parallelisation of passes.
Furthermore they could be much more easily marked with memory effects
than function calls if that ever proved useful.

This builds on top of llvm#107879.

Resolves llvm#108016

Author: tblah

flang/include/flang/Optimizer/Builder/FIRBuilder.h

@@ -29,6 +29,7 @@
 #include <utility>
 
 namespace mlir {
+class DataLayout;
 class SymbolTable;
 }
 
@@ -253,6 +254,15 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
                       mlir::ValueRange lenParams = {},
                       llvm::ArrayRef<mlir::NamedAttribute> attrs = {});
 
+  /// Create an LLVM stack save intrinsic op. Returns the saved stack pointer.
+  /// The stack address space is fetched from the data layout of the current
+  /// module.
+  mlir::Value genStackSave(mlir::Location loc);
+
+  /// Create an LLVM stack restore intrinsic op. stackPointer should be a value
+  /// previously returned from genStackSave.
+  void genStackRestore(mlir::Location loc, mlir::Value stackPointer);
+
   /// Create a global value.
   fir::GlobalOp createGlobal(mlir::Location loc, mlir::Type type,
                              llvm::StringRef name,
@@ -523,6 +533,9 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
     setCommonAttributes(op);
   }
 
+  /// Construct a data layout on demand and return it
+  mlir::DataLayout &getDataLayout();
+
 private:
   /// Set attributes (e.g. FastMathAttr) to \p op operation
   /// based on the current attributes setting.
@@ -537,6 +550,11 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
   /// fir::GlobalOp and func::FuncOp symbol table to speed-up
   /// lookups.
   mlir::SymbolTable *symbolTable = nullptr;
+
+  /// DataLayout constructed on demand. Access via getDataLayout().
+  /// Stored via a unique_ptr rather than an optional so as not to bloat this
+  /// class when most instances won't ever need a data layout.
+  std::unique_ptr<mlir::DataLayout> dataLayout = nullptr;
 };
 
 } // namespace fir
@@ -729,6 +747,9 @@ elideExtentsAlreadyInType(mlir::Type type, mlir::ValueRange shape);
 llvm::SmallVector<mlir::Value>
 elideLengthsAlreadyInType(mlir::Type type, mlir::ValueRange lenParams);
 
+/// Get the address space which should be used for allocas
+uint64_t getAllocaAddressSpace(mlir::DataLayout *dataLayout);
+
 } // namespace fir::factory
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H

flang/include/flang/Optimizer/Builder/LowLevelIntrinsics.h

@@ -42,12 +42,6 @@ mlir::func::FuncOp getLlvmGetRounding(FirOpBuilder &builder);
 /// Get the `llvm.set.rounding` intrinsic.
 mlir::func::FuncOp getLlvmSetRounding(FirOpBuilder &builder);
 
-/// Get the `llvm.stacksave` intrinsic.
-mlir::func::FuncOp getLlvmStackSave(FirOpBuilder &builder);
-
-/// Get the `llvm.stackrestore` intrinsic.
-mlir::func::FuncOp getLlvmStackRestore(FirOpBuilder &builder);
-
 /// Get the `llvm.init.trampoline` intrinsic.
 mlir::func::FuncOp getLlvmInitTrampoline(FirOpBuilder &builder);
 

flang/include/flang/Optimizer/Support/DataLayout.h

@@ -45,7 +45,6 @@ void setMLIRDataLayoutFromAttributes(mlir::ModuleOp mlirModule,
 /// std::nullopt.
 std::optional<mlir::DataLayout>
 getOrSetDataLayout(mlir::ModuleOp mlirModule, bool allowDefaultLayout = false);
-
 } // namespace fir::support
 
 #endif // FORTRAN_OPTIMIZER_SUPPORT_DATALAYOUT_H

flang/lib/Lower/Bridge.cpp

@@ -3257,15 +3257,10 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         const Fortran::parser::CharBlock &endPosition =
             eval.getLastNestedEvaluation().position;
         localSymbols.pushScope();
-        mlir::func::FuncOp stackSave = fir::factory::getLlvmStackSave(*builder);
-        mlir::func::FuncOp stackRestore =
-            fir::factory::getLlvmStackRestore(*builder);
-        mlir::Value stackPtr =
-            builder->create<fir::CallOp>(toLocation(), stackSave).getResult(0);
+        mlir::Value stackPtr = builder->genStackSave(toLocation());
         mlir::Location endLoc = genLocation(endPosition);
-        stmtCtx.attachCleanup([=]() {
-          builder->create<fir::CallOp>(endLoc, stackRestore, stackPtr);
-        });
+        stmtCtx.attachCleanup(
+            [=]() { builder->genStackRestore(endLoc, stackPtr); });
         Fortran::semantics::Scope &scope =
             bridge.getSemanticsContext().FindScope(endPosition);
         scopeBlockIdMap.try_emplace(&scope, ++blockId);

flang/lib/Lower/ConvertCall.cpp

@@ -368,22 +368,9 @@ std::pair<fir::ExtendedValue, bool> Fortran::lower::genCallOpAndResult(
 
     if (!extents.empty() || !lengths.empty()) {
       auto *bldr = &converter.getFirOpBuilder();
-      auto stackSaveFn = fir::factory::getLlvmStackSave(builder);
-      auto stackSaveSymbol = bldr->getSymbolRefAttr(stackSaveFn.getName());
-      mlir::Value sp;
-      fir::CallOp call = bldr->create<fir::CallOp>(
-          loc, stackSaveSymbol, stackSaveFn.getFunctionType().getResults(),
-          mlir::ValueRange{});
-      if (call.getNumResults() != 0)
-        sp = call.getResult(0);
-      stmtCtx.attachCleanup([bldr, loc, sp]() {
-        auto stackRestoreFn = fir::factory::getLlvmStackRestore(*bldr);
-        auto stackRestoreSymbol =
-            bldr->getSymbolRefAttr(stackRestoreFn.getName());
-        bldr->create<fir::CallOp>(loc, stackRestoreSymbol,
-                                  stackRestoreFn.getFunctionType().getResults(),
-                                  mlir::ValueRange{sp});
-      });
+      mlir::Value sp = bldr->genStackSave(loc);
+      stmtCtx.attachCleanup(
+          [bldr, loc, sp]() { bldr->genStackRestore(loc, sp); });
     }
     mlir::Value temp =
         builder.createTemporary(loc, type, ".result", extents, resultLengths);

flang/lib/Optimizer/Builder/FIRBuilder.cpp

@@ -18,6 +18,7 @@
 #include "flang/Optimizer/Dialect/FIRAttr.h"
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Support/DataLayout.h"
 #include "flang/Optimizer/Support/FatalError.h"
 #include "flang/Optimizer/Support/InternalNames.h"
 #include "flang/Optimizer/Support/Utils.h"
@@ -328,6 +329,17 @@ mlir::Value fir::FirOpBuilder::createHeapTemporary(
                                  name, dynamicLength, dynamicShape, attrs);
 }
 
+mlir::Value fir::FirOpBuilder::genStackSave(mlir::Location loc) {
+  mlir::Type voidPtr = mlir::LLVM::LLVMPointerType::get(
+      getContext(), fir::factory::getAllocaAddressSpace(&getDataLayout()));
+  return create<mlir::LLVM::StackSaveOp>(loc, voidPtr);
+}
+
+void fir::FirOpBuilder::genStackRestore(mlir::Location loc,
+                                        mlir::Value stackPointer) {
+  create<mlir::LLVM::StackRestoreOp>(loc, stackPointer);
+}
+
 /// Create a global variable in the (read-only) data section. A global variable
 /// must have a unique name to identify and reference it.
 fir::GlobalOp fir::FirOpBuilder::createGlobal(
@@ -791,6 +803,15 @@ void fir::FirOpBuilder::setFastMathFlags(
   setFastMathFlags(arithFMF);
 }
 
+// Construction of an mlir::DataLayout is expensive so only do it on demand and
+// memoise it in the builder instance
+mlir::DataLayout &fir::FirOpBuilder::getDataLayout() {
+  if (dataLayout)
+    return *dataLayout;
+  dataLayout = std::make_unique<mlir::DataLayout>(getModule());
+  return *dataLayout;
+}
+
 //===--------------------------------------------------------------------===//
 // ExtendedValue inquiry helper implementation
 //===--------------------------------------------------------------------===//
@@ -1664,3 +1685,10 @@ void fir::factory::setInternalLinkage(mlir::func::FuncOp func) {
       mlir::LLVM::LinkageAttr::get(func->getContext(), internalLinkage);
   func->setAttr("llvm.linkage", linkage);
 }
+
+uint64_t fir::factory::getAllocaAddressSpace(mlir::DataLayout *dataLayout) {
+  if (dataLayout)
+    if (mlir::Attribute addrSpace = dataLayout->getAllocaMemorySpace())
+      return mlir::cast<mlir::IntegerAttr>(addrSpace).getUInt();
+  return 0;
+}

flang/lib/Optimizer/Builder/LowLevelIntrinsics.cpp

@@ -76,25 +76,6 @@ fir::factory::getLlvmSetRounding(fir::FirOpBuilder &builder) {
                                 funcTy);
 }
 
-mlir::func::FuncOp fir::factory::getLlvmStackSave(fir::FirOpBuilder &builder) {
-  // FIXME: This should query the target alloca address space
-  auto ptrTy = builder.getRefType(builder.getIntegerType(8));
-  auto funcTy =
-      mlir::FunctionType::get(builder.getContext(), std::nullopt, {ptrTy});
-  return builder.createFunction(builder.getUnknownLoc(), "llvm.stacksave.p0",
-                                funcTy);
-}
-
-mlir::func::FuncOp
-fir::factory::getLlvmStackRestore(fir::FirOpBuilder &builder) {
-  // FIXME: This should query the target alloca address space
-  auto ptrTy = builder.getRefType(builder.getIntegerType(8));
-  auto funcTy =
-      mlir::FunctionType::get(builder.getContext(), {ptrTy}, std::nullopt);
-  return builder.createFunction(builder.getUnknownLoc(), "llvm.stackrestore.p0",
-                                funcTy);
-}
-
 mlir::func::FuncOp
 fir::factory::getLlvmInitTrampoline(fir::FirOpBuilder &builder) {
   auto ptrTy = builder.getRefType(builder.getIntegerType(8));

flang/lib/Optimizer/CodeGen/TargetRewrite.cpp

@@ -1236,25 +1236,18 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
 
   inline void clearMembers() { setMembers(nullptr, nullptr, nullptr); }
 
-  uint64_t getAllocaAddressSpace() const {
-    if (dataLayout)
-      if (mlir::Attribute addrSpace = dataLayout->getAllocaMemorySpace())
-        return llvm::cast<mlir::IntegerAttr>(addrSpace).getUInt();
-    return 0;
-  }
-
   // Inserts a call to llvm.stacksave at the current insertion
   // point and the given location. Returns the call's result Value.
   inline mlir::Value genStackSave(mlir::Location loc) {
-    mlir::Type voidPtr = mlir::LLVM::LLVMPointerType::get(
-        rewriter->getContext(), getAllocaAddressSpace());
-    return rewriter->create<mlir::LLVM::StackSaveOp>(loc, voidPtr);
+    fir::FirOpBuilder builder(*rewriter, getModule());
+    return builder.genStackSave(loc);
   }
 
   // Inserts a call to llvm.stackrestore at the current insertion
   // point and the given location and argument.
   inline void genStackRestore(mlir::Location loc, mlir::Value sp) {
-    rewriter->create<mlir::LLVM::StackRestoreOp>(loc, sp);
+    fir::FirOpBuilder builder(*rewriter, getModule());
+    return builder.genStackRestore(loc, sp);
   }
 
   fir::CodeGenSpecifics *specifics = nullptr;

flang/lib/Optimizer/Transforms/StackArrays.cpp

@@ -734,28 +734,12 @@ void AllocMemConversion::insertStackSaveRestore(
   auto mod = oldAlloc->getParentOfType<mlir::ModuleOp>();
   fir::FirOpBuilder builder{rewriter, mod};
 
-  mlir::func::FuncOp stackSaveFn = fir::factory::getLlvmStackSave(builder);
-  mlir::SymbolRefAttr stackSaveSym =
-      builder.getSymbolRefAttr(stackSaveFn.getName());
-
   builder.setInsertionPoint(oldAlloc);
-  mlir::Value sp =
-      builder
-          .create<fir::CallOp>(oldAlloc.getLoc(), stackSaveSym,
-                               stackSaveFn.getFunctionType().getResults(),
-                               mlir::ValueRange{})
-          .getResult(0);
-
-  mlir::func::FuncOp stackRestoreFn =
-      fir::factory::getLlvmStackRestore(builder);
-  mlir::SymbolRefAttr stackRestoreSym =
-      builder.getSymbolRefAttr(stackRestoreFn.getName());
+  mlir::Value sp = builder.genStackSave(oldAlloc.getLoc());
 
   auto createStackRestoreCall = [&](mlir::Operation *user) {
     builder.setInsertionPoint(user);
-    builder.create<fir::CallOp>(user->getLoc(), stackRestoreSym,
-                                stackRestoreFn.getFunctionType().getResults(),
-                                mlir::ValueRange{sp});
+    builder.genStackRestore(user->getLoc(), sp);
   };
 
   for (mlir::Operation *user : oldAlloc->getUsers()) {

flang/lib/Optimizer/Transforms/StackReclaim.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "flang/Common/Fortran.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Transforms/Passes.h"
@@ -31,34 +32,20 @@ class StackReclaimPass : public fir::impl::StackReclaimBase<StackReclaimPass> {
 };
 } // namespace
 
-uint64_t getAllocaAddressSpace(Operation *op) {
-  mlir::ModuleOp module = mlir::dyn_cast_or_null<mlir::ModuleOp>(op);
-  if (!module)
-    module = op->getParentOfType<mlir::ModuleOp>();
-
-  if (mlir::Attribute addrSpace =
-          mlir::DataLayout(module).getAllocaMemorySpace())
-    return llvm::cast<mlir::IntegerAttr>(addrSpace).getUInt();
-  return 0;
-}
-
 void StackReclaimPass::runOnOperation() {
   auto *op = getOperation();
-  auto *context = &getContext();
-  mlir::OpBuilder builder(context);
-  mlir::Type voidPtr =
-      mlir::LLVM::LLVMPointerType::get(context, getAllocaAddressSpace(op));
+  fir::FirOpBuilder builder(op, fir::getKindMapping(op));
 
   op->walk([&](fir::DoLoopOp loopOp) {
     mlir::Location loc = loopOp.getLoc();
 
     if (!loopOp.getRegion().getOps<fir::AllocaOp>().empty()) {
       builder.setInsertionPointToStart(&loopOp.getRegion().front());
-      auto stackSaveOp = builder.create<LLVM::StackSaveOp>(loc, voidPtr);
+      mlir::Value sp = builder.genStackSave(loc);
 
       auto *terminator = loopOp.getRegion().back().getTerminator();
       builder.setInsertionPoint(terminator);
-      builder.create<LLVM::StackRestoreOp>(loc, stackSaveOp);
+      builder.genStackRestore(loc, sp);
     }
   });
 }