From fb385903dc0ce602db3fdb1a1c0b0df9c8c395c1 Mon Sep 17 00:00:00 2001 From: mixu Date: Wed, 10 Jun 2026 15:32:16 -0400 Subject: [PATCH] Preserve full shadow memtransfers Ensure reverse-mode shadow copies keep the full memtransfer when Enzyme also emits sliced shadow updates, so inactive pointer fields in copied aggregates remain initialized for later reverse use. --- enzyme/Enzyme/AdjointGenerator.h | 4 +++ .../Enzyme/ReverseMode/full-shadow-memcpy.ll | 32 +++++++++++++++++++ 2 files changed, 36 insertions(+) create mode 100644 enzyme/test/Enzyme/ReverseMode/full-shadow-memcpy.ll diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h index 5dbeb27c909..ea7ed83dbce 100644 --- a/enzyme/Enzyme/AdjointGenerator.h +++ b/enzyme/Enzyme/AdjointGenerator.h @@ -3653,6 +3653,10 @@ class AdjointGenerator : public llvm::InstVisitor { } } + if (size != 0) { + toIterate.emplace_back(nullptr, 0, size); + } + for (auto &&[floatTy_ref, seg_start_ref, seg_size_ref] : toIterate) { auto floatTy = floatTy_ref; auto seg_start = seg_start_ref; diff --git a/enzyme/test/Enzyme/ReverseMode/full-shadow-memcpy.ll b/enzyme/test/Enzyme/ReverseMode/full-shadow-memcpy.ll new file mode 100644 index 00000000000..1f1269b54c4 --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/full-shadow-memcpy.ll @@ -0,0 +1,32 @@ +; RUN: %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="enzyme" -S | FileCheck %s + +%State = type { ptr, double, ptr, double } + +define void @scatter(ptr byval(%State) align 8 %arg) { +entry: + %local = alloca %State, align 8 + call void @llvm.memcpy.p0.p0.i64(ptr align 8 %local, ptr align 8 %arg, i64 32, i1 false) + %out.gep = getelementptr inbounds %State, ptr %local, i32 0, i32 0 + %out = load ptr, ptr %out.gep, align 8 + %scale.gep = getelementptr inbounds %State, ptr %local, i32 0, i32 1 + %scale = load double, ptr %scale.gep, align 8 + %in.gep = getelementptr inbounds %State, ptr %local, i32 0, i32 2 + %in = load ptr, ptr %in.gep, align 8 + %x = load double, ptr %in, align 8 + %res = fmul double %x, %scale + store double %res, ptr %out, align 8 + ret void +} + +declare void @llvm.memcpy.p0.p0.i64(ptr noalias nocapture writeonly, ptr noalias nocapture readonly, i64, i1) + +define void @caller(ptr %arg, ptr %darg) { +entry: + call void (...) @__enzyme_autodiff(void (ptr)* @scatter, ptr %arg, ptr %darg) + ret void +} + +declare void @__enzyme_autodiff(...) + +; CHECK: define internal void @diffescatter(ptr byval(%State) align 8 %arg, ptr align 8 %"arg'") +; CHECK: call void @llvm.memcpy.p0.p0.i64(ptr align 8 %{{.*}}, ptr align 8 %"arg'", i64 32, i1 false)