Fix data race in BSpline and QuadraticSpline evaluation (rebase of #533)

Project.toml

@@ -10,6 +10,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -73,11 +74,10 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "AllocCheck", "BenchmarkTools", "SafeTestsets", "SciMLTesting", "ChainRulesCore", "Mooncake", "Optim", "Test", "StableRNGs", "FiniteDifferences", "QuadGK", "ForwardDiff", "StaticArrays", "Symbolics", "Unitful", "Zygote", "SparseConnectivityTracer"]
+test = ["Aqua", "AllocCheck", "BenchmarkTools", "SafeTestsets", "SciMLTesting", "ChainRulesCore", "Mooncake", "Optim", "Test", "StableRNGs", "FiniteDifferences", "QuadGK", "ForwardDiff", "Symbolics", "Unitful", "Zygote", "SparseConnectivityTracer"]

src/DataInterpolations.jl

@@ -9,6 +9,7 @@ using RecipesBase: RecipesBase, @recipe, @series
 using PrettyTables: PrettyTables, pretty_table
 using ForwardDiff: ForwardDiff
 using EnumX: EnumX, @enumx
+using StaticArrays: MVector, SVector
 import FindFirstFunctions
 import FindFirstFunctions: Guesser
 

src/derivatives.jl

@@ -256,14 +256,16 @@ function _derivative(A::BSplineInterpolation{<:AbstractVector{<:Number}}, t::Num
     n = length(A.t)
     scale = (A.p[idx + 1] - A.p[idx]) / (A.t[idx + 1] - A.t[idx])
     t_ = A.p[idx] + (t - A.t[idx]) * scale
-    sc = t isa ForwardDiff.Dual ? zeros(eltype(t), n) : A.sc
-    spline_coefficients!(sc, A.d - 1, A.k, t_)
     ducum = zero(eltype(A.u))
     if t == A.t[1]
         ducum = (A.c[2] - A.c[1]) / (A.k[A.d + 2])
     else
-        for i in 1:(n - 1)
-            ducum += sc[i + 1] * (A.c[i + 1] - A.c[i]) / (A.k[i + A.d + 1] - A.k[i + 1])
+        # Stack-allocated basis window: differentiation must be reentrant (#532)
+        vals, offset, m = bspline_nonzero_coefficients(A.d - 1, A.k, t_, n)
+        @inbounds for l in 1:m
+            i = offset + l - 1
+            (1 <= i <= n - 1) || continue
+            ducum += vals[l] * (A.c[i + 1] - A.c[i]) / (A.k[i + A.d + 1] - A.k[i + 1])
         end
     end
     return ducum * A.d * scale
@@ -280,15 +282,17 @@ function _derivative(
     n = length(A.t)
     scale = (A.p[idx + 1] - A.p[idx]) / (A.t[idx + 1] - A.t[idx])
     t_ = A.p[idx] + (t - A.t[idx]) * scale
-    sc = t isa ForwardDiff.Dual ? zeros(eltype(t), n) : A.sc
-    spline_coefficients!(sc, A.d - 1, A.k, t_)
     ducum = zeros(size(A.u)[1:(end - 1)]...)
     if t == A.t[1]
         ducum = (A.c[ax_u..., 2] - A.c[ax_u..., 1]) / (A.k[A.d + 2])
     else
-        for i in 1:(n - 1)
+        # Stack-allocated basis window: differentiation must be reentrant (#532)
+        vals, offset, m = bspline_nonzero_coefficients(A.d - 1, A.k, t_, n)
+        @inbounds for l in 1:m
+            i = offset + l - 1
+            (1 <= i <= n - 1) || continue
             ducum = ducum +
-                sc[i + 1] * (A.c[ax_u..., i + 1] - A.c[ax_u..., i]) /
+                vals[l] * (A.c[ax_u..., i + 1] - A.c[ax_u..., i]) /
                 (A.k[i + A.d + 1] - A.k[i + 1])
         end
     end
@@ -302,14 +306,16 @@ function _derivative(A::BSplineApprox{<:AbstractVector{<:Number}}, t::Number, ig
     idx = get_idx(A, t, iguess)
     scale = (A.p[idx + 1] - A.p[idx]) / (A.t[idx + 1] - A.t[idx])
     t_ = A.p[idx] + (t - A.t[idx]) * scale
-    sc = t isa ForwardDiff.Dual ? zeros(eltype(t), A.h) : A.sc
-    spline_coefficients!(sc, A.d - 1, A.k, t_)
     ducum = zero(eltype(A.u))
     if t == A.t[1]
         ducum = (A.c[2] - A.c[1]) / (A.k[A.d + 2])
     else
-        for i in 1:(A.h - 1)
-            ducum += sc[i + 1] * (A.c[i + 1] - A.c[i]) / (A.k[i + A.d + 1] - A.k[i + 1])
+        # Stack-allocated basis window: differentiation must be reentrant (#532)
+        vals, offset, m = bspline_nonzero_coefficients(A.d - 1, A.k, t_, A.h)
+        @inbounds for l in 1:m
+            i = offset + l - 1
+            (1 <= i <= A.h - 1) || continue
+            ducum += vals[l] * (A.c[i + 1] - A.c[i]) / (A.k[i + A.d + 1] - A.k[i + 1])
         end
     end
     return ducum * A.d * scale
@@ -325,14 +331,16 @@ function _derivative(
     idx = get_idx(A, t, iguess)
     scale = (A.p[idx + 1] - A.p[idx]) / (A.t[idx + 1] - A.t[idx])
     t_ = A.p[idx] + (t - A.t[idx]) * scale
-    sc = t isa ForwardDiff.Dual ? zeros(eltype(t), A.h) : A.sc
-    spline_coefficients!(sc, A.d - 1, A.k, t_)
     ducum = zeros(size(A.u)[1:(end - 1)]...)
     if t == A.t[1]
         ducum = (A.c[ax_u..., 2] - A.c[ax_u..., 1]) / (A.k[A.d + 2])
     else
-        for i in 1:(A.h - 1)
-            ducum += sc[i + 1] * (A.c[ax_u..., i + 1] - A.c[ax_u..., i]) /
+        # Stack-allocated basis window: differentiation must be reentrant (#532)
+        vals, offset, m = bspline_nonzero_coefficients(A.d - 1, A.k, t_, A.h)
+        @inbounds for l in 1:m
+            i = offset + l - 1
+            (1 <= i <= A.h - 1) || continue
+            ducum = ducum + vals[l] * (A.c[ax_u..., i + 1] - A.c[ax_u..., i]) /
                 (A.k[i + A.d + 1] - A.k[i + 1])
         end
     end

src/interpolation_caches.jl

@@ -668,7 +668,7 @@ function QuadraticSpline(
     k, A = quadratic_spline_params(t, sc)
     c = A \ u
 
-    p = QuadraticSplineParameterCache(u, t, k, c, sc, cache_parameters)
+    p = QuadraticSplineParameterCache(u, t, k, c, cache_parameters)
     A = QuadraticSpline(
         u, t, nothing, p, k, c, sc, extrapolation_left,
         extrapolation_right, cache_parameters, t_props
@@ -710,7 +710,7 @@ function QuadraticSpline(
         end
     end
 
-    p = QuadraticSplineParameterCache(u, t, k, c, sc, cache_parameters)
+    p = QuadraticSplineParameterCache(u, t, k, c, cache_parameters)
     A = QuadraticSpline(
         u, t, nothing, p, k, c, sc, extrapolation_left,
         extrapolation_right, cache_parameters, t_props
@@ -1020,6 +1020,8 @@ function BSplineInterpolation(
     t_props = something(search_properties, FindFirstFunctions.SearchProperties(t))
     n = length(t)
     n < d + 1 && error("BSplineInterpolation needs at least d + 1, i.e. $(d + 1) points.")
+    d ≥ BSPLINE_STACK_MAXLEN &&
+        error("BSplineInterpolation supports degree d < $(BSPLINE_STACK_MAXLEN); got d = $d.")
     s = zero(eltype(u))
     p = zero(t)
     k = zeros(eltype(t), n + d + 1)
@@ -1101,6 +1103,8 @@ function BSplineInterpolation(
     t_props = something(search_properties, FindFirstFunctions.SearchProperties(t))
     n = length(t)
     n < d + 1 && error("BSplineInterpolation needs at least d + 1, i.e. $(d + 1) points.")
+    d ≥ BSPLINE_STACK_MAXLEN &&
+        error("BSplineInterpolation supports degree d < $(BSPLINE_STACK_MAXLEN); got d = $d.")
     s = zero(eltype(u))
     p = zero(t)
     k = zeros(eltype(t), n + d + 1)
@@ -1274,6 +1278,8 @@ function BSplineApprox(
     t_props = something(search_properties, FindFirstFunctions.SearchProperties(t))
     n = length(t)
     h < d + 1 && error("BSplineApprox needs at least d + 1, i.e. $(d + 1) control points.")
+    d ≥ BSPLINE_STACK_MAXLEN &&
+        error("BSplineApprox supports degree d < $(BSPLINE_STACK_MAXLEN); got d = $d.")
     s = zero(eltype(u))
     p = zero(t)
     k = zeros(eltype(t), h + d + 1)
@@ -1376,6 +1382,8 @@ function BSplineApprox(
     t_props = something(search_properties, FindFirstFunctions.SearchProperties(t))
     n = length(t)
     h < d + 1 && error("BSplineApprox needs at least d + 1, i.e. $(d + 1) control points.")
+    d ≥ BSPLINE_STACK_MAXLEN &&
+        error("BSplineApprox supports degree d < $(BSPLINE_STACK_MAXLEN); got d = $d.")
     s = zero(eltype(u))
     p = zero(t)
     k = zeros(eltype(t), h + d + 1)

src/interpolation_methods.jl

@@ -1264,11 +1264,11 @@ function _interpolate(
     idx = get_idx(A, t, iguess)
     t = A.p[idx] + (t - A.t[idx]) / (A.t[idx + 1] - A.t[idx]) * (A.p[idx + 1] - A.p[idx])
     n = length(A.t)
-    sc = t isa ForwardDiff.Dual ? zeros(eltype(t), n) : A.sc
-    nonzero_coefficient_idxs = spline_coefficients!(sc, A.d, A.k, t)
+    # Stack-allocated basis window: evaluation must be reentrant for thread safety (#532)
+    vals, offset, m = bspline_nonzero_coefficients(A.d, A.k, t, n)
     ucum = zero(eltype(A.u))
-    for i in nonzero_coefficient_idxs
-        ucum += sc[i] * A.c[i]
+    @inbounds for l in 1:m
+        ucum += vals[l] * A.c[offset + l]
     end
     return ucum
 end
@@ -1285,11 +1285,11 @@ function _interpolate(
     idx = get_idx(A, t, iguess)
     t = A.p[idx] + (t - A.t[idx]) / (A.t[idx + 1] - A.t[idx]) * (A.p[idx + 1] - A.p[idx])
     n = length(A.t)
-    sc = t isa ForwardDiff.Dual ? zeros(eltype(t), n) : A.sc
-    nonzero_coefficient_idxs = spline_coefficients!(sc, A.d, A.k, t)
+    # Stack-allocated basis window: evaluation must be reentrant for thread safety (#532)
+    vals, offset, m = bspline_nonzero_coefficients(A.d, A.k, t, n)
     ucum = zeros(eltype(A.u), size(A.u)[1:(end - 1)]...)
-    for i in nonzero_coefficient_idxs
-        ucum = ucum + (sc[i] * A.c[ax_u..., i])
+    @inbounds for l in 1:m
+        ucum = ucum + (vals[l] * A.c[ax_u..., offset + l])
     end
     return ucum
 end
@@ -1301,11 +1301,11 @@ function _interpolate(A::BSplineApprox{<:AbstractVector{<:Number}}, t::Number, i
     # change t into param [0 1]
     idx = get_idx(A, t, iguess)
     t = A.p[idx] + (t - A.t[idx]) / (A.t[idx + 1] - A.t[idx]) * (A.p[idx + 1] - A.p[idx])
-    sc = t isa ForwardDiff.Dual ? zeros(eltype(t), A.h) : A.sc
-    nonzero_coefficient_idxs = spline_coefficients!(sc, A.d, A.k, t)
+    # Stack-allocated basis window: evaluation must be reentrant for thread safety (#532)
+    vals, offset, m = bspline_nonzero_coefficients(A.d, A.k, t, A.h)
     ucum = zero(eltype(A.u))
-    for i in nonzero_coefficient_idxs
-        ucum += sc[i] * A.c[i]
+    @inbounds for l in 1:m
+        ucum += vals[l] * A.c[offset + l]
     end
     return ucum
 end
@@ -1319,11 +1319,11 @@ function _interpolate(
     # change t into param [0 1]
     idx = get_idx(A, t, iguess)
     t = A.p[idx] + (t - A.t[idx]) / (A.t[idx + 1] - A.t[idx]) * (A.p[idx + 1] - A.p[idx])
-    sc = t isa ForwardDiff.Dual ? zeros(eltype(t), A.h) : A.sc
-    nonzero_coefficient_idxs = spline_coefficients!(sc, A.d, A.k, t)
+    # Stack-allocated basis window: evaluation must be reentrant for thread safety (#532)
+    vals, offset, m = bspline_nonzero_coefficients(A.d, A.k, t, A.h)
     ucum = zeros(eltype(A.u), size(A.u)[1:(end - 1)]...)
-    for i in nonzero_coefficient_idxs
-        ucum = ucum + (sc[i] * A.c[ax_u..., i])
+    @inbounds for l in 1:m
+        ucum = ucum + (vals[l] * A.c[ax_u..., offset + l])
     end
     return ucum
 end

src/interpolation_utils.jl

@@ -67,6 +67,46 @@ function spline_coefficients!(N, d, k, u::AbstractVector)
     return nothing
 end
 
+# A degree-`d` B-spline has only `d + 1` nonzero basis functions at any point, so
+# evaluation needs a scratch buffer of just `d + 1` entries — not one sized to the
+# full knot/control-point vector. `bspline_nonzero_coefficients` computes those
+# values into a stack-allocated `MVector` (no heap allocation, hence reentrant /
+# thread-safe, #532) for degrees up to `BSPLINE_STACK_MAXLEN - 1`; callers fall back
+# to the heap `spline_coefficients!` for larger degrees, which essentially never occur.
+const BSPLINE_STACK_MAXLEN = 16
+
+# Returns `(vals, offset, m)`: `vals[l]` is the basis weight of control point
+# `offset + l` for `l in 1:m`. `ncp` is the number of control points (needed only to
+# place the single nonzero weight at the right boundary). Mirrors the locator and
+# recurrence of `spline_coefficients!`, but writes into a local window indexed `1:m`
+# instead of absolute knot indices.
+@inline function bspline_nonzero_coefficients(d::Integer, k, u::T, ncp::Integer) where {T}
+    N = zero(MVector{BSPLINE_STACK_MAXLEN, T})
+    if u == k[1]
+        @inbounds N[1] = one(u)
+        return SVector(N), 0, 1
+    elseif u == k[end]
+        @inbounds N[1] = one(u)
+        return SVector(N), ncp - 1, 1
+    end
+    i = searchsortedlast(k, u)
+    # Local index of global knot index `g` is `g + off` (so `i - d → 1`, `i → d + 1`).
+    off = d + 1 - i
+    @inbounds begin
+        N[i + off] = one(u)
+        for deg in 1:d
+            N[i - deg + off] = (k[i + 1] - u) / (k[i + 1] - k[i - deg + 1]) *
+                N[i - deg + 1 + off]
+            for j in (i - deg + 1):(i - 1)
+                N[j + off] = (u - k[j]) / (k[j + deg] - k[j]) * N[j + off] +
+                    (k[j + deg + 1] - u) / (k[j + deg + 1] - k[j + 1]) * N[j + 1 + off]
+            end
+            N[i + off] = (u - k[i]) / (k[i + deg] - k[i]) * N[i + off]
+        end
+    end
+    return SVector(N), i - d - 1, d + 1
+end
+
 function quadratic_spline_params(t::AbstractVector, sc::AbstractVector)
     # Duplicate time points make the collocation system singular
     if any(i -> t[i] == t[i + 1], 1:(length(t) - 1))
@@ -301,7 +341,7 @@ function get_parameters(A::QuadraticSpline, idx)
     return if A.cache_parameters
         A.p.α[idx], A.p.β[idx]
     else
-        quadratic_spline_parameters(A.u, A.t, A.k, A.c, A.sc, idx)
+        quadratic_spline_parameters(A.u, A.t, A.k, A.c, idx)
     end
 end
 

src/parameter_caches.jl

@@ -131,33 +131,46 @@ struct QuadraticSplineParameterCache{pType}
     β::pType
 end
 
-function QuadraticSplineParameterCache(u, t, k, c, sc, cache_parameters)
+function QuadraticSplineParameterCache(u, t, k, c, cache_parameters)
     return if cache_parameters
         parameters = quadratic_spline_parameters.(
-            Ref(u), Ref(t), Ref(k), Ref(c), Ref(sc), 1:(length(t) - 1)
+            Ref(u), Ref(t), Ref(k), Ref(c), 1:(length(t) - 1)
         )
         α, β = collect.(eachrow(stack(collect.(parameters))))
         QuadraticSplineParameterCache(α, β)
     else
         # Compute parameters once to infer types
-        α, β = quadratic_spline_parameters(u, t, k, c, sc, 1)
+        α, β = quadratic_spline_parameters(u, t, k, c, 1)
         QuadraticSplineParameterCache(typeof(α)[], typeof(β)[])
     end
 end
 
-function quadratic_spline_parameters(u, t, k, c, sc, idx)
+function quadratic_spline_parameters(u, t, k, c, idx)
     uᵢ₊ = if length(t) == 2
         # For 2 data points the knot vector has boundary multiplicity 2 < degree + 1,
         # so the B-spline basis cannot be evaluated at interior points; the spline
         # degenerates to the linear interpolant (α = 0, β = u₂ - u₁).
         (u[1] + u[2]) / 2
     else
         tᵢ₊ = (t[idx] + t[idx + 1]) / 2
-        nonzero_coefficient_idxs = spline_coefficients!(sc, 2, k, tᵢ₊)
+        # Value of the spline at the segment midpoint via the (degree 2) B-spline basis.
+        # The three nonzero basis values are evaluated into local variables rather than a
+        # shared scratch buffer so that evaluation is reentrant / thread-safe (#532).
+        # `tᵢ₊` is always interior, so only the non-boundary branch of the Cox-de Boor
+        # recursion is needed (cf. `spline_coefficients!`).
+        i = findfirst(x -> x > tᵢ₊, k)::Int - 1
+        w₁ = (k[i + 1] - tᵢ₊) / (k[i + 1] - k[i])
+        w₂ = (tᵢ₊ - k[i]) / (k[i + 1] - k[i])
+        N₁ = (k[i + 1] - tᵢ₊) / (k[i + 1] - k[i - 1]) * w₁
+        N₂ = (tᵢ₊ - k[i - 1]) / (k[i + 1] - k[i - 1]) * w₁ +
+            (k[i + 2] - tᵢ₊) / (k[i + 2] - k[i]) * w₂
+        N₃ = (tᵢ₊ - k[i]) / (k[i + 2] - k[i]) * w₂
+        # Seed the accumulator with `zero(first(u))` (as the buffer-based version did) so
+        # `uᵢ₊` keeps the element type of `u` for vector-valued data.
         uᵢ₊ = zero(first(u))
-        for j in nonzero_coefficient_idxs
-            uᵢ₊ += sc[j] * c[j]
-        end
+        uᵢ₊ += N₁ * c[i - 2]
+        uᵢ₊ += N₂ * c[i - 1]
+        uᵢ₊ += N₃ * c[i]
         uᵢ₊
     end
     α = 2 * (u[idx + 1] + u[idx]) - 4uᵢ₊