TBAAForest.h

//===-- TBAAForest.h - A TBAA tree for each function -----------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#ifndef FORTRAN_OPTIMIZER_ANALYSIS_TBAA_FOREST_H
#define FORTRAN_OPTIMIZER_ANALYSIS_TBAA_FOREST_H
 
#include "flang/Optimizer/Dialect/FIROpsSupport.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/MLIRContext.h"
#include "llvm/ADT/DenseMap.h"
#include <string>
 
namespace fir {
 
//===----------------------------------------------------------------------===//
// TBAATree
//===----------------------------------------------------------------------===//
struct TBAATree {
  //===----------------------------------------------------------------------===//
  // TBAAForrest::TBAATree::SubtreeState
  //===----------------------------------------------------------------------===//
  class SubtreeState {
    friend TBAATree; // only allow construction by TBAATree
  public:
    SubtreeState() = delete;
    SubtreeState(const SubtreeState &) = delete;
    SubtreeState(SubtreeState &&) = default;
 
    mlir::LLVM::TBAATagAttr getTag(llvm::StringRef uniqueId) const;

    mlir::LLVM::TBAATagAttr getTag() const;
 
    mlir::LLVM::TBAATypeDescriptorAttr getRoot() const { return parent; }

    SubtreeState &getOrCreateNamedSubtree(mlir::StringAttr name);
 
  private:
    SubtreeState(mlir::MLIRContext *ctx, std::string name,
                 mlir::LLVM::TBAANodeAttr grandParent)
        : parentId{std::move(name)}, context(ctx) {
      parent = mlir::LLVM::TBAATypeDescriptorAttr::get(
          context, parentId, mlir::LLVM::TBAAMemberAttr::get(grandParent, 0));
    }
 
    const std::string parentId;
    mlir::MLIRContext *const context;
    mlir::LLVM::TBAATypeDescriptorAttr parent;
    // A map of named sub-trees, e.g. sub-trees of the COMMON blocks
    // placed under the "global data" root.
    llvm::DenseMap<mlir::StringAttr, SubtreeState> namedSubtrees;
  };

  SubtreeState targetDataTree;
  SubtreeState globalDataTree;
  SubtreeState allocatedDataTree;
  SubtreeState dummyArgDataTree;
  SubtreeState directDataTree;
  mlir::LLVM::TBAATypeDescriptorAttr anyAccessDesc;
  mlir::LLVM::TBAATypeDescriptorAttr boxMemberTypeDesc;
  mlir::LLVM::TBAATypeDescriptorAttr anyDataTypeDesc;
 
  // Structure of the created tree:
  //   Function root
  //   |
  //   "any access"
  //   |
  //   |- "descriptor member"
  //   |- "any data access"
  //      |
  //      |- "dummy arg data"
  //        |
  //        |- <dummy arg name 1>
  //        |- <dummy arg name 2>
  //      |- "target data" <-- Any POINTER variable or TARGET dummy arg
  //        |
  //        |- <target name 1> <--- any TARGET variable which isn't a dummy arg
  //        |- <target name 2>
  //      |- "allocated data"
  //        |
  //        |- <allocated name 1>
  //        |- <allocated name 2>
  //      |- "direct data"
  //        |
  //        |- <direct name 1>
  //        |- <direct name 2>
  //      |- "global data"
  //        |
  //        |- <global name 1>
  //        |- <global name 2>
  static TBAATree buildTree(mlir::StringAttr functionName);
 
private:
  TBAATree(mlir::LLVM::TBAATypeDescriptorAttr anyAccess,
           mlir::LLVM::TBAATypeDescriptorAttr dataRoot,
           mlir::LLVM::TBAATypeDescriptorAttr boxMemberTypeDesc);
};
 
//===----------------------------------------------------------------------===//
// TBAAForrest
//===----------------------------------------------------------------------===//
class TBAAForrest {
public:
  explicit TBAAForrest(bool separatePerFunction = true)
      : separatePerFunction{separatePerFunction} {}
 
  inline const TBAATree &operator[](mlir::func::FuncOp func) {
    return getFuncTree(getInternalOrSymbolName(func));
  }
  inline const TBAATree &operator[](mlir::LLVM::LLVMFuncOp func) {
    return getFuncTree(getInternalOrSymbolName(func));
  }
  // Returns the TBAA tree associated with the scope enclosed
  // within the given function. With MLIR inlining, there may
  // be multiple scopes within a single function. It is the caller's
  // responsibility to provide unique name for the scope.
  // If the scope string is empty, returns the TBAA tree for the
  // "root" scope of the given function.
  inline TBAATree &getMutableFuncTreeWithScope(mlir::func::FuncOp func,
                                               llvm::StringRef scope) {
    mlir::StringAttr name = getInternalOrSymbolName(func);
    if (!scope.empty())
      name = mlir::StringAttr::get(name.getContext(),
                                   llvm::Twine(name) + " - " + scope);
    return getFuncTree(name);
  }
 
  inline const TBAATree &getFuncTreeWithScope(mlir::func::FuncOp func,
                                              llvm::StringRef scope) {
    return getMutableFuncTreeWithScope(func, scope);
  }
 
private:
  template <typename FuncOp>
  static mlir::StringAttr getInternalOrSymbolName(FuncOp func) {
    // External name conversion may rename a function before TBAA construction.
    // Use its original name consistently so tags created at different pipeline
    // stages belong to the same tree.
    if (auto name = func->template getAttrOfType<mlir::StringAttr>(
            getInternalFuncNameAttrName()))
      return name;
    return func.getSymNameAttr();
  }
 
  TBAATree &getFuncTree(mlir::StringAttr symName) {
    if (!separatePerFunction)
      symName = mlir::StringAttr::get(symName.getContext(), "");
    if (!trees.contains(symName))
      trees.insert({symName, TBAATree::buildTree(symName)});
    auto it = trees.find(symName);
    assert(it != trees.end());
    return it->second;
  }
 
  // Should each function use a different tree?
  const bool separatePerFunction;
  // TBAA tree per function
  llvm::DenseMap<mlir::StringAttr, TBAATree> trees;
};
 
} // namespace fir
 
#endif // FORTRAN_OPTIMIZER_ANALYSIS_TBAA_FOREST_H