EnzymeAD / EnzymeAD/Enzyme

Incorrect gradient

Open
#3,174 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
LLVM
Stars
1.7k
Forks
188
Avg merge
1d 22h
Merged PRs (30d)
26

Description

```llvm
; Zeroing memset of a stack slot whose element type TypeAnalysis cannot determine.
;;
; visitMemSetCommon splits on whether the element type is known. With a null secret type
; it emitted the shadow memset only into the forward sweep, so the shadow accumulator
; was never re-zeroed between reverse iterations and any gradient accumulated into it
; grew with the trip count. A zeroing memset kills every prior value in the region, so
; no adjoint can flow past it and the shadow must be zeroed in the reverse sweep too.
;
; Written with opaque pointers, so it is restricted to LLVM 16 and later.

; RUN: if [ %llvmver -ge 16 ]; then %opt < %s %newLoadEnzyme -enzyme-preopt=false -enzyme-strict-aliasing=0 -passes="enzyme" -S | FileCheck %s; fi

declare void @__enzyme_autodiff(...)

declare void @llvm.memset.p0.i64(ptr nocapture writeonly, i8, i64, i1 immarg)

define void @f(ptr %x, i64 %n) {
entry:
%acc = alloca [32 x i8], align 4
br label %loop

loop:
%i = phi i64 [ 0, %entry ], [ %inext, %latch ]
call void @llvm.memset.p0.i64(ptr align 4 %acc, i8 0, i64 32, i1 false)
%v = load float, ptr %x, align 4
store float %v, ptr %acc, align 4
%a = load float, ptr %acc, align 4
store float %a, ptr %x, align 4
br label %latch

latch:
%inext = add nuw i64 %i, 1
%done = icmp eq i64 %inext, %n
br i1 %done, label %exit, label %loop

exit:
ret void
}

define void @df(ptr %x, ptr %dx, i64 %n) {
call void (...) @__enzyme_autodiff(ptr @f, metadata !"enzyme_dup", ptr %x, ptr %dx, metadata !"enzyme_const", i64 %n)
ret void
}

; The shadow accumulator must be re-zeroed in the reverse sweep. Before a fix no
; memset is emitted into the invert blocks at all.
```

Contributor guide

Open the contributing guide

Research direction

Start at visitMemSetCommon and reproduce the issue with the provided LLVM IR using the RUN command and FileCheck. Verify that a zeroing memset for the shadow accumulator is emitted in the reverse or invert blocks, so the gradient does not accumulate across reverse iterations.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.