[flang][FIR] add FIR TBAA pass

See RFC at
https://discourse.llvm.org/t/rfc-propagate-fir-alias-analysis-information-using-tbaa/73755

This pass adds TBAA tags to all accesses to non-pointer/target dummy
arguments. These TBAA tags tell LLVM that these accesses cannot alias:
allowing better dead code elimination, hoisting out of loops, and
vectorization.

Each function has its own TBAA tree so that accesses between funtions
MayAlias after inlining.

I also included code for adding tags for local allocations and for
global variables. Enabling all three kinds of tag is known to produce a
miscompile and so these are disabled by default. But it isn't much
code and I thought it could be interesting to play with these later if
one is looking at a benchmark which looks like it would benefit from
more alias information. I'm open to removing this code too.

TBAA tags are also added separately by TBAABuilder during CodeGen.
TBAABuilder has to run during CodeGen because it adds tags to box
accesses, many of which are implicit in FIR. This pass cannot (easily)
run in CodeGen because fir::AliasAnalysis has difficulty tracing values
between blocks, and by the time CodeGen runs, structured control flow
has already been lowered.

Coming in follow up patches
  - Change CodeGen/TBAABuilder to use TBAAForest to add tags within the
    same per-function trees as are used here (delayed to a later patch
    to make it easier to revert)
  - Command line argument processing to actually enable the pass

Author: tblah

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

@@ -0,0 +1,104 @@
+//===-- 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 "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/MLIRContext.h"
+#include "llvm/ADT/DenseMap.h"
+#include <string>
+
+namespace fir {
+
+//===----------------------------------------------------------------------===//
+// TBAATree
+//===----------------------------------------------------------------------===//
+/// Per-function TBAA tree. Each tree contains branches for data (of various
+/// kinds) and descriptor access
+struct TBAATree {
+  //===----------------------------------------------------------------------===//
+  // TBAAForrest::TBAATree::SubtreeState
+  //===----------------------------------------------------------------------===//
+  /// This contains a TBAA subtree based on some parent. New tags can be added
+  /// under the parent using getTag.
+  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;
+
+  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;
+    llvm::DenseMap<llvm::StringRef, mlir::LLVM::TBAATagAttr> tagDedup;
+  };
+
+  SubtreeState globalDataTree;
+  SubtreeState allocatedDataTree;
+  SubtreeState dummyArgDataTree;
+  mlir::LLVM::TBAATypeDescriptorAttr anyAccessDesc;
+  mlir::LLVM::TBAATypeDescriptorAttr boxMemberTypeDesc;
+  mlir::LLVM::TBAATypeDescriptorAttr anyDataTypeDesc;
+
+  static TBAATree buildTree(mlir::StringAttr functionName);
+
+private:
+  TBAATree(mlir::LLVM::TBAATypeDescriptorAttr anyAccess,
+           mlir::LLVM::TBAATypeDescriptorAttr dataRoot,
+           mlir::LLVM::TBAATypeDescriptorAttr boxMemberTypeDesc);
+};
+
+//===----------------------------------------------------------------------===//
+// TBAAForrest
+//===----------------------------------------------------------------------===//
+/// Collection of TBAATrees, usually indexed by function (so that each function
+/// has a different TBAATree)
+class TBAAForrest {
+public:
+  explicit TBAAForrest(bool separatePerFunction = true)
+      : separatePerFunction{separatePerFunction} {}
+
+  inline const TBAATree &operator[](mlir::func::FuncOp func) {
+    return getFuncTree(func.getSymNameAttr());
+  }
+  inline const TBAATree &operator[](mlir::LLVM::LLVMFuncOp func) {
+    return getFuncTree(func.getSymNameAttr());
+  }
+
+private:
+  const 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);
+  }
+
+  // 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

flang/include/flang/Optimizer/Transforms/Passes.h

@@ -61,6 +61,7 @@ std::unique_ptr<mlir::Pass> createMemDataFlowOptPass();
 std::unique_ptr<mlir::Pass> createPromoteToAffinePass();
 std::unique_ptr<mlir::Pass> createMemoryAllocationPass();
 std::unique_ptr<mlir::Pass> createStackArraysPass();
+std::unique_ptr<mlir::Pass> createAliasTagsPass();
 std::unique_ptr<mlir::Pass> createSimplifyIntrinsicsPass();
 std::unique_ptr<mlir::Pass> createAddDebugFoundationPass();
 std::unique_ptr<mlir::Pass> createLoopVersioningPass();

