Add Mooncake extension marking searches inactive

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"]

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

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.zero_tangent(g)), (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

test/test_groups.toml

@@ -21,3 +21,6 @@ versions = ["lts", "1"]
 
 [Enzyme]
 versions = ["1"]
+
+[Mooncake]
+versions = ["1"]