diff --git a/include/clad/Differentiator/BuiltinDerivatives.h b/include/clad/Differentiator/BuiltinDerivatives.h index d13f6d81e..4c4e499ae 100644 --- a/include/clad/Differentiator/BuiltinDerivatives.h +++ b/include/clad/Differentiator/BuiltinDerivatives.h @@ -12,6 +12,8 @@ #include #include +#include +#include #include #define elidable_reverse_forw __attribute__((annotate("elidable_reverse_forw"))) @@ -1600,6 +1602,30 @@ void constructor_pullback(ValueAndPushforward rhs, } } // namespace class_functions } // namespace custom_derivatives + +// Reverse-mode helper for an in-place realloc: resize `ptr` back to the byte +// size it had before the forward realloc. When that regrows a buffer the +// forward pass shrank, zero the re-grown tail for adjoint buffers so fresh +// derivatives start at 0; primal buffers pass zeroGrownTail=false, their tail +// being overwritten by value restores. Returns the possibly-moved pointer. +// NOLINTBEGIN(cppcoreguidelines-no-malloc) +// NOLINTBEGIN(cppcoreguidelines-owning-memory) +// NOLINTBEGIN(cppcoreguidelines-pro-bounds-pointer-arithmetic) +inline void* reverse_realloc(void* ptr, size_t oldBytes, size_t newBytes, + bool zeroGrownTail) { + void* resized = ::realloc(ptr, oldBytes); + if (!resized) // realloc failed: keep the original buffer instead of leaking + // it + return ptr; + ptr = resized; + if (zeroGrownTail && oldBytes > newBytes) + ::memset(static_cast(ptr) + newBytes, 0, oldBytes - newBytes); + return ptr; +} +// NOLINTEND(cppcoreguidelines-pro-bounds-pointer-arithmetic) +// NOLINTEND(cppcoreguidelines-owning-memory) +// NOLINTEND(cppcoreguidelines-no-malloc) + } // namespace clad // FIXME: These math functions depend on promote_2 just like pow: diff --git a/include/clad/Differentiator/ReverseModeVisitor.h b/include/clad/Differentiator/ReverseModeVisitor.h index 8d14453ad..14dd4b40c 100644 --- a/include/clad/Differentiator/ReverseModeVisitor.h +++ b/include/clad/Differentiator/ReverseModeVisitor.h @@ -591,6 +591,18 @@ namespace clad { /// @returns The call to memset if the condition is met, otherwise nullptr. clang::Expr* CheckAndBuildCallToMemset(clang::Expr* LHS, clang::Expr* RHS); + /// The byte-size operand of a malloc/calloc/realloc call, in the primal's + /// terms: malloc(n) -> n, calloc(k, s) -> k*s, realloc(p, n) -> n. Returns + /// null if `allocExpr` is not one of those. Used to track allocation sizes + /// so a realloc can be undone in the reverse sweep. + clang::Expr* buildAllocByteSize(clang::Expr* allocExpr); + + /// Pointer variables that are the target of an in-place `p = realloc(p, n)` + /// somewhere in the function -- only these need an allocation-size shadow; + /// one for any other allocated pointer would be dead. Filled by a pre-scan + /// of the body before differentiation, then read at each allocation. + llvm::SmallPtrSet m_InPlaceReallocPtrs; + static DeclDiff DifferentiateStaticAssertDecl(const clang::StaticAssertDecl* SAD); diff --git a/include/clad/Differentiator/VisitorBase.h b/include/clad/Differentiator/VisitorBase.h index 8fe7341bd..c46de80b9 100644 --- a/include/clad/Differentiator/VisitorBase.h +++ b/include/clad/Differentiator/VisitorBase.h @@ -242,6 +242,11 @@ namespace clad { struct AdjointInfo { clang::VarDecl* Decl = nullptr; enum WrapKind : std::uint8_t { Plain, Deref, ParenDeref } Wrap = Plain; + /// For a pointer allocated with malloc/calloc/realloc, the `size_t` + /// shadow holding its current allocation size in bytes; null otherwise. + /// Reverse mode uses it to undo an in-place realloc. See + /// ReverseModeVisitor's VisitBinaryOperator/DifferentiateVarDecl. + clang::VarDecl* AllocSize = nullptr; }; /// Map used to keep track of variable declarations and match them /// with their derivatives. diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp index 086a263d9..38cdfd0b7 100644 --- a/lib/Differentiator/ReverseModeVisitor.cpp +++ b/lib/Differentiator/ReverseModeVisitor.cpp @@ -269,6 +269,9 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { return AllocCallInfo(k, CE); } + [[nodiscard]] Kind getKind() const { return m_Kind; } + [[nodiscard]] clang::CallExpr* getCall() const { return m_Call; } + // The number-of-bytes operand that a following memset must zero: // malloc(n) -> n, realloc(p, n) -> n. calloc self-zeroes and needs no // memset, so it (and free/none) report null here. @@ -318,6 +321,21 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { return nullptr; } + Expr* ReverseModeVisitor::buildAllocByteSize(Expr* allocExpr) { + AllocCallInfo ac = AllocCallInfo::recognize(allocExpr); + CallExpr* call = ac.getCall(); + switch (ac.getKind()) { + case AllocCallInfo::Kind::Malloc: + return call->getArg(0); + case AllocCallInfo::Kind::Realloc: + return call->getArg(1); + case AllocCallInfo::Kind::Calloc: + return BuildOp(BO_Mul, call->getArg(0), call->getArg(1)); + default: + return nullptr; + } + } + ReverseModeVisitor::ReverseModeVisitor(DerivativeBuilder& builder, const DiffRequest& request) : VisitorBase(builder, request) {} @@ -552,6 +570,26 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { } } + // Pre-scan the body for in-place `p = realloc(p, n)`. Only those pointers + // get an allocation-size shadow at their malloc/calloc/realloc; a shadow + // for any other allocated pointer would be dead code. + for (llvm::SmallVector worklist = {m_DiffReq->getBody()}; + !worklist.empty();) { + const Stmt* S = worklist.pop_back_val(); + if (!S) + continue; + if (const auto* BO = dyn_cast(S)) + if (BO->getOpcode() == BO_Assign) { + AllocCallInfo ac = AllocCallInfo::recognize(BO->getRHS()); + if (ac.isInPlaceRealloc(BO->getLHS())) + m_InPlaceReallocPtrs.insert(cast( + cast(BO->getLHS()->IgnoreParenCasts()) + ->getDecl())); + } + for (const Stmt* C : S->children()) + worklist.push_back(C); + } + // Start the visitation process which outputs the statements in the // current block. StmtDiff BodyDiff = Visit(m_DiffReq->getBody()); @@ -3225,15 +3263,60 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { // For an in-place `p = realloc(p, sz)`, realloc frees the old block, so // saving p to restore it in the reverse sweep would dangle (and be // double-freed at cleanup). Keep the reallocated pointer instead, for - // both p and its adjoint _d_p. This is sound for a growing or same-size - // realloc, where every reverse access stays in bounds. A shrinking - // in-place realloc is not handled -- its reverse sweep may read the - // freed tail out of bounds, as it did before this change; no test - // exercises it. + // both p and its adjoint _d_p. bool isReallocAssignment = opCode == BO_Assign && AllocCallInfo::recognize(R).isInPlaceRealloc(L); + // For an in-place realloc whose buffer size we tracked at allocation, + // undo the resize in the reverse sweep: restore p and its adjoint _d_p to + // their pre-realloc byte size so statements that ran before the realloc + // access valid indices again (clad::reverse_realloc also zeroes the + // re-grown adjoint tail). Without a tracked size we fall back to keeping + // the reallocated pointer, which is correct only for grow/same-size. + if (isReallocAssignment) { + const auto* pDRE = cast(L->IgnoreParenCasts()); + // Resolve the pointer to its adjoint record via the primal clone (the + // key m_Variables uses), then read the allocation-size shadow off it. + auto replIt = m_DeclReplacements.find(cast(pDRE->getDecl())); + auto varIt = replIt == m_DeclReplacements.end() + ? m_Variables.end() + : m_Variables.find(replIt->second); + if (varIt != m_Variables.end() && varIt->second.AllocSize) { + VarDecl* sizeVar = varIt->second.AllocSize; + Expr* newBytes = AllocCallInfo::recognize(R).getCall()->getArg(1); + // Capture both sizes in the forward pass. The reverse sweep must use + // the sizes as they were at the realloc, but the operands may refer + // to variables whose values differ by the time the sweep runs, so + // neither can be re-evaluated there. oldBytes reads the shadow before + // it advances; newBytes is stored once and reused for the update. + Expr* oldBytes = + StoreAndRef(BuildDeclRef(sizeVar), direction::forward, "_t", + /*forceDeclCreation=*/true); + Expr* newBytesRef = StoreAndRef(Clone(newBytes), direction::forward, + "_t", /*forceDeclCreation=*/true); + addToCurrentBlock( + BuildOp(BO_Assign, BuildDeclRef(sizeVar), Clone(newBytesRef)), + direction::forward); + auto buildReverseRealloc = [&](Expr* ptr, bool zeroTail) { + Expr* zt = new (m_Context) + CXXBoolLiteralExpr(zeroTail, m_Context.BoolTy, noLoc); + llvm::SmallVector args = {Clone(ptr), Clone(oldBytes), + Clone(newBytesRef), zt}; + Expr* call = GetFunctionCall("reverse_realloc", "clad", args); + // reverse_realloc returns void*; cast back to the pointer's type. + TypeSourceInfo* TSI = + m_Context.getTrivialTypeSourceInfo(ptr->getType(), noLoc); + call = m_Sema.BuildCStyleCastExpr(noLoc, TSI, noLoc, call).get(); + return BuildOp(BO_Assign, Clone(ptr), call); + }; + addToCurrentBlock(buildReverseRealloc(L, /*zeroTail=*/false), + direction::reverse); + addToCurrentBlock(buildReverseRealloc(ResultRef, /*zeroTail=*/true), + direction::reverse); + } + } + // Store the value of the LHS of the assignment in the forward pass // and restore it in the reverse pass if (m_DiffReq.shouldBeRecorded(L) && !isReallocAssignment) { @@ -3830,6 +3913,18 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { BuildDeclRef(VDDerived), VDDerived->getInit()->IgnoreCasts())) memsetCalls.push_back(memsetCall); + // Track this pointer's allocation size in bytes so an in-place + // realloc of it can be undone in the reverse sweep. The shadow is + // recorded on the pointer's adjoint entry, keyed like every other + // m_Variables record by the primal clone. + if (m_InPlaceReallocPtrs.contains(VD)) + if (Expr* bytes = buildAllocByteSize(VDDerived->getInit())) { + VarDecl* sizeVar = BuildVarDecl( + m_Context.getSizeType(), + "_" + VD->getNameAsString() + "_size", Clone(bytes)); + memsetCalls.push_back(BuildDeclStmt(sizeVar)); + m_Variables[VDDiff.getDecl()].AllocSize = sizeVar; + } } } } else if (auto* SAD = dyn_cast(D)) { diff --git a/test/Gradient/Pointers.C b/test/Gradient/Pointers.C index 5e09e6f66..6bb58a12c 100644 --- a/test/Gradient/Pointers.C +++ b/test/Gradient/Pointers.C @@ -331,31 +331,39 @@ double cStyleMemoryAlloc(double x, size_t n) { // CHECK-NEXT: double _t0 = t->x; // CHECK-NEXT: t->x = x; // CHECK-NEXT: double *_d_p = (double *)calloc(1, sizeof(double)); +// CHECK-NEXT: {{__size_t|size_t|unsigned long long|unsigned long|unsigned int}} _p_size0 = 1 * sizeof(double); // CHECK-NEXT: double *p = (double *)calloc(1, sizeof(double)); // CHECK-NEXT: double _t1 = *p; // CHECK-NEXT: *p = x; // CHECK-NEXT: double _d_res = 0.; // CHECK-NEXT: double res = t->x + *p; +// CHECK-NEXT: {{__size_t|size_t|unsigned long long|unsigned long|unsigned int}} _t2 = _p_size0; +// CHECK-NEXT: {{__size_t|size_t|unsigned long long|unsigned long|unsigned int}} _t3 = 2 * sizeof(double); +// CHECK-NEXT: _p_size0 = _t3; // CHECK-NEXT: _d_p = (double *)realloc(_d_p, 2 * sizeof(double)); // CHECK-NEXT: memset(_d_p, 0, 2 * sizeof(double)); // CHECK-NEXT: p = (double *)realloc(p, 2 * sizeof(double)); -// CHECK-NEXT: double _t2 = p[1]; +// CHECK-NEXT: double _t4 = p[1]; // CHECK-NEXT: p[1] = 2 * x; -// CHECK-NEXT: double _t3 = res; +// CHECK-NEXT: double _t5 = res; // CHECK-NEXT: res += p[1]; // CHECK-NEXT: _d_res += 1; // CHECK-NEXT: { -// CHECK-NEXT: res = _t3; +// CHECK-NEXT: res = _t5; // CHECK-NEXT: double _r_d3 = _d_res; // CHECK-NEXT: _d_p[1] += _r_d3; // CHECK-NEXT: } // CHECK-NEXT: { -// CHECK-NEXT: p[1] = _t2; +// CHECK-NEXT: p[1] = _t4; // CHECK-NEXT: double _r_d2 = _d_p[1]; // CHECK-NEXT: _d_p[1] = 0.; // CHECK-NEXT: *_d_x += 2 * _r_d2; // CHECK-NEXT: } // CHECK-NEXT: { +// CHECK-NEXT: p = (double *)clad::reverse_realloc(p, _t2, _t3, false); +// CHECK-NEXT: _d_p = (double *)clad::reverse_realloc(_d_p, _t2, _t3, true); +// CHECK-NEXT: } +// CHECK-NEXT: { // CHECK-NEXT: _d_t->x += _d_res; // CHECK-NEXT: *_d_p += _d_res; // CHECK-NEXT: } diff --git a/test/Gradient/ReallocShrink.C b/test/Gradient/ReallocShrink.C new file mode 100644 index 000000000..9833ac34c --- /dev/null +++ b/test/Gradient/ReallocShrink.C @@ -0,0 +1,29 @@ +// RUN: %cladclang %s -I%S/../../include -oReallocShrink.out 2>&1 | %filecheck %s +// RUN: ./ReallocShrink.out | %filecheck_exec %s + +// Shrinking in-place realloc: the reverse sweep must resize the buffer (and its +// adjoint) back to the pre-realloc size so the pre-realloc reads stay in bounds +// and the adjoint accumulates correctly. Runs Memcheck-clean under the valgrind +// CI row. + +#include "clad/Differentiator/Differentiator.h" +#include + +// res = (x + x*x) + x -> dres/dx = 2 + 2x ; at x = 3 -> 8 +double shrink_realloc(double x) { + double* p = (double*)malloc(2 * sizeof(double)); + p[0] = x; + p[1] = x * x; + double res = p[0] + p[1]; + p = (double*)realloc(p, 1 * sizeof(double)); // shrink 2 -> 1 + res += p[0]; + free(p); + return res; +} + +int main() { + auto g = clad::gradient(shrink_realloc, "x"); + double dx = 0; + g.execute(3.0, &dx); + printf("{%.2f}\n", dx); // CHECK-EXEC: {8.00} +}