From 12cc0d6195f2a627ca74dfd3c46984675eea6118 Mon Sep 17 00:00:00 2001
From: Jacob Peng <jacobmpeng@gmail.com>
Date: Mon, 4 Mar 2024 15:55:46 -0500
Subject: [PATCH] Add AssumeOp to no-op list, LogOp to fully active ops

---
 .../MLIR/Analysis/ActivityAnnotations.cpp     | 29 +++++++++++++++----
 enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp |  9 ++++--
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/enzyme/Enzyme/MLIR/Analysis/ActivityAnnotations.cpp b/enzyme/Enzyme/MLIR/Analysis/ActivityAnnotations.cpp
index 047394896abe..1bf1442c9854 100644
--- a/enzyme/Enzyme/MLIR/Analysis/ActivityAnnotations.cpp
+++ b/enzyme/Enzyme/MLIR/Analysis/ActivityAnnotations.cpp
@@ -66,16 +66,21 @@ void enzyme::ForwardActivityAnnotationAnalysis::setToEntryState(
 
 /// True iff all results differentially depend on all operands
 // TODO: differential dependency/activity interface
+// TODO: Select cond is not fully active
 static bool isFullyActive(Operation *op) {
   return isa<LLVM::FMulOp, LLVM::FAddOp, LLVM::FDivOp, LLVM::FSubOp,
              LLVM::FNegOp, LLVM::FAbsOp, LLVM::SqrtOp, LLVM::SinOp, LLVM::CosOp,
-             LLVM::Exp2Op, LLVM::ExpOp, LLVM::InsertValueOp,
-             LLVM::ExtractValueOp, LLVM::BitcastOp>(op);
+             LLVM::Exp2Op, LLVM::ExpOp, LLVM::LogOp, LLVM::InsertValueOp,
+             LLVM::ExtractValueOp, LLVM::BitcastOp, LLVM::SelectOp>(op);
 }
 
 static bool isNoOp(Operation *op) {
   return isa<LLVM::NoAliasScopeDeclOp, LLVM::LifetimeStartOp,
-             LLVM::LifetimeEndOp>(op);
+             LLVM::LifetimeEndOp, LLVM::AssumeOp>(op);
+}
+
+static bool isPossiblyActive(Type type) {
+  return isa<FloatType, ComplexType>(type);
 }
 
 void enzyme::ForwardActivityAnnotationAnalysis::visitOperation(
@@ -139,8 +144,10 @@ void enzyme::ForwardActivityAnnotationAnalysis::processMemoryRead(
   // those origins to the read results.
   for (DistinctAttr srcClass : srcClasses->getAliasClasses()) {
     for (ForwardOriginsLattice *result : results) {
-      propagateIfChanged(result,
-                         result->merge(originsMap->getOrigins(srcClass)));
+      if (isPossiblyActive(result->getPoint().getType())) {
+        propagateIfChanged(result,
+                           result->merge(originsMap->getOrigins(srcClass)));
+      }
     }
   }
 }
@@ -189,6 +196,12 @@ void enzyme::ForwardActivityAnnotationAnalysis::visitExternalCall(
                                          results);
     }
   }
+
+  // In the absence of a summary attribute, assume all results differentially
+  // depend on all operands
+  for (ForwardOriginsLattice *result : results)
+    for (const ForwardOriginsLattice *operand : operands)
+      join(result, *operand);
 }
 
 void enzyme::ForwardActivityAnnotationAnalysis::processCallToSummarizedFunc(
@@ -291,6 +304,12 @@ void enzyme::BackwardActivityAnnotationAnalysis::visitExternalCall(
                                          results);
     }
   }
+
+  // In the absence of a summary attribute, assume all results differentially
+  // depend on all operands
+  for (BackwardOriginsLattice *operand : operands)
+    for (const BackwardOriginsLattice *result : results)
+      meet(operand, *result);
 }
 
 void enzyme::BackwardActivityAnnotationAnalysis::processCallToSummarizedFunc(
diff --git a/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp b/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp
index 660e2df11140..4687e9270851 100644
--- a/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp
+++ b/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp
@@ -268,6 +268,12 @@ static bool isMustStore(Operation *op, Value pointer) {
   return false; // isa<LLVM::StoreOp>(op);
 }
 
+// TODO: This should be integrated into an interface somewhere
+static bool isNoOp(Operation *op) {
+  return isa<LLVM::NoAliasScopeDeclOp, LLVM::LifetimeStartOp,
+             LLVM::LifetimeEndOp, LLVM::AssumeOp>(op);
+}
+
 void enzyme::PointsToPointerAnalysis::visitOperation(Operation *op,
                                                      const PointsToSets &before,
                                                      PointsToSets *after) {
@@ -277,8 +283,7 @@ void enzyme::PointsToPointerAnalysis::visitOperation(Operation *op,
   // fixpoint and bail.
   auto memory = dyn_cast<MemoryEffectOpInterface>(op);
   if (!memory) {
-    if (isa<LLVM::NoAliasScopeDeclOp, LLVM::LifetimeStartOp,
-            LLVM::LifetimeEndOp>(op)) {
+    if (isNoOp(op)) {
       // Treat this as a no-op
       return;
     }