diff --git a/enzyme/Enzyme/MLIR/Passes/CMakeLists.txt b/enzyme/Enzyme/MLIR/Passes/CMakeLists.txt
index f5297b5fd531..620a3f0e1542 100644
--- a/enzyme/Enzyme/MLIR/Passes/CMakeLists.txt
+++ b/enzyme/Enzyme/MLIR/Passes/CMakeLists.txt
@@ -6,6 +6,7 @@ add_mlir_doc(Passes EnzymePasses ./ -gen-pass-doc)
 add_mlir_dialect_library(MLIREnzymeTransforms
     LowerAffineAtomicRmwPass.cpp
     EnzymeMLIRPass.cpp
+    EnzymeAcitivityOpt.cpp
     ExpandImpulsePass.cpp
     EnzymeBatchPass.cpp
     EnzymeBatchDiffPass.cpp
@@ -50,6 +51,7 @@ add_mlir_dialect_library(MLIREnzymeTransforms
   MLIRMemRefDialect
   MLIRNVVMDialect
   MLIRPass
+  MLIREnzymeAnalysis
   MLIRSideEffectInterfaces
   MLIRControlFlowInterfaces
   MLIRSCFToControlFlow
diff --git a/enzyme/Enzyme/MLIR/Passes/EnzymeAcitivityOpt.cpp b/enzyme/Enzyme/MLIR/Passes/EnzymeAcitivityOpt.cpp
new file mode 100644
index 000000000000..f54ecc150940
--- /dev/null
+++ b/enzyme/Enzyme/MLIR/Passes/EnzymeAcitivityOpt.cpp
@@ -0,0 +1,558 @@
+//===- EnzymeAcitivtyOpt.cpp - Optimize activity for differentiation ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "Analysis/DataFlowActivityAnalysis.h"
+#include "Analysis/DataFlowAliasAnalysis.h"
+#include "Dialect/Ops.h"
+#include "Interfaces/AutoDiffTypeInterface.h"
+#include "Interfaces/EnzymeLogic.h"
+#include "PassDetails.h"
+#include "Passes/Passes.h"
+
+#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlowFramework.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "llvm/Support/raw_ostream.h"
+
+#include <Analysis/ActivityAnalysis.h>
+#include <deque>
+
+#define DEBUG_TYPE "enzyme-activity-opt"
+#define ENZYME_DBGS llvm::dbgs() << "[" << DEBUG_TYPE << "]"
+
+using namespace mlir;
+using namespace mlir::enzyme;
+using namespace mlir::dataflow;
+
+namespace mlir {
+namespace enzyme {
+#define GEN_PASS_DEF_ENZYMEACTIVITYOPT
+#include "Passes/Passes.h.inc"
+} // namespace enzyme
+} // namespace mlir
+
+namespace {
+
+struct TaggedActivity {
+  Attribute tag;
+  bool isConstant;
+};
+
+struct ActivityOptimization {
+  SmallVector<Activity> argumentActivity;
+  SmallVector<Activity> resultActivity;
+  SmallVector<TaggedActivity> taggedActivities;
+};
+
+static void
+printTaggedActivityAnalysis(FunctionOpInterface callee,
+                            ArrayRef<TaggedActivity> taggedActivities) {
+  llvm::errs() << FlatSymbolRefAttr::get(callee) << ":\n";
+  for (const TaggedActivity &activity : taggedActivities) {
+    llvm::errs() << "  " << activity.tag << ": "
+                 << (activity.isConstant ? "Constant" : "Active") << "\n";
+  }
+}
+
+// NOTE: Currently, we only optimize for return activities in the classical
+// activity analyzer. This is because the ActivityAnalyzer computes and runs
+// activity analysis based on a fixed assignment to argument activities
+static FailureOr<ActivityOptimization> computeClassicalActivityOptimization(
+    Operation *diffOp, FunctionOpInterface callee,
+    ArrayRef<Activity> argumentActivity, ArrayRef<Activity> resultActivity,
+    SymbolTableCollection &symbolTable) {
+  SmallPtrSet<Block *, 4> blocksNotForAnalysis;
+  mlir::enzyme::MTypeResults TR;
+  SmallPtrSet<mlir::Value, 1> constargvals;
+  SmallPtrSet<mlir::Value, 1> activeargvals;
+
+  // TODO: check if we need to specialize here for fwddiff and autodiff
+  for (auto &&[arg, act] :
+       llvm::zip(callee.getFunctionBody().getArguments(), argumentActivity)) {
+    if (act == enzyme::Activity::enzyme_const)
+      constargvals.insert(arg);
+    else
+      activeargvals.insert(arg);
+  }
+
+  SmallVector<DIFFE_TYPE> ReturnActivity;
+  for (auto act : resultActivity) {
+    if (act != enzyme::Activity::enzyme_const)
+      ReturnActivity.push_back(DIFFE_TYPE::DUP_ARG);
+    else
+      ReturnActivity.push_back(DIFFE_TYPE::CONSTANT);
+  }
+
+  DenseMap<Operation *, bool> readOnlyCache;
+  enzyme::ActivityAnalyzer activityAnalyzer(blocksNotForAnalysis, readOnlyCache,
+                                            constargvals, activeargvals,
+                                            ReturnActivity);
+
+  // walk thru all return ops and set activity
+  SmallVector<bool> isRetValueConstant(resultActivity.size(), true);
+  callee.walk([&](Operation *op) {
+    if (!op->hasTrait<OpTrait::ReturnLike>() ||
+        op->getParentOp() != callee.getOperation())
+      return;
+
+    for (const auto &[idx, opret] : llvm::enumerate(op->getOperands())) {
+      bool icv = activityAnalyzer.isConstantValue(TR, opret);
+      isRetValueConstant[idx] &= icv;
+    }
+  });
+
+  // Initialize activity results to be the same as that of the callsite
+  ActivityOptimization optimized;
+  optimized.argumentActivity.assign(argumentActivity.begin(),
+                                    argumentActivity.end());
+  optimized.resultActivity.assign(resultActivity.begin(), resultActivity.end());
+
+  for (auto &&[idx, ract] : llvm::enumerate(optimized.resultActivity)) {
+    if (isRetValueConstant[idx] && (ract != Activity::enzyme_const ||
+                                    ract != Activity::enzyme_constnoneed)) {
+      if (ract == Activity::enzyme_dup || ract == Activity::enzyme_active) {
+        ract = Activity::enzyme_const;
+      } else if (ract == Activity::enzyme_activenoneed ||
+                 ract == Activity::enzyme_dupnoneed) {
+        ract = Activity::enzyme_constnoneed;
+      }
+    }
+  }
+
+  // set return activity based on final inferred activity
+  return optimized;
+}
+
+static FailureOr<ActivityOptimization> computeDataFlowActivityOptimization(
+    Operation *diffOp, FunctionOpInterface callee,
+    ArrayRef<Activity> argumentActivity, ArrayRef<Activity> resultActivity,
+    SymbolTableCollection &symbolTable, bool collectActivityAnalysis = false) {
+
+  Region &body = callee.getFunctionBody();
+  if (body.empty() || body.front().empty())
+    return failure();
+  if (callee.getArguments().size() != argumentActivity.size()) {
+    diffOp->emitError()
+        << "callee argument count does not match activity count";
+    return failure();
+  }
+
+  DataFlowSolver solver;
+  solver.load<enzyme::PointsToPointerAnalysis>();
+  solver.load<enzyme::AliasAnalysis>(callee.getContext());
+  solver.load<enzyme::SparseForwardActivityAnalysis>();
+  solver.load<enzyme::DenseForwardActivityAnalysis>(&body.front(),
+                                                    argumentActivity);
+  solver.load<enzyme::SparseBackwardActivityAnalysis>(symbolTable);
+  solver.load<enzyme::DenseBackwardActivityAnalysis>(
+      symbolTable, callee, argumentActivity, resultActivity);
+  solver.load<DeadCodeAnalysis>();
+  solver.load<SparseConstantPropagation>();
+
+  for (const auto &[arg, activity] :
+       llvm::zip(callee.getArguments(), argumentActivity)) {
+    // we should only initialize for enzyme_active, enzyme_dup, enzyme_dupnoneed
+    if (activity != Activity::enzyme_active &&
+        activity != Activity::enzyme_dup &&
+        activity != Activity::enzyme_dupnoneed)
+      continue;
+    auto *argLattice =
+        solver.getOrCreateState<enzyme::ForwardValueActivity>(arg);
+    (void)argLattice->join(ValueActivity::getActiveVal());
+  }
+
+  SmallVector<Operation *> returnOps;
+  bool invalidReturn = false;
+  callee->walk(
+      [&](Operation *op) {
+        if (!op->hasTrait<OpTrait::ReturnLike>() ||
+            op->getParentOp() != callee.getOperation())
+          return;
+
+        if (op->getNumOperands() != resultActivity.size()) {
+          invalidReturn = true;
+          diffOp->emitError()
+              << "return operand count does not match ret_activity count";
+          return;
+        }
+
+        returnOps.push_back(op);
+        for (const auto &[operand, activity] :
+             llvm::zip(op->getOperands(), resultActivity)) {
+          // Seed scalar return activity here; dense memory activity handles
+          // pointed-to data for duplicated memory returns.
+          auto *returnLattice =
+              solver.getOrCreateState<enzyme::BackwardValueActivity>(operand);
+          if (activity == Activity::enzyme_active ||
+              activity == Activity::enzyme_activenoneed ||
+              activity == Activity::enzyme_dup ||
+              activity == Activity::enzyme_dupnoneed) {
+            (void)returnLattice->meet(ValueActivity::getActiveVal());
+          }
+        }
+      });
+
+  if (invalidReturn)
+    return failure();
+
+  if (returnOps.empty())
+    return failure();
+
+  Operation *analysisRoot = callee->getParentOfType<ModuleOp>();
+  if (!analysisRoot)
+    analysisRoot = callee.getOperation();
+  if (failed(solver.initializeAndRun(analysisRoot))) {
+    diffOp->emitError() << "dataflow activity analysis failed";
+    return failure();
+  }
+
+  auto isActiveData = [&](Value value) {
+    auto *forward = solver.lookupState<enzyme::ForwardValueActivity>(value);
+    auto *backward = solver.lookupState<enzyme::BackwardValueActivity>(value);
+    bool forwardActive = forward && forward->getValue().isActiveVal();
+    bool backwardActive = backward && backward->getValue().isActiveVal();
+    return forwardActive && backwardActive;
+  };
+
+  auto isConstantPointer = [&](Value value) {
+    auto *backwardMemory = solver.lookupState<enzyme::BackwardMemoryActivity>(
+        solver.getProgramPointBefore(&body.front().front()));
+    auto *aliasClassLattice =
+        solver.lookupState<enzyme::AliasClassLattice>(value);
+
+    if (!backwardMemory || !aliasClassLattice ||
+        aliasClassLattice->isUndefined() || aliasClassLattice->isUnknown())
+      return false;
+
+    auto scheduleVisit = [](std::deque<DistinctAttr> &frontier,
+                            DenseSet<DistinctAttr> &visited,
+                            const enzyme::AliasClassSet &aliasClasses) {
+      bool known = true;
+      (void)aliasClasses.foreachElement(
+          [&](DistinctAttr neighbor, enzyme::AliasClassSet::State state) {
+            if (state != enzyme::AliasClassSet::State::Defined || !neighbor) {
+              known = false;
+              return ChangeResult::NoChange;
+            }
+            if (!visited.contains(neighbor)) {
+              visited.insert(neighbor);
+              frontier.push_back(neighbor);
+            }
+            return ChangeResult::NoChange;
+          });
+      return known;
+    };
+
+    for (Operation *returnOp : returnOps) {
+      auto *forwardMemory = solver.lookupState<enzyme::ForwardMemoryActivity>(
+          solver.getProgramPointAfter(returnOp));
+      const enzyme::PointsToSets *pointsToSets =
+          solver.lookupState<enzyme::PointsToSets>(
+              solver.getProgramPointAfter(returnOp));
+      if (!forwardMemory || !pointsToSets)
+        return false;
+
+      std::deque<DistinctAttr> frontier;
+      DenseSet<DistinctAttr> visited;
+      if (!scheduleVisit(frontier, visited,
+                         aliasClassLattice->getAliasClassesObject()))
+        return false;
+
+      while (!frontier.empty()) {
+        DistinctAttr aliasClass = frontier.front();
+        frontier.pop_front();
+
+        if (forwardMemory->hasActiveData(aliasClass) &&
+            backwardMemory->activeDataFlowsOut(aliasClass))
+          return false;
+
+        const enzyme::AliasClassSet &pointsTo =
+            pointsToSets->getPointsTo(aliasClass);
+        if (pointsTo.isUndefined() || pointsTo.isUnknown())
+          return false;
+        if (!scheduleVisit(frontier, visited, pointsTo))
+          return false;
+      }
+    }
+
+    return true;
+  };
+
+  auto isConstantValue = [&](Value value) {
+    if (isa<LLVM::LLVMPointerType, MemRefType, UnrankedMemRefType>(
+            value.getType()))
+      return isConstantPointer(value);
+    return !isActiveData(value);
+  };
+
+  // Initialize activity results to be the same as that of the callsite
+  ActivityOptimization optimized;
+  optimized.argumentActivity.assign(argumentActivity.begin(),
+                                    argumentActivity.end());
+  optimized.resultActivity.assign(resultActivity.begin(), resultActivity.end());
+
+  for (const auto &[index, arg, activity] :
+       llvm::enumerate(callee.getArguments(), argumentActivity)) {
+    if (isConstantValue(arg)) {
+      if (activity == Activity::enzyme_active ||
+          activity == Activity::enzyme_dup ||
+          activity == Activity::enzyme_dupnoneed)
+        optimized.argumentActivity[index] = Activity::enzyme_const;
+    }
+  }
+
+  // Conservatively combine all the activities of all ReturnLike ops as a
+  // summary for ret_activity originating from the callsite autodiff op
+  SmallVector<bool> allReturnValuesConstant(resultActivity.size(), true);
+  for (Operation *returnOp : returnOps) {
+    for (const auto &[index, operand] :
+         llvm::enumerate(returnOp->getOperands())) {
+      allReturnValuesConstant[index] =
+          allReturnValuesConstant[index] && isConstantValue(operand);
+    }
+  }
+
+  for (const auto &[index, activity] : llvm::enumerate(resultActivity)) {
+    if (allReturnValuesConstant[index]) {
+      if (activity == Activity::enzyme_active ||
+          activity == Activity::enzyme_dup)
+        optimized.resultActivity[index] = Activity::enzyme_const;
+      else if (activity == Activity::enzyme_activenoneed ||
+               activity == Activity::enzyme_dupnoneed)
+        optimized.resultActivity[index] = Activity::enzyme_constnoneed;
+    }
+  }
+
+  if (collectActivityAnalysis) {
+    for (BlockArgument arg : callee.getArguments()) {
+      if (Attribute tag = callee.getArgAttr(arg.getArgNumber(), "enzyme.tag"))
+        optimized.taggedActivities.push_back({tag, isConstantValue(arg)});
+    }
+
+    callee.walk([&](Operation *op) {
+      Attribute tag = op->getAttr("tag");
+      if (!tag)
+        return;
+      for (OpResult result : op->getResults())
+        optimized.taggedActivities.push_back({tag, isConstantValue(result)});
+    });
+  }
+
+  return optimized;
+}
+
+class MinimizeForwardActivity : public OpRewritePattern<ForwardDiffOp> {
+public:
+  MinimizeForwardActivity(MLIRContext *context, bool dataflow,
+                          bool printActivityAnalysis,
+                          PatternBenefit benefit = 1)
+      : OpRewritePattern(context, benefit), dataflow(dataflow),
+        printActivityAnalysis(printActivityAnalysis) {}
+
+  LogicalResult matchAndRewrite(ForwardDiffOp op,
+                                PatternRewriter &rewriter) const override {
+    auto *ctx = rewriter.getContext();
+    SymbolTableCollection symbolTable;
+    FunctionOpInterface callee = dyn_cast_or_null<FunctionOpInterface>(
+        symbolTable.lookupNearestSymbolFrom(op, op.getFnAttr()));
+    if (!callee)
+      return failure();
+
+    SmallVector<Activity> argumentActivity =
+        llvm::to_vector(op.getActivity().getAsValueRange<ActivityAttr>());
+    SmallVector<Activity> resultActivity =
+        llvm::to_vector(op.getRetActivity().getAsValueRange<ActivityAttr>());
+
+    FailureOr<ActivityOptimization> optimized =
+        dataflow
+            ? computeDataFlowActivityOptimization(op, callee, argumentActivity,
+                                                  resultActivity, symbolTable,
+                                                  printActivityAnalysis)
+            : computeClassicalActivityOptimization(op, callee, argumentActivity,
+                                                   resultActivity, symbolTable);
+
+    if (failed(optimized))
+      return failure();
+
+    // debug printing activity
+    if (printActivityAnalysis && dataflow)
+      printTaggedActivityAnalysis(callee, optimized->taggedActivities);
+
+    if (llvm::equal(argumentActivity, optimized->argumentActivity) &&
+        llvm::equal(resultActivity, optimized->resultActivity))
+      return failure();
+
+    // reconstruct op according to modified activity
+    SmallVector<Value> in_args;
+    auto in_idx = 0;
+    for (const auto &[idx, act, nact] :
+         llvm::enumerate(argumentActivity, optimized->argumentActivity)) {
+      Value inp = op.getInputs()[in_idx];
+      if (act == nact) {
+        in_args.push_back(inp);
+        in_idx++;
+        if (act == Activity::enzyme_dup || act == Activity::enzyme_dupnoneed) {
+          inp = op.getInputs()[in_idx];
+          in_args.push_back(inp);
+          in_idx++;
+        }
+      } else {
+        // Handle const demotion for enzyme_dup/enzyme_dupnoneed/enzyme_active
+        in_args.push_back(inp);
+        in_idx++;
+        // we skip the duplicated derivative arg
+        if ((act == Activity::enzyme_dup ||
+             act == Activity::enzyme_dupnoneed) &&
+            nact == Activity::enzyme_const)
+          in_idx++;
+      }
+    }
+
+    SmallVector<Value> out_args;
+    auto out_idx = 0;
+    for (const auto &[idx, act, nact] :
+         llvm::enumerate(resultActivity, optimized->resultActivity)) {
+      if (act == Activity::enzyme_constnoneed)
+        continue;
+
+      Value out = op.getOutputs()[out_idx];
+      if (act == nact) {
+        switch (act) {
+        case Activity::enzyme_active:
+        case Activity::enzyme_activenoneed:
+        case Activity::enzyme_const:
+          out_args.push_back(out);
+          out_idx++;
+          break;
+
+        case Activity::enzyme_dup:
+        case Activity::enzyme_dupnoneed:
+          out_args.push_back(out);
+          out_idx++;
+
+          // handle derivative
+          out = op.getOutputs()[out_idx];
+          out_args.push_back(out);
+          out_idx++;
+          break;
+
+        default:
+          llvm_unreachable("unexpected return activity arg");
+        }
+      } else {
+        // NOTE: These are the only possible activity demotions
+        // enzyme_active -> enzyme_const
+        // enzyme_dup -> enzyme_const
+        // enzyme_activenoneed -> enzyme_constnoneed
+        // enzyme_dupnoneed -> enzyme_constnoneed
+        switch (act) {
+        case Activity::enzyme_active:
+        case Activity::enzyme_dup:
+          out_args.push_back(out);
+          // we skip the derivative
+          out_idx += 2;
+          break;
+
+        case Activity::enzyme_dupnoneed:
+        case Activity::enzyme_activenoneed:
+          // the primal is already skipped. We also skip the derivative.
+          out_idx++;
+          break;
+
+        default:
+          llvm_unreachable("unexpected demotion for return activity");
+        }
+      }
+    }
+
+    TypeRange out_ty = ValueRange(out_args).getTypes();
+    ArrayAttr newInActivity = ArrayAttr::get(
+        ctx, llvm::map_to_vector(optimized->argumentActivity,
+                                 [&](enzyme::Activity act) -> Attribute {
+                                   return ActivityAttr::get(ctx, act);
+                                 }));
+
+    ArrayAttr newOutActivity = ArrayAttr::get(
+        ctx, llvm::map_to_vector(optimized->resultActivity,
+                                 [&](enzyme::Activity act) -> Attribute {
+                                   return ActivityAttr::get(ctx, act);
+                                 }));
+
+    ForwardDiffOp newOp = ForwardDiffOp::create(
+        rewriter, op->getLoc(), out_ty, op.getFnAttr(), in_args, newInActivity,
+        newOutActivity, op.getWidthAttr(), op.getStrongZeroAttr());
+
+    // Map old uses of op to newOp
+    auto oldIdx = 0;
+    auto newIdx = 0;
+    for (auto [idx, old_val, new_val] :
+         llvm::enumerate(resultActivity, optimized->resultActivity)) {
+
+      if (old_val == new_val) {
+        // don't index into op if its a const_noneed
+        if (old_val == Activity::enzyme_constnoneed) {
+          continue;
+        }
+        // replace use
+        op.getOutputs()[oldIdx++].replaceAllUsesWith(
+            newOp.getOutputs()[newIdx++]);
+        if (old_val == Activity::enzyme_dup) {
+          // 2nd replacement for derivative
+          op.getOutputs()[oldIdx++].replaceAllUsesWith(
+              newOp.getOutputs()[newIdx++]);
+        }
+      } else {
+        // handle all substitutions
+        if (new_val == Activity::enzyme_constnoneed &&
+            old_val == Activity::enzyme_dupnoneed) {
+          ++oldIdx; // skip gradient too
+        } else if (new_val == Activity::enzyme_const &&
+                   old_val == Activity::enzyme_dup) {
+          op.getOutputs()[oldIdx++].replaceAllUsesWith(
+              newOp.getOutputs()[newIdx++]);
+          ++oldIdx; // skip derivative
+        } else {
+          llvm_unreachable("unexpected demotion for return activity");
+        }
+      }
+    }
+
+    rewriter.eraseOp(op);
+    return success();
+  }
+
+private:
+  bool dataflow;
+  bool printActivityAnalysis;
+};
+
+struct EnzymeActivityOptPass
+    : public enzyme::impl::EnzymeActivityOptBase<EnzymeActivityOptPass> {
+  using EnzymeActivityOptBase::EnzymeActivityOptBase;
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    patterns.add<MinimizeForwardActivity>(&getContext(), dataflow,
+                                          printActivityAnalysis);
+
+    GreedyRewriteConfig config;
+    config.enableFolding();
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns),
+                                     config))) {
+      signalPassFailure();
+    }
+  }
+};
+
+} // namespace
diff --git a/enzyme/Enzyme/MLIR/Passes/Passes.td b/enzyme/Enzyme/MLIR/Passes/Passes.td
index f835b8002182..475a98cf3e53 100644
--- a/enzyme/Enzyme/MLIR/Passes/Passes.td
+++ b/enzyme/Enzyme/MLIR/Passes/Passes.td
@@ -81,6 +81,29 @@ def BatchDiffPass : Pass<"enzyme-diff-batch">{
   ];
 }
 
