Initialize DenseBackwardActivityAnalysis with ret_activity (#2874)

Author: vimarsh6739

enzyme/Enzyme/MLIR/Analysis/DataFlowActivityAnalysis.cpp

@@ -582,21 +582,28 @@ LogicalResult enzyme::DenseBackwardActivityAnalysis::visitOperation(
       propagateIfChanged(before, changed);
     }
 
+    assert(op->getNumOperands() == returnActivity.size() &&
+           "the number of returned arguments must match the number of return "
+           "activity entries");
+
     // Initialize the return activity of the operands
-    for (Value operand : op->getOperands()) {
-      if (isa<MemRefType, LLVM::LLVMPointerType>(operand.getType())) {
-        auto *retAliasClasses = getOrCreateFor<AliasClassLattice>(
-            getProgramPointBefore(op), operand);
-        ChangeResult changed =
-            retAliasClasses->getAliasClassesObject().foreachElement(
-                [before](DistinctAttr retAliasClass,
-                         enzyme::AliasClassSet::State state) {
-                  if (state == enzyme::AliasClassSet::State::Undefined)
-                    return ChangeResult::NoChange;
-                  return before->setActiveOut(retAliasClass);
-                });
-        propagateIfChanged(before, changed);
-      }
+    for (const auto &[operand, operandActivity] :
+         llvm::zip(op->getOperands(), returnActivity)) {
+      if (operandActivity != enzyme::Activity::enzyme_dup &&
+          operandActivity != enzyme::Activity::enzyme_dupnoneed)
+        continue;
+
+      auto *retAliasClasses =
+          getOrCreateFor<AliasClassLattice>(getProgramPointBefore(op), operand);
+      ChangeResult changed =
+          retAliasClasses->getAliasClassesObject().foreachElement(
+              [before](DistinctAttr retAliasClass,
+                       enzyme::AliasClassSet::State state) {
+                if (state == enzyme::AliasClassSet::State::Undefined)
+                  return ChangeResult::NoChange;
+                return before->setActiveOut(retAliasClass);
+              });
+      propagateIfChanged(before, changed);
     }
   }
 
@@ -906,7 +913,8 @@ void printActivityAnalysisResults(DataFlowSolver &solver,
 
 void enzyme::runDataFlowActivityAnalysis(
     FunctionOpInterface callee, ArrayRef<enzyme::Activity> argumentActivity,
-    bool print, bool verbose, bool annotate) {
+    ArrayRef<enzyme::Activity> returnActivity, bool print, bool verbose,
+    bool annotate) {
   SymbolTableCollection symbolTable;
   DataFlowSolver solver;
 
@@ -917,7 +925,7 @@ void enzyme::runDataFlowActivityAnalysis(
                                             argumentActivity);
   solver.load<SparseBackwardActivityAnalysis>(symbolTable);
   solver.load<DenseBackwardActivityAnalysis>(symbolTable, callee,
-                                             argumentActivity);
+                                             argumentActivity, returnActivity);
 
   // Required for the dataflow framework to traverse region-based control flow
   solver.load<DeadCodeAnalysis>();

enzyme/Enzyme/MLIR/Analysis/DataFlowActivityAnalysis.h

@@ -299,9 +299,10 @@ class DenseBackwardActivityAnalysis
   DenseBackwardActivityAnalysis(DataFlowSolver &solver,
                                 SymbolTableCollection &symbolTable,
                                 FunctionOpInterface parentOp,
-                                ArrayRef<enzyme::Activity> argumentActivity)
+                                ArrayRef<enzyme::Activity> argumentActivity,
+                                ArrayRef<enzyme::Activity> returnActivity)
       : DenseBackwardDataFlowAnalysis(solver, symbolTable), parentOp(parentOp),
-        argumentActivity(argumentActivity) {}
+        argumentActivity(argumentActivity), returnActivity(returnActivity) {}
 
   LogicalResult visitOperation(Operation *op,
                                const BackwardMemoryActivity &after,
@@ -319,10 +320,12 @@ class DenseBackwardActivityAnalysis
 private:
   FunctionOpInterface parentOp;
   SmallVector<enzyme::Activity> argumentActivity;
+  SmallVector<enzyme::Activity> returnActivity;
 };
 
 void runDataFlowActivityAnalysis(FunctionOpInterface callee,
                                  ArrayRef<enzyme::Activity> argumentActivity,
+                                 ArrayRef<enzyme::Activity> returnActivity,
                                  bool print = false, bool verbose = false,
                                  bool annotate = false);
 

enzyme/Enzyme/MLIR/Passes/PrintActivityAnalysis.cpp

@@ -145,6 +145,7 @@ struct PrintActivityAnalysisPass
       enzyme::runActivityAnnotations(callee, argActivities, config);
     } else if (config.dataflow) {
       enzyme::runDataFlowActivityAnalysis(callee, argActivities,
+                                          resultActivities,
                                           /*print=*/true, verbose, annotate);
     } else {