Skip to content

Commit 238dde5

Browse files
xys-syxxys-syx
committed
unified alloca_scope
1 parent fa61965 commit 238dde5

2 files changed

Lines changed: 132 additions & 31 deletions

File tree

‎enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp‎

Lines changed: 79 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -220,19 +220,41 @@ struct AllocaScopeOpInterfaceReverse
220220
}
221221
}
222222

223-
auto revScope =
224-
memref::AllocaScopeOp::create(builder, op->getLoc(), TypeRange());
225-
Block *revBody = builder.createBlock(&revScope.getBodyRegion());
226-
OpBuilder bodyBuilder(revBody, revBody->end());
227-
memref::AllocaScopeReturnOp::create(bodyBuilder, op->getLoc(),
228-
ValueRange());
229-
bodyBuilder.setInsertionPoint(revBody->getTerminator());
223+
Region &scopeRegion = scopeOp.getBodyRegion();
224+
SmallVector<Value> capturedInputs;
225+
scopeRegion.walk([&](Operation *inner) {
226+
for (Value operand : inner->getOperands()) {
227+
Region *defRegion = operand.getParentRegion();
228+
if (!defRegion || scopeRegion.isAncestor(defRegion))
229+
continue;
230+
if (gutils->isConstantValue(operand))
231+
continue;
232+
auto iface = dyn_cast<AutoDiffTypeInterface>(operand.getType());
233+
if (!iface || iface.isMutable())
234+
continue;
235+
if (!llvm::is_contained(capturedInputs, operand))
236+
capturedInputs.push_back(operand);
237+
}
238+
});
239+
240+
SmallVector<Value> capturedPre;
241+
capturedPre.reserve(capturedInputs.size());
242+
for (Value v : capturedInputs) {
243+
capturedPre.push_back(gutils->diffe(v, builder));
244+
gutils->zeroDiffe(v, builder);
245+
}
246+
247+
auto newScope =
248+
cast<memref::AllocaScopeOp>(gutils->getNewFromOriginal(op));
249+
newScope->moveBefore(builder.getInsertionBlock(),
250+
builder.getInsertionPoint());
251+
252+
Block &newBody = newScope.getBodyRegion().front();
253+
OpBuilder bodyBuilder(newBody.getTerminator());
230254

231255
Block &oldBody = scopeOp.getBodyRegion().front();
232256
bool valid = true;
233257

