Recognize Numba NRT_decref as a deallocation function#2930
Recognize Numba NRT_decref as a deallocation function#2930ludgerpaehler wants to merge 6 commits into
Conversation
Numba-Enzyme differentiates functions that allocate arrays via the Numba runtime (NRT). NRT_MemInfo_alloc* is registered as an allocation (shadow handler), but its paired release, NRT_decref, was not recognized as a deallocation. Enzyme then duplicated the (inactive-marked) refcount op across the augmented-forward and reverse passes and separately freed the allocation, multi-freeing it and corrupting the heap. Recognizing NRT_decref in isDeallocationFunction lets Enzyme cache the allocation and emit a single, correctly-placed free, matching how it handles malloc / jl_alloc_array. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Removed comments explaining NRT_decref handling.
Isolates the isDeallocationFunction(NRT_decref) recognition with a built-in-recognized allocation (malloc) whose only release is NRT_decref -- the malloc/free-nouse pattern of alloctomallocnouse.ll with free replaced by NRT_decref. With the recognition the pair is treated exactly as malloc/free (empty augmented pass; reverse re-allocates a shadow and frees it once); before it, differentiation fails with "No augmented forward pass found for NRT_decref". Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015Jbd9RehxJeX7BvDvSiPiQ
|
Test added. CI still needs to pass. |
The test used opaque-pointer (`ptr`) CHECK lines, which only match the IR
Enzyme emits on LLVM >= 17. On LLVM 15 (typed pointers by default) and
LLVM 16 (typed forced via `-opaque-pointers=0` in the lit config) the
output uses `double*`/`i8*`, so FileCheck failed the "Enzyme CI LLVM 15/16"
jobs while >= 17 passed.
Match pointers loosely so the single CHECK set works under both forms: the
augmented function is checked for emptiness without pinning its parameter
types, and the shadow malloc/free lines use `{{.*}}` in place of the pointer
type. Also add the old-PM RUN line gated `%llvmver < 16` (mirroring
alloctomallocnouse.ll, the sibling this test is modeled on) for parity.
The assertion is unchanged: recognizing NRT_decref as a deallocation makes
the augmented pass empty and frees the shadow once, with no NRT_decref
duplicated into the derivative.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015Jbd9RehxJeX7BvDvSiPiQ
|
Tests look good afaict, let me know if you need more for the PR |
wsmoses
left a comment
There was a problem hiding this comment.
can you put an actual numba ir case here?
Replace the synthetic malloc stand-in with the real Numba runtime functions: an array is allocated with NRT_MemInfo_alloc_aligned and released with NRT_decref (the IR Numba emits). The allocation is registered self-contained with __enzyme_allocation_like (in Numba-Enzyme this is a downstream C-API shadow handler), run via the preserve-nvvm pass; its shadow-free is registered as @free -- NOT NRT_decref -- so recognition of the SOURCE NRT_decref as a deallocation depends solely on isDeallocationFunction (the change under test). With the recognition the augmented pass is empty and the reverse pass re-allocates a shadow via NRT_MemInfo_alloc_aligned and frees it once, with no NRT_decref duplicated into the derivative; renaming the source NRT_decref makes differentiation fail ("No augmented forward pass found for NRT_decref"). Pointer spellings stay matched loosely and the input uses typed pointers, so the single new-PM RUN line works across LLVM 15-20. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015Jbd9RehxJeX7BvDvSiPiQ
Remove the reference to how the NRT allocation is shadowed downstream (the C-API shadow handler) -- not relevant to the upstream test -- and fix the stale "malloc" phrasing now that the allocation is NRT_MemInfo_alloc. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015Jbd9RehxJeX7BvDvSiPiQ
Put in, CI looks on track to me. Let me know if you want further changes. |
wsmoses
left a comment
There was a problem hiding this comment.
sorry this isn't the right way of doing thing and looks kind of like claude just spewing a bit of ai slop [e.g. this should be symmetrically registered in isAllocationFunction like the free function was added]
Point taken, I am sorry that you had the read the incoherent IR. I'll sit down in calm on the weekend and hand-write a proper example without running between multiple things s.t. it is semantically coherent and is in line with the style of e.g. the rust deallocator example |
To cleanly differentiate via the Numba runtime (NRT), NRT_MemInfo_alloc* is registered as an allocation (shadow handler), but its paired release, NRT_decref, was not recognized as a deallocation.
Without this addition Enzyme duplicates the (inactive-marked) refcount op across the augmented-forward and reverse passes and separately frees the allocation, multi-freeing it and corrupting the heap.
Recognizing NRT_decref in isDeallocationFunction lets Enzyme cache the allocation and emit a single, correctly-placed free, matching how it handles malloc / jl_alloc_array.