[flang] Use fir.declare/fir.dummy_scope for TBAA tags attachments. (#92472)

vzakhari · web-flow · commit 6cd86d0fae8c · 2024-06-04T08:33:40.000-07:00
With MLIR inlining (e.g. `flang-new -mmlir -inline-all=true`)
the current TBAA tags attachment is suboptimal, because
we may lose information about the callee's dummy arguments
(by bypassing fir.declare in AliasAnalysis::getSource).
This is a conservative first step to improve the situation.
This patch makes AddAliasTagsPass to account for fir.dummy_scope
hierarchy after MLIR inlining and use it to place the TBAA tags
into TBAA trees corresponding to different function scopes.
The pass uses special mode of AliasAnalysis to find the instantiation
point of a Fortran variable (a [hl]fir.decalre) when searching
for the source of a memory reference. In this mode, AliasAnalysis
will always stop at fir.declare operations that have dummy_scope
operands - there should not be a reason to past throught it
for the purpose of TBAA tags attachment.
diff --git a/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h b/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h
@@ -120,6 +120,17 @@ struct AliasAnalysis {
       /// Source definition of a value.
       SourceUnion u;
 
+      /// A value definition denoting the place where the corresponding
+      /// source variable was instantiated by the front-end.
+      /// Currently, it is the result of [hl]fir.declare of the source,
+      /// if we can reach it.
+      /// It helps to identify the scope where the corresponding variable
+      /// was defined in the original Fortran source, e.g. when MLIR
+      /// inlining happens an inlined fir.declare of the callee's
+      /// dummy argument identifies the scope where the source
+      /// may be treated as a dummy argument.
+      mlir::Value instantiationPoint;
+
       /// Whether the source was reached following data or box reference
       bool isData{false};
     };
@@ -168,7 +179,10 @@ struct AliasAnalysis {
   mlir::ModRefResult getModRef(mlir::Operation *op, mlir::Value location);
 
   /// Return the memory source of a value.
-  Source getSource(mlir::Value);
+  /// If getInstantiationPoint is true, the search for the source
+  /// will stop at [hl]fir.declare if it represents a dummy
+  /// argument declaration (i.e. it has the dummy_scope operand).
+  Source getSource(mlir::Value, bool getInstantiationPoint = false);
 };
 
 inline bool operator==(const AliasAnalysis::Source::SourceOrigin &lhs,
diff --git a/flang/include/flang/Optimizer/Analysis/TBAAForest.h b/flang/include/flang/Optimizer/Analysis/TBAAForest.h
@@ -92,6 +92,20 @@ class TBAAForrest {
     }
     return getFuncTree(func.getSymNameAttr());
   }
+  // 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 const TBAATree &getFuncTreeWithScope(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);
+  }
 
 private:
   const TBAATree &getFuncTree(mlir::StringAttr symName) {
diff --git a/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp b/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
@@ -95,6 +95,9 @@ bool AliasAnalysis::Source::isRecordWithPointerComponent() const {
 }
 
 AliasResult AliasAnalysis::alias(Value lhs, Value rhs) {
+  // TODO: alias() has to be aware of the function scopes.
+  // After MLIR inlining, the current implementation may
+  // not recognize non-aliasing entities.
   auto lhsSrc = getSource(lhs);
   auto rhsSrc = getSource(rhs);
   bool approximateSource = lhsSrc.approximateSource || rhsSrc.approximateSource;
@@ -232,7 +235,8 @@ getAttrsFromVariable(fir::FortranVariableOpInterface var) {
   return attrs;
 }
 
-AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v) {
+AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v,
+                                               bool getInstantiationPoint) {
   auto *defOp = v.getDefiningOp();
   SourceKind type{SourceKind::Unknown};
   mlir::Type ty;
@@ -244,6 +248,7 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v) {
   bool followingData = !isBoxRef;
   mlir::SymbolRefAttr global;
   Source::Attributes attributes;
+  mlir::Value instantiationPoint;
   while (defOp && !breakFromLoop) {
     ty = defOp->getResultTypes()[0];
     llvm::TypeSwitch<Operation *>(defOp)
@@ -334,6 +339,21 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v) {
             breakFromLoop = true;
             return;
           }
+          if (getInstantiationPoint) {
+            // Fetch only the innermost instantiation point.
+            if (!instantiationPoint)
+              instantiationPoint = op->getResult(0);
+
+            if (op.getDummyScope()) {
+              // Do not track past DeclareOp that has the dummy_scope
+              // operand. This DeclareOp is known to represent
+              // a dummy argument for some runtime instantiation
+              // of a procedure.
+              type = SourceKind::Argument;
+              breakFromLoop = true;
+              return;
+            }
+          }
           // TODO: Look for the fortran attributes present on the operation
           // Track further through the operand
           v = op.getMemref();
@@ -372,9 +392,17 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v) {
     }
 
   if (type == SourceKind::Global) {
-    return {{global, followingData}, type, ty, attributes, approximateSource};
+    return {{global, instantiationPoint, followingData},
+            type,
+            ty,
+            attributes,
+            approximateSource};
   }
-  return {{v, followingData}, type, ty, attributes, approximateSource};
+  return {{v, instantiationPoint, followingData},
+          type,
+          ty,
+          attributes,
+          approximateSource};
 }
 
 } // namespace fir