flang/include/flang/Optimizer/Transforms/Passes.td

@@ -252,6 +252,26 @@ def StackArrays : Pass<"stack-arrays", "mlir::ModuleOp"> {
   let constructor = "::fir::createStackArraysPass()";
 }
 
+def AddAliasTags : Pass<"fir-add-alias-tags", "mlir::ModuleOp"> {
+  let summary = "Add tbaa tags to operations that implement FirAliasAnalysisOpInterface";
+  let description = [{
+    TBAA (type based alias analysis) is one method to pass pointer alias information
+    from language frontends to LLVM. This pass uses fir::AliasAnalysis to add this
+    information to fir.load and fir.store operations.
+    Additional tags are added during codegen. See fir::TBAABuilder.
+    This needs to be a separate pass so that it happens before structured control
+    flow operations are lowered to branches and basic blocks (this makes tracing
+    the source of values much eaiser). The other TBAA tags need to be applied to
+    box loads and stores which are implicit in FIR and so cannot be annotated
+    until codegen.
+    TODO: this is currently a pass on mlir::ModuleOp to avoid parallelism. In
+    theory, each operation could be considered in prallel, so long as there
+    aren't races adding new tags to the mlir context.
+  }];
+  let dependentDialects = [ "fir::FIROpsDialect" ];
+  let constructor = "::fir::createAliasTagsPass()";
+}
+
 def SimplifyRegionLite : Pass<"simplify-region-lite", "mlir::ModuleOp"> {
   let summary = "Region simplification";
   let description = [{

flang/lib/Optimizer/Analysis/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_flang_library(FIRAnalysis
   AliasAnalysis.cpp
+  TBAAForest.cpp
 
   DEPENDS
   FIRDialect

flang/lib/Optimizer/Analysis/TBAAForest.cpp

@@ -0,0 +1,60 @@
+//===- TBAAForest.cpp - Per-functon TBAA Trees ----------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Analysis/TBAAForest.h"
+#include <mlir/Dialect/LLVMIR/LLVMAttrs.h>
+
+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(
+          context, id, mlir::LLVM::TBAAMemberAttr::get(parent, 0));
+  return mlir::LLVM::TBAATagAttr::get(type, type, 0);
+  // return tag;
+}
+
+fir::TBAATree fir::TBAATree::buildTree(mlir::StringAttr func) {
+  llvm::StringRef funcName = func.getValue();
+  std::string rootId = ("Flang function root " + funcName).str();
+  mlir::MLIRContext *ctx = func.getContext();
+  mlir::LLVM::TBAARootAttr funcRoot =
+      mlir::LLVM::TBAARootAttr::get(ctx, mlir::StringAttr::get(ctx, rootId));
+
+  static constexpr llvm::StringRef anyAccessTypeDescId = "any access";
+  mlir::LLVM::TBAATypeDescriptorAttr anyAccess =
+      mlir::LLVM::TBAATypeDescriptorAttr::get(
+          ctx, anyAccessTypeDescId,
+          mlir::LLVM::TBAAMemberAttr::get(funcRoot, 0));
+
+  static constexpr llvm::StringRef anyDataAccessTypeDescId = "any data access";
+  mlir::LLVM::TBAATypeDescriptorAttr dataRoot =
+      mlir::LLVM::TBAATypeDescriptorAttr::get(
+          ctx, anyDataAccessTypeDescId,
+          mlir::LLVM::TBAAMemberAttr::get(anyAccess, 0));
+
+  static constexpr llvm::StringRef boxMemberTypeDescId = "descriptor member";
+  mlir::LLVM::TBAATypeDescriptorAttr boxMemberTypeDesc =
+      mlir::LLVM::TBAATypeDescriptorAttr::get(
+          ctx, boxMemberTypeDescId,
+          mlir::LLVM::TBAAMemberAttr::get(anyAccess, 0));
+
+  return TBAATree{anyAccess, dataRoot, boxMemberTypeDesc};
+}
+
+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),
+      dummyArgDataTree(dataRoot.getContext(), "dummy arg data", dataRoot),
+      anyAccessDesc(anyAccess), boxMemberTypeDesc(boxMemberTypeDesc),
+      anyDataTypeDesc(dataRoot) {}