diff --git a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp index 1a47acf6513..8d9630fbdb3 100644 --- a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp +++ b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp @@ -18,13 +18,10 @@ #include "Interfaces/GradientUtilsReverse.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/DialectRegistry.h" #include "mlir/Support/LogicalResult.h" -// TODO: We need a way to zero out a memref (which linalg.fill does), but -// ideally we wouldn't depend on the linalg dialect. -#include "mlir/Dialect/Linalg/IR/Linalg.h" - using namespace mlir; using namespace mlir::enzyme; @@ -203,6 +200,147 @@ struct SubViewOpInterfaceReverse } }; +struct AllocaScopeOpInterfaceReverse + : public ReverseAutoDiffOpInterface::ExternalModel< + AllocaScopeOpInterfaceReverse, memref::AllocaScopeOp> { + LogicalResult createReverseModeAdjoint(Operation *op, OpBuilder &builder, + MGradientUtilsReverse *gutils, + SmallVector caches) const { + auto scopeOp = cast(op); + + SmallVector resultsActive(scopeOp->getNumResults(), false); + SmallVector incomingGradients; + incomingGradients.reserve(scopeOp->getNumResults()); + for (OpResult result : scopeOp->getResults()) { + bool active = !gutils->isConstantValue(result); + resultsActive[result.getResultNumber()] = active; + if (active) { + incomingGradients.push_back(gutils->diffe(result, builder)); + gutils->zeroDiffe(result, builder); + } + } + + Region &scopeRegion = scopeOp.getBodyRegion(); + SmallVector capturedInputs; + scopeRegion.walk([&](Operation *inner) { + for (Value operand : inner->getOperands()) { + Region *defRegion = operand.getParentRegion(); + if (!defRegion || scopeRegion.isAncestor(defRegion)) + continue; + if (gutils->isConstantValue(operand)) + continue; + auto iface = dyn_cast(operand.getType()); + if (!iface || iface.isMutable()) + continue; + if (!llvm::is_contained(capturedInputs, operand)) + capturedInputs.push_back(operand); + } + }); + + SmallVector capturedPre; + capturedPre.reserve(capturedInputs.size()); + for (Value v : capturedInputs) { + capturedPre.push_back(gutils->diffe(v, builder)); + gutils->zeroDiffe(v, builder); + } + + auto newScope = cast(gutils->getNewFromOriginal(op)); + newScope->moveBefore(builder.getInsertionBlock(), + builder.getInsertionPoint()); + + Block &newBody = newScope.getBodyRegion().front(); + OpBuilder bodyBuilder(newBody.getTerminator()); + + Block &oldBody = scopeOp.getBodyRegion().front(); + bool valid = true; + + for (Operation &innerOp : oldBody.getOperations()) { + for (Value result : innerOp.getResults()) { + if (!gutils->isConstantValue(result)) { + auto iface = dyn_cast(result.getType()); + if (iface && !iface.isMutable()) + gutils->zeroDiffe(result, bodyBuilder); + } + } + } + + Operation *term = oldBody.getTerminator(); + unsigned incomingIdx = 0; + for (auto indexedOperand : llvm::enumerate(term->getOperands())) { + if (!resultsActive[indexedOperand.index()]) + continue; + + Value operand = indexedOperand.value(); + if (!gutils->isConstantValue(operand)) + gutils->addToDiffe(operand, incomingGradients[incomingIdx], + bodyBuilder); + incomingIdx++; + } + + auto first = oldBody.rbegin(); + first++; + + for (auto it = first; it != oldBody.rend(); it++) { + valid &= gutils->Logic.visitChild(&*it, bodyBuilder, gutils).succeeded(); + } + + if (capturedInputs.empty()) + return success(valid); + + SmallVector contributions; + contributions.reserve(capturedInputs.size()); + for (Value v : capturedInputs) + contributions.push_back(gutils->diffe(v, bodyBuilder)); + + unsigned numPrimal = newScope->getNumResults(); + Operation *bodyTerm = newBody.getTerminator(); + SmallVector retVals(bodyTerm->getOperands().begin(), + bodyTerm->getOperands().end()); + retVals.append(contributions.begin(), contributions.end()); + + SmallVector newResultTypes(newScope->getResultTypes().begin(), + newScope->getResultTypes().end()); + for (Value v : capturedInputs) + newResultTypes.push_back(gutils->getShadowType(v.getType())); + + OpBuilder scopeBuilder(newScope); + auto extScope = memref::AllocaScopeOp::create(scopeBuilder, op->getLoc(), + newResultTypes); + extScope.getBodyRegion().takeBody(newScope.getBodyRegion()); + + Block &extBody = extScope.getBodyRegion().front(); + Operation *movedTerm = extBody.getTerminator(); + OpBuilder termBuilder(movedTerm); + memref::AllocaScopeReturnOp::create(termBuilder, movedTerm->getLoc(), + retVals); + gutils->erase(movedTerm); + + for (unsigned i = 0; i < numPrimal; ++i) { + newScope->getResult(i).replaceAllUsesWith(extScope.getResult(i)); + gutils->originalToNewFn.map(scopeOp->getResult(i), extScope.getResult(i)); + } + gutils->erase(newScope); + + builder.setInsertionPointAfter(extScope); + for (auto indexed : llvm::enumerate(capturedInputs)) { + Value v = indexed.value(); + gutils->setDiffe(v, capturedPre[indexed.index()], builder); + gutils->addToDiffe(v, extScope.getResult(numPrimal + indexed.index()), + builder); + } + + return success(valid); + } + + SmallVector cacheValues(Operation *op, + MGradientUtilsReverse *gutils) const { + return SmallVector(); + } + + void createShadowValues(Operation *op, OpBuilder &builder, + MGradientUtilsReverse *gutils) const {} +}; + class MemRefClonableTypeInterface : public ClonableTypeInterface::ExternalModel { @@ -273,14 +411,33 @@ class MemRefAutoDiffTypeInterface LogicalResult zeroInPlace(Type self, OpBuilder &builder, Location loc, Value val) const { auto MT = cast(self); - if (auto iface = dyn_cast(MT.getElementType())) { - if (!iface.isMutable()) { - Value zero = iface.createNullValue(builder, loc); - linalg::FillOp::create(builder, loc, zero, val); - } - } else { + auto eltIface = dyn_cast(MT.getElementType()); + if (!eltIface || eltIface.isMutable()) return failure(); + Value zero = eltIface.createNullValue(builder, loc); + + if (MT.getRank() == 0) { + memref::StoreOp::create(builder, loc, zero, val, ValueRange{}); + return success(); } + + Value c0 = arith::ConstantIndexOp::create(builder, loc, 0); + Value c1 = arith::ConstantIndexOp::create(builder, loc, 1); + + SmallVector lbs(MT.getRank(), c0); + SmallVector steps(MT.getRank(), c1); + SmallVector ubs; + for (auto [i, d] : llvm::enumerate(MT.getShape())) { + ubs.push_back( + d == ShapedType::kDynamic + ? memref::DimOp::create(builder, loc, val, i).getResult() + : arith::ConstantIndexOp::create(builder, loc, d).getResult()); + } + + scf::ParallelOp::create(builder, loc, lbs, ubs, steps, + [&](OpBuilder &b, Location l, ValueRange ivs) { + memref::StoreOp::create(b, l, zero, val, ivs); + }); return success(); } @@ -299,5 +456,7 @@ void mlir::enzyme::registerMemRefDialectAutoDiffInterface( memref::LoadOp::attachInterface(*context); memref::StoreOp::attachInterface(*context); memref::SubViewOp::attachInterface(*context); + memref::AllocaScopeOp::attachInterface( + *context); }); } diff --git a/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp b/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp index f0c06429385..77684116748 100644 --- a/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp +++ b/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp @@ -223,8 +223,12 @@ void mlir::enzyme::MGradientUtils::setDiffe(mlir::Value val, mlir::Value toset, } assert(!isConstantValue(val)); + bool isMutable = false; + if (auto iface = dyn_cast(val.getType())) + isMutable = iface.isMutable(); + if (mode == DerivativeMode::ForwardMode || - mode == DerivativeMode::ForwardModeSplit) { + mode == DerivativeMode::ForwardModeSplit || isMutable) { setInvertedPointer(val, toset); } /* diff --git a/enzyme/test/MLIR/ReverseMode/alloca.mlir b/enzyme/test/MLIR/ReverseMode/alloca.mlir new file mode 100644 index 00000000000..76b47c8c4da --- /dev/null +++ b/enzyme/test/MLIR/ReverseMode/alloca.mlir @@ -0,0 +1,62 @@ +// RUN: %eopt --enzyme %s | FileCheck %s + +func.func @foo_flat(%x : f64) -> f64 { + %buf = memref.alloca() : memref + memref.store %x, %buf[] : memref + %y = memref.load %buf[] : memref + return %y : f64 +} + +func.func @dfoo_flat(%x: f64, %dout : f64) -> f64 { + %dx = enzyme.autodiff @foo_flat(%x, %dout) { + activity = [#enzyme], + ret_activity = [#enzyme] + } : (f64, f64) -> (f64) + return %dx : f64 +} + +// CHECK-LABEL: func.func private @diffefoo_flat( +// CHECK-SAME: %[[X:[^,]+]]: f64, +// CHECK-SAME: %[[DOUT:[^)]+]]: f64) -> f64 { + +// CHECK: %[[GX:.+]] = "enzyme.init"() : () -> !enzyme.Gradient +// CHECK: "enzyme.set"(%[[GX]], %{{.*}}) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[CS:.+]] = "enzyme.init"() : () -> !enzyme.Cache> +// CHECK: %[[CL:.+]] = "enzyme.init"() : () -> !enzyme.Cache> +// CHECK: %[[GY:.+]] = "enzyme.init"() : () -> !enzyme.Gradient +// CHECK: "enzyme.set"(%[[GY]], %{{.*}}) : (!enzyme.Gradient, f64) -> () + +// CHECK: %[[DBUF:.+]] = memref.alloca() : memref +// CHECK: %[[ZERO_INIT:.+]] = arith.constant 0.000000e+00 : f64 +// CHECK: memref.store %[[ZERO_INIT]], %[[DBUF]][] : memref + +// CHECK: %[[BUF:.+]] = memref.alloca() : memref +// CHECK: "enzyme.push"(%[[CS]], %[[DBUF]]) : (!enzyme.Cache>, memref) -> () +// CHECK: memref.store %[[X]], %[[BUF]][] : memref +// CHECK: "enzyme.push"(%[[CL]], %[[DBUF]]) : (!enzyme.Cache>, memref) -> () +// CHECK: memref.load %[[BUF]][] : memref +// CHECK: cf.br ^bb1 +// CHECK: ^bb1: + +// CHECK: %{{.+}} = "enzyme.get"(%[[GY]]) : (!enzyme.Gradient) -> f64 +// CHECK: arith.addf %{{.+}}, %[[DOUT]] : f64 +// CHECK: "enzyme.set"(%[[GY]], %{{.*}}) : (!enzyme.Gradient, f64) -> () + +// CHECK: %{{.+}} = "enzyme.get"(%[[GY]]) : (!enzyme.Gradient) -> f64 +// CHECK: %[[POPL:.+]] = "enzyme.pop"(%[[CL]]) : (!enzyme.Cache>) -> memref +// CHECK: memref.load %[[POPL]][] : memref +// CHECK: arith.addf +// CHECK: memref.store %{{.+}}, %[[POPL]][] : memref + +// CHECK: %[[POPS:.+]] = "enzyme.pop"(%[[CS]]) : (!enzyme.Cache>) -> memref +// CHECK: memref.load %[[POPS]][] : memref +// CHECK: %{{.+}} = "enzyme.get"(%[[GX]]) : (!enzyme.Gradient) -> f64 +// CHECK: arith.addf +// CHECK: "enzyme.set"(%[[GX]], %{{.*}}) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[ZERO_CLEAR:.+]] = arith.constant 0.000000e+00 : f64 +// CHECK: memref.store %[[ZERO_CLEAR]], %[[POPS]][] : memref + +// CHECK: %{{.+}} = "enzyme.get"(%[[GX]]) : (!enzyme.Gradient) -> f64 +// CHECK: return %{{.+}} : f64 + +// CHECK-NOT: enzyme.placeholder diff --git a/enzyme/test/MLIR/ReverseMode/alloca_scope.mlir b/enzyme/test/MLIR/ReverseMode/alloca_scope.mlir new file mode 100644 index 00000000000..5f863756d35 --- /dev/null +++ b/enzyme/test/MLIR/ReverseMode/alloca_scope.mlir @@ -0,0 +1,94 @@ +// RUN: %eopt --enzyme --canonicalize --remove-unnecessary-enzyme-ops --enzyme-simplify-math %s | FileCheck %s + +func.func private @scope_overwrite(%x: memref) { + memref.alloca_scope { + %val = memref.load %x[] : memref + %cos = math.cos %val : f32 + memref.store %cos, %x[] : memref + } + return +} + +func.func @dscope_overwrite(%x: memref, %dx: memref) { + enzyme.autodiff @scope_overwrite(%x, %dx) { + activity = [#enzyme], + ret_activity = [] + } : (memref, memref) -> () + return +} + +func.func private @scope_load_mul(%x: memref, %c: f64) -> f64 { + %r = memref.alloca_scope -> (f64) { + %v = memref.load %x[] : memref + %p = arith.mulf %v, %c : f64 + memref.alloca_scope.return %p : f64 + } + return %r : f64 +} + +func.func @dscope_load_mul(%x: memref, %dx: memref, + %c: f64, %dout: f64) -> f64 { + %res = enzyme.autodiff @scope_load_mul(%x, %dx, %c, %dout) { + activity = [#enzyme, #enzyme], + ret_activity = [#enzyme] + } : (memref, memref, f64, f64) -> f64 + return %res : f64 +} + +// CHECK-LABEL: func.func private @diffescope_overwrite( +// CHECK-SAME: %[[X:.*]]: memref, +// CHECK-SAME: %[[DX:.*]]: memref) { +// CHECK: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 +// CHECK: %[[CDX0:.*]] = "enzyme.init"() : () -> !enzyme.Cache> +// CHECK: %[[CVAL:.*]] = "enzyme.init"() : () -> !enzyme.Cache +// CHECK: %[[CDX1:.*]] = "enzyme.init"() : () -> !enzyme.Cache> +// CHECK: memref.alloca_scope { +// CHECK: "enzyme.push"(%[[CDX0]], %[[DX]]) +// CHECK: %[[V:.*]] = memref.load %[[X]][] : memref +// CHECK: "enzyme.push"(%[[CVAL]], %[[V]]) +// CHECK: %[[COS:.*]] = math.cos %[[V]] : f32 +// CHECK: "enzyme.push"(%[[CDX1]], %[[DX]]) +// CHECK: memref.store %[[COS]], %[[X]][] : memref +// CHECK: } +// CHECK: %[[DX_R1:.*]] = "enzyme.pop"(%[[CDX1]]) +// CHECK: %[[DOUT:.*]] = memref.load %[[DX_R1]][] : memref +// CHECK: memref.store %[[ZERO]], %[[DX_R1]][] : memref +// CHECK: %[[VAL:.*]] = "enzyme.pop"(%[[CVAL]]) +// CHECK: %[[SIN:.*]] = math.sin %[[VAL]] : f32 +// CHECK: %[[NEG:.*]] = arith.negf %[[SIN]] : f32 +// CHECK: %[[GRAD:.*]] = arith.mulf %[[DOUT]], %[[NEG]] : f32 +// CHECK: %[[DX_R0:.*]] = "enzyme.pop"(%[[CDX0]]) +// CHECK: %[[OLD:.*]] = memref.load %[[DX_R0]][] : memref +// CHECK: %[[NEW:.*]] = arith.addf %[[OLD]], %[[GRAD]] : f32 +// CHECK: memref.store %[[NEW]], %[[DX_R0]][] : memref +// CHECK: return + +// CHECK-LABEL: func.func private @diffescope_load_mul( +// CHECK-SAME: %[[X:[^,]+]]: memref, +// CHECK-SAME: %[[DX:[^,]+]]: memref, +// CHECK-SAME: %[[C:[^,]+]]: f64, +// CHECK-SAME: %[[DOUT:[^)]+]]: f64) -> f64 { +// CHECK: %[[CDX:.*]] = "enzyme.init"() : () -> !enzyme.Cache> +// CHECK: %[[GC:.*]] = "enzyme.init"() : () -> !enzyme.Gradient +// CHECK: %[[CC:.*]] = "enzyme.init"() : () -> !enzyme.Cache +// CHECK: %[[CV:.*]] = "enzyme.init"() : () -> !enzyme.Cache +// CHECK: memref.alloca_scope -> (f64) { +// CHECK: "enzyme.push"(%[[CDX]], %[[DX]]) +// CHECK: %[[V:.*]] = memref.load %[[X]][] : memref +// CHECK: "enzyme.push"(%[[CV]], %[[V]]) +// CHECK: "enzyme.push"(%[[CC]], %[[C]]) +// CHECK: %[[P:.*]] = arith.mulf %[[V]], %[[C]] : f64 +// CHECK: memref.alloca_scope.return %[[P]] : f64 +// CHECK: } +// CHECK: memref.alloca_scope { +// CHECK: %[[V_R:.*]] = "enzyme.pop"(%[[CV]]) +// CHECK: %[[C_R:.*]] = "enzyme.pop"(%[[CC]]) +// CHECK-DAG: %[[DC:.*]] = arith.mulf %[[DOUT]], %[[V_R]] : f64 +// CHECK-DAG: %[[DV:.*]] = arith.mulf %[[DOUT]], %[[C_R]] : f64 +// CHECK: %[[DX_R:.*]] = "enzyme.pop"(%[[CDX]]) +// CHECK: %[[OLD:.*]] = memref.load %[[DX_R]][] : memref +// CHECK: %[[NEW:.*]] = arith.addf %[[OLD]], %[[DV]] : f64 +// CHECK: memref.store %[[NEW]], %[[DX_R]][] : memref +// CHECK: } +// CHECK: %[[OUT:.*]] = "enzyme.get"(%[[GC]]) : (!enzyme.Gradient) -> f64 +// CHECK: return %[[OUT]] : f64 diff --git a/enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir b/enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir new file mode 100644 index 00000000000..1a7f50c5553 --- /dev/null +++ b/enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir @@ -0,0 +1,72 @@ +// RUN: %eopt --enzyme %s | FileCheck %s + +func.func @foo(%x : f64) -> f64{ + %out = memref.alloca_scope -> (f64) { + %buf = memref.alloca() : memref + memref.store %x, %buf[] : memref + %y = memref.load %buf[] : memref + memref.alloca_scope.return %y : f64 + } + return %out : f64 +} + +func.func @dfoo(%x: f64, %dout : f64) -> f64 { + %dx = enzyme.autodiff @foo(%x, %dout) { + activity = [#enzyme], + ret_activity = [#enzyme] + } : (f64, f64) -> (f64) + return %dx : f64 +} + +// CHECK-LABEL: func.func private @diffefoo( +// CHECK-SAME: %[[X:[^,]+]]: f64, +// CHECK-SAME: %[[DOUT:[^)]+]]: f64) -> f64 { +// CHECK: %[[C0:.*]] = "enzyme.init"() : () -> !enzyme.Cache> +// CHECK: %[[C1:.*]] = "enzyme.init"() : () -> !enzyme.Cache> +// CHECK: %[[G0:.*]] = "enzyme.init"() : () -> !enzyme.Gradient +// CHECK: "enzyme.set"(%[[G0]], %{{.*}}) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[G1:.*]] = "enzyme.init"() : () -> !enzyme.Gradient +// CHECK: "enzyme.set"(%[[G1]], %{{.*}}) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[G2:.*]] = "enzyme.init"() : () -> !enzyme.Gradient +// CHECK: "enzyme.set"(%[[G2]], %{{.*}}) : (!enzyme.Gradient, f64) -> () +// CHECK: cf.br ^bb1 +// CHECK: ^bb1: +// CHECK: %[[T5:.*]] = "enzyme.get"(%[[G2]]) : (!enzyme.Gradient) -> f64 +// CHECK: %[[T6:.*]] = arith.addf %[[T5]], %[[DOUT]] : f64 +// CHECK: "enzyme.set"(%[[G2]], %[[T6]]) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[T7:.*]] = "enzyme.get"(%[[G2]]) : (!enzyme.Gradient) -> f64 +// CHECK: "enzyme.set"(%[[G2]], %{{.*}}) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[T8:.*]] = "enzyme.get"(%[[G1]]) : (!enzyme.Gradient) -> f64 +// CHECK: "enzyme.set"(%[[G1]], %{{.*}}) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[SCOPE:.*]]:2 = memref.alloca_scope -> (f64, f64) { +// CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref +// CHECK: memref.store %{{.*}}, %[[ALLOCA]][] : memref +// CHECK: %[[ALLOCA5:.*]] = memref.alloca() : memref +// CHECK: "enzyme.push"(%[[C0]], %[[ALLOCA]]) : (!enzyme.Cache>, memref) -> () +// CHECK: memref.store %[[X]], %[[ALLOCA5]][] : memref +// CHECK: "enzyme.push"(%[[C1]], %[[ALLOCA]]) : (!enzyme.Cache>, memref) -> () +// CHECK: %[[LD:.*]] = memref.load %[[ALLOCA5]][] : memref +// CHECK: "enzyme.set"(%[[G0]], %{{.*}}) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[T14:.*]] = "enzyme.get"(%[[G0]]) : (!enzyme.Gradient) -> f64 +// CHECK: %[[T15:.*]] = arith.addf %[[T14]], %[[T7]] : f64 +// CHECK: "enzyme.set"(%[[G0]], %[[T15]]) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[T16:.*]] = "enzyme.get"(%[[G0]]) : (!enzyme.Gradient) -> f64 +// CHECK: %[[POP1:.*]] = "enzyme.pop"(%[[C1]]) : (!enzyme.Cache>) -> memref +// CHECK: %[[LD18:.*]] = memref.load %[[POP1]][] : memref +// CHECK: %[[T19:.*]] = arith.addf %[[LD18]], %[[T16]] : f64 +// CHECK: memref.store %[[T19]], %[[POP1]][] : memref +// CHECK: %[[POP0:.*]] = "enzyme.pop"(%[[C0]]) : (!enzyme.Cache>) -> memref +// CHECK: %[[LD21:.*]] = memref.load %[[POP0]][] : memref +// CHECK: %[[T22:.*]] = "enzyme.get"(%[[G1]]) : (!enzyme.Gradient) -> f64 +// CHECK: %[[T23:.*]] = arith.addf %[[T22]], %[[LD21]] : f64 +// CHECK: "enzyme.set"(%[[G1]], %[[T23]]) : (!enzyme.Gradient, f64) -> () +// CHECK: memref.store %{{.*}}, %[[POP0]][] : memref +// CHECK: %[[T24:.*]] = "enzyme.get"(%[[G1]]) : (!enzyme.Gradient) -> f64 +// CHECK: memref.alloca_scope.return %[[LD]], %[[T24]] : f64, f64 +// CHECK: } +// CHECK: "enzyme.set"(%[[G1]], %[[T8]]) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[T10:.*]] = "enzyme.get"(%[[G1]]) : (!enzyme.Gradient) -> f64 +// CHECK: %[[T11:.*]] = arith.addf %[[T10]], %[[SCOPE]]#1 : f64 +// CHECK: "enzyme.set"(%[[G1]], %[[T11]]) : (!enzyme.Gradient, f64) -> () +// CHECK: %[[T12:.*]] = "enzyme.get"(%[[G1]]) : (!enzyme.Gradient) -> f64 +// CHECK: return %[[T12]] : f64