From f7a46fd41562e13a622613f63b12fa67418213bf Mon Sep 17 00:00:00 2001
From: "Oleksandr \"Alex\" Zinenko" <zinenko@google.com>
Date: Tue, 27 Feb 2024 23:46:11 +0100
Subject: [PATCH] [mlir] move alias lattice update to relevant transfer
 function (#1573)

---
 enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp | 21 ++++++++++++-------
 .../MLIR/AliasAnalysis/func_attributes.mlir   |  2 +-
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp b/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp
index 767a60a3ed2b..9af864e592d6 100644
--- a/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp
+++ b/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp
@@ -366,6 +366,19 @@ void enzyme::PointsToPointerAnalysis::visitOperation(Operation *op,
 
   SmallVector<MemoryEffects::EffectInstance> effects;
   memory.getEffects(effects);
+
+  // If the operation allocates fresh memory and doesn't write into it, that
+  // memory is known not to point to any known alias class.
+  if (effects.size() == 1 &&
+      isa<MemoryEffects::Allocate>(effects.front().getEffect()) &&
+      effects.front().getValue()) {
+    const auto *destClasses =
+        getOrCreateFor<AliasClassLattice>(op, effects.front().getValue());
+    propagateIfChanged(
+        after, after->setPointingToEmpty(destClasses->getAliasClassesObject()));
+    return;
+  }
+
   for (const auto &effect : effects) {
     if (!isa<MemoryEffects::Write>(effect.getEffect()))
       continue;
@@ -860,14 +873,6 @@ void enzyme::AliasAnalysis::transfer(
               result->getPoint(),
               originalClasses.getOriginalClass(result->getPoint(), debugLabel));
           propagateIfChanged(result, result->join(fresh));
-
-          // The pointer to freshly allocated memory is known not to point to
-          // anything.
-          // TODO(zinenko): this is a bit strange to update _another_ lattice
-          // here.
-          auto *pointsTo = getOrCreate<PointsToSets>(op);
-          propagateIfChanged(pointsTo, pointsTo->setPointingToEmpty(
-                                           fresh.getAliasClassesObject()));
         }
       }
     } else if (isa<MemoryEffects::Read>(effect.getEffect())) {
diff --git a/enzyme/test/MLIR/AliasAnalysis/func_attributes.mlir b/enzyme/test/MLIR/AliasAnalysis/func_attributes.mlir
index 87c7e4105a2b..e1f4b74a21f4 100644
--- a/enzyme/test/MLIR/AliasAnalysis/func_attributes.mlir
+++ b/enzyme/test/MLIR/AliasAnalysis/func_attributes.mlir
@@ -115,7 +115,7 @@ func.func private @callee(%ptr : !llvm.ptr {llvm.readonly}) attributes {
 }
 
 // CHECK: points-to-pointer sets
-// CHECK-NEXT: <empty>
+// CHECK: <empty>
 // CHECK-LABEL @call_other_none_arg_rw_readonly
 func.func @call_other_none_arg_rw_readonly(%input: !llvm.ptr {enzyme.tag = "input"}) {
   call @callee(%input) : (!llvm.ptr) -> ()