From 193db38dfefce90a0a301574e0c7617224b063b7 Mon Sep 17 00:00:00 2001 From: AstitvaAggarwal Date: Fri, 18 Sep 2026 15:56:43 +0100 Subject: [PATCH 1/2] Add Mooncake extension marking searches inactive Mirrors the existing EnzymeCore extension: searchsorted_last, searchsorted_first, searchsortedlast!, searchsortedfirst!, searchsortedrange, findequal, findfirstequal, and findfirstsortedequal all return a discrete index with no meaningful derivative, so mark them @zero_derivative for Mooncake too. This also works around a real Mooncake limitation: KIND_SIMD_LINEAR_SCAN's kernel embeds a raw llvmcall SIMD intrinsic that Mooncake's intrinsic-lifting machinery cannot translate. Marking the dispatcher itself inactive stops Mooncake before it ever descends into that (or any other strategy's) kernel. Surfaced via native Hessian-vector products through OrdinaryDiffEq's dense-output interpolation: SciML/SciMLSensitivity.jl#1648. --- Project.toml | 6 +- ext/FindFirstFunctionsMooncakeExt.jl | 65 +++++++++++++++++++ test/Mooncake/mooncake_tests.jl | 94 ++++++++++++++++++++++++++++ test/test_groups.toml | 3 + 4 files changed, 167 insertions(+), 1 deletion(-) create mode 100644 ext/FindFirstFunctionsMooncakeExt.jl create mode 100644 test/Mooncake/mooncake_tests.jl diff --git a/Project.toml b/Project.toml index ba40a16..a77950d 100644 --- a/Project.toml +++ b/Project.toml @@ -8,14 +8,17 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" [weakdeps] EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" [extensions] FindFirstFunctionsEnzymeCoreExt = "EnzymeCore" +FindFirstFunctionsMooncakeExt = "Mooncake" [compat] Enzyme = "0.13" EnzymeCore = "0.8" ForwardDiff = "1" +Mooncake = "0.5" PrecompileTools = "1.0.1" SafeTestsets = "0.1" SciMLTesting = "2.4" @@ -26,10 +29,11 @@ julia = "1.10" [extras] Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Enzyme", "ForwardDiff", "SafeTestsets", "SciMLTesting", "StableRNGs", "Test"] +test = ["Enzyme", "ForwardDiff", "Mooncake", "SafeTestsets", "SciMLTesting", "StableRNGs", "Test"] diff --git a/ext/FindFirstFunctionsMooncakeExt.jl b/ext/FindFirstFunctionsMooncakeExt.jl new file mode 100644 index 0000000..26aa0cc --- /dev/null +++ b/ext/FindFirstFunctionsMooncakeExt.jl @@ -0,0 +1,65 @@ +module FindFirstFunctionsMooncakeExt + +using FindFirstFunctions +using Mooncake: Mooncake, @is_primitive, @zero_derivative, CoDual, Dual, MinimalCtx, NoPullback, primal, zero_dual, zero_fcodual +import Mooncake: frule!!, rrule!! + +# These functions all return a discrete index (or a batch of indices) locating a value +# within a sorted/searched collection: the result has no meaningful derivative. Same +# position as `EnzymeRules.inactive` already takes for this exact list of functions -- see +# `FindFirstFunctionsEnzymeCoreExt`. +# +# For Mooncake specifically, this also works around a real translation limitation: +# `searchsorted_first`/`searchsorted_last`'s `KIND_SIMD_LINEAR_SCAN` strategy embeds a +# hand-written SIMD kernel via a raw `Core.Intrinsics.llvmcall`, which Mooncake's +# intrinsic-lifting machinery cannot translate (arbitrary embedded LLVM IR text has no +# general translation). Marking the dispatcher itself `@zero_derivative` -- rather than +# patching the SIMD kernel -- stops Mooncake before it ever descends into that (or any +# other strategy's) kernel. Surfaced via native Hessian-vector products through +# `OrdinaryDiffEq`'s dense-output interpolation: see +# https://github.com/SciML/SciMLSensitivity.jl/issues/1648. +const _INACTIVE_FNS = ( + FindFirstFunctions.searchsorted_last, + FindFirstFunctions.searchsorted_first, + FindFirstFunctions.searchsortedlast!, + FindFirstFunctions.searchsortedfirst!, + FindFirstFunctions.searchsortedrange, + FindFirstFunctions.findequal, + FindFirstFunctions.findfirstequal, + FindFirstFunctions.findfirstsortedequal, +) + +for f in _INACTIVE_FNS + @eval @zero_derivative MinimalCtx Tuple{typeof($f),Vararg} +end + +const _InactiveFn = Union{map(typeof, _INACTIVE_FNS)...} + +# Keyword arguments (`order`, `strategy`, `queries_sorted`, ...) are reached through the +# compiler-generated keyword-body method via `Core.kwcall`, not the plain signature above, +# so give that entry point the same treatment. The primal call must still go through +# `Core.kwcall` itself so that a non-default keyword is honoured -- only the derivative is +# zero here. +@is_primitive MinimalCtx Tuple{typeof(Core.kwcall),NamedTuple,F,Vararg} where {F<:_InactiveFn} + +function frule!!( + ::Dual{typeof(Core.kwcall)}, + kwargs::Dual{<:NamedTuple}, + f::Dual{<:_InactiveFn}, + args::Vararg{Dual}, +) + y = Core.kwcall(primal(kwargs), primal(f), map(primal, args)...) + return zero_dual(y) +end + +function rrule!!( + kwcall::CoDual{typeof(Core.kwcall)}, + kwargs::CoDual{<:NamedTuple}, + f::CoDual{<:_InactiveFn}, + args::Vararg{CoDual}, +) + y = Core.kwcall(primal(kwargs), primal(f), map(primal, args)...) + return zero_fcodual(y), NoPullback(kwcall, kwargs, f, args...) +end + +end diff --git a/test/Mooncake/mooncake_tests.jl b/test/Mooncake/mooncake_tests.jl new file mode 100644 index 0000000..f99791d --- /dev/null +++ b/test/Mooncake/mooncake_tests.jl @@ -0,0 +1,94 @@ +@safetestset "Mooncake extension: searches are inactive" begin + using FindFirstFunctions, Mooncake, Test + + function interp(u, knots, x) + i0 = searchsorted_last(KIND_BINARY_BRACKET, knots, x) + i = clamp(i0, 1, length(knots) - 1) + w = (x - knots[i]) / (knots[i + 1] - knots[i]) + return (1 - w) * u[i] + w * u[i + 1] + end + + knots = collect(0.0:0.5:10.0) + u = sin.(knots) + x = 3.3 + + iref = searchsorted_last(KIND_BINARY_BRACKET, knots, x) + wref = (x - knots[iref]) / (knots[iref + 1] - knots[iref]) + + @testset "reverse mode, gradient w.r.t. the knot values" begin + f(u) = interp(u, knots, x) + _, (_, du) = Mooncake.value_and_gradient!!(Mooncake.build_rrule(f, u), f, u) + expected = zero(u) + expected[iref] = 1 - wref + expected[iref + 1] = wref + @test du ≈ expected + end + + @testset "reverse mode, derivative w.r.t. the query point" begin + g(x) = interp(u, knots, x) + _, (_, dx) = Mooncake.value_and_gradient!!(Mooncake.build_rrule(g, x), g, x) + @test dx ≈ (u[iref + 1] - u[iref]) / (knots[iref + 1] - knots[iref]) + end + + @testset "forward mode agrees" begin + g(x) = interp(u, knots, x) + cache = Mooncake.prepare_derivative_cache(g, x) + _, dx = Mooncake.value_and_derivative!!(cache, (g, Mooncake.NoTangent()), (x, 1.0)) + @test dx ≈ (u[iref + 1] - u[iref]) / (knots[iref + 1] - knots[iref]) + end + + # The motivating case (SciMLSensitivity.jl#1648): native forward-over-reverse HVP + # needs both directions to work through the same search, since it forward- + # differentiates the whole reverse pass. + @testset "native HVP w.r.t. the knot values" begin + f(u) = interp(u, knots, x) + hcache = Mooncake.prepare_hvp_cache(f, u) + v = zero(u) + v[iref] = 1.0 + _, du, hvp_u = Mooncake.value_and_hvp!!(hcache, f, v, u) + expected = zero(u) + expected[iref] = 1 - wref + expected[iref + 1] = wref + @test du ≈ expected + # `interp` is linear in `u`, so the Hessian w.r.t. `u` is exactly zero. + @test all(iszero, hvp_u) + end + + # Regression test for the actual bug (SciMLSensitivity.jl#1648): `KIND_SIMD_LINEAR_SCAN` + # specifically is the strategy whose kernel embeds a raw `llvmcall` that Mooncake cannot + # translate. `KIND_BINARY_BRACKET` above never reaches that kernel at all, so it alone + # would not catch a regression here. + @testset "KIND_SIMD_LINEAR_SCAN specifically" begin + function interp_simd(u, knots, x) + i0 = searchsorted_last(KIND_SIMD_LINEAR_SCAN, knots, x) + i = clamp(i0, 1, length(knots) - 1) + w = (x - knots[i]) / (knots[i + 1] - knots[i]) + return (1 - w) * u[i] + w * u[i + 1] + end + f(u) = interp_simd(u, knots, x) + _, (_, du) = Mooncake.value_and_gradient!!(Mooncake.build_rrule(f, u), f, u) + hcache = Mooncake.prepare_hvp_cache(f, u) + v = zero(u) + v[iref] = 1.0 + _, du2, hvp_u = Mooncake.value_and_hvp!!(hcache, f, v, u) + expected = zero(u) + expected[iref] = 1 - wref + expected[iref + 1] = wref + @test du ≈ expected + @test du2 ≈ expected + @test all(iszero, hvp_u) + end + + @testset "batched and equality searches accept the rule" begin + function g(u, knots, xs) + idx = similar(xs, Int) + searchsortedlast!(idx, knots, xs) + j = findfirstsortedequal(knots[3], knots) + return sum(u) + zero(eltype(u)) * (idx[1] + something(j, 0)) + end + xs = [0.7, 2.2, 9.1] + h(u) = g(u, knots, xs) + _, (_, du) = Mooncake.value_and_gradient!!(Mooncake.build_rrule(h, u), h, u) + @test all(==(1.0), du) + end +end diff --git a/test/test_groups.toml b/test/test_groups.toml index 0c58956..ab35aee 100644 --- a/test/test_groups.toml +++ b/test/test_groups.toml @@ -21,3 +21,6 @@ versions = ["lts", "1"] [Enzyme] versions = ["1"] + +[Mooncake] +versions = ["1"] From 63a6f2e7f60e1fd2c386f236d0a15041b6810c81 Mon Sep 17 00:00:00 2001 From: AstitvaAggarwal Date: Fri, 18 Sep 2026 16:10:09 +0100 Subject: [PATCH 2/2] Fix formatting and a forward-mode test bug Runic wants spaces after commas and around <: in where clauses. The forward-mode test passed a NoTangent() for the closure, but inside @safetestset closures capture their free variables as real fields (unlike a plain top-level script), so it needs the closure's actual zero tangent instead. --- ext/FindFirstFunctionsMooncakeExt.jl | 24 ++++++++++++------------ test/Mooncake/mooncake_tests.jl | 2 +- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/ext/FindFirstFunctionsMooncakeExt.jl b/ext/FindFirstFunctionsMooncakeExt.jl index 26aa0cc..5fb77a8 100644 --- a/ext/FindFirstFunctionsMooncakeExt.jl +++ b/ext/FindFirstFunctionsMooncakeExt.jl @@ -30,7 +30,7 @@ const _INACTIVE_FNS = ( ) for f in _INACTIVE_FNS - @eval @zero_derivative MinimalCtx Tuple{typeof($f),Vararg} + @eval @zero_derivative MinimalCtx Tuple{typeof($f), Vararg} end const _InactiveFn = Union{map(typeof, _INACTIVE_FNS)...} @@ -40,24 +40,24 @@ const _InactiveFn = Union{map(typeof, _INACTIVE_FNS)...} # so give that entry point the same treatment. The primal call must still go through # `Core.kwcall` itself so that a non-default keyword is honoured -- only the derivative is # zero here. -@is_primitive MinimalCtx Tuple{typeof(Core.kwcall),NamedTuple,F,Vararg} where {F<:_InactiveFn} +@is_primitive MinimalCtx Tuple{typeof(Core.kwcall), NamedTuple, F, Vararg} where {F <: _InactiveFn} function frule!!( - ::Dual{typeof(Core.kwcall)}, - kwargs::Dual{<:NamedTuple}, - f::Dual{<:_InactiveFn}, - args::Vararg{Dual}, -) + ::Dual{typeof(Core.kwcall)}, + kwargs::Dual{<:NamedTuple}, + f::Dual{<:_InactiveFn}, + args::Vararg{Dual}, + ) y = Core.kwcall(primal(kwargs), primal(f), map(primal, args)...) return zero_dual(y) end function rrule!!( - kwcall::CoDual{typeof(Core.kwcall)}, - kwargs::CoDual{<:NamedTuple}, - f::CoDual{<:_InactiveFn}, - args::Vararg{CoDual}, -) + kwcall::CoDual{typeof(Core.kwcall)}, + kwargs::CoDual{<:NamedTuple}, + f::CoDual{<:_InactiveFn}, + args::Vararg{CoDual}, + ) y = Core.kwcall(primal(kwargs), primal(f), map(primal, args)...) return zero_fcodual(y), NoPullback(kwcall, kwargs, f, args...) end diff --git a/test/Mooncake/mooncake_tests.jl b/test/Mooncake/mooncake_tests.jl index f99791d..58b5055 100644 --- a/test/Mooncake/mooncake_tests.jl +++ b/test/Mooncake/mooncake_tests.jl @@ -33,7 +33,7 @@ @testset "forward mode agrees" begin g(x) = interp(u, knots, x) cache = Mooncake.prepare_derivative_cache(g, x) - _, dx = Mooncake.value_and_derivative!!(cache, (g, Mooncake.NoTangent()), (x, 1.0)) + _, dx = Mooncake.value_and_derivative!!(cache, (g, Mooncake.zero_tangent(g)), (x, 1.0)) @test dx ≈ (u[iref + 1] - u[iref]) / (knots[iref + 1] - knots[iref]) end