Recognize Numba NRT_decref as a deallocation function#2930
Open
ludgerpaehler wants to merge 6 commits into
Open
Recognize Numba NRT_decref as a deallocation function#2930ludgerpaehler wants to merge 6 commits into
ludgerpaehler 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
Collaborator
Author
|
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
Collaborator
Author
|
Tests look good afaict, let me know if you need more for the PR |
wsmoses
requested changes
Jul 8, 2026
wsmoses
left a comment
wsmoses
left a comment
Member
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
Collaborator
Author
Put in, CI looks on track to me. Let me know if you want further changes. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.