[mlir] correctly handle "unknown" state in activity analysis (#1571)

This requires injecting ModRef information about library functions.

enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp

@@ -456,6 +456,51 @@ static bool mayWriteArg(FunctionOpInterface callee, unsigned argNo,
   return !hasReadOnlyAttr && !hasReadNoneAttr && funcMayWrite;
 }
 
+/// Returns information indicating whether the function may read or write into
+/// the memory pointed to by its arguments. When unknown, returns `nullopt`.
+static std::optional<LLVM::ModRefInfo>
+getFunctionArgModRef(FunctionOpInterface func) {
+  // First, handle some library functions with statically known behavior.
+  StringRef name = cast<SymbolOpInterface>(func.getOperation()).getName();
+  auto hardcoded = llvm::StringSwitch<std::optional<LLVM::ModRefInfo>>(name)
+                       // printf: only reads from arguments.
+                       .Case("printf", LLVM::ModRefInfo::Ref)
+                       // operator delete(void *) doesn't read from arguments.
+                       .Case("_ZdlPv", LLVM::ModRefInfo::NoModRef)
+                       .Default(std::nullopt);
+  if (hardcoded)
+    return hardcoded;
+
+  if (auto memoryAttr =
+          func->getAttrOfType<LLVM::MemoryEffectsAttr>(kLLVMMemoryAttrName))
+    return memoryAttr.getArgMem();
+  return std::nullopt;
+}
+
+/// Returns information indicating whether the function may read or write into
+/// the memory other than that pointed to by its arguments, though still
+/// accessible from (any) calling context. When unknown, returns `nullopt`.
+static std::optional<LLVM::ModRefInfo>
+getFunctionOtherModRef(FunctionOpInterface func) {
+  // First, handle some library functions with statically known behavior.
+  StringRef name = cast<SymbolOpInterface>(func.getOperation()).getName();
+  auto hardcoded =
+      llvm::StringSwitch<std::optional<LLVM::ModRefInfo>>(name)
+          // printf: doesn't access other (technically, stdout is pointer-like,
+          // but we cannot flow information through it since it is write-only.
+          .Case("printf", LLVM::ModRefInfo::NoModRef)
+          // operator delete(void *) doesn't access other.
+          .Case("_ZdlPv", LLVM::ModRefInfo::NoModRef)
+          .Default(std::nullopt);
+  if (hardcoded)
+    return hardcoded;
+
+  if (auto memoryAttr =
+          func->getAttrOfType<LLVM::MemoryEffectsAttr>(kLLVMMemoryAttrName))
+    return memoryAttr.getOther();
+  return std::nullopt;
+}
+
 void enzyme::PointsToPointerAnalysis::visitCallControlFlowTransfer(
     CallOpInterface call, CallControlFlowAction action,
     const PointsToSets &before, PointsToSets *after) {
@@ -497,13 +542,10 @@ void enzyme::PointsToPointerAnalysis::visitCallControlFlowTransfer(
     //     into pointers that are non-arguments.
     if (auto callee = SymbolTable::lookupNearestSymbolFrom<FunctionOpInterface>(
             call, symbol.getLeafReference())) {
-      auto memoryAttr =
-          callee->getAttrOfType<LLVM::MemoryEffectsAttr>(kLLVMMemoryAttrName);
-      std::optional<LLVM::ModRefInfo> argModRef =
-          memoryAttr ? std::make_optional(memoryAttr.getArgMem())
-                     : std::nullopt;
+      std::optional<LLVM::ModRefInfo> argModRef = getFunctionArgModRef(callee);
       std::optional<LLVM::ModRefInfo> otherModRef =
-          memoryAttr ? std::make_optional(memoryAttr.getOther()) : std::nullopt;
+          getFunctionOtherModRef(callee);
+
       SmallVector<int> pointerLikeOperands;
       for (auto &&[i, operand] : llvm::enumerate(call.getArgOperands())) {
         if (isPointerLike(operand.getType()))

enzyme/Enzyme/MLIR/Analysis/DataFlowActivityAnalysis.cpp

@@ -794,6 +794,20 @@ void traverseCallGraph(FunctionOpInterface root,
   }
 }
 
+static const enzyme::AliasClassSet &
+getDefaultPointsTo(const enzyme::PointsToSets &pointsToSets) {
+  // Get the default points-to alias class set, which is where the
+  // "unknown" and any other unlisted class set points to.
+  const enzyme::AliasClassSet &defaultPointsTo =
+      pointsToSets.getPointsTo(nullptr);
+  // Unknown class can point to unknown or nothing, unless further
+  // refined.
+  assert((defaultPointsTo.isUnknown() ||
+          defaultPointsTo.getAliasClasses().empty()) &&
+         "new case introduced for AliasClassSet?");
+  return defaultPointsTo;
+}
+
 void printActivityAnalysisResults(const DataFlowSolver &solver,
                                   FunctionOpInterface callee,
                                   const SmallPtrSet<Operation *, 2> &returnOps,
@@ -815,19 +829,33 @@ void printActivityAnalysisResults(const DataFlowSolver &solver,
       auto *bma = solver.lookupState<BackwardMemoryActivity>(
           &callee.getFunctionBody().front().front());
 
-      auto *pointsToSets =
+      const enzyme::PointsToSets *pointsToSets =
           solver.lookupState<enzyme::PointsToSets>(*returnOps.begin());
       auto *aliasClassLattice = solver.lookupState<AliasClassLattice>(value);
       // Traverse the points-to sets in a simple BFS
       std::deque<DistinctAttr> frontier;
       DenseSet<DistinctAttr> visited;
-      // TODO(zinenko): FIXME, handle unknown...
-      if (!aliasClassLattice->isUnknown()) {
+      auto scheduleVisit = [&](auto range) {
+        for (DistinctAttr neighbor : range) {
+          if (!visited.contains(neighbor)) {
+            visited.insert(neighbor);
+            frontier.push_back(neighbor);
+          }
+        }
+      };
+
+      if (aliasClassLattice->isUnknown()) {
+        // If this pointer is in unknown alias class, it may point to active
+        // data if the unknown alias class is known to point to something and
+        // may not point to active data if the unknown alias class is known not
+        // to point to anything.
+        auto &defaultPointsTo = getDefaultPointsTo(*pointsToSets);
+        return !defaultPointsTo.isUnknown() &&
+               defaultPointsTo.getAliasClasses().empty();
+      } else {
         const DenseSet<DistinctAttr> &aliasClasses =
             aliasClassLattice->getAliasClasses();
-        frontier.insert(frontier.end(), aliasClasses.begin(),
-                        aliasClasses.end());
-        visited.insert(aliasClasses.begin(), aliasClasses.end());
+        scheduleVisit(aliasClasses);
       }
       while (!frontier.empty()) {
         DistinctAttr aliasClass = frontier.front();
@@ -841,19 +869,19 @@ void printActivityAnalysisResults(const DataFlowSolver &solver,
 
         // Or if it points to a pointer that points to active data.
         if (pointsToSets->getPointsTo(aliasClass).isUnknown()) {
-          // TODO(zinenko): FIXME handle unknown. Conservative assumption here
-          // is to assume the value is active (or unknown if we can return
-          // that). Is there a less conservative option?
+          // If a pointer points to an unknown alias set, query the default
+          // points-to alias set (which also applies to the unknown alias set).
+          auto &defaultPointsTo = getDefaultPointsTo(*pointsToSets);
+          // If it is in turn unknown, conservatively assume the pointer may be
+          // pointing to some active data.
+          if (defaultPointsTo.isUnknown())
+            return false;
+          // Otherwise look at classes pointed to by unknown (which can only be
+          // an empty set as of time of writing).
+          scheduleVisit(defaultPointsTo.getAliasClasses());
           continue;
         }
-        const DenseSet<DistinctAttr> &neighbors =
-            pointsToSets->getPointsTo(aliasClass).getAliasClasses();
-        for (DistinctAttr neighbor : neighbors) {
-          if (!visited.contains(neighbor)) {
-            visited.insert(neighbor);
-            frontier.push_back(neighbor);
-          }
-        }
+        scheduleVisit(pointsToSets->getPointsTo(aliasClass).getAliasClasses());
       }
       // Otherwise, it's constant
       return true;