[flang] Create TBAA subtree for COMMON block variables.

In order to help LLVM disambiguate accesses to the COMMON
block variables, this patch creates a TBAA sub-tree for each
COMMON block, and the places all variables belonging to this
COMMON block into this sub-tree. The structure looks like this:
```
  common /blk/ a, b, c

  "global data"
  |
  |- "blk_"
     |
     |- "blk_/a"
     |- "blk_/b"
     |- "blk_/c"
```

The TBAA tag for "a" is created in "blk_/a" root, etc.
If, for some reason, we cannot identify a variable's name,
but we know that it belongs to COMMON "blk", the TBAA tag
will be created in "blk_" root - this tag indicates that this access
can overlap with any accesses of a/b/c.

I measured 10% speed-up on 434.zeusmp and 20% speed-up on 200.sixtrack
on Zen4. I expect around the same speed-ups on ARM.

Author: vzakhari

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

@@ -46,6 +46,12 @@ struct TBAATree {
 
     mlir::LLVM::TBAATypeDescriptorAttr getRoot() const { return parent; }
 
+    /// For the given name, get or create a subtree in the current
+    /// subtree. For example, this is used for creating subtrees
+    /// inside the "global data" subtree for the COMMON block variables
+    /// belonging to the same COMMON block.
+    SubtreeState &getOrCreateNamedSubtree(mlir::StringAttr name);
+
   private:
     SubtreeState(mlir::MLIRContext *ctx, std::string name,
                  mlir::LLVM::TBAANodeAttr grandParent)
@@ -57,6 +63,9 @@ struct TBAATree {
     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;
   };
 
   /// A subtree for POINTER/TARGET variables data.
@@ -131,22 +140,29 @@ class TBAAForrest {
   // 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 const TBAATree &getFuncTreeWithScope(mlir::func::FuncOp func,
-                                              llvm::StringRef scope) {
+  inline TBAATree &getMutableFuncTreeWithScope(mlir::func::FuncOp func,
+                                               llvm::StringRef scope) {
     mlir::StringAttr name = func.getSymNameAttr();
     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:
-  const TBAATree &getFuncTree(mlir::StringAttr symName) {
+  TBAATree &getFuncTree(mlir::StringAttr symName) {
     if (!separatePerFunction)
       symName = mlir::StringAttr::get(symName.getContext(), "");
     if (!trees.contains(symName))
       trees.insert({symName, TBAATree::buildTree(symName)});
-    return trees.at(symName);
+    auto it = trees.find(symName);
+    assert(it != trees.end());
+    return it->second;
   }
 
   // Should each function use a different tree?

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

@@ -365,7 +365,12 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
   // Linkage helpers (inline). The default linkage is external.
   //===--------------------------------------------------------------------===//
 
-  mlir::StringAttr createCommonLinkage() { return getStringAttr("common"); }
+  static mlir::StringAttr createCommonLinkage(mlir::MLIRContext *context) {
+    return mlir::StringAttr::get(context, "common");
+  }
+  mlir::StringAttr createCommonLinkage() {
+    return createCommonLinkage(getContext());
+  }
 
   mlir::StringAttr createInternalLinkage() { return getStringAttr("internal"); }
 

flang/lib/Optimizer/Analysis/TBAAForest.cpp

@@ -11,12 +11,21 @@
 
 mlir::LLVM::TBAATagAttr
 fir::TBAATree::SubtreeState::getTag(llvm::StringRef uniqueName) const {
-  std::string id = (parentId + "/" + uniqueName).str();
+  std::string id = (parentId + '/' + uniqueName).str();
   mlir::LLVM::TBAATypeDescriptorAttr type =
       mlir::LLVM::TBAATypeDescriptorAttr::get(
           context, id, mlir::LLVM::TBAAMemberAttr::get(parent, 0));
   return mlir::LLVM::TBAATagAttr::get(type, type, 0);
-  // return tag;
+}
+
+fir::TBAATree::SubtreeState &
+fir::TBAATree::SubtreeState::getOrCreateNamedSubtree(mlir::StringAttr name) {
+  if (!namedSubtrees.contains(name))
+    namedSubtrees.insert(
+        {name, SubtreeState(context, parentId + '/' + name.str(), parent)});
+  auto it = namedSubtrees.find(name);
+  assert(it != namedSubtrees.end());
+  return it->second;
 }
 
 mlir::LLVM::TBAATagAttr fir::TBAATree::SubtreeState::getTag() const {

flang/lib/Optimizer/Transforms/AddAliasTags.cpp

@@ -14,6 +14,7 @@
 
 #include "flang/Optimizer/Analysis/AliasAnalysis.h"
 #include "flang/Optimizer/Analysis/TBAAForest.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/Dialect/FirAliasTagOpInterface.h"
 #include "flang/Optimizer/Transforms/Passes.h"
@@ -61,8 +62,10 @@ namespace {
 class PassState {
 public:
   PassState(mlir::DominanceInfo &domInfo,
-            std::optional<unsigned> localAllocsThreshold)
-      : domInfo(domInfo), localAllocsThreshold(localAllocsThreshold) {}
+            std::optional<unsigned> localAllocsThreshold,
+            const mlir::SymbolTable &symTab)
+      : domInfo(domInfo), localAllocsThreshold(localAllocsThreshold),
+        symTab(symTab) {}
   /// memoised call to fir::AliasAnalysis::getSource
   inline const fir::AliasAnalysis::Source &getSource(mlir::Value value) {
     if (!analysisCache.contains(value))
@@ -72,13 +75,14 @@ class PassState {
   }
 
   /// get the per-function TBAATree for this function
-  inline const fir::TBAATree &getFuncTree(mlir::func::FuncOp func) {
-    return forrest[func];
+  inline fir::TBAATree &getMutableFuncTreeWithScope(mlir::func::FuncOp func,
+                                                    fir::DummyScopeOp scope) {
+    auto &scopeMap = scopeNames.at(func);
+    return forrest.getMutableFuncTreeWithScope(func, scopeMap.lookup(scope));
   }
   inline const fir::TBAATree &getFuncTreeWithScope(mlir::func::FuncOp func,
                                                    fir::DummyScopeOp scope) {
-    auto &scopeMap = scopeNames.at(func);
-    return forrest.getFuncTreeWithScope(func, scopeMap.lookup(scope));
+    return getMutableFuncTreeWithScope(func, scope);
   }
 
   void processFunctionScopes(mlir::func::FuncOp func);
@@ -98,8 +102,14 @@ class PassState {
   // attachment.
   bool attachLocalAllocTag();
 
+  fir::GlobalOp getGlobalDefiningOp(mlir::StringAttr name) const {
+    return symTab.lookup<fir::GlobalOp>(name);
+  }
+
 private:
   mlir::DominanceInfo &domInfo;
+  std::optional<unsigned> localAllocsThreshold;
+  const mlir::SymbolTable &symTab;
   fir::AliasAnalysis analysis;
   llvm::DenseMap<mlir::Value, fir::AliasAnalysis::Source> analysisCache;
   fir::TBAAForrest forrest;
@@ -117,8 +127,6 @@ class PassState {
   // Local pass cache for derived types that contain descriptor
   // member(s), to avoid the cost of isRecordWithDescriptorMember().
   llvm::DenseSet<mlir::Type> typesContainingDescriptors;
-
-  std::optional<unsigned> localAllocsThreshold;
 };
 
 // Process fir.dummy_scope operations in the given func:
@@ -310,14 +318,55 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
              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");
-    if (source.isPointer())
+    mlir::StringAttr name = glbl.getRootReference();
+    LLVM_DEBUG(llvm::dbgs().indent(2) << "Found reference to global "
+                                      << name.str() << " at " << *op << "\n");
+    if (source.isPointer()) {
       tag = state.getFuncTreeWithScope(func, scopeOp).targetDataTree.getTag();
-    else
-      tag =
-          state.getFuncTreeWithScope(func, scopeOp).globalDataTree.getTag(name);
+    } else {
+      // In general, place the tags under the "global data" root.
+      fir::TBAATree::SubtreeState *subTree =
+          &state.getMutableFuncTreeWithScope(func, scopeOp).globalDataTree;
+
+      // The COMMON blocks have their own sub-tree root under the "global data"
+      // root, which is named after the name of the COMMON block.
+      // If we can identify the name of the member variable, then
+      // we create a sub-tree under the root of the COMMON block
+      // and place the tag there. If we cannot identify the name
+      // of the member variable (e.g. for whatever reason there is no
+      // fir.declare for it), then we place the tag under the root
+      // of the COMMON block.
+      auto globalOp = state.getGlobalDefiningOp(name);
+      // TODO: this is a subtle identification of the fact that
+      // the variable belongs to a COMMON block.
+      // Should we have an attribute on [hl]fir.declare
+      // that specifies the name of the COMMON block the variable
+      // belongs to?
+      if (globalOp &&
+          globalOp.getLinkName() ==
+              fir::FirOpBuilder::createCommonLinkage(globalOp->getContext())) {
+        // Get or create a sub-tree for the COMMON block.
+        subTree = &subTree->getOrCreateNamedSubtree(name);
+
+        auto declOp = mlir::dyn_cast_or_null<fir::DeclareOp>(
+            source.origin.instantiationPoint);
+        mlir::StringAttr varName;
+        if (declOp) {
+          // The tag for the variable will be placed under its own
+          // root in the COMMON sub-tree.
+          varName = declOp.getUniqName();
+          tag = subTree->getTag(varName.str());
+        } else {
+          tag = subTree->getTag();
+        }
+        LLVM_DEBUG(llvm::dbgs().indent(2)
+                   << "Variable named '"
+                   << (varName ? varName.str() : "<unknown>")
+                   << "' is from COMMON block '" << name.str() << "'\n");
+      } else {
+        tag = subTree->getTag(name.str());
+      }
+    }
 
     // TBAA for global variables with descriptors
   } else if (enableDirect &&
@@ -401,11 +450,14 @@ void AddAliasTagsPass::runOnOperation() {
   // thinks the pass operates on), then the real work of the pass is done in
   // runOnAliasInterface
   auto &domInfo = getAnalysis<mlir::DominanceInfo>();
-  PassState state(domInfo, localAllocsThreshold.getPosition()
-                               ? std::optional<unsigned>(localAllocsThreshold)
-                               : std::nullopt);
-
   mlir::ModuleOp mod = getOperation();
+  mlir::SymbolTable symTab(mod);
+  PassState state(domInfo,
+                  localAllocsThreshold.getPosition()
+                      ? std::optional<unsigned>(localAllocsThreshold)
+                      : std::nullopt,
+                  symTab);
+
   mod.walk(
       [&](fir::FirAliasTagOpInterface op) { runOnAliasInterface(op, state); });