diff --git a/flang/lib/Optimizer/CodeGen/TBAABuilder.cpp b/flang/lib/Optimizer/CodeGen/TBAABuilder.cpp
@@ -52,6 +52,38 @@ TBAABuilder::TBAABuilder(MLIRContext *context, bool applyTBAA,
                          bool forceUnifiedTree)
     : enableTBAA(applyTBAA && !disableTBAA),
       trees(/*separatePerFunction=*/perFunctionTBAATrees && !forceUnifiedTree) {
+  // TODO: the TBAA tags created here are rooted in the root scope
+  // of the enclosing function. This does not work best with MLIR inlining.
+  // A better approach is to root them according to the scopes they belong to
+  // and that were used by AddAliasTagsPass to create TBAA tags before
+  // the CodeGen. For example:
+  //   subroutine caller(a, b, ptr)
+  //     real, target :: a(:), b(:)
+  //     integer, pointer :: ptr(:)
+  //     call callee(a, b, ptr)
+  //   end
+  //   subroutine callee(a, b, ptr)
+  //     real :: a(:), b(:)
+  //     integer, pointer :: ptr(:)
+  //     do i=...
+  //       a(ptr(i)) = b(ptr(i))
+  //     end do
+  //   end
+  //
+  // When callee is inlined, the dummy arguments 'a' and 'b' will
+  // be rooted in TBAA tree corresponding to the `call callee` call site,
+  // saying that the references to 'a' and 'b' cannot alias each other.
+  // These tags will be created by AddAliasTagsPass, but it will not be able
+  // to create any tags for 'ptr' references.
+  // During the CodeGen, we create 'any data access' tags for the
+  // 'ptr' acceses. If they are rooted within the root scope of `caller`,
+  // they end up in a different TBAA tree with the 'a' and 'b' access
+  // tags, so 'ptr', 'a' and 'b' references MayAlias. Moreover,
+  // the box access of 'ptr' will also be in a different TBAA tree
+  // with 'a' and 'b' tags, meaning they can also alias.
+  // This will prevent LLVM vectorization even with memory conflict checks.
+  // It seems that we'd better move all TBAA tags assignment to
+  // AddAliasTagsPass, which can at least rely on the dummy arguments scopes.
   if (!enableTBAA)
     return;
 }
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -3910,8 +3910,15 @@ std::optional<std::int64_t> fir::getIntIfConstant(mlir::Value value) {
 
 bool fir::isDummyArgument(mlir::Value v) {
   auto blockArg{mlir::dyn_cast<mlir::BlockArgument>(v)};
-  if (!blockArg)
+  if (!blockArg) {
+    auto defOp = v.getDefiningOp();
+    if (defOp) {
+      if (auto declareOp = mlir::dyn_cast<fir::DeclareOp>(defOp))
+        if (declareOp.getDummyScope())
+          return true;
+    }
     return false;
+  }
 
   auto *owner{blockArg.getOwner()};
   return owner->isEntryBlock() &&
diff --git a/flang/lib/Optimizer/Transforms/AddAliasTags.cpp b/flang/lib/Optimizer/Transforms/AddAliasTags.cpp
@@ -17,9 +17,11 @@
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/Dialect/FirAliasTagOpInterface.h"
 #include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/IR/Dominance.h"
 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/Twine.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -54,24 +56,85 @@ namespace {
 /// Shared state per-module
 class PassState {
 public:
+  PassState(mlir::DominanceInfo &domInfo) : domInfo(domInfo) {}
   /// memoised call to fir::AliasAnalysis::getSource
   inline const fir::AliasAnalysis::Source &getSource(mlir::Value value) {
     if (!analysisCache.contains(value))
-      analysisCache.insert({value, analysis.getSource(value)});
+      analysisCache.insert(
+          {value, analysis.getSource(value, /*getInstantiationPoint=*/true)});
     return analysisCache[value];
   }
 
   /// get the per-function TBAATree for this function
   inline const fir::TBAATree &getFuncTree(mlir::func::FuncOp func) {
     return forrest[func];
   }
+  inline const fir::TBAATree &getFuncTreeWithScope(mlir::func::FuncOp func,
+                                                   fir::DummyScopeOp scope) {
+    auto &scopeMap = scopeNames.at(func);
+    return forrest.getFuncTreeWithScope(func, scopeMap.lookup(scope));
+  }
+
+  void processFunctionScopes(mlir::func::FuncOp func);
+  fir::DummyScopeOp getDeclarationScope(fir::DeclareOp declareOp);
 
 private:
+  mlir::DominanceInfo &domInfo;
   fir::AliasAnalysis analysis;
   llvm::DenseMap<mlir::Value, fir::AliasAnalysis::Source> analysisCache;
   fir::TBAAForrest forrest;
+  // Unique names for fir.dummy_scope operations within
+  // the given function.
+  llvm::DenseMap<mlir::func::FuncOp,
+                 llvm::DenseMap<fir::DummyScopeOp, std::string>>
+      scopeNames;
+  // A map providing a vector of fir.dummy_scope operations
+  // for the given function. The vectors are sorted according
+  // to the dominance information.
+  llvm::DenseMap<mlir::func::FuncOp, llvm::SmallVector<fir::DummyScopeOp, 16>>
+      sortedScopeOperations;
 };
 
+// Process fir.dummy_scope operations in the given func:
+// sort them according to the dominance information, and
+// associate a unique (within the current function) scope name
+// with each of them.
+void PassState::processFunctionScopes(mlir::func::FuncOp func) {
+  if (scopeNames.contains(func))
+    return;
+
+  auto &scopeMap = scopeNames.getOrInsertDefault(func);
+  auto &scopeOps = sortedScopeOperations.getOrInsertDefault(func);
+  func.walk([&](fir::DummyScopeOp op) { scopeOps.push_back(op); });
+  llvm::stable_sort(scopeOps, [&](const fir::DummyScopeOp &op1,
+                                  const fir::DummyScopeOp &op2) {
+    return domInfo.properlyDominates(&*op1, &*op2);
+  });
+  unsigned scopeId = 0;
+  for (auto scope : scopeOps) {
+    if (scopeId != 0) {
+      std::string name = (llvm::Twine("Scope ") + llvm::Twine(scopeId)).str();
+      LLVM_DEBUG(llvm::dbgs() << "Creating scope '" << name << "':\n"
+                              << scope << "\n");
+      scopeMap.insert({scope, std::move(name)});
+    }
+    ++scopeId;
+  }
+}
+
+// For the given fir.declare returns the dominating fir.dummy_scope
+// operation.
+fir::DummyScopeOp PassState::getDeclarationScope(fir::DeclareOp declareOp) {
+  auto func = declareOp->getParentOfType<mlir::func::FuncOp>();
+  assert(func && "fir.declare does not have parent func.func");
+  auto &scopeOps = sortedScopeOperations.at(func);
+  for (auto II = scopeOps.rbegin(), IE = scopeOps.rend(); II != IE; ++II) {
+    if (domInfo.dominates(&**II, &*declareOp))
+      return *II;
+  }
+  return nullptr;
+}
+
 class AddAliasTagsPass : public fir::impl::AddAliasTagsBase<AddAliasTagsPass> {
 public:
   void runOnOperation() override;
@@ -85,6 +148,9 @@ class AddAliasTagsPass : public fir::impl::AddAliasTagsBase<AddAliasTagsPass> {
 } // namespace
 
 static fir::DeclareOp getDeclareOp(mlir::Value arg) {
+  if (auto declare =
+          mlir::dyn_cast_or_null<fir::DeclareOp>(arg.getDefiningOp()))
+    return declare;
   for (mlir::Operation *use : arg.getUsers())
     if (fir::DeclareOp declare = mlir::dyn_cast<fir::DeclareOp>(use))
       return declare;
@@ -94,7 +160,7 @@ static fir::DeclareOp getDeclareOp(mlir::Value arg) {
 /// Get the name of a function argument using the "fir.bindc_name" attribute,
 /// or ""
 static std::string getFuncArgName(mlir::Value arg) {
-  // first try getting the name from the hlfir.declare
+  // first try getting the name from the fir.declare
   if (fir::DeclareOp declare = getDeclareOp(arg))
     return declare.getUniqName().str();
 
@@ -139,6 +205,23 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
     return;
   }
 
+  // Process the scopes, if not processed yet.
+  state.processFunctionScopes(func);
+
+  fir::DummyScopeOp scopeOp;
+  if (auto declVal = source.origin.instantiationPoint) {
+    // If the source is a dummy argument within some fir.dummy_scope,
+    // then find the corresponding innermost scope to be used for finding
+    // the right TBAA tree.
+    auto declareOp =
+        mlir::dyn_cast_or_null<fir::DeclareOp>(declVal.getDefiningOp());
+    assert(declareOp && "Instantiation point must be fir.declare");
+    if (auto dummyScope = declareOp.getDummyScope())
+      scopeOp = mlir::cast<fir::DummyScopeOp>(dummyScope.getDefiningOp());
+    if (!scopeOp)
+      scopeOp = state.getDeclarationScope(declareOp);
+  }
+
   mlir::LLVM::TBAATagAttr tag;
   // TBAA for dummy arguments
   if (enableDummyArgs &&
@@ -147,7 +230,8 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
                << "Found reference to dummy argument at " << *op << "\n");
     std::string name = getFuncArgName(source.origin.u.get<mlir::Value>());
     if (!name.empty())
-      tag = state.getFuncTree(func).dummyArgDataTree.getTag(name);
+      tag = state.getFuncTreeWithScope(func, scopeOp)
+                .dummyArgDataTree.getTag(name);
     else
       LLVM_DEBUG(llvm::dbgs().indent(2)
                  << "WARN: couldn't find a name for dummy argument " << *op
@@ -161,7 +245,7 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
     const char *name = glbl.getRootReference().data();
     LLVM_DEBUG(llvm::dbgs().indent(2) << "Found reference to global " << name
                                       << " at " << *op << "\n");
-    tag = state.getFuncTree(func).globalDataTree.getTag(name);
+    tag = state.getFuncTreeWithScope(func, scopeOp).globalDataTree.getTag(name);
 
     // TBAA for SourceKind::Direct
   } else if (enableDirect &&
@@ -172,7 +256,8 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
       const char *name = glbl.getRootReference().data();
       LLVM_DEBUG(llvm::dbgs().indent(2) << "Found reference to direct " << name
                                         << " at " << *op << "\n");
-      tag = state.getFuncTree(func).directDataTree.getTag(name);
+      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
@@ -193,7 +278,8 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
     if (name) {
       LLVM_DEBUG(llvm::dbgs().indent(2) << "Found reference to allocation "
                                         << name << " at " << *op << "\n");
-      tag = state.getFuncTree(func).allocatedDataTree.getTag(*name);
+      tag = state.getFuncTreeWithScope(func, scopeOp)
+                .allocatedDataTree.getTag(*name);
     } else {
       LLVM_DEBUG(llvm::dbgs().indent(2)
                  << "WARN: couldn't find a name for allocation " << *op
@@ -219,7 +305,8 @@ void AddAliasTagsPass::runOnOperation() {
   // Instead this pass stores state per mlir::ModuleOp (which is what MLIR
   // thinks the pass operates on), then the real work of the pass is done in
   // runOnAliasInterface
-  PassState state;
+  auto &domInfo = getAnalysis<mlir::DominanceInfo>();
+  PassState state(domInfo);
 
   mlir::ModuleOp mod = getOperation();
   mod.walk(
diff --git a/flang/lib/Optimizer/Transforms/AddDebugInfo.cpp b/flang/lib/Optimizer/Transforms/AddDebugInfo.cpp
@@ -92,8 +92,13 @@ void AddDebugInfoPass::handleDeclareOp(fir::cg::XDeclareOp declOp,
 
   // FIXME: There may be cases where an argument is processed a bit before
   // DeclareOp is generated. In that case, DeclareOp may point to an
-  // intermediate op and not to BlockArgument. We need to find those cases and
-  // walk the chain to get to the actual argument.
+  // intermediate op and not to BlockArgument.
+  // Moreover, with MLIR inlining we cannot use the BlockArgument
+  // position to identify the original number of the dummy argument.
+  // If we want to keep running AddDebugInfoPass late, the dummy argument
+  // position in the argument list has to be expressed in FIR (e.g. as a
+  // constant attribute of [hl]fir.declare/fircg.ext_declare operation that has
+  // a dummy_scope operand).
   unsigned argNo = 0;
   if (fir::isDummyArgument(declOp.getMemref())) {
     auto arg = llvm::cast<mlir::BlockArgument>(declOp.getMemref());
diff --git a/flang/test/Transforms/tbaa-with-dummy-scope.fir b/flang/test/Transforms/tbaa-with-dummy-scope.fir
diff --git a/flang/test/Transforms/tbaa.fir b/flang/test/Transforms/tbaa.fir
diff --git a/flang/test/Transforms/tbaa2.fir b/flang/test/Transforms/tbaa2.fir
