From 5ddbd10a3c90e1c4bc46331d746d64ee6e918c0a Mon Sep 17 00:00:00 2001 From: Yuansui Xu Date: Wed, 13 May 2026 15:05:24 -0500 Subject: [PATCH 01/11] using scf.parallel+memref.store to zero out a memref --- .../MemRefAutoDiffOpInterfaceImpl.cpp | 35 +++++++++++++------ 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp index 1a47acf65136..a2837a6ba48a 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; @@ -273,14 +270,32 @@ 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(); } From 3a65f7ff1a9bc68617364e6d7ed54942b7c69932 Mon Sep 17 00:00:00 2001 From: Yuansui Xu Date: Wed, 13 May 2026 15:10:22 -0500 Subject: [PATCH 02/11] add test for memref.alloca --- .../MemRefAutoDiffOpInterfaceImpl.cpp | 15 +++++----- enzyme/test/MLIR/ReverseMode/alloca.mlir | 29 +++++++++++++++++++ 2 files changed, 37 insertions(+), 7 deletions(-) create mode 100644 enzyme/test/MLIR/ReverseMode/alloca.mlir diff --git a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp index a2837a6ba48a..62ee51ba125d 100644 --- a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp +++ b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp @@ -274,7 +274,7 @@ class MemRefAutoDiffTypeInterface 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(); @@ -287,15 +287,16 @@ class MemRefAutoDiffTypeInterface 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()); + 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); - }); + [&](OpBuilder &b, Location l, ValueRange ivs) { + memref::StoreOp::create(b, l, zero, val, ivs); + }); return success(); } diff --git a/enzyme/test/MLIR/ReverseMode/alloca.mlir b/enzyme/test/MLIR/ReverseMode/alloca.mlir new file mode 100644 index 000000000000..a61fb3eaea0d --- /dev/null +++ b/enzyme/test/MLIR/ReverseMode/alloca.mlir @@ -0,0 +1,29 @@ +// RUN: %eopt --split-input-file --enzyme --canonicalize --remove-unnecessary-enzyme-ops --enzyme-simplify-math %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 { +// A shadow memref.alloca must be created and zero-initialized so the +// reverse-mode adjoint can accumulate gradients into it. +// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f64 +// CHECK-DAG: %[[DBUF:.*]] = memref.alloca() : memref +// CHECK: memref.store %[[ZERO]], %[[DBUF]][] : memref +// No leftover placeholders should survive the differentiation. +// CHECK-NOT: enzyme.placeholder +// The function ultimately returns the gradient w.r.t. x (= dout). +// CHECK: return %{{.*}} : f64 \ No newline at end of file From 7615eee96dd68671c980450ab09ffe589a177f24 Mon Sep 17 00:00:00 2001 From: Yuansui Xu Date: Thu, 14 May 2026 00:12:53 -0500 Subject: [PATCH 03/11] revert --- .../Enzyme/MLIR/Interfaces/GradientUtils.cpp | 18 ++++--- enzyme/test/MLIR/ReverseMode/alloca.mlir | 51 +++++++++++++++---- 2 files changed, 54 insertions(+), 15 deletions(-) diff --git a/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp b/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp index f0c064293856..5e58e6e0b6c9 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); } /* @@ -240,11 +244,13 @@ void mlir::enzyme::MGradientUtils::setDiffe(mlir::Value val, mlir::Value toset, void mlir::enzyme::MGradientUtils::setInvertedPointer(Value val, Value toset) { assert(getShadowType(val.getType()) == toset.getType()); - auto found = invertedPointers.lookupOrNull(val); - assert(found != nullptr); - auto placeholder = found.getDefiningOp(); - placeholder.replaceAllUsesWith(toset); - erase(placeholder); + + if (auto found = invertedPointers.lookupOrNull(val)) { + if (auto placeholder = found.getDefiningOp()) { + placeholder.replaceAllUsesWith(toset); + erase(placeholder); + } + } invertedPointers.map(val, toset); } diff --git a/enzyme/test/MLIR/ReverseMode/alloca.mlir b/enzyme/test/MLIR/ReverseMode/alloca.mlir index a61fb3eaea0d..76b47c8c4da3 100644 --- a/enzyme/test/MLIR/ReverseMode/alloca.mlir +++ b/enzyme/test/MLIR/ReverseMode/alloca.mlir @@ -1,4 +1,4 @@ -// RUN: %eopt --split-input-file --enzyme --canonicalize --remove-unnecessary-enzyme-ops --enzyme-simplify-math %s | FileCheck %s +// RUN: %eopt --enzyme %s | FileCheck %s func.func @foo_flat(%x : f64) -> f64 { %buf = memref.alloca() : memref @@ -18,12 +18,45 @@ func.func @dfoo_flat(%x: f64, %dout : f64) -> f64 { // CHECK-LABEL: func.func private @diffefoo_flat( // CHECK-SAME: %[[X:[^,]+]]: f64, // CHECK-SAME: %[[DOUT:[^)]+]]: f64) -> f64 { -// A shadow memref.alloca must be created and zero-initialized so the -// reverse-mode adjoint can accumulate gradients into it. -// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f64 -// CHECK-DAG: %[[DBUF:.*]] = memref.alloca() : memref -// CHECK: memref.store %[[ZERO]], %[[DBUF]][] : memref -// No leftover placeholders should survive the differentiation. + +// 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 -// The function ultimately returns the gradient w.r.t. x (= dout). -// CHECK: return %{{.*}} : f64 \ No newline at end of file From 1e885ec01df57b3543c6a534543e476e0291f512 Mon Sep 17 00:00:00 2001 From: Yuansui Xu Date: Thu, 14 May 2026 00:16:56 -0500 Subject: [PATCH 04/11] fmt --- enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp b/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp index 5e58e6e0b6c9..a3f9cb8099da 100644 --- a/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp +++ b/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp @@ -244,7 +244,7 @@ void mlir::enzyme::MGradientUtils::setDiffe(mlir::Value val, mlir::Value toset, void mlir::enzyme::MGradientUtils::setInvertedPointer(Value val, Value toset) { assert(getShadowType(val.getType()) == toset.getType()); - + if (auto found = invertedPointers.lookupOrNull(val)) { if (auto placeholder = found.getDefiningOp()) { placeholder.replaceAllUsesWith(toset); From e94d56898846c8cef52a2ca1f7bc64e50ac51f37 Mon Sep 17 00:00:00 2001 From: Yuansui Xu Date: Wed, 20 May 2026 02:30:05 -0500 Subject: [PATCH 05/11] revert --- enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp b/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp index a3f9cb8099da..776841167486 100644 --- a/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp +++ b/enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp @@ -244,13 +244,11 @@ void mlir::enzyme::MGradientUtils::setDiffe(mlir::Value val, mlir::Value toset, void mlir::enzyme::MGradientUtils::setInvertedPointer(Value val, Value toset) { assert(getShadowType(val.getType()) == toset.getType()); - - if (auto found = invertedPointers.lookupOrNull(val)) { - if (auto placeholder = found.getDefiningOp()) { - placeholder.replaceAllUsesWith(toset); - erase(placeholder); - } - } + auto found = invertedPointers.lookupOrNull(val); + assert(found != nullptr); + auto placeholder = found.getDefiningOp(); + placeholder.replaceAllUsesWith(toset); + erase(placeholder); invertedPointers.map(val, toset); } From 1d901892bba7a961e1fbcaffe36f5369f300f831 Mon Sep 17 00:00:00 2001 From: Yuansui Xu Date: Thu, 30 Apr 2026 14:02:40 -0500 Subject: [PATCH 06/11] register reverse mode AD for memref.alloca_scope --- .../MemRefAutoDiffOpInterfaceImpl.cpp | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp index 62ee51ba125d..78b970b3d8ad 100644 --- a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp +++ b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp @@ -200,6 +200,81 @@ 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); + } + } + + auto revScope = + memref::AllocaScopeOp::create(builder, op->getLoc(), TypeRange()); + Block *revBody = builder.createBlock(&revScope.getBodyRegion()); + OpBuilder bodyBuilder(revBody, revBody->end()); + memref::AllocaScopeReturnOp::create(bodyBuilder, op->getLoc(), + ValueRange()); + bodyBuilder.setInsertionPoint(revBody->getTerminator()); + + Block &oldBody = scopeOp.getBodyRegion().front(); + bool valid = true; + + // Values defined in the scoped region cannot be used outside it. Reset their + // adjoints before propagating gradients through the scoped body. + 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(); + } + + 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 { @@ -315,5 +390,7 @@ void mlir::enzyme::registerMemRefDialectAutoDiffInterface( memref::LoadOp::attachInterface(*context); memref::StoreOp::attachInterface(*context); memref::SubViewOp::attachInterface(*context); + memref::AllocaScopeOp::attachInterface( + *context); }); } From 9417dff27333a6336d571f234f5d3f4e2491260f Mon Sep 17 00:00:00 2001 From: Yuansui Xu Date: Thu, 30 Apr 2026 14:06:42 -0500 Subject: [PATCH 07/11] fmt --- .../MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp index 78b970b3d8ad..425456471d61 100644 --- a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp +++ b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp @@ -231,8 +231,8 @@ struct AllocaScopeOpInterfaceReverse Block &oldBody = scopeOp.getBodyRegion().front(); bool valid = true; - // Values defined in the scoped region cannot be used outside it. Reset their - // adjoints before propagating gradients through the scoped body. + // Values defined in the scoped region cannot be used outside it. Reset + // their adjoints before propagating gradients through the scoped body. for (Operation &innerOp : oldBody.getOperations()) { for (Value result : innerOp.getResults()) { if (!gutils->isConstantValue(result)) { From 1ad6a68d6e4ce8783e94ff719d1514175ef1dc20 Mon Sep 17 00:00:00 2001 From: Yuansui Xu Date: Wed, 6 May 2026 21:56:08 -0500 Subject: [PATCH 08/11] add lit tests --- .../test/MLIR/ReverseMode/alloca_scope.mlir | 94 +++++++++++++++++++ 1 file changed, 94 insertions(+) create mode 100644 enzyme/test/MLIR/ReverseMode/alloca_scope.mlir diff --git a/enzyme/test/MLIR/ReverseMode/alloca_scope.mlir b/enzyme/test/MLIR/ReverseMode/alloca_scope.mlir new file mode 100644 index 000000000000..5f863756d358 --- /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 From fa61965a4c4630465dcce31d2eeed4a48de1beda Mon Sep 17 00:00:00 2001 From: Vimarsh Sathia Date: Thu, 7 May 2026 14:53:50 -0500 Subject: [PATCH 09/11] Add a failing testcase --- .../test/MLIR/ReverseMode/alloca_scope2.mlir | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir diff --git a/enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir b/enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir new file mode 100644 index 000000000000..c18e0f713e6b --- /dev/null +++ b/enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir @@ -0,0 +1,38 @@ +// 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 +} + +func.func @foo2(%x : f64) -> f64{ + %cst = arith.constant 0.0000e+00 : f64 + %buf = memref.alloca() : memref + memref.store %x, %buf[] : memref + %y = memref.load %buf[] : memref + %out = arith.addf %cst, %y : f64 + return %out : f64 +} + + +func.func @dfoo2(%x: f64, %dout : f64) -> f64 { + %dx = enzyme.autodiff @foo2(%x, %dout) { + activity = [#enzyme], + ret_activity = [#enzyme] + } : (f64, f64) -> (f64) + return %dx : f64 +} From 238dde5059d45f5a6f301ef23ba29fef04e6ce6b Mon Sep 17 00:00:00 2001 From: Yuansui Xu Date: Thu, 21 May 2026 01:28:27 -0500 Subject: [PATCH 10/11] unified alloca_scope --- .../MemRefAutoDiffOpInterfaceImpl.cpp | 91 ++++++++++++++++--- .../test/MLIR/ReverseMode/alloca_scope2.mlir | 72 +++++++++++---- 2 files changed, 132 insertions(+), 31 deletions(-) diff --git a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp index 425456471d61..34a369db7425 100644 --- a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp +++ b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp @@ -220,19 +220,41 @@ struct AllocaScopeOpInterfaceReverse } } - auto revScope = - memref::AllocaScopeOp::create(builder, op->getLoc(), TypeRange()); - Block *revBody = builder.createBlock(&revScope.getBodyRegion()); - OpBuilder bodyBuilder(revBody, revBody->end()); - memref::AllocaScopeReturnOp::create(bodyBuilder, op->getLoc(), - ValueRange()); - bodyBuilder.setInsertionPoint(revBody->getTerminator()); + 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; - // Values defined in the scoped region cannot be used outside it. Reset - // their adjoints before propagating gradients through the scoped body. for (Operation &innerOp : oldBody.getOperations()) { for (Value result : innerOp.getResults()) { if (!gutils->isConstantValue(result)) { @@ -253,16 +275,61 @@ struct AllocaScopeOpInterfaceReverse if (!gutils->isConstantValue(operand)) gutils->addToDiffe(operand, incomingGradients[incomingIdx], bodyBuilder); - ++incomingIdx; + incomingIdx++; } auto first = oldBody.rbegin(); - ++first; + first++; - for (auto it = first; it != oldBody.rend(); ++it) { + 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); } diff --git a/enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir b/enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir index c18e0f713e6b..1a7f50c55539 100644 --- a/enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir +++ b/enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir @@ -10,29 +10,63 @@ func.func @foo(%x : f64) -> 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 -} - -func.func @foo2(%x : f64) -> f64{ - %cst = arith.constant 0.0000e+00 : f64 - %buf = memref.alloca() : memref - memref.store %x, %buf[] : memref - %y = memref.load %buf[] : memref - %out = arith.addf %cst, %y : f64 - return %out : f64 -} - - -func.func @dfoo2(%x: f64, %dout : f64) -> f64 { - %dx = enzyme.autodiff @foo2(%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 From 6f780f5c47cdf8320ddeef238e61bf2605993529 Mon Sep 17 00:00:00 2001 From: Yuansui Xu Date: Thu, 21 May 2026 01:38:38 -0500 Subject: [PATCH 11/11] fmt --- .../MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp index 34a369db7425..8d9630fbdb3a 100644 --- a/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp +++ b/enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp @@ -244,8 +244,7 @@ struct AllocaScopeOpInterfaceReverse gutils->zeroDiffe(v, builder); } - auto newScope = - cast(gutils->getNewFromOriginal(op)); + auto newScope = cast(gutils->getNewFromOriginal(op)); newScope->moveBefore(builder.getInsertionBlock(), builder.getInsertionPoint());