Skip to content

Commit 54a896e

Browse files
xys-syxvimarsh6739
xys-syx
authored and
vimarsh6739
committed
[MLIR] Fix reverse-mode shadows for mutable allocations
Allow setDiffe to replace placeholder shadows for mutable values in reverse mode, while preserving existing user-provided shadows for mutable arguments. Add memref.alloca coverage for both internal allocation shadows and memref argument shadows.
1 parent d7dd2e0 commit 54a896e

2 files changed

Lines changed: 123 additions & 1 deletion

File tree

‎enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,8 +225,20 @@ void mlir::enzyme::MGradientUtils::setDiffe(mlir::Value val, mlir::Value toset,
225225
}
226226
assert(!isConstantValue(val));
227227

228+
// Reverse mode also uses invertedPointers for mutable values. For allocation
229+
// results, forceAugmentedReturns() installs a placeholder shadow; the
230+
// allocation derivative later calls setDiffe with the real shadow, so replace
231+
// it here.
232+
bool isMutable = cast<AutoDiffTypeInterface>(val.getType()).isMutable();
228233
if (mode == DerivativeMode::ForwardMode ||
229-
mode == DerivativeMode::ForwardModeSplit) {
234+
mode == DerivativeMode::ForwardModeSplit || isMutable) {
235+
if (isMutable) {
236+
// Mutable arguments may already be mapped to a real user-provided shadow
237+
// rather than a placeholder. Only replace placeholders here.
238+
auto existing = invertedPointers.lookupOrNull(val);
239+
if (existing && !existing.getDefiningOp<enzyme::PlaceholderOp>())
240+
return;
241+
}
230242
setInvertedPointer(val, toset);
231243
}
232244
/*

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

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
// RUN: %eopt --enzyme %s | FileCheck --implicit-check-not=enzyme.placeholder %s
2+
3+
func.func @foo_flat(%x : f64) -> f64 {
4+
%buf = memref.alloca() : memref<f64>
5+
memref.store %x, %buf[] : memref<f64>
6+
%y = memref.load %buf[] : memref<f64>
7+
return %y : f64
8+
}
9+
10+
func.func @dfoo_flat(%x: f64, %dout : f64) -> f64 {
11+
%dx = enzyme.autodiff @foo_flat(%x, %dout) {
12+
activity = [#enzyme<activity enzyme_active>],
13+
ret_activity = [#enzyme<activity enzyme_activenoneed>]
14+
} : (f64, f64) -> (f64)
15+
return %dx : f64
16+
}
17+
18+
// CHECK-LABEL: func.func private @diffefoo_flat(
19+
// CHECK-SAME: %[[X:[^,]+]]: f64,
20+
// CHECK-SAME: %[[DOUT:[^)]+]]: f64) -> f64 {
21+
22+
// CHECK: %[[GX:.+]] = "enzyme.init"() : () -> !enzyme.Gradient<f64>
23+
// CHECK: "enzyme.set"(%[[GX]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
24+
// CHECK: %[[CS:.+]] = "enzyme.init"() : () -> !enzyme.Cache<memref<f64>>
25+
// CHECK: %[[CL:.+]] = "enzyme.init"() : () -> !enzyme.Cache<memref<f64>>
26+
// CHECK: %[[GY:.+]] = "enzyme.init"() : () -> !enzyme.Gradient<f64>
27+
// CHECK: "enzyme.set"(%[[GY]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
28+
29+
// CHECK: %[[DBUF:.+]] = memref.alloca() : memref<f64>
30+
// CHECK: %[[ZERO_INIT:.+]] = arith.constant 0.000000e+00 : f64
31+
// CHECK: linalg.fill ins(%[[ZERO_INIT]] : f64) outs(%[[DBUF]] : memref<f64>)
32+
33+
// CHECK: %[[BUF:.+]] = memref.alloca() : memref<f64>
34+
// CHECK: "enzyme.push"(%[[CS]], %[[DBUF]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
35+
// CHECK: memref.store %[[X]], %[[BUF]][] : memref<f64>
36+
// CHECK: "enzyme.push"(%[[CL]], %[[DBUF]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
37+
// CHECK: memref.load %[[BUF]][] : memref<f64>
38+
// CHECK: cf.br ^bb1
39+
// CHECK: ^bb1:
40+
41+
// CHECK: %{{.+}} = "enzyme.get"(%[[GY]]) : (!enzyme.Gradient<f64>) -> f64
42+
// CHECK: arith.addf %{{.+}}, %[[DOUT]] : f64
43+
// CHECK: "enzyme.set"(%[[GY]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
44+
45+
// CHECK: %{{.+}} = "enzyme.get"(%[[GY]]) : (!enzyme.Gradient<f64>) -> f64
46+
// CHECK: %[[POPL:.+]] = "enzyme.pop"(%[[CL]]) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
47+
// CHECK: memref.load %[[POPL]][] : memref<f64>
48+
// CHECK: arith.addf
49+
// CHECK: memref.store %{{.+}}, %[[POPL]][] : memref<f64>
50+
51+
// CHECK: %[[POPS:.+]] = "enzyme.pop"(%[[CS]]) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
52+
// CHECK: memref.load %[[POPS]][] : memref<f64>
53+
// CHECK: %{{.+}} = "enzyme.get"(%[[GX]]) : (!enzyme.Gradient<f64>) -> f64
54+
// CHECK: arith.addf
55+
// CHECK: "enzyme.set"(%[[GX]], %{{.*}}) : (!enzyme.Gradient<f64>, f64) -> ()
56+
// CHECK: %[[ZERO_CLEAR:.+]] = arith.constant 0.000000e+00 : f64
57+
// CHECK: memref.store %[[ZERO_CLEAR]], %[[POPS]][] : memref<f64>
58+
59+
// CHECK: %{{.+}} = "enzyme.get"(%[[GX]]) : (!enzyme.Gradient<f64>) -> f64
60+
// CHECK: return %{{.+}} : f64
61+
62+
63+
// -----
64+
65+
func.func @load_arg(%mem: memref<f64>) -> f64 {
66+
%buf = memref.alloca() : memref<f64>
67+
%x = memref.load %mem[] : memref<f64>
68+
memref.store %x, %buf[] : memref<f64>
69+
%y = memref.load %buf[] : memref<f64>
70+
return %y : f64
71+
}
72+
73+
func.func @dload_arg(%mem: memref<f64>, %dmem: memref<f64>, %dout : f64) {
74+
enzyme.autodiff @load_arg(%mem, %dmem, %dout) {
75+
activity = [#enzyme<activity enzyme_dup>],
76+
ret_activity = [#enzyme<activity enzyme_activenoneed>]
77+
} : (memref<f64>, memref<f64>, f64) -> ()
78+
return
79+
}
80+
81+
// CHECK-LABEL: func.func private @diffeload_arg(
82+
// CHECK-SAME: %[[MEM:[^,]+]]: memref<f64>,
83+
// CHECK-SAME: %[[DMEM:[^,]+]]: memref<f64>,
84+
// CHECK-SAME: %[[DOUT:[^)]+]]: f64) {
85+
// CHECK: %[[DBUF:.+]] = memref.alloca() : memref<f64>
86+
// CHECK: linalg.fill ins(%{{.+}} : f64) outs(%[[DBUF]] : memref<f64>)
87+
// CHECK: %[[BUF:.+]] = memref.alloca() : memref<f64>
88+
// CHECK: "enzyme.push"(%{{.+}}, %[[DMEM]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
89+
// CHECK: memref.load %[[MEM]][] : memref<f64>
90+
// CHECK: "enzyme.push"(%{{.+}}, %[[DBUF]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
91+
// CHECK: memref.store %{{.+}}, %[[BUF]][] : memref<f64>
92+
// CHECK: "enzyme.push"(%{{.+}}, %[[DBUF]]) : (!enzyme.Cache<memref<f64>>, memref<f64>) -> ()
93+
// CHECK: memref.load %[[BUF]][] : memref<f64>
94+
// CHECK: arith.addf %{{.+}}, %[[DOUT]] : f64
95+
// CHECK: %[[POPL:.+]] = "enzyme.pop"(%{{.+}}) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
96+
// CHECK: memref.load %[[POPL]][] : memref<f64>
97+
// CHECK: arith.addf
98+
// CHECK: memref.store %{{.+}}, %[[POPL]][] : memref<f64>
99+
// CHECK: %[[POPS:.+]] = "enzyme.pop"(%{{.+}}) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
100+
// CHECK: memref.load %[[POPS]][] : memref<f64>
101+
// CHECK: "enzyme.get"
102+
// CHECK: arith.addf
103+
// CHECK: "enzyme.set"
104+
// CHECK: memref.store %{{.+}}, %[[POPS]][] : memref<f64>
105+
// CHECK: "enzyme.get"
106+
// CHECK: %[[POPD:.+]] = "enzyme.pop"(%{{.+}}) : (!enzyme.Cache<memref<f64>>) -> memref<f64>
107+
// CHECK: memref.load %[[POPD]][] : memref<f64>
108+
// CHECK: arith.addf
109+
// CHECK: memref.store %{{.+}}, %[[POPD]][] : memref<f64>
110+
// CHECK: return

0 commit comments

Comments
 (0)