234-
// Values defined in the scoped region cannot be used outside it. Reset
235-
// their adjoints before propagating gradients through the scoped body.
236258
for (Operation &innerOp : oldBody.getOperations()) {
237259
for (Value result : innerOp.getResults()) {
238260
if (!gutils->isConstantValue(result)) {
@@ -253,16 +275,61 @@ struct AllocaScopeOpInterfaceReverse
253275
if (!gutils->isConstantValue(operand))
254276
gutils->addToDiffe(operand, incomingGradients[incomingIdx],
255277
bodyBuilder);
256-
++incomingIdx;
278+
incomingIdx++;
257279
}
258280

259281
auto first = oldBody.rbegin();
260-
++first;
282+
first++;
261283

262-
for (auto it = first; it != oldBody.rend(); ++it) {
284+
for (auto it = first; it != oldBody.rend(); it++) {
263285
valid &= gutils->Logic.visitChild(&*it, bodyBuilder, gutils).succeeded();
264286
}
265287

288+
if (capturedInputs.empty())
289+
return success(valid);
290+
291+
SmallVector<Value> contributions;
292+
contributions.reserve(capturedInputs.size());
293+
for (Value v : capturedInputs)
294+
contributions.push_back(gutils->diffe(v, bodyBuilder));
295+
296+
unsigned numPrimal = newScope->getNumResults();
297+
Operation *bodyTerm = newBody.getTerminator();
298+
SmallVector<Value> retVals(bodyTerm->getOperands().begin(),
299+
bodyTerm->getOperands().end());
300+
retVals.append(contributions.begin(), contributions.end());
301+
302+
SmallVector<Type> newResultTypes(newScope->getResultTypes().begin(),
303+
newScope->getResultTypes().end());
304+
for (Value v : capturedInputs)
305+
newResultTypes.push_back(gutils->getShadowType(v.getType()));
306+
307+
OpBuilder scopeBuilder(newScope);
308+
auto extScope = memref::AllocaScopeOp::create(scopeBuilder, op->getLoc(),
309+
newResultTypes);
310+
extScope.getBodyRegion().takeBody(newScope.getBodyRegion());
311+
312+
Block &extBody = extScope.getBodyRegion().front();
313+
Operation *movedTerm = extBody.getTerminator();
314+
OpBuilder termBuilder(movedTerm);
315+
memref::AllocaScopeReturnOp::create(termBuilder, movedTerm->getLoc(),
316+
retVals);
317+
gutils->erase(movedTerm);
318+
319+
for (unsigned i = 0; i < numPrimal; ++i) {
320+
newScope->getResult(i).replaceAllUsesWith(extScope.getResult(i));
321+
gutils->originalToNewFn.map(scopeOp->getResult(i), extScope.getResult(i));
322+
}
323+
gutils->erase(newScope);
324+
325+
builder.setInsertionPointAfter(extScope);
326+
for (auto indexed : llvm::enumerate(capturedInputs)) {
327+
Value v = indexed.value();
328+
gutils->setDiffe(v, capturedPre[indexed.index()], builder);
329+
gutils->addToDiffe(v, extScope.getResult(numPrimal + indexed.index()),
330+
builder);
331+
}
332+
266333
return success(valid);
267334
}
268335

‎enzyme/test/MLIR/ReverseMode/alloca_scope2.mlir‎

Lines changed: 53 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -10,29 +10,63 @@ func.func @foo(%x : f64) -> f64{
1010
return %out : f64
1111
}
1212

13-
1413
func.func @dfoo(%x: f64, %dout : f64) -> f64 {
1514
%dx = enzyme.autodiff @foo(%x, %dout) {
16-
activity = [#enzyme<activity enzyme_active>],
17-
ret_activity = [#enzyme<acitivity enzyme_activenoneed>]
18-
} : (f64, f64) -> (f64)
19-
return %dx : f64
20-
}
21-
22-
func.func @foo2(%x : f64) -> f64{
23-
%cst = arith.constant 0.0000e+00 : f64
24-
%buf = memref.alloca() : memref<f64>
25-
memref.store %x, %buf[] : memref<f64>
26-
%y = memref.load %buf[] : memref<f64>
27-
%out = arith.addf %cst, %y : f64
28-
return %out : f64
29-
}
30-
31-
32-
func.func @dfoo2(%x: f64, %dout : f64) -> f64 {
33-
%dx = enzyme.autodiff @foo2(%x, %dout) {
3415
activity = [#enzyme<activity enzyme_active>],
3516
ret_activity = [#enzyme<activity enzyme_activenoneed>]
3617
} : (f64, f64) -> (f64)
3718
return %dx : f64
3819
}
20+
21+
// CHECK-LABEL: func.func private @diffefoo(
22+
// CHECK-SAME: %[[X:[^,]+]]: f64,
23+
// CHECK-SAME: %[[DOUT:[^)]+]]: f64) -> f64 {
24+
// CHECK: %[[C0:.*]] = "enzyme.init"() : () -> !enzyme.Cache<memref<f64>>
25+
// CHECK: %[[C1:.*]] = "enzyme.init"() : () -> !enzyme.Cache<memref<f64>>
26+
// CHECK: %[[G0:.*]] = "enzyme.init"() : () -> !enzyme.Gradient<f64>
27+
// CHECK: "enzyme.set"(%[[G0]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
28+
// CHECK: %[[G1:.*]] = "enzyme.init"() : () -> !enzyme.Gradient<f64>
29+
// CHECK: "enzyme.set"(%[[G1]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
30+
// CHECK: %[[G2:.*]] = "enzyme.init"() : () -> !enzyme.Gradient<f64>
31+
// CHECK: "enzyme.set"(%[[G2]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
32+
// CHECK: cf.br ^bb1
33+
// CHECK: ^bb1:
34+
// CHECK: %[[T5:.*]] = "enzyme.get"(%[[G2]]) : (!enzyme.Gradient<f64>) -> f64
35+
// CHECK: %[[T6:.*]] = arith.addf %[[T5]], %[[DOUT]] : f64
36+
// CHECK: "enzyme.set"(%[[G2]], %[[T6]]) : (!enzyme.Gradient<f64>, f64) -> ()
37+
// CHECK: %[[T7:.*]] = "enzyme.get"(%[[G2]]) : (!enzyme.Gradient<f64>) -> f64
38+
// CHECK: "enzyme.set"(%[[G2]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
39+
// CHECK: %[[T8:.*]] = "enzyme.get"(%[[G1]]) : (!enzyme.Gradient<f64>) -> f64
40+
// CHECK: "enzyme.set"(%[[G1]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
41+
// CHECK: %[[SCOPE:.*]]:2 = memref.alloca_scope -> (f64, f64) {
42+
// CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<f64>
43+
// CHECK: memref.store %{{.*}}, %[[ALLOCA]][] : memref<f64>
44+
// CHECK: %[[ALLOCA5:.*]] = memref.alloca() : memref<f64>
45+
// CHECK: "enzyme.push"(%[[C0]], %[[ALLOCA]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
46+
// CHECK: memref.store %[[X]], %[[ALLOCA5]][] : memref<f64>
47+
// CHECK: "enzyme.push"(%[[C1]], %[[ALLOCA]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
48+
// CHECK: %[[LD:.*]] = memref.load %[[ALLOCA5]][] : memref<f64>
49+
// CHECK: "enzyme.set"(%[[G0]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
50+
// CHECK: %[[T14:.*]] = "enzyme.get"(%[[G0]]) : (!enzyme.Gradient<f64>) -> f64
51+
// CHECK: %[[T15:.*]] = arith.addf %[[T14]], %[[T7]] : f64
52+
// CHECK: "enzyme.set"(%[[G0]], %[[T15]]) : (!enzyme.Gradient<f64>, f64) -> ()
53+
// CHECK: %[[T16:.*]] = "enzyme.get"(%[[G0]]) : (!enzyme.Gradient<f64>) -> f64
54+
// CHECK: %[[POP1:.*]] = "enzyme.pop"(%[[C1]]) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
55+
// CHECK: %[[LD18:.*]] = memref.load %[[POP1]][] : memref<f64>
56+
// CHECK: %[[T19:.*]] = arith.addf %[[LD18]], %[[T16]] : f64
57+
// CHECK: memref.store %[[T19]], %[[POP1]][] : memref<f64>
58+
// CHECK: %[[POP0:.*]] = "enzyme.pop"(%[[C0]]) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
59+
// CHECK: %[[LD21:.*]] = memref.load %[[POP0]][] : memref<f64>
60+
// CHECK: %[[T22:.*]] = "enzyme.get"(%[[G1]]) : (!enzyme.Gradient<f64>) -> f64
61+
// CHECK: %[[T23:.*]] = arith.addf %[[T22]], %[[LD21]] : f64
62+
// CHECK: "enzyme.set"(%[[G1]], %[[T23]]) : (!enzyme.Gradient<f64>, f64) -> ()
63+
// CHECK: memref.store %{{.*}}, %[[POP0]][] : memref<f64>
64+
// CHECK: %[[T24:.*]] = "enzyme.get"(%[[G1]]) : (!enzyme.Gradient<f64>) -> f64
65+
// CHECK: memref.alloca_scope.return %[[LD]], %[[T24]] : f64, f64
66+
// CHECK: }
67+
// CHECK: "enzyme.set"(%[[G1]], %[[T8]]) : (!enzyme.Gradient<f64>, f64) -> ()
68+
// CHECK: %[[T10:.*]] = "enzyme.get"(%[[G1]]) : (!enzyme.Gradient<f64>) -> f64
69+
// CHECK: %[[T11:.*]] = arith.addf %[[T10]], %[[SCOPE]]#1 : f64
70+
// CHECK: "enzyme.set"(%[[G1]], %[[T11]]) : (!enzyme.Gradient<f64>, f64) -> ()
71+
// CHECK: %[[T12:.*]] = "enzyme.get"(%[[G1]]) : (!enzyme.Gradient<f64>) -> f64
72+
// CHECK: return %[[T12]] : f64

0 commit comments

Comments
 (0)