[MLIR] Add StoreLikeInterface to decouple activity analysis from stores (#2978)
* [MLIR] Add ActiveStoreOpInterface to decouple activity analysis from stores
Activity analysis hard-codes which ops are stores (dyn_cast<LLVM::StoreOp> /
memref::StoreOp) when reasoning about stored-value / pointer activity. That
prevents out-of-tree dialects from participating and duplicates logic per
dialect.
Add enzyme::ActiveStoreOpInterface exposing (getStoredValue, getStoredPointer),
attach it to memref.store and llvm.store, and use it at the potential-active-
store site that already handled both dialects -- now a single dialect-agnostic
branch. Any dialect (e.g. an out-of-tree fir.store / hlfir.assign) can opt in by
attaching the interface.
Behavior-preserving: llvm.store/memref.store return the same value/pointer
through the interface. The MLIR test suite is unchanged (the two
ReverseMode failures are pre-existing on this LLVM build, independent of this
change). The remaining LLVM-specific store sites (which walk llvm.alloca /
llvm.load origins) are intentionally left as-is.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* [MLIR] Rename ActiveStoreOpInterface to StoreLikeInterface
Address review: rename the interface and clarify getStoredPointer's
description to note it returns the base pointer before any in-op offsets.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* [MLIR] Finish dialect-agnostic activity analysis for stores/mutable refs
Extends the StoreLikeInterface generalization to the remaining store sites
in ActivityAnalysis and gates the pointer-like type check on the type
interface, so out-of-tree dialects (e.g. FIR/HLFIR) participate fully in
activity analysis without hard-coded dyn_casts. These are the
Fortran-independent pieces of EnzymeAD/Enzyme#2969.
- isConstantValue: besides LLVM ptr / memref, treat any type whose
AutoDiffTypeInterface reports isMutable() (e.g. !fir.ref) as a reference
that carries active memory. Behavior-preserving in-tree: among core
dialects only memref and LLVM pointer report isMutable(), and both are
already matched by the existing isa<> check.
- isOperationInactiveFromOrigin: dialect-agnostic store branch mirroring
the LLVM::StoreOp case (inactive iff stored value or pointer is constant).
- isValueActivelyStoredOrReturned: dialect-agnostic store branch mirroring
the LLVM::StoreOp case.
Deliberately excludes the Fortran-coupled / behavior-changing parts of
#2969 (FIR/HLFIR models, flang plugin, the CoreDialects registration TU
split, and dropping tensor/linalg from the pass dependent dialects).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* Cleanup PR a bit
* apply clang-format
* Update enzyme/Enzyme/MLIR/Interfaces/AutoDiffOpInterface.td
Co-authored-by: Paul Berg <naydex.mc+github@gmail.com>
* [MLIR] Attach StoreLikeInterface to affine.store
Lets activity analysis reason about affine.store's stored-value/pointer
activity generically, the same way memref.store and llvm.store now do,
instead of requiring a hard-coded dyn_cast. getStoredPointer returns the
base memref (the affine map's indices are applied within the op).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* [MLIR] Add forward-mode conditional-store tests for memref/llvm
The affine.store `@if_then` test in ForwardMode/affine.mlir covers a memref
that is unconditionally initialized with a constant and then conditionally
overwritten with an active value: forward mode must zero-initialize the shadow
before the conditional store, otherwise the not-taken path reads uninitialized
shadow memory and returns a garbage tangent instead of 0.
No equivalent coverage existed for memref.store or llvm.store. Add memref_if.mlir
and llvm_if.mlir as the direct analogues. Both are marked XFAIL: forward-mode
differentiation of memref.store / llvm.store does not currently emit the shadow
zero-initialization, so the tests document the known bug and will XPASS once it
is fixed.
Verified with a local enzymemlir-opt build: FileCheck fails on exactly the
zero-init CHECK-DAG line and matches everywhere else.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* fix constant handling
* clang format
* [MLIR] Zero shadow memory for inactive stores in forward mode
Restore the original store-inactivity condition in activity analysis: a
store is inactive if *either* the stored value or the pointer is inactive,
matching LLVM Enzyme's isInstructionInactiveFromOrigin.
That condition alone made the shadow store disappear, because
MGradientUtils::visitChild skips any constant operation in forward mode, so
memoryIdentityForwardHandler never ran and never emitted the null value it
already knows how to produce for a constant operand. LLVM Enzyme has no such
gate: its forward-mode visitCommonStore returns early only when the pointer
is constant. Mirror that by continuing to visit store-like ops whose pointer
is active.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
* [MLIR] Address review: generalize the forward-mode skip condition
Replace the store-specific storesThroughActivePointer check with the general
condition wsmoses suggested: only skip a constant operation in forward mode if
it is pure, or if all of its operands are constant. A side-effecting op with an
active operand may still need to touch shadow memory, which now also covers
llvm.memcpy / llvm.memset rather than only StoreLikeInterface ops. This also
makes the stale /*iface.hasNoEffect()*/ note redundant.
Retain the "if we are being stored into, not storing this value" comment in
isValueActivelyStoredOrReturned.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Paul Berg <naydex.mc+github@gmail.com>
diff --git a/enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp b/enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp
index fdbbeeb..ff97b21 100644
--- a/enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp
+++ b/enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp
@@ -1814,9 +1814,10 @@
containsPointer = false;
// if (!TR.intType(1, Val, /*errIfNotFound*/ false).isPossiblePointer())
- // TODO: this should be an MLIR type interface connected to type analysis.
- if (!isa<LLVM::LLVMPointerType, MemRefType>(Val.getType()))
+ auto typeIface = dyn_cast<AutoDiffTypeInterface>(Val.getType());
+ if (!typeIface || !typeIface.isMutable()) {
containsPointer = false;
+ }
if (containsPointer && !isValuePotentiallyUsedAsPointer(Val)) {
containsPointer = false;
@@ -2454,21 +2455,11 @@
if (EnzymePrintActivity)
llvm::errs() << "potential active store: " << *op << " Val=" << Val
<< "\n";
- if (auto SI = dyn_cast<LLVM::StoreOp>(op)) {
- bool cop = !Hypothesis->isConstantValue(TR, SI.getValue());
+ if (auto SI = dyn_cast<enzyme::StoreLikeInterface>(op)) {
+ bool cop = !Hypothesis->isConstantValue(TR, SI.getStoredValue());
if (EnzymePrintActivity)
llvm::errs() << " -- store potential activity: " << (int)cop
- << " - " << *SI << " of "
- << " Val=" << Val << "\n";
- potentialStore = true;
- if (cop)
- potentiallyActiveStore = true;
- } else if (auto SI = dyn_cast<memref::StoreOp>(op)) {
- // FIXME: this is a copy-pasta form above to work with MLIR memrefs.
- bool cop = !Hypothesis->isConstantValue(TR, SI.getValueToStore());
- if (EnzymePrintActivity)
- llvm::errs() << " -- store potential activity: " << (int)cop
- << " - " << *SI << " of "
+ << " - " << *op << " of "
<< " Val=" << Val << "\n";
potentialStore = true;
if (cop)
@@ -2790,17 +2781,19 @@
if (EnzymePrintActivity)
llvm::errs() << " < UPSEARCH" << (int)directions << ">" << *op << "\n";
- if (auto store = dyn_cast<LLVM::StoreOp>(op)) {
- if (isConstantValue(TR, store.getValue()) ||
- isConstantValue(TR, store.getAddr())) {
+ // if either src or dst is inactive, there cannot be a transfer of active
+ // values and thus the store is inactive
+ if (auto store = dyn_cast<enzyme::StoreLikeInterface>(op)) {
+ if (isConstantValue(TR, store.getStoredValue()) ||
+ isConstantValue(TR, store.getStoredPointer())) {
if (EnzymePrintActivity)
llvm::errs() << " constant instruction as store operand is inactive"
<< *op << "\n";
return true;
}
if (inactArg) {
- inactArg->insert(store.getValue());
- inactArg->insert(store.getAddr());
+ inactArg->insert(store.getStoredValue());
+ inactArg->insert(store.getStoredPointer());
}
return false;
}
@@ -3734,13 +3727,13 @@
}
}
- if (auto SI = dyn_cast<LLVM::StoreOp>(a)) {
+ if (auto SI = dyn_cast<enzyme::StoreLikeInterface>(a)) {
// If we are being stored into, not storing this value
// this case can be skipped
- if (SI.getValue() != val) {
+ if (SI.getStoredValue() != val) {
if (!ignoreStoresInto) {
- // Storing into active value, return true
- if (!isConstantValue(TR, SI.getValue())) {
+ // Active value stored into `val` (the pointer): `val` is active.
+ if (!isConstantValue(TR, SI.getStoredValue())) {
StoredOrReturnedCache[key] = true;
if (EnzymePrintActivity)
llvm::errs() << " </ASOR" << (int)directions
@@ -3752,8 +3745,8 @@
}
continue;
} else {
- // Storing into active memory, return true
- if (!isConstantValue(TR, SI.getAddr())) {
+ // `val` is stored into active memory: active.
+ if (!isConstantValue(TR, SI.getStoredPointer())) {
StoredOrReturnedCache[key] = true;
if (EnzymePrintActivity)
llvm::errs() << " </ASOR" << (int)directions
diff --git a/enzyme/Enzyme/MLIR/Implementations/AffineAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/AffineAutoDiffOpInterfaceImpl.cpp
index da8b531..dbf2185 100644
--- a/enzyme/Enzyme/MLIR/Implementations/AffineAutoDiffOpInterfaceImpl.cpp
+++ b/enzyme/Enzyme/MLIR/Implementations/AffineAutoDiffOpInterfaceImpl.cpp
@@ -590,6 +590,20 @@
}
};
+// Lets activity analysis treat affine.store generically via StoreLikeInterface.
+// getStoredPointer returns the base memref; the affine map indices apply within
+// the op.
+struct AffineStoreLike
+ : public StoreLikeInterface::ExternalModel<AffineStoreLike,
+ affine::AffineStoreOp> {
+ Value getStoredValue(Operation *op) const {
+ return cast<affine::AffineStoreOp>(op).getValueToStore();
+ }
+ Value getStoredPointer(Operation *op) const {
+ return cast<affine::AffineStoreOp>(op).getMemRef();
+ }
+};
+
struct AffineStoreOpInterfaceReverse
: public ReverseAutoDiffOpInterface::ExternalModel<
AffineStoreOpInterfaceReverse, affine::AffineStoreOp> {
@@ -844,6 +858,7 @@
registerInterfaces(context);
affine::AffineLoadOp::attachInterface<AffineLoadOpInterfaceReverse>(
*context);
+ affine::AffineStoreOp::attachInterface<AffineStoreLike>(*context);
affine::AffineStoreOp::attachInterface<AffineStoreOpInterfaceReverse>(
*context);
affine::AffineForOp::attachInterface<AffineForOpInterfaceReverse>(*context);
diff --git a/enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffRegistration.cpp b/enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffRegistration.cpp
index e4b8f37..0613250 100644
--- a/enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffRegistration.cpp
+++ b/enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffRegistration.cpp
@@ -6,16 +6,11 @@
//
//===----------------------------------------------------------------------===//
//
-// The aggregate registration entry point that attaches Enzyme's autodiff
-// external models for every core upstream dialect. It is deliberately isolated
-// in its own translation unit: because it references every per-dialect
-// registration (Linalg and NVVM included), linking it forces those models --
-// and their dialect symbols -- into the consumer. Tools that want them all
-// (enzymemlir-opt, fir-enzyme-opt, ...) call this; a consumer that only needs a
-// subset (the lean `flang -fc1` plugin, which must not pull in Linalg symbols
-// flang does not export) instead links the individual
-// register*DialectAutoDiffInterface functions it wants and never references
-// this TU, so the Linalg/NVVM models are not linked at all.
+// Aggregate entry point registering autodiff external models for every core
+// upstream dialect. Isolated in its own TU so that linking it pulls in all the
+// per-dialect models (Linalg, NVVM, ...); consumers needing only a subset (e.g.
+// the flang -fc1 plugin, which must not pull in Linalg) link the individual
+// register*DialectAutoDiffInterface functions instead and skip this TU.
//
//===----------------------------------------------------------------------===//
diff --git a/enzyme/Enzyme/MLIR/Implementations/LLVMAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/LLVMAutoDiffOpInterfaceImpl.cpp
index 3b45d63..d76469c 100644
--- a/enzyme/Enzyme/MLIR/Implementations/LLVMAutoDiffOpInterfaceImpl.cpp
+++ b/enzyme/Enzyme/MLIR/Implementations/LLVMAutoDiffOpInterfaceImpl.cpp
@@ -27,6 +27,17 @@
namespace {
#include "Implementations/LLVMDerivatives.inc"
+// Lets activity analysis treat llvm.store generically via StoreLikeInterface.
+struct LLVMStoreLike
+ : public StoreLikeInterface::ExternalModel<LLVMStoreLike, LLVM::StoreOp> {
+ Value getStoredValue(Operation *op) const {
+ return cast<LLVM::StoreOp>(op).getValue();
+ }
+ Value getStoredPointer(Operation *op) const {
+ return cast<LLVM::StoreOp>(op).getAddr();
+ }
+};
+
struct InlineAsmActivityInterface
: public ActivityOpInterface::ExternalModel<InlineAsmActivityInterface,
LLVM::InlineAsmOp> {
@@ -564,6 +575,7 @@
LLVM::LLVMArrayType::attachInterface<ArrayTypeInterface>(*context);
LLVM::SelectOp::attachInterface<SelectActivityInterface>(*context);
+ LLVM::StoreOp::attachInterface<LLVMStoreLike>(*context);
LLVM::LoadOp::attachInterface<LoadOpInterfaceReverse>(*context);
LLVM::StoreOp::attachInterface<StoreOpInterfaceReverse>(*context);
LLVM::GEPOp::attachInterface<GEPOpInterfaceReverse>(*context);
diff --git a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp
index ffbf395..2427bf9 100644
--- a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp
+++ b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp
@@ -27,6 +27,18 @@
namespace {
#include "Implementations/MemRefDerivatives.inc"
+// Lets activity analysis treat memref.store generically via StoreLikeInterface.
+struct MemRefStoreLike
+ : public StoreLikeInterface::ExternalModel<MemRefStoreLike,
+ memref::StoreOp> {
+ Value getStoredValue(Operation *op) const {
+ return cast<memref::StoreOp>(op).getValueToStore();
+ }
+ Value getStoredPointer(Operation *op) const {
+ return cast<memref::StoreOp>(op).getMemRef();
+ }
+};
+
struct LoadOpInterfaceReverse
: public ReverseAutoDiffOpInterface::ExternalModel<LoadOpInterfaceReverse,
memref::LoadOp> {
@@ -312,6 +324,7 @@
MemRefType::attachInterface<MemRefAutoDiffTypeInterface>(*context);
MemRefType::attachInterface<MemRefClonableTypeInterface>(*context);
+ memref::StoreOp::attachInterface<MemRefStoreLike>(*context);
memref::LoadOp::attachInterface<LoadOpInterfaceReverse>(*context);
memref::StoreOp::attachInterface<StoreOpInterfaceReverse>(*context);
memref::SubViewOp::attachInterface<SubViewOpInterfaceReverse>(*context);
diff --git a/enzyme/Enzyme/MLIR/Interfaces/AutoDiffOpInterface.td b/enzyme/Enzyme/MLIR/Interfaces/AutoDiffOpInterface.td
index 81c76f6..f3df11c 100644
--- a/enzyme/Enzyme/MLIR/Interfaces/AutoDiffOpInterface.td
+++ b/enzyme/Enzyme/MLIR/Interfaces/AutoDiffOpInterface.td
@@ -126,10 +126,32 @@
];
}
+def StoreLikeInterface : OpInterface<"StoreLikeInterface"> {
+ let description = [{
+ A store-like operation that writes one operand (the value) through another
+ (the pointer/reference). Lets activity analysis reason about stored-value and
+ pointer activity without hard-coding specific dialects (memref.store,
+ LLVM.store, ...), which keeps the analysis dialect-agnostic and lets
+ out-of-tree dialects participate by attaching this interface.
+ }];
+ let cppNamespace = "::mlir::enzyme";
+
+ let methods = [InterfaceMethod<
+ /*desc=*/"Returns the value written by this store.",
+ /*retTy=*/"::mlir::Value",
+ /*methodName=*/"getStoredValue">,
+ InterfaceMethod<
+ /*desc=*/"Returns the base pointer/reference this store "
+ "writes through (before any potential offsets "
+ "within the op are applied).",
+ /*retTy=*/"::mlir::Value",
+ /*methodName=*/"getStoredPointer">];
+}
+
def ActivityOpInterface
: OpInterface<"ActivityOpInterface"> {
let cppNamespace = "::mlir::enzyme";
-
+
let methods = [
InterfaceMethod<
/*desc=*/[{
diff --git a/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp b/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp
index 9c98b68..a7a1088 100644
--- a/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp
+++ b/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp
@@ -317,10 +317,18 @@
LogicalResult MGradientUtils::visitChild(Operation *op) {
if (mode == DerivativeMode::ForwardMode) {
+ // An op with side effects may still need to touch shadow memory even when
+ // it is constant: a store of an inactive value into active memory has to
+ // zero the shadow, or a later load reads a stale tangent. Only skip it if
+ // it is pure, or if every operand is constant and there is no shadow to
+ // write through.
if ((op->getBlock()->getTerminator() != op) &&
llvm::all_of(op->getResults(),
[this](Value v) { return isConstantValue(v); }) &&
- /*iface.hasNoEffect()*/ activityAnalyzer->isConstantOperation(TR, op)) {
+ (isPure(op) ||
+ llvm::all_of(op->getOperands(),
+ [this](Value v) { return isConstantValue(v); })) &&
+ activityAnalyzer->isConstantOperation(TR, op)) {
return success();
}
// }
diff --git a/enzyme/test/MLIR/ForwardMode/llvm_if.mlir b/enzyme/test/MLIR/ForwardMode/llvm_if.mlir
new file mode 100644
index 0000000..b869d02
--- /dev/null
+++ b/enzyme/test/MLIR/ForwardMode/llvm_if.mlir
@@ -0,0 +1,51 @@
+// RUN: %eopt --enzyme %s | FileCheck %s
+
+// An llvm.alloca that is unconditionally initialized with a constant and then
+// conditionally overwritten with an active value. In forward mode the shadow
+// allocation must be zero-initialized before the conditional store, otherwise
+// the load on the not-taken path reads uninitialized shadow memory and the
+// returned tangent is garbage instead of 0.
+//
+// This is the llvm.store analogue of the affine.store `@if_then` test.
+
+module {
+ func.func @if_then(%x : f64, %c : i1) -> f64 {
+ %c1_i64 = arith.constant 1 : i64
+ %c2 = arith.constant 2.000000e+00 : f64
+ %mem = llvm.alloca %c1_i64 x f64 : (i64) -> !llvm.ptr
+ llvm.store %c2, %mem : f64, !llvm.ptr
+ scf.if %c {
+ %mul = arith.mulf %x, %x : f64
+ llvm.store %mul, %mem : f64, !llvm.ptr
+ }
+ %r = llvm.load %mem : !llvm.ptr -> f64
+ %res = arith.mulf %c2, %r : f64
+ return %res : f64
+ }
+ func.func @dif_then(%x : f64, %dx : f64, %c : i1) -> f64 {
+ %r = enzyme.fwddiff @if_then(%x, %dx, %c) { activity=[#enzyme<activity enzyme_dup>, #enzyme<activity enzyme_const>], ret_activity=[#enzyme<activity enzyme_dupnoneed>] } : (f64, f64, i1) -> (f64)
+ return %r : f64
+ }
+}
+
+// CHECK: @fwddiffeif_then
+// CHECK: (%[[arg0:.+]]: f64, %[[arg1:.+]]: f64, %[[arg2:.+]]: i1) -> f64 {
+// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : i64
+// CHECK-DAG: %[[cst2:.+]] = arith.constant 2.000000e+00 : f64
+// CHECK: %[[alloc:.+]] = llvm.alloca %[[c1]] x f64 : (i64) -> !llvm.ptr
+// CHECK: %[[alloc_0:.+]] = llvm.alloca %[[c1]] x f64 : (i64) -> !llvm.ptr
+// CHECK-DAG: %[[cst0:.+]] = arith.constant 0.000000e+00 : f64
+// CHECK: llvm.store %[[cst0]], %[[alloc]] : f64, !llvm.ptr
+// CHECK: llvm.store %[[cst2]], %[[alloc_0]] : f64, !llvm.ptr
+// CHECK: scf.if %[[arg2]] {
+// CHECK: %[[v4:.+]] = arith.mulf %[[arg1]], %[[arg0]] : f64
+// CHECK: %[[v5:.+]] = arith.mulf %[[arg1]], %[[arg0]] : f64
+// CHECK: %[[v6:.+]] = arith.addf %[[v4]], %[[v5]] : f64
+// CHECK: %[[v7:.+]] = arith.mulf %[[arg0]], %[[arg0]] : f64
+// CHECK: llvm.store %[[v6]], %[[alloc]] : f64, !llvm.ptr
+// CHECK: llvm.store %[[v7]], %[[alloc_0]] : f64, !llvm.ptr
+// CHECK: }
+// CHECK: %[[v0:.+]] = llvm.load %[[alloc]] : !llvm.ptr -> f64
+// CHECK: %[[v1:.+]] = llvm.load %[[alloc_0]] : !llvm.ptr -> f64
+// CHECK: %[[v2:.+]] = arith.mulf %[[v0]], %[[cst2]] : f64
+// CHECK: return %[[v2]] : f64
diff --git a/enzyme/test/MLIR/ForwardMode/memref_if.mlir b/enzyme/test/MLIR/ForwardMode/memref_if.mlir
new file mode 100644
index 0000000..ceff151
--- /dev/null
+++ b/enzyme/test/MLIR/ForwardMode/memref_if.mlir
@@ -0,0 +1,51 @@
+// RUN: %eopt --enzyme %s | FileCheck %s
+
+// A memref that is unconditionally initialized with a constant and then
+// conditionally overwritten with an active value. In forward mode the shadow
+// memref must be zero-initialized before the conditional store, otherwise the
+// load on the not-taken path reads uninitialized shadow memory and the returned
+// tangent is garbage instead of 0.
+//
+// This is the memref analogue of the affine.store `@if_then` test.
+
+module {
+ func.func @if_then(%x : f64, %c : i1) -> f64 {
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2.000000e+00 : f64
+ %mem = memref.alloc() : memref<1xf64>
+ memref.store %c2, %mem[%c0] : memref<1xf64>
+ scf.if %c {
+ %mul = arith.mulf %x, %x : f64
+ memref.store %mul, %mem[%c0] : memref<1xf64>
+ }
+ %r = memref.load %mem[%c0] : memref<1xf64>
+ %res = arith.mulf %c2, %r : f64
+ return %res : f64
+ }
+ func.func @dif_then(%x : f64, %dx : f64, %c : i1) -> f64 {
+ %r = enzyme.fwddiff @if_then(%x, %dx, %c) { activity=[#enzyme<activity enzyme_dup>, #enzyme<activity enzyme_const>], ret_activity=[#enzyme<activity enzyme_dupnoneed>] } : (f64, f64, i1) -> (f64)
+ return %r : f64
+ }
+}
+
+// CHECK: @fwddiffeif_then
+// CHECK: (%[[arg0:.+]]: f64, %[[arg1:.+]]: f64, %[[arg2:.+]]: i1) -> f64 {
+// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[cst2:.+]] = arith.constant 2.000000e+00 : f64
+// CHECK: %[[alloc:.+]] = memref.alloc() : memref<1xf64>
+// CHECK: %[[alloc_0:.+]] = memref.alloc() : memref<1xf64>
+// CHECK-DAG: %[[cst0:.+]] = arith.constant 0.000000e+00 : f64
+// CHECK: memref.store %[[cst0]], %[[alloc]][%[[c0]]] : memref<1xf64>
+// CHECK: memref.store %[[cst2]], %[[alloc_0]][%[[c0]]] : memref<1xf64>
+// CHECK: scf.if %[[arg2]] {
+// CHECK: %[[v4:.+]] = arith.mulf %[[arg1]], %[[arg0]] : f64
+// CHECK: %[[v5:.+]] = arith.mulf %[[arg1]], %[[arg0]] : f64
+// CHECK: %[[v6:.+]] = arith.addf %[[v4]], %[[v5]] : f64
+// CHECK: %[[v7:.+]] = arith.mulf %[[arg0]], %[[arg0]] : f64
+// CHECK: memref.store %[[v6]], %[[alloc]][%[[c0]]] : memref<1xf64>
+// CHECK: memref.store %[[v7]], %[[alloc_0]][%[[c0]]] : memref<1xf64>
+// CHECK: }
+// CHECK: %[[v0:.+]] = memref.load %[[alloc]][%[[c0]]] : memref<1xf64>
+// CHECK: %[[v1:.+]] = memref.load %[[alloc_0]][%[[c0]]] : memref<1xf64>
+// CHECK: %[[v2:.+]] = arith.mulf %[[v0]], %[[cst2]] : f64
+// CHECK: return %[[v2]] : f64
