From d99c1a32aeee7c03f5172a7e22721c16b59ac4d7 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Sun, 5 Jul 2026 06:49:24 -0700 Subject: [PATCH 1/6] Recognize Numba NRT_decref as a deallocation function 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) --- enzyme/Enzyme/LibraryFuncs.h | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/enzyme/Enzyme/LibraryFuncs.h b/enzyme/Enzyme/LibraryFuncs.h index 27697bc41a53..9aed9ab2e0bc 100644 --- a/enzyme/Enzyme/LibraryFuncs.h +++ b/enzyme/Enzyme/LibraryFuncs.h @@ -135,6 +135,13 @@ static inline bool isDeallocationFunction(const llvm::StringRef name, return true; if (name == "swift_release") return true; + // Numba runtime (NRT) reference-count release / deallocation. Recognizing + // NRT_decref as a deallocation lets Enzyme cache the paired + // NRT_MemInfo_alloc* allocation and emit a single, correctly-placed free + // instead of duplicating the refcount op across the augmented/reverse + // passes (which multi-frees and corrupts the heap). + if (name == "NRT_decref") + return true; return false; } From 6584a5756b04d9a1a2a13a0f7ccfcbd5ce12d2c5 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Tue, 7 Jul 2026 21:54:13 -0700 Subject: [PATCH 2/6] Clean up excessive comment Removed comments explaining NRT_decref handling. --- enzyme/Enzyme/LibraryFuncs.h | 5 ----- 1 file changed, 5 deletions(-) diff --git a/enzyme/Enzyme/LibraryFuncs.h b/enzyme/Enzyme/LibraryFuncs.h index 9aed9ab2e0bc..b22372c2fdf7 100644 --- a/enzyme/Enzyme/LibraryFuncs.h +++ b/enzyme/Enzyme/LibraryFuncs.h @@ -135,11 +135,6 @@ static inline bool isDeallocationFunction(const llvm::StringRef name, return true; if (name == "swift_release") return true; - // Numba runtime (NRT) reference-count release / deallocation. Recognizing - // NRT_decref as a deallocation lets Enzyme cache the paired - // NRT_MemInfo_alloc* allocation and emit a single, correctly-placed free - // instead of duplicating the refcount op across the augmented/reverse - // passes (which multi-frees and corrupts the heap). if (name == "NRT_decref") return true; return false; From 72e073125a6da7f279b168e4a4b514c3df5548f0 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Wed, 8 Jul 2026 10:17:42 -0700 Subject: [PATCH 3/6] Add reverse-mode regression test for NRT_decref deallocation recognition 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 Claude-Session: https://claude.ai/code/session_015Jbd9RehxJeX7BvDvSiPiQ --- enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll | 86 ++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll diff --git a/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll b/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll new file mode 100644 index 000000000000..79353a767efb --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll @@ -0,0 +1,86 @@ +; RUN: %opt < %s %newLoadEnzyme -passes="enzyme,function(mem2reg,instsimplify,adce,loop(loop-deletion),correlated-propagation,%simplifycfg)" -enzyme-preopt=false -S | FileCheck %s + +; Regression test: the Numba runtime's reference-count release, NRT_decref, +; must be recognized as a deallocation function (isDeallocationFunction), +; the same way swift_release / __rust_dealloc / _mlir_memref_to_llvm_free are. +; +; Numba frees a heap allocation with a paired NRT_decref. Unless Enzyme knows +; NRT_decref frees its argument, it cannot pair the release with the allocation +; and cannot manage the allocation's lifetime across the augmented/reverse +; split (downstream, where NRT_decref is additionally marked enzyme_inactive, +; the unrecognized release was instead duplicated across both passes, separately +; freeing the allocation and corrupting the heap). +; +; This test isolates the deallocation recognition using 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, Enzyme treats the pair exactly as malloc/free: the value is dead +; in the forward pass, so the augmented function is empty and the reverse pass +; re-allocates a shadow and frees it once. (Enzyme emits malloc's canonical +; deallocator, @free, for its own shadow allocation; NRT_decref itself is a +; Numba allocation's deallocator and is paired downstream, so it does not +; appear in this self-contained module -- what is tested here is that the +; source NRT_decref is understood as a free.) +; +; Before recognizing NRT_decref, differentiating @sum fails outright with +; "Enzyme: No augmented forward pass found for NRT_decref" (the release has no +; known derivative and is not a deallocation); with the recognition it compiles +; to the IR checked below. + +declare i8* @malloc(i64) +declare void @NRT_decref(i8*) + +define dso_local double @subsum(double* nocapture readonly %x, i64 %n) { +entry: + %m = call i8* @malloc(i64 8) + %v = bitcast i8* %m to double* + br label %for.body + +for.cond.cleanup: ; preds = %for.body + %res = load double, double* %v + call void @NRT_decref(i8* %m) + ret double %res + +for.body: ; preds = %entry, %for.body + %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ] + %total.07 = phi double [ 0.000000e+00, %entry ], [ %add, %for.body ] + %arrayidx = getelementptr inbounds double, double* %x, i64 %indvars.iv + %0 = load double, double* %arrayidx, align 8 + %add = fadd fast double %0, %total.07 + store double %add, double* %v + %indvars.iv.next = add nuw i64 %indvars.iv, 1 + %exitcond = icmp eq i64 %indvars.iv, %n + br i1 %exitcond, label %for.cond.cleanup, label %for.body +} + +define dso_local double @sum(double* nocapture readonly %x, i64 %n) { +entry: + %res = call double @subsum(double* %x, i64 %n) + store double 0.000000e+00, double* %x + ret double %res +} + +define dso_local void @dsum(double* %x, double* %xp, i64 %n) { +entry: + %0 = tail call double (double (double*, i64)*, ...) @__enzyme_autodiff(double (double*, i64)* nonnull @sum, double* %x, double* %xp, i64 %n) + ret void +} + +declare double @__enzyme_autodiff(double (double*, i64)*, ...) + +; The recognized malloc/NRT_decref pair is dead in the forward pass, so the +; augmented function caches nothing and is empty -- impossible unless the +; NRT_decref is understood to free the malloc. +; CHECK: define internal void @augmented_subsum(ptr nocapture readonly %x, ptr %"x'", i64 %n) +; CHECK-NEXT: entry: +; CHECK-NEXT: ret void +; CHECK-NEXT: } + +; The reverse pass re-allocates a shadow and frees it exactly once (via +; malloc's canonical deallocator). NRT_decref is never duplicated into the +; derivative. +; CHECK: define internal void @diffesubsum(ptr nocapture readonly %x, ptr %"x'", i64 %n, double %differeturn) +; CHECK-NEXT: entry: +; CHECK-NEXT: %"m'mi" = call {{.*}}ptr @malloc(i64 8) +; CHECK: call void @free(ptr nonnull %"m'mi") +; CHECK-NOT: NRT_decref From 00588074e48ba21446900e9d505be15210acc2da Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Wed, 8 Jul 2026 12:10:06 -0700 Subject: [PATCH 4/6] nrtdealloc.ll: make CHECKs pointer-spelling-agnostic (fix LLVM 15/16 CI) 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 Claude-Session: https://claude.ai/code/session_015Jbd9RehxJeX7BvDvSiPiQ --- enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll b/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll index 79353a767efb..b3d334b481b8 100644 --- a/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll +++ b/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll @@ -1,3 +1,4 @@ +; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -adce -loop-deletion -correlated-propagation -simplifycfg -S | FileCheck %s; fi ; RUN: %opt < %s %newLoadEnzyme -passes="enzyme,function(mem2reg,instsimplify,adce,loop(loop-deletion),correlated-propagation,%simplifycfg)" -enzyme-preopt=false -S | FileCheck %s ; Regression test: the Numba runtime's reference-count release, NRT_decref, @@ -25,7 +26,9 @@ ; Before recognizing NRT_decref, differentiating @sum fails outright with ; "Enzyme: No augmented forward pass found for NRT_decref" (the release has no ; known derivative and is not a deallocation); with the recognition it compiles -; to the IR checked below. +; to the IR checked below. Pointer spellings are matched loosely so the test +; passes under both the typed-pointer (LLVM <= 16) and opaque-pointer +; (LLVM >= 17) forms of the emitted IR. declare i8* @malloc(i64) declare void @NRT_decref(i8*) @@ -71,16 +74,15 @@ declare double @__enzyme_autodiff(double (double*, i64)*, ...) ; The recognized malloc/NRT_decref pair is dead in the forward pass, so the ; augmented function caches nothing and is empty -- impossible unless the ; NRT_decref is understood to free the malloc. -; CHECK: define internal void @augmented_subsum(ptr nocapture readonly %x, ptr %"x'", i64 %n) +; CHECK: define internal void @augmented_subsum( ; CHECK-NEXT: entry: -; CHECK-NEXT: ret void +; CHECK-NEXT: ret void ; CHECK-NEXT: } ; The reverse pass re-allocates a shadow and frees it exactly once (via -; malloc's canonical deallocator). NRT_decref is never duplicated into the -; derivative. -; CHECK: define internal void @diffesubsum(ptr nocapture readonly %x, ptr %"x'", i64 %n, double %differeturn) -; CHECK-NEXT: entry: -; CHECK-NEXT: %"m'mi" = call {{.*}}ptr @malloc(i64 8) -; CHECK: call void @free(ptr nonnull %"m'mi") +; malloc's canonical deallocator @free). NRT_decref is never duplicated into +; the derivative. +; CHECK: define internal void @diffesubsum( +; CHECK: %"m'mi" = call {{.*}}@malloc(i64 8) +; CHECK: call void @free({{.*}}%"m'mi") ; CHECK-NOT: NRT_decref From ef8998cbccef137095292f16dcb16f43c515b532 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Wed, 8 Jul 2026 18:50:54 -0700 Subject: [PATCH 5/6] nrtdealloc.ll: use the actual Numba NRT allocation IR 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 Claude-Session: https://claude.ai/code/session_015Jbd9RehxJeX7BvDvSiPiQ --- enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll | 65 ++++++++++---------- 1 file changed, 32 insertions(+), 33 deletions(-) diff --git a/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll b/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll index b3d334b481b8..a1c08cc02cc5 100644 --- a/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll +++ b/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll @@ -1,41 +1,41 @@ -; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -adce -loop-deletion -correlated-propagation -simplifycfg -S | FileCheck %s; fi -; RUN: %opt < %s %newLoadEnzyme -passes="enzyme,function(mem2reg,instsimplify,adce,loop(loop-deletion),correlated-propagation,%simplifycfg)" -enzyme-preopt=false -S | FileCheck %s +; RUN: %opt < %s %newLoadEnzyme -passes="preserve-nvvm,enzyme,function(mem2reg,instsimplify,adce,loop(loop-deletion),correlated-propagation,%simplifycfg)" -enzyme-preopt=false -S | FileCheck %s ; Regression test: the Numba runtime's reference-count release, NRT_decref, ; must be recognized as a deallocation function (isDeallocationFunction), ; the same way swift_release / __rust_dealloc / _mlir_memref_to_llvm_free are. ; -; Numba frees a heap allocation with a paired NRT_decref. Unless Enzyme knows -; NRT_decref frees its argument, it cannot pair the release with the allocation -; and cannot manage the allocation's lifetime across the augmented/reverse -; split (downstream, where NRT_decref is additionally marked enzyme_inactive, -; the unrecognized release was instead duplicated across both passes, separately -; freeing the allocation and corrupting the heap). +; This uses the actual Numba NRT functions: an array is allocated with +; NRT_MemInfo_alloc_aligned and released with NRT_decref (the IR Numba emits). +; The allocation is registered with __enzyme_allocation_like so Enzyme knows +; how to shadow it self-contained (in Numba-Enzyme this is done downstream by a +; C-API shadow handler); its shadow-free is registered as @free, NOT NRT_decref +; -- so the recognition of the SOURCE NRT_decref as a deallocation depends +; solely on isDeallocationFunction (the change under test), not on the +; allocation/free pairing. ; -; This test isolates the deallocation recognition using 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, Enzyme treats the pair exactly as malloc/free: the value is dead -; in the forward pass, so the augmented function is empty and the reverse pass -; re-allocates a shadow and frees it once. (Enzyme emits malloc's canonical -; deallocator, @free, for its own shadow allocation; NRT_decref itself is a -; Numba allocation's deallocator and is paired downstream, so it does not -; appear in this self-contained module -- what is tested here is that the -; source NRT_decref is understood as a free.) -; -; Before recognizing NRT_decref, differentiating @sum fails outright with -; "Enzyme: No augmented forward pass found for NRT_decref" (the release has no -; known derivative and is not a deallocation); with the recognition it compiles -; to the IR checked below. Pointer spellings are matched loosely so the test -; passes under both the typed-pointer (LLVM <= 16) and opaque-pointer -; (LLVM >= 17) forms of the emitted IR. +; With the recognition, the malloc/NRT_decref-style pair is dead in the forward +; pass, so the augmented function 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. Before recognizing NRT_decref, differentiating +; @sum fails outright with "Enzyme: No augmented forward pass found for +; NRT_decref". Pointer spellings are matched loosely so the test passes under +; both the typed-pointer (LLVM <= 16) and opaque-pointer (LLVM >= 17) forms. + +@.dealloc = private unnamed_addr constant [3 x i8] c"-1\00" +@__enzyme_allocation_like = global [4 x i8*] [ + i8* bitcast (i8* (i64, i32)* @NRT_MemInfo_alloc_aligned to i8*), + i8* inttoptr (i64 0 to i8*), + i8* getelementptr inbounds ([3 x i8], [3 x i8]* @.dealloc, i64 0, i64 0), + i8* bitcast (void (i8*)* @free to i8*) +] -declare i8* @malloc(i64) +declare noalias i8* @NRT_MemInfo_alloc_aligned(i64, i32) declare void @NRT_decref(i8*) +declare void @free(i8*) define dso_local double @subsum(double* nocapture readonly %x, i64 %n) { entry: - %m = call i8* @malloc(i64 8) + %m = call i8* @NRT_MemInfo_alloc_aligned(i64 8, i32 32) %v = bitcast i8* %m to double* br label %for.body @@ -71,18 +71,17 @@ entry: declare double @__enzyme_autodiff(double (double*, i64)*, ...) -; The recognized malloc/NRT_decref pair is dead in the forward pass, so the +; The recognized allocation/NRT_decref pair is dead in the forward pass, so the ; augmented function caches nothing and is empty -- impossible unless the -; NRT_decref is understood to free the malloc. +; NRT_decref is understood to free the allocation. ; CHECK: define internal void @augmented_subsum( ; CHECK-NEXT: entry: ; CHECK-NEXT: ret void ; CHECK-NEXT: } -; The reverse pass re-allocates a shadow and frees it exactly once (via -; malloc's canonical deallocator @free). NRT_decref is never duplicated into -; the derivative. +; The reverse pass re-allocates a shadow with the real NRT allocator and frees +; it exactly once. NRT_decref is never duplicated into the derivative. ; CHECK: define internal void @diffesubsum( -; CHECK: %"m'mi" = call {{.*}}@malloc(i64 8) +; CHECK: %"m'mi" = call {{.*}}@NRT_MemInfo_alloc_aligned(i64 8, i32 32) ; CHECK: call void @free({{.*}}%"m'mi") ; CHECK-NOT: NRT_decref From 42a38369094a88293742faca11543587b5c25b4c Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Wed, 8 Jul 2026 18:58:21 -0700 Subject: [PATCH 6/6] nrtdealloc.ll: drop downstream-project detail from the comment 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 Claude-Session: https://claude.ai/code/session_015Jbd9RehxJeX7BvDvSiPiQ --- enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll b/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll index a1c08cc02cc5..5bdfe9464c86 100644 --- a/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll +++ b/enzyme/test/Enzyme/ReverseMode/nrtdealloc.ll @@ -7,13 +7,12 @@ ; This uses the actual Numba NRT functions: an array is allocated with ; NRT_MemInfo_alloc_aligned and released with NRT_decref (the IR Numba emits). ; The allocation is registered with __enzyme_allocation_like so Enzyme knows -; how to shadow it self-contained (in Numba-Enzyme this is done downstream by a -; C-API shadow handler); its shadow-free is registered as @free, NOT NRT_decref -; -- so the recognition of the SOURCE NRT_decref as a deallocation depends -; solely on isDeallocationFunction (the change under test), not on the +; how to shadow it self-contained; its shadow-free is registered as @free, NOT +; NRT_decref -- so the recognition of the SOURCE NRT_decref as a deallocation +; depends solely on isDeallocationFunction (the change under test), not on the ; allocation/free pairing. ; -; With the recognition, the malloc/NRT_decref-style pair is dead in the forward +; With the recognition, the allocation/NRT_decref pair is dead in the forward ; pass, so the augmented function 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. Before recognizing NRT_decref, differentiating