+def EnzymeActivityOpt : Pass<"enzyme-activity-opt">{
+  let summary = "Optimize activity for differentiation";
+  let dependentDialects = [
+    "enzyme::EnzymeDialect",
+  ];
+  let options = [
+    Option<
+      /*C++ variable name=*/"dataflow",
+      /*CLI argument=*/"dataflow",
+      /*type=*/"bool",
+      /*default=*/"false",
+      /*description=*/"Whether to use the new Dataflow activity analysis"
+    >,
+    Option<
+      /*C++ variable name=*/"printActivityAnalysis",
+      /*CLI argument=*/"print",
+      /*type=*/"bool",
+      /*default=*/"false",
+      /*description=*/"Print the dataflow activity analysis results used by the optimizer"
+    >,
+  ];
+}
+
 def DifferentiateWrapperPass : Pass<"enzyme-wrap"> {
   let summary = "Add wrapper function to be differentiated";
   let dependentDialects = [
diff --git a/enzyme/test/MLIR/OptimizeAD/optimize_activity.mlir b/enzyme/test/MLIR/OptimizeAD/optimize_activity.mlir
new file mode 100644
index 000000000000..c2225c98933e
--- /dev/null
+++ b/enzyme/test/MLIR/OptimizeAD/optimize_activity.mlir
@@ -0,0 +1,44 @@
+// RUN: %eopt --enzyme-activity-opt --split-input-file %s | FileCheck %s --check-prefix=CLASSIC
+// RUN: %eopt --enzyme-activity-opt="dataflow=true" --split-input-file %s | FileCheck %s --check-prefix=DATAFLOW
+
+module {
+  func.func @square2(%x: f32 {enzyme.tag = "x"},
+                     %y: f32 {enzyme.tag = "y"}) -> (f32, f32) {
+    %p = arith.mulf %x, %x {tag = "p"} : f32
+    %q = arith.mulf %y, %y {tag = "q"} : f32
+    return %p, %q : f32, f32
+  }
+
+  func.func @fwd_return_activity(%x: f32, %y: f32, %dy: f32) -> f32 {
+    %p, %dp = enzyme.fwddiff @square2(%x, %y, %dy)
+        {activity = [#enzyme<activity enzyme_const>,
+                     #enzyme<activity enzyme_dup>],
+         ret_activity = [#enzyme<activity enzyme_dup>,
+                         #enzyme<activity enzyme_constnoneed>]}
+        : (f32, f32, f32) -> (f32, f32)
+    return %p : f32
+  }
+
+  func.func @fwd_argument_activity(%x: f32, %y: f32, %dx: f32, %dy: f32)
+      -> (f32, f32, f32) {
+    %p, %q, %dq = enzyme.fwddiff @square2(%x, %dx, %y, %dy)
+        {activity = [#enzyme<activity enzyme_dup>,
+                     #enzyme<activity enzyme_dup>],
+         ret_activity = [#enzyme<activity enzyme_const>,
+                         #enzyme<activity enzyme_dup>]}
+        : (f32, f32, f32, f32) -> (f32, f32, f32)
+    return %p, %q, %dq : f32, f32, f32
+  }
+}
+
+// CLASSIC-LABEL: func.func @fwd_return_activity
+// CLASSIC: enzyme.fwddiff @square2(%arg0, %arg1, %arg2) {{.*}}activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_dup>]{{.*}}ret_activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_constnoneed>]
+
+// CLASSIC-LABEL: func.func @fwd_argument_activity
+// CLASSIC: enzyme.fwddiff @square2(%arg0, %arg2, %arg1, %arg3) {{.*}}activity = [#enzyme<activity enzyme_dup>, #enzyme<activity enzyme_dup>]{{.*}}ret_activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_dup>]
+
+// DATAFLOW-LABEL: func.func @fwd_return_activity
+// DATAFLOW: enzyme.fwddiff @square2(%arg0, %arg1) {{.*}}activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_const>]{{.*}}ret_activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_constnoneed>]
+
+// DATAFLOW-LABEL: func.func @fwd_argument_activity
+// DATAFLOW: enzyme.fwddiff @square2(%arg0, %arg1, %arg3) {{.*}}activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_dup>]{{.*}}ret_activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_dup>]