[flang] Restructured TBAA trees in AddAliasTags pass. (#136725)

This patch produces the following TBAA tree for a function:
```
  Function root
  |
  "any access"
  |
  |- "descriptor member"
  |- "any data access"
     |
     |- "dummy arg data"
     |- "target data"
        |
        |- "allocated data"
        |- "direct data"
        |- "global data"
```
The TBAA tags are assigned using the following logic:
  * All POINTER variables point to the root of "target data".
  * Dummy arguments without POINTER/TARGET point to their
    leafs under "dummy arg data".
  * Dummy arguments with TARGET point to the root of "target data".
  * Global variables without descriptors point to their leafs under
    "global data" (including the ones with TARGET).
  * Global variables with descriptors point to their leafs under
    "direct data" (including the ones with TARGET).
  * Locally allocated variables point to their leafs under
   "allocated data" (including the ones with TARGET).

This change makes it possible to disambiguate globals like:
```
  module data
    real, allocatable :: a(:)
    real, allocatable, target :: b(:)
  end
```

Indeed, two direct references to global vars cannot alias
even if any/both of them have TARGET attribute.
In addition, the dummy arguments without POINTER/TARGET cannot alias
any other variable even with POINTER/TARGET. This was not expressed
in TBAA before this change.

As before, any "unknown" memory references (such as with Indirect
source, as classified by FIR alias analysis) may alias with
anything, as long as they point to the root of "any access".

Please think of the counterexamples for which this structure
may not work.

Author: vzakhari

flang/include/flang/Optimizer/Analysis/AliasAnalysis.h

@@ -155,6 +155,12 @@ struct AliasAnalysis {
     /// Return true, if Target or Pointer attribute is set.
     bool isTargetOrPointer() const;
 
+    /// Return true, if Target attribute is set.
+    bool isTarget() const;
+
+    /// Return true, if Pointer attribute is set.
+    bool isPointer() const;
+
     bool isDummyArgument() const;
     bool isData() const;
     bool isBoxData() const;

flang/include/flang/Optimizer/Analysis/TBAAForest.h

@@ -40,6 +40,12 @@ struct TBAATree {
 
     mlir::LLVM::TBAATagAttr getTag(llvm::StringRef uniqueId) const;
 
+    /// Create a TBAA tag pointing to the root of this subtree,
+    /// i.e. all the children tags will alias with this tag.
+    mlir::LLVM::TBAATagAttr getTag() const;
+
+    mlir::LLVM::TBAATypeDescriptorAttr getRoot() const { return parent; }
+
   private:
     SubtreeState(mlir::MLIRContext *ctx, std::string name,
                  mlir::LLVM::TBAANodeAttr grandParent)
@@ -51,17 +57,44 @@ struct TBAATree {
     const std::string parentId;
     mlir::MLIRContext *const context;
     mlir::LLVM::TBAATypeDescriptorAttr parent;
-    llvm::DenseMap<llvm::StringRef, mlir::LLVM::TBAATagAttr> tagDedup;
   };
 
+  /// A subtree for POINTER/TARGET variables data.
+  /// Any POINTER variable must use a tag that points
+  /// to the root of this subtree.
+  /// A TARGET dummy argument must also point to this root.
+  SubtreeState targetDataTree;
+  /// A subtree for global variables data (e.g. user module variables).
   SubtreeState globalDataTree;
+  /// A subtree for variables allocated via fir.alloca or fir.allocmem.
   SubtreeState allocatedDataTree;
+  /// A subtree for subprogram's dummy arguments.
+  /// It only contains children for the dummy arguments
+  /// that are not POINTER/TARGET. They all do not conflict
+  /// with each other and with any other data access, except
+  /// with unknown data accesses (FIR alias analysis uses
+  /// SourceKind::Indirect for sources of such accesses).
   SubtreeState dummyArgDataTree;
+  /// A subtree for global variables descriptors.
   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"
+  //      |- "target data"
+  //         |
+  //         |- "allocated data"
+  //         |- "direct data"
+  //         |- "global data"
   static TBAATree buildTree(mlir::StringAttr functionName);
 
 private:

flang/lib/Optimizer/Analysis/AliasAnalysis.cpp

@@ -110,6 +110,14 @@ bool AliasAnalysis::Source::isTargetOrPointer() const {
          attributes.test(Attribute::Target);
 }
 
+bool AliasAnalysis::Source::isTarget() const {
+  return attributes.test(Attribute::Target);
+}
+
+bool AliasAnalysis::Source::isPointer() const {
+  return attributes.test(Attribute::Pointer);
+}
+
 bool AliasAnalysis::Source::isDummyArgument() const {
   if (auto v = origin.u.dyn_cast<mlir::Value>()) {
     return fir::isDummyArgument(v);

flang/lib/Optimizer/Analysis/TBAAForest.cpp

@@ -11,9 +11,6 @@
 
 mlir::LLVM::TBAATagAttr
 fir::TBAATree::SubtreeState::getTag(llvm::StringRef uniqueName) const {
-  // mlir::LLVM::TBAATagAttr &tag = tagDedup[uniqueName];
-  // if (tag)
-  //   return tag;
   std::string id = (parentId + "/" + uniqueName).str();
   mlir::LLVM::TBAATypeDescriptorAttr type =
       mlir::LLVM::TBAATypeDescriptorAttr::get(
@@ -22,6 +19,10 @@ fir::TBAATree::SubtreeState::getTag(llvm::StringRef uniqueName) const {
   // return tag;
 }
 
+mlir::LLVM::TBAATagAttr fir::TBAATree::SubtreeState::getTag() const {
+  return mlir::LLVM::TBAATagAttr::get(parent, parent, 0);
+}
+
 fir::TBAATree fir::TBAATree::buildTree(mlir::StringAttr func) {
   llvm::StringRef funcName = func.getValue();
   std::string rootId = ("Flang function root " + funcName).str();
@@ -53,9 +54,13 @@ fir::TBAATree fir::TBAATree::buildTree(mlir::StringAttr func) {
 fir::TBAATree::TBAATree(mlir::LLVM::TBAATypeDescriptorAttr anyAccess,
                         mlir::LLVM::TBAATypeDescriptorAttr dataRoot,
                         mlir::LLVM::TBAATypeDescriptorAttr boxMemberTypeDesc)
-    : globalDataTree(dataRoot.getContext(), "global data", dataRoot),
-      allocatedDataTree(dataRoot.getContext(), "allocated data", dataRoot),
+    : targetDataTree(dataRoot.getContext(), "target data", dataRoot),
+      globalDataTree(dataRoot.getContext(), "global data",
+                     targetDataTree.getRoot()),
+      allocatedDataTree(dataRoot.getContext(), "allocated data",
+                        targetDataTree.getRoot()),
       dummyArgDataTree(dataRoot.getContext(), "dummy arg data", dataRoot),
-      directDataTree(dataRoot.getContext(), "direct data", dataRoot),
+      directDataTree(dataRoot.getContext(), "direct data",
+                     targetDataTree.getRoot()),
       anyAccessDesc(anyAccess), boxMemberTypeDesc(boxMemberTypeDesc),
       anyDataTypeDesc(dataRoot) {}

flang/lib/Optimizer/Transforms/AddAliasTags.cpp

@@ -198,12 +198,6 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
   LLVM_DEBUG(llvm::dbgs() << "Analysing " << op << "\n");
 
   const fir::AliasAnalysis::Source &source = state.getSource(memref);
-  if (source.isTargetOrPointer()) {
-    LLVM_DEBUG(llvm::dbgs().indent(2) << "Skipping TARGET/POINTER\n");
-    // These will get an "any data access" tag in TBAABuilder (CodeGen): causing
-    // them to "MayAlias" with all non-box accesses
-    return;
-  }
 
   // Process the scopes, if not processed yet.
   state.processFunctionScopes(func);
@@ -228,37 +222,48 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
     LLVM_DEBUG(llvm::dbgs().indent(2)
                << "Found reference to dummy argument at " << *op << "\n");
     std::string name = getFuncArgName(llvm::cast<mlir::Value>(source.origin.u));
-    if (!name.empty())
+    // If it is a TARGET or POINTER, then we do not care about the name,
+    // because the tag points to the root of the subtree currently.
+    if (source.isTargetOrPointer()) {
+      tag = state.getFuncTreeWithScope(func, scopeOp).targetDataTree.getTag();
+    } else if (!name.empty()) {
       tag = state.getFuncTreeWithScope(func, scopeOp)
                 .dummyArgDataTree.getTag(name);
-    else
+    } else {
       LLVM_DEBUG(llvm::dbgs().indent(2)
                  << "WARN: couldn't find a name for dummy argument " << *op
                  << "\n");
+      tag = state.getFuncTreeWithScope(func, scopeOp).dummyArgDataTree.getTag();
+    }
 
-    // TBAA for global variables
+    // TBAA for global variables without descriptors
   } else if (enableGlobals &&
              source.kind == fir::AliasAnalysis::SourceKind::Global &&
              !source.isBoxData()) {
     mlir::SymbolRefAttr glbl = llvm::cast<mlir::SymbolRefAttr>(source.origin.u);
     const char *name = glbl.getRootReference().data();
     LLVM_DEBUG(llvm::dbgs().indent(2) << "Found reference to global " << name
                                       << " at " << *op << "\n");
-    tag = state.getFuncTreeWithScope(func, scopeOp).globalDataTree.getTag(name);
+    if (source.isPointer())
+      tag = state.getFuncTreeWithScope(func, scopeOp).targetDataTree.getTag();
+    else
+      tag =
+          state.getFuncTreeWithScope(func, scopeOp).globalDataTree.getTag(name);
 
-    // TBAA for SourceKind::Direct
+    // TBAA for global variables with descriptors
   } else if (enableDirect &&
              source.kind == fir::AliasAnalysis::SourceKind::Global &&
              source.isBoxData()) {
     if (auto glbl = llvm::dyn_cast<mlir::SymbolRefAttr>(source.origin.u)) {
       const char *name = glbl.getRootReference().data();
       LLVM_DEBUG(llvm::dbgs().indent(2) << "Found reference to direct " << name
                                         << " at " << *op << "\n");
-      tag =
-          state.getFuncTreeWithScope(func, scopeOp).directDataTree.getTag(name);
+      if (source.isPointer())
+        tag = state.getFuncTreeWithScope(func, scopeOp).targetDataTree.getTag();
+      else
+        tag = state.getFuncTreeWithScope(func, scopeOp)
+                  .directDataTree.getTag(name);
     } else {
-      // SourceKind::Direct is likely to be extended to cases which are not a
-      // SymbolRefAttr in the future
       LLVM_DEBUG(llvm::dbgs().indent(2) << "Can't get name for direct "
                                         << source << " at " << *op << "\n");
     }
@@ -269,11 +274,23 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
     std::optional<llvm::StringRef> name;
     mlir::Operation *sourceOp =
         llvm::cast<mlir::Value>(source.origin.u).getDefiningOp();
+    bool unknownAllocOp = false;
     if (auto alloc = mlir::dyn_cast_or_null<fir::AllocaOp>(sourceOp))
       name = alloc.getUniqName();
     else if (auto alloc = mlir::dyn_cast_or_null<fir::AllocMemOp>(sourceOp))
       name = alloc.getUniqName();
-    if (name) {
+    else
+      unknownAllocOp = true;
+
+    if (unknownAllocOp) {
+      LLVM_DEBUG(llvm::dbgs().indent(2)
+                 << "WARN: unknown defining op for SourceKind::Allocate " << *op
+                 << "\n");
+    } else if (source.isPointer()) {
+      LLVM_DEBUG(llvm::dbgs().indent(2)
+                 << "Found reference to allocation at " << *op << "\n");
+      tag = state.getFuncTreeWithScope(func, scopeOp).targetDataTree.getTag();
+    } else if (name) {
       LLVM_DEBUG(llvm::dbgs().indent(2) << "Found reference to allocation "
                                         << name << " at " << *op << "\n");
       tag = state.getFuncTreeWithScope(func, scopeOp)
@@ -282,6 +299,8 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
       LLVM_DEBUG(llvm::dbgs().indent(2)
                  << "WARN: couldn't find a name for allocation " << *op
                  << "\n");
+      tag =
+          state.getFuncTreeWithScope(func, scopeOp).allocatedDataTree.getTag();
     }
   } else {
     if (source.kind != fir::AliasAnalysis::SourceKind::Argument &&