Skip to content

Commit 2fe294e

Browse files
vimarsh6739vimarsh6739
authored
Fix indexing bug in EnzymeBatchDiffPass (#2944)
The op's index was being used to index into primal inputs
1 parent 7be5247 commit 2fe294e

3 files changed

Lines changed: 62 additions & 86 deletions

File tree

‎enzyme/Enzyme/MLIR/Passes/EnzymeBatchDiffPass.cpp‎

Lines changed: 22 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,9 @@ struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
5353
OpBuilder builder(op);
5454

5555
op->walk([&](Block *blk) {
56-
// map tracking batchable AD calls
5756
std::map<enzyme::batchutils::BatchDiffCacheKey,
5857
SmallVector<enzyme::ForwardDiffOp>>
59-
toMerge;
58+
diffMergeSet;
6059

6160
for (auto fwdOp : blk->getOps<enzyme::ForwardDiffOp>()) {
6261
auto fnOp = dyn_cast_or_null<FunctionOpInterface>(
@@ -67,10 +66,10 @@ struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
6766
batchutils::BatchDiffCacheKey key =
6867
batchutils::createDiffCacheKey(fwdOp, fnOp);
6968

70-
toMerge[key].push_back(fwdOp);
69+
diffMergeSet[key].push_back(fwdOp);
7170
}
7271

73-
for (auto &pair : toMerge) {
72+
for (auto &pair : diffMergeSet) {
7473
auto key = pair.first;
7574
auto allDiffs = pair.second;
7675
if (allDiffs.size() < 2)
@@ -185,15 +184,12 @@ struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
185184
SmallVector<ForwardDiffOp> legalMerge = batchutils::pruneMemoryEffects(
186185
symbolTable, key, prunedSources, callerEffectMap, innerEffectCache);
187186

188-
// go ahead and actually do the merge now
189187
{
190188
SmallVector<enzyme::ForwardDiffOp> &allOps = legalMerge;
191189
int64_t width = allOps.size();
192-
193190
if (width < 2)
194191
continue;
195192

196-
// We will insert the merged op before the first fwddiff call
197193
auto firstDiffOp = allOps.front();
198194
IRRewriter::InsertionGuard insertGuard(builder);
199195
builder.setInsertionPoint(firstDiffOp);
@@ -205,14 +201,12 @@ struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
205201
SmallVector<ActivityAttr, 2> retActivityAttrs;
206202
SmallVector<mlir::Type, 2> out_ty;
207203
auto in_idx = 0;
208-
209-
// process input, d<input>
210204
for (auto [idx, act] : llvm::enumerate(key.inActivity)) {
211205
ActivityAttr iattr = ActivityAttr::get(context, act);
212206
inActivityAttrs.push_back(iattr);
213-
in_args.push_back(key.inputs[in_idx]);
207+
in_args.push_back(key.inputs[idx]);
214208
in_idx++;
215-
209+
// batched input derivative
216210
SmallVector<mlir::Value> derivList;
217211
if (act == Activity::enzyme_dup ||
218212
act == Activity::enzyme_dupnoneed) {
@@ -227,104 +221,90 @@ struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
227221
}
228222
}
229223

230-
// process out, d<out> (only need types)
231224
auto out_idx = 0;
232225
for (auto [idx, ract] : llvm::enumerate(key.retActivity)) {
233226
ActivityAttr iattr = ActivityAttr::get(context, ract);
234-
235227
retActivityAttrs.push_back(iattr);
236228
switch (ract) {
237-
238229
case Activity::enzyme_active: {
239230
mlir::Value res = firstDiffOp.getOutputs()[out_idx];
240231
out_ty.push_back(res.getType());
241232
++out_idx;
242233
break;
243234
}
244-
245235
case Activity::enzyme_const: {
246236
mlir::Value res = firstDiffOp.getOutputs()[out_idx];
247237
out_ty.push_back(res.getType());
248238
++out_idx;
249239
break;
250240
}
251-
252241
case Activity::enzyme_dupnoneed: {
253-
// derivative
254-
242+
// batched output derivative
255243
mlir::Value dres = firstDiffOp.getOutputs()[out_idx];
256244
out_ty.push_back(getConcatType(dres, width));
257245
++out_idx;
258246
break;
259247
}
260-
261248
case Activity::enzyme_dup: {
262249
mlir::Value res = firstDiffOp.getOutputs()[out_idx];
263250
out_ty.push_back(res.getType());
264-
265251
++out_idx;
266-
267-
// derivative
252+
// batched output derivative
268253
mlir::Value dres = firstDiffOp.getOutputs()[out_idx];
269254
out_ty.push_back(getConcatType(dres, width));
270255
++out_idx;
271256
break;
272257
}
273-
274258
case Activity::enzyme_constnoneed: {
275259
break;
276260
}
277-
278261
case Activity::enzyme_activenoneed: {
279262
mlir::Value res = firstDiffOp.getOutputs()[out_idx];
280263
out_ty.push_back(res.getType());
281264
++out_idx;
282265
break;
283266
}
284-
285267
default:
286268
llvm_unreachable(
287269
"unknown activity value encountered for ret_activity");
288270
}
289271
}
290272

291-
// create new FwdDiffOp
292273
ArrayAttr newInActivity = ArrayAttr::get(
293274
context, llvm::ArrayRef<Attribute>(inActivityAttrs.begin(),
294275
inActivityAttrs.end()));
295-
296276
ArrayAttr newRetActivity = ArrayAttr::get(
297277
context, llvm::ArrayRef<Attribute>(retActivityAttrs.begin(),
298278
retActivityAttrs.end()));
299-
300279
IntegerAttr newWidthAttr =
301280
IntegerAttr::get(firstDiffOp.getWidthAttr().getType(), width);
302-
303281
auto newDiffOp = ForwardDiffOp::create(
304282
builder, loc, out_ty, firstDiffOp.getFnAttr(), in_args,
305283
newInActivity, newRetActivity, newWidthAttr,
306284
firstDiffOp.getStrongZeroAttr());
307285

308-
// Rename old users of out,d<out> to new users
286+
// Rename uses
287+
// out -> primal
288+
// dout -> derivative
309289
out_idx = 0;
310290
for (auto [idx, ract] : llvm::enumerate(key.retActivity)) {
311291
switch (ract) {
312292
case Activity::enzyme_constnoneed:
313-
// no-op
314293
break;
294+
295+
case Activity::enzyme_active:
296+
case Activity::enzyme_activenoneed:
315297
case Activity::enzyme_const: {
316298
auto new_out = newDiffOp.getOutputs()[out_idx];
317-
318299
for (auto dop : allOps) {
319300
dop.getOutputs()[out_idx].replaceAllUsesWith(new_out);
320301
}
321-
322302
out_idx++;
323303
break;
324304
}
325305

326306
case Activity::enzyme_dupnoneed: {
327-
// derivative
307+
// batched derivative
328308
auto batch_dout = newDiffOp.getOutputs()[out_idx];
329309
for (auto [dop_idx, dop] : llvm::enumerate(allOps)) {
330310
auto old_dout = dop.getOutputs()[out_idx];
@@ -340,45 +320,27 @@ struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
340320

341321
case Activity::enzyme_dup: {
342322
mlir::Value new_out = newDiffOp.getOutputs()[out_idx];
343-
344323
for (ForwardDiffOp dop : allOps) {
345324
dop.getOutputs()[out_idx].replaceAllUsesWith(new_out);
346325
}
347326
out_idx++;
348327

349-
// derivative
328+
// batched derivative
350329
auto batch_dout = newDiffOp.getOutputs()[out_idx];
351330
for (auto [dop_idx, dop] : llvm::enumerate(allOps)) {
352-
353331
auto old_dout = dop.getOutputs()[out_idx];
354332
auto doutTy = old_dout.getType();
355333
auto new_dout =
356334
getExtractValue(builder, loc, doutTy, batch_dout, dop_idx);
357335

358336
old_dout.replaceAllUsesWith(new_dout);
359337
}
360-
361338
++out_idx;
362339
break;
363340
}
364-
case Activity::enzyme_active: {
365-
auto new_out = newDiffOp.getOutputs()[out_idx];
366-
367-
for (ForwardDiffOp dop : allOps) {
368-
dop.getOutputs()[out_idx].replaceAllUsesWith(new_out);
369-
}
370-
out_idx++;
371-
break;
372-
}
373-
case Activity::enzyme_activenoneed: {
374-
auto new_out = newDiffOp.getOutputs()[out_idx];
375-
376-
for (ForwardDiffOp dop : allOps) {
377-
dop.getOutputs()[out_idx].replaceAllUsesWith(new_out);
378-
}
379-
out_idx++;
380-
break;
381-
}
341+
default:
342+
llvm_unreachable(
343+
"unknown activity value encountered for ret_activity");
382344
}
383345
}
384346

@@ -408,7 +370,7 @@ struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
408370
// map tracking batchable AD calls
409371
std::map<enzyme::batchutils::BatchDiffCacheKey,
410372
SmallVector<enzyme::AutoDiffOp>>
411-
toMerge;
373+
diffMergeSet;
412374

413375
for (auto revOp : blk->getOps<enzyme::AutoDiffOp>()) {
414376
auto fnOp = dyn_cast_or_null<FunctionOpInterface>(
@@ -419,10 +381,10 @@ struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
419381
batchutils::BatchDiffCacheKey key =
420382
batchutils::createDiffCacheKey(revOp, fnOp);
421383

422-
toMerge[key].push_back(revOp);
384+
diffMergeSet[key].push_back(revOp);
423385
}
424386

425-
for (auto &pair : toMerge) {
387+
for (auto &pair : diffMergeSet) {
426388
auto key = pair.first;
427389
auto allDiffs = pair.second;
428390
if (allDiffs.size() < 2)
@@ -552,7 +514,7 @@ struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
552514
for (auto [idx, act] : llvm::enumerate(key.inActivity)) {
553515
auto iattr = ActivityAttr::get(context, act);
554516
inActivityAttrs.push_back(iattr);
555-
in_args.push_back(key.inputs[call_idx]);
517+
in_args.push_back(key.inputs[idx]);
556518
call_idx++;
557519

558520
if (act == Activity::enzyme_dup ||
@@ -564,7 +526,6 @@ struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
564526
}
565527

566528
mlir::Value b_din = getConcatValue(builder, loc, derivList);
567-
568529
in_args.push_back(b_din);
569530
call_idx++;
570531
}

‎enzyme/Enzyme/MLIR/Passes/EnzymeBatchDiffPass.h‎

Lines changed: 3 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -64,31 +64,9 @@ struct BatchDiffCacheKey {
6464
template <typename SourceOp>
6565
BatchDiffCacheKey createDiffCacheKey(SourceOp uop, FunctionOpInterface fn) {
6666
// extract in_activity, ret_activity, in_args
67-
SmallVector<Activity> inActivity;
68-
SmallVector<Activity> retActivity;
69-
SmallVector<Value> in_args;
70-
71-
auto in_idx = 0;
72-
73-
for (auto [idx, act] : llvm::enumerate(uop.getActivity())) {
74-
auto iattr = cast<ActivityAttr>(act);
75-
auto val = iattr.getValue();
76-
inActivity.push_back(val);
77-
78-
in_args.push_back(uop.getInputs()[in_idx]);
79-
++in_idx;
80-
81-
if (val == Activity::enzyme_dup || val == Activity::enzyme_dupnoneed) {
82-
++in_idx;
83-
}
84-
}
85-
86-
for (auto [idx, ract] : llvm::enumerate(uop.getRetActivity())) {
87-
auto iattr = cast<ActivityAttr>(ract);
88-
auto val = iattr.getValue();
89-
retActivity.push_back(val);
90-
}
91-
67+
SmallVector<Activity> inActivity(uop.getActivity().template getAsValueRange<ActivityAttr>());
68+
SmallVector<Activity> retActivity(uop.getRetActivity().template getAsValueRange<ActivityAttr>());
69+
SmallVector<Value> in_args = uop.getPrimalInputs();
9270
batchutils::BatchDiffCacheKey key{fn, in_args, inActivity, retActivity,
9371
uop->getBlock()};
9472
return key;

‎enzyme/test/MLIR/OptimizeAD/fwd_batch.mlir‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,43 @@ module {
3333

3434
// -----
3535

36+
// Multiple duplicated inputs followed by a constant input
37+
module {
38+
func.func @muladd(%x : f64, %y : f64, %c : f64) -> f64 {
39+
%xy = arith.mulf %x, %y : f64
40+
%z = arith.addf %xy, %c : f64
41+
return %z : f64
42+
}
43+
func.func @test3(%x : f64, %dx1 : f64, %dx2 : f64,
44+
%y : f64, %dy1 : f64, %dy2 : f64, %c : f64) -> (f64, f64) {
45+
%r1 = enzyme.fwddiff @muladd(%x, %dx1, %y, %dy1, %c) { activity=[#enzyme<activity enzyme_dup>, #enzyme<activity enzyme_dup>, #enzyme<activity enzyme_const>], ret_activity=[#enzyme<activity enzyme_dupnoneed>] } : (f64, f64, f64, f64, f64) -> f64
46+
%r2 = enzyme.fwddiff @muladd(%x, %dx2, %y, %dy2, %c) { activity=[#enzyme<activity enzyme_dup>, #enzyme<activity enzyme_dup>, #enzyme<activity enzyme_const>], ret_activity=[#enzyme<activity enzyme_dupnoneed>] } : (f64, f64, f64, f64, f64) -> f64
47+
return %r1, %r2 : f64, f64
48+
}
49+
}
50+
51+
// CHECK-LABEL: func.func @test3
52+
// CHECK-SAME: (%[[X:.*]]: f64, %[[DX1:.*]]: f64, %[[DX2:.*]]: f64, %[[Y:.*]]: f64, %[[DY1:.*]]: f64, %[[DY2:.*]]: f64, %[[C:.*]]: f64)
53+
// CHECK: %[[DX:.*]] = enzyme.concat(%[[DX1]], %[[DX2]]) : (f64, f64) -> tensor<2xf64>
54+
// CHECK: %[[DY:.*]] = enzyme.concat(%[[DY1]], %[[DY2]]) : (f64, f64) -> tensor<2xf64>
55+
// CHECK: %[[BATCHED_RES:.*]] = enzyme.fwddiff @muladd(%[[X]], %[[DX]], %[[Y]], %[[DY]], %[[C]]) {{.*}} width = 2 {{.*}} : (f64, tensor<2xf64>, f64, tensor<2xf64>, f64) -> tensor<2xf64>
56+
// CHECK: %[[RES0:.*]] = enzyme.extract %[[BATCHED_RES]][0] : (tensor<2xf64>) -> f64
57+
// CHECK-NEXT: %[[RES1:.*]] = enzyme.extract %[[BATCHED_RES]][1] : (tensor<2xf64>) -> f64
58+
// CHECK-NEXT: return %[[RES0]], %[[RES1]]
59+
60+
// LEGAL-LABEL: func.func @test3
61+
// LEGAL-SAME: (%[[X:.*]]: f64, %[[DX1:.*]]: f64, %[[DX2:.*]]: f64, %[[Y:.*]]: f64, %[[DY1:.*]]: f64, %[[DY2:.*]]: f64, %[[C:.*]]: f64)
62+
// LEGAL: %[[DX:.*]] = tensor.from_elements %[[DX1]], %[[DX2]] : tensor<2xf64>
63+
// LEGAL: %[[DY:.*]] = tensor.from_elements %[[DY1]], %[[DY2]] : tensor<2xf64>
64+
// LEGAL: %[[BATCHED_RES:.*]] = enzyme.fwddiff @muladd(%[[X]], %[[DX]], %[[Y]], %[[DY]], %[[C]]) {{.*}} width = 2 {{.*}} : (f64, tensor<2xf64>, f64, tensor<2xf64>, f64) -> tensor<2xf64>
65+
// LEGAL: %[[C0:.*]] = arith.constant 0 : index
66+
// LEGAL-NEXT: %[[RES0:.*]] = tensor.extract %[[BATCHED_RES]][%[[C0]]] : tensor<2xf64>
67+
// LEGAL-NEXT: %[[C1:.*]] = arith.constant 1 : index
68+
// LEGAL-NEXT: %[[RES1:.*]] = tensor.extract %[[BATCHED_RES]][%[[C1]]] : tensor<2xf64>
69+
// LEGAL-NEXT: return %[[RES0]], %[[RES1]]
70+
71+
// -----
72+
3673
// 2. Tensor test
3774
module {
3875
func.func @square(%x : tensor<10xf64>) -> tensor<10xf64>{

0 commit comments

Comments
 (0)