[flang-commits] [flang] [flang][NFC] Extract StackArrays analysis and rewrite into a header - memory passes unification [1/5] (PR #210721)

[Tom Eccles via flang-commits]

================
@@ -0,0 +1,154 @@
+//===- StackArrays.h ------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
+//
+//===----------------------------------------------------------------------===//
+//
+// This header exposes the reusable pieces of the StackArrays pass: the
+// analysis that determines which fir.allocmem operations can safely be moved
+// to the stack (and where the replacement fir.alloca should be inserted), and
+// the pattern that performs the heap-to-stack rewrite. They are shared so that
+// other passes can reuse the same "is this heap allocation safely stackifiable,
+// and where" logic.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_OPTIMIZER_TRANSFORMS_STACKARRAYS_H
+#define FORTRAN_OPTIMIZER_TRANSFORMS_STACKARRAYS_H
+
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/Support/KindMapping.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Support/TypeID.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/PointerUnion.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include <optional>
+
+namespace fir {
+
+/// Stores where an alloca should be inserted. If the PointerUnion is an
+/// Operation the alloca should be inserted /after/ the operation. If it is a
+/// block, the alloca can be placed anywhere in that block.
+class InsertionPoint {
+  llvm::PointerUnion<mlir::Operation *, mlir::Block *> location;
+  bool saveRestoreStack;
+
+  /// Get contained pointer type or nullptr
+  template <class T>
+  T *tryGetPtr() const {
+    // Use llvm::dyn_cast_if_present because location may be null here.
+    if (T *ptr = llvm::dyn_cast_if_present<T *>(location))
+      return ptr;
+    return nullptr;
+  }
+
+public:
+  template <class T>
+  InsertionPoint(T *ptr, bool saveRestoreStack = false)
+      : location(ptr), saveRestoreStack{saveRestoreStack} {}
+  InsertionPoint(std::nullptr_t null)
+      : location(null), saveRestoreStack{false} {}
+
+  /// Get contained operation, or nullptr
+  mlir::Operation *tryGetOperation() const {
+    return tryGetPtr<mlir::Operation>();
+  }
+
+  /// Get contained block, or nullptr
+  mlir::Block *tryGetBlock() const { return tryGetPtr<mlir::Block>(); }
+
+  /// Get whether the stack should be saved/restored. If yes, an llvm.stacksave
+  /// intrinsic should be added before the alloca, and an llvm.stackrestore
+  /// intrinsic should be added where the freemem is
+  bool shouldSaveRestoreStack() const { return saveRestoreStack; }
+
+  operator bool() const { return tryGetOperation() || tryGetBlock(); }
+
+  bool operator==(const InsertionPoint &rhs) const {
+    return (location == rhs.location) &&
+           (saveRestoreStack == rhs.saveRestoreStack);
+  }
+
+  bool operator!=(const InsertionPoint &rhs) const { return !(*this == rhs); }
+};
+
+/// Drives analysis to find candidate fir.allocmem operations which could be
+/// moved to the stack. Intended to be used with mlir::Pass::getAnalysis
+class StackArraysAnalysisWrapper {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(StackArraysAnalysisWrapper)
+
+  // Maps fir.allocmem -> place to insert alloca
+  using AllocMemMap = llvm::DenseMap<mlir::Operation *, InsertionPoint>;
+
+  StackArraysAnalysisWrapper(mlir::Operation *op) {}
+
+  // returns nullptr if analysis failed
+  const AllocMemMap *getCandidateOps(mlir::Operation *func);
----------------
tblah wrote:

nit: for a public interface this needs a warning about pointer invalidation every time `funcMaps` grows

https://github.com/llvm/llvm-project/pull/210721