Repository: FluxML/NNlib.jl
Pull request #633: Don't spawn if: only one iteration or no threads or threading disabled
Author: login: IanButterworth
Created: 2025-03-12T14:13:28Z | Merged: 2025-03-18T22:12:27Z

A single pass of ObjectDetector.jl creates 150 new tasks, all of which were within for loops with length-1 iterators.. i.e. the new task can be entirely avoided.

With this there are 0 spawned tasks.

Avoiding constructing the iterators entirely should be possible too, but I thought I'd open a PR with these changes before looking into that.

### PR Checklist

- [x] Tests are added
- [x] Documentation, if applicable



Diff:
diff --git a/Project.toml b/Project.toml
index 026eeb17..bad42524 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "NNlib"
 uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-version = "0.9.28"
+version = "0.9.29"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -10,6 +10,7 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [weakdeps]
@@ -43,6 +44,7 @@ GPUArraysCore = "0.1, 0.2"
 KernelAbstractions = "0.9.2"
 LinearAlgebra = "<0.0.1, 1"
 Random = "<0.0.1, 1"
+ScopedValues = "1.3.0"
 SpecialFunctions = "2"
 Statistics = "1"
 cuDNN = "1"
diff --git a/docs/src/index.md b/docs/src/index.md
index 46958da1..78d09df4 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -12,4 +12,13 @@ for CUDA support, or
 ```julia
 using NNlib, AMDGPU
 ```
-for AMDGPU support.
\ No newline at end of file
+for AMDGPU support.
+
+## Threading
+
+Various `NNlib` functions utilize available julia threads on divisible workloads. To disable this use
+the `ScopedValue`-backed switch `NNlib.@disallow_spawns`
+i.e.
+```julia
+NNlib.@disallow_spawns function_that_uses_nnlib()
+```
diff --git a/src/NNlib.jl b/src/NNlib.jl
index ad5c7441..5ea90778 100644
--- a/src/NNlib.jl
+++ b/src/NNlib.jl
@@ -13,11 +13,26 @@ using LinearAlgebra
 using LinearAlgebra.BLAS: @blasfunc, BlasInt
 using LinearAlgebra: AdjOrTransAbsMat, Adjoint, BlasFloat, Transpose
 using Random
+using ScopedValues
 using Statistics
 using Statistics: mean
 
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 
+# internal. TODO: change to an approach where amount of threading is controlled, not just on/off
+const ALLOW_SPAWNS = ScopedValue(true)
+should_use_spawn() = Threads.nthreads(:default) > 1 && ALLOW_SPAWNS[]
+"""
+    @disallow_spawns ex
+
+Disallow NNlib to use `@spawn` on divisible workloads. i.e. within `conv` etc.
+"""
+macro disallow_spawns(ex)
+    quote
+        @with ALLOW_SPAWNS => false $(esc(ex))
+    end
+end
+
 # Include APIs
 include("dim_helpers.jl")
 export ConvDims, DenseConvDims, PoolDims, DepthwiseConvDims
@@ -35,7 +50,7 @@ include("dropout.jl")
 export dropout, dropout!
 
 include("softmax.jl")
-export softmax, softmax!, ∇softmax, ∇softmax!, logsoftmax, 
+export softmax, softmax!, ∇softmax, ∇softmax!, logsoftmax,
     logsoftmax!, ∇logsoftmax, ∇logsoftmax!, logsumexp
 
 include("batched/batchedadjtrans.jl")
@@ -47,9 +62,9 @@ include("gemm.jl")
 export grid_sample, ∇grid_sample
 
 include("conv.jl")
-export conv, conv!, ∇conv_data, ∇conv_data!, ∇conv_filter, 
-    ∇conv_filter!, depthwiseconv, depthwiseconv!, 
-    ∇depthwiseconv_data, ∇depthwiseconv_data!, 
+export conv, conv!, ∇conv_data, ∇conv_data!, ∇conv_filter,
+    ∇conv_filter!, depthwiseconv, depthwiseconv!,
+    ∇depthwiseconv_data, ∇depthwiseconv_data!,
     ∇depthwiseconv_filter, ∇depthwiseconv_filter!
 
 include("conv_bias_act.jl")
diff --git a/src/conv.jl b/src/conv.jl
index f4050bdb..26b3d036 100644
--- a/src/conv.jl
+++ b/src/conv.jl
@@ -181,7 +181,7 @@ for (front_name, backend, signature) in (
 )
     # We only define 3d conv primitives, we reshape lower down to get 1d and 2d convolution
     @eval begin
-        
+
         function $(Symbol("$(front_name)!"))(
                         out::AbstractArray{$(signature[1][1]), $(signature[1][2])},
                         in1::AbstractArray{$(signature[2][1]), $(signature[1][2])},
@@ -202,11 +202,21 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (xc, wc) in zip(x_cs, w_cs)
+            function conv_group(xc, wc)
                 x = @view in1[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
                 w = @view in2[ntuple(i -> i == 5 ? wc : Colon(), 5)...]
                 y = @view out[ntuple(i -> i == 4 ? wc : Colon(), 5)...]
-                Threads.@spawn $(Symbol("$(front_name)_$(backend)!"))(y, x, w, cdims2; kwargs...)
+                $(Symbol("$(front_name)_$(backend)!"))(y, x, w, cdims2; kwargs...)
+            end
+
+            if should_use_spawn() && length(x_cs) > 1
+                Threads.@sync for (xc, wc) in zip(x_cs, w_cs)
+                    Threads.@spawn conv_group(xc, wc)
+                end
+            else
+                for (xc, wc) in zip(x_cs, w_cs)
+                    conv_group(xc, wc)
+                end
             end
 
             return out
@@ -246,11 +256,21 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (xc, yc, wc) in zip(dx_cs, dy_cs, w_cs)
+            function ∇conv_data_group(xc, yc, wc)
                 dxv = @view out[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
                 dyv = @view in1[ntuple(i -> i == 4 ? yc : Colon(), 5)...]
                 wv = @view in2[ntuple(i -> i == 5  ? wc : Colon(), 5)...]
-                Threads.@spawn $(Symbol("$(front_name)_$(backend)!"))(dxv, dyv, wv, cdims2; kwargs...)
+                $(Symbol("$(front_name)_$(backend)!"))(dxv, dyv, wv, cdims2; kwargs...)
+            end
+
+            if should_use_spawn() && length(dx_cs) > 1
+                Threads.@sync for (xc, yc, wc) in zip(dx_cs, dy_cs, w_cs)
+                    Threads.@spawn ∇conv_data_group(xc, yc, wc)
+                end
+            else
+                for (xc, yc, wc) in zip(dx_cs, dy_cs, w_cs)
+                    ∇conv_data_group(xc, yc, wc)
+                end
             end
 
             return out
@@ -288,11 +308,21 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (wc, xc, yc) in zip(dw_cs, x_cs, dy_cs)
+            function ∇conv_filter_group(wc, xc, yc)
                 x = @view in1[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
                 dy = @view in2[ntuple(i -> i == 4 ? yc : Colon(), 5)...]
-                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...]
-                Threads.@spawn $(Symbol("$(front_name)_$(backend)!"))(dw, x, dy, cdims2; kwargs...)
+                dw = @view out[ntuple(i -> i == 5 ? wc : Colon(), 5)...]
+                $(Symbol("$(front_name)_$(backend)!"))(dw, x, dy, cdims2; kwargs...)
+            end
+
+            if should_use_spawn() && length(dw_cs) > 1
+                Threads.@sync for (wc, xc, yc) in zip(dw_cs, x_cs, dy_cs)
+                    Threads.@spawn ∇conv_filter_group(wc, xc, yc)
+                end
+            else
+                for (wc, xc, yc) in zip(dw_cs, x_cs, dy_cs)
+                    ∇conv_filter_group(wc, xc, yc)
+                end
             end
 
             return out
@@ -306,10 +336,10 @@ for (front_name, backend, signature) in (
     # (frontend, backend, (out Array signature, in1 Array signature, in2 Array signature, (parametric Types)))
     (:depthwiseconv, :im2col, ((:T, 5), (:T, 5), (:T, 5), :C, (:(T <: G), :(C <: ConvDims)))),
     (:depthwiseconv, :direct, ((:yT, :N), (:T1, :N), (:T2, :N), :C, (:yT, :T1, :T2, :N, :(C <: ConvDims)))),
-    
+
     (:∇depthwiseconv_data, :im2col, ((:T, 5), (:T, 5), (:T, 5), :C, (:(T <: G), :(C <: ConvDims)))),
     (:∇depthwiseconv_data, :direct, ((:yT, :N), (:T1, :N), (:T2, :N), :C, (:yT, :T1, :T2, :N, :(C <: ConvDims)))),
-    
+
     (:∇depthwiseconv_filter, :im2col, ((:T, 5), (:T, 5), (:T, 5), :C, (:(T <: G), :(C <: ConvDims)))),
     (:∇depthwiseconv_filter, :direct, ((:yT, :N), (:T1, :N), (:T2, :N), :C, (:yT, :T1, :T2, :N, :(C <: ConvDims)))),
 )
diff --git a/src/gemm.jl b/src/gemm.jl
index 9cbe5019..7ee407a9 100644
--- a/src/gemm.jl
+++ b/src/gemm.jl
@@ -104,22 +104,34 @@ for (gemm, elt) in gemm_datatype_mappings
 
                 old_threads = get_num_threads()
                 set_num_threads(1)
-                Threads.@sync for ks in Iterators.partition(1:size(C, 3), cld(size(C, 3), n_threads))
-                    Threads.@spawn for k in ks
+
+                parts = Iterators.partition(1:size(C, 3), cld(size(C, 3), n_threads))
+
+                function gemm!_part(ks)
+                    for k in ks
 
                         ptrAk = ptrA + (k-1) * strA * sizeof($elt)
                         ptrBk = ptrB + (k-1) * strB * sizeof($elt)
                         ptrCk = ptrC + (k-1) * strC * sizeof($elt)
 
                         ccall((@blasfunc($(gemm)), libblas), Nothing,
-                              (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ref{BlasInt},
-                               Ref{BlasInt}, Ref{$elt}, Ptr{$elt}, Ref{BlasInt},
-                               Ptr{$elt}, Ref{BlasInt}, Ref{$elt}, Ptr{$elt},
-                               Ref{BlasInt}),
-                              transA, transB, m, n,
-                              ka, alpha, ptrAk, max(1,Base.stride(A,2)),
-                              ptrBk, max(1,Base.stride(B,2)), beta, ptrCk,
-                              max(1,Base.stride(C,2)))
+                            (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ref{BlasInt},
+                            Ref{BlasInt}, Ref{$elt}, Ptr{$elt}, Ref{BlasInt},
+                            Ptr{$elt}, Ref{BlasInt}, Ref{$elt}, Ptr{$elt},
+                            Ref{BlasInt}),
+                            transA, transB, m, n,
+                            ka, alpha, ptrAk, max(1,Base.stride(A,2)),
+                            ptrBk, max(1,Base.stride(B,2)), beta, ptrCk,
+                            max(1,Base.stride(C,2)))
+                    end
+                end
+                if should_use_spawn() && length(parts) > 1
+                    Threads.@sync for ks in parts
+                        Threads.@spawn gemm!_part(ks)
+                    end
+                else
+                    for ks in parts
+                        gemm!_part(ks)
                     end
                 end
                 set_num_threads(old_threads)
diff --git a/src/impl/conv_im2col.jl b/src/impl/conv_im2col.jl
index 8c497954..1893822d 100644
--- a/src/impl/conv_im2col.jl
+++ b/src/impl/conv_im2col.jl
@@ -47,20 +47,28 @@ function conv_im2col!(
 
     parts = Iterators.partition(axes(x, 5), ceil(Int, size(x, 5) / ntasks))
 
-    @sync for (task_n, part) in enumerate(parts)
-        Threads.@spawn begin
-            col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
-            for batch_idx in part
-                im2col!(col_slice, view(x, :, :, :, :, batch_idx), cdims)
-                GC.@preserve col_slice w y begin
-                    col_ptr = pointer(col_slice)
-                    w_ptr = pointer(w)
-                    y_ptr = pointer(y, (batch_idx - 1)*M*N + 1)
-                    gemm!(Val(false), Val(false), M, N, K, alpha, col_ptr, w_ptr, beta, y_ptr)
-                end
+    function conv_part(task_n, part)
+        col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
+        for batch_idx in part
+            im2col!(col_slice, view(x, :, :, :, :, batch_idx), cdims)
+            GC.@preserve col_slice w y begin
+                col_ptr = pointer(col_slice)
+                w_ptr = pointer(w)
+                y_ptr = pointer(y, (batch_idx - 1)*M*N + 1)
+                gemm!(Val(false), Val(false), M, N, K, alpha, col_ptr, w_ptr, beta, y_ptr)
             end
         end
     end
+
+    if should_use_spawn() && length(parts) > 1
+        @sync for (task_n, part) in enumerate(parts)
+            Threads.@spawn conv_part(task_n, part)
+        end
+    else
+        for (task_n, part) in enumerate(parts)
+            conv_part(task_n, part)
+        end
+    end
     return y
 end
 
@@ -152,18 +160,25 @@ function ∇conv_data_im2col!(
 
     parts = Iterators.partition(axes(dx, 5), ceil(Int, size(dx, 5) / ntasks))
 
-    @sync for (task_n, part) in enumerate(parts)
-        Threads.@spawn begin
-            col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
-            for batch_idx in part
-                GC.@preserve col_slice w dy begin
-                    dy_ptr = pointer(dy, (batch_idx - 1)*M*K + 1)
-                    w_ptr = pointer(w)
-                    col_ptr = pointer(col_slice)
-                    gemm!(Val(false), Val(true), M, N, K, alpha, dy_ptr, w_ptr, T(0), col_ptr)
-                end
-                col2im!(view(dx, :, :, :, :, batch_idx), col_slice, cdims, beta)
+    function ∇conv_data_part(task_n, part)
+        col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
+        for batch_idx in part
+            GC.@preserve col_slice w dy begin
+                dy_ptr = pointer(dy, (batch_idx - 1)*M*K + 1)
+                w_ptr = pointer(w)
+                col_ptr = pointer(col_slice)
+                gemm!(Val(false), Val(true), M, N, K, alpha, dy_ptr, w_ptr, T(0), col_ptr)
             end
+            col2im!(view(dx, :, :, :, :, batch_idx), col_slice, cdims, beta)
+        end
+    end
+    if should_use_spawn() && length(parts) > 1
+        @sync for (task_n, part) in enumerate(parts)
+            Threads.@spawn ∇conv_data_part(task_n, part)
+        end
+    else
+        for (task_n, part) in enumerate(parts)
+            ∇conv_data_part(task_n, part)
         end
     end
     return dx
diff --git a/src/impl/depthwiseconv_im2col.jl b/src/impl/depthwiseconv_im2col.jl
index a909dc05..60e40a9d 100644
--- a/src/impl/depthwiseconv_im2col.jl
+++ b/src/impl/depthwiseconv_im2col.jl
@@ -30,25 +30,32 @@ function depthwiseconv_im2col!(
 
     dcdims = DenseConvDims(cdims)
 
-    @sync for (task_n, part) in enumerate(parts)
-        Threads.@spawn begin
-            col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
-            for batch_idx in part
-                im2col!(col_slice, view(x, :, :, :, :, batch_idx), dcdims)
-
-                # We do a separate convolution for each channel in x, as we must
-                for c_in in 1:channels_in(cdims)
-                    # Walk each pointer forward as we process each input channel
-                    GC.@preserve col_slice w y begin
-                        col_ptr = pointer(col_slice, (c_in-1)*M*K+1)
-                        w_ptr = pointer(w, (c_in-1)*K*N+1)
-                        y_ptr = pointer(y, ((batch_idx - 1)*channels_in(cdims) + c_in - 1)*M*N + 1)
-                        gemm!(Val(false), Val(false), M, N, K, alpha, col_ptr, w_ptr, beta, y_ptr)
-                    end
+    function depthwiseconv_part(task_n, part)
+        col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
+        for batch_idx in part
+            im2col!(col_slice, view(x, :, :, :, :, batch_idx), dcdims)
+
+            # We do a separate convolution for each channel in x, as we must
+            for c_in in 1:channels_in(cdims)
+                # Walk each pointer forward as we process each input channel
+                GC.@preserve col_slice w y begin
+                    col_ptr = pointer(col_slice, (c_in-1)*M*K+1)
+                    w_ptr = pointer(w, (c_in-1)*K*N+1)
+                    y_ptr = pointer(y, ((batch_idx - 1)*channels_in(cdims) + c_in - 1)*M*N + 1)
+                    gemm!(Val(false), Val(false), M, N, K, alpha, col_ptr, w_ptr, beta, y_ptr)
                 end
             end
         end
     end
+    if should_use_spawn() && length(parts) > 1
+        @sync for (task_n, part) in enumerate(parts)
+            Threads.@spawn depthwiseconv_part(task_n, part)
+        end
+    else
+        for (task_n, part) in enumerate(parts)
+            depthwiseconv_part(task_n, part)
+        end
+    end
     return y
 end
 
@@ -117,22 +124,29 @@ function ∇depthwiseconv_data_im2col!(
 
     parts = Iterators.partition(axes(dx)[end], ceil(Int, size(dx, 5) / ntasks))
 
-    @sync for (task_n, part) in enumerate(parts)
-        Threads.@spawn begin
-            col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
-            for batch_idx in part
-                # We do a separate convolution for each channel in x, as we must
-                for cidx in 1:channels_in(cdims)
-                    GC.@preserve col_slice w dy begin
-                        # Walk each pointer forward as we process each input channel
-                        dy_ptr = pointer(dy, (batch_idx - 1)*M*K*channels_in(cdims)+(cidx - 1)*K*M + 1)
-                        w_ptr = pointer(w, (cidx - 1)*K*N + 1)
-                        col_ptr = pointer(col_slice, (cidx - 1)*M*N + 1)
-                        gemm!(Val(false), Val(true), M, N, K, alpha, dy_ptr, w_ptr, T(0), col_ptr)
-                    end
+    function ∇depthwiseconv_data_part(task_n, part)
+        col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
+        for batch_idx in part
+            # We do a separate convolution for each channel in x, as we must
+            for cidx in 1:channels_in(cdims)
+                GC.@preserve col_slice w dy begin
+                    # Walk each pointer forward as we process each input channel
+                    dy_ptr = pointer(dy, (batch_idx - 1)*M*K*channels_in(cdims)+(cidx - 1)*K*M + 1)
+                    w_ptr = pointer(w, (cidx - 1)*K*N + 1)
+                    col_ptr = pointer(col_slice, (cidx - 1)*M*N + 1)
+                    gemm!(Val(false), Val(true), M, N, K, alpha, dy_ptr, w_ptr, T(0), col_ptr)
                 end
-                col2im!(view(dx, :, :, :, :, batch_idx), col_slice, cdims, beta)
             end
+            col2im!(view(dx, :, :, :, :, batch_idx), col_slice, cdims, beta)
+        end
+    end
+    if should_use_spawn() && length(parts) > 1
+        @sync for (task_n, part) in enumerate(parts)
+            Threads.@spawn ∇depthwiseconv_data_part(task_n, part)
+        end
+    else
+        for (task_n, part) in enumerate(parts)
+            ∇depthwiseconv_data_part(task_n, part)
         end
     end
     return dx
diff --git a/test/conv.jl b/test/conv.jl
index badb2c5f..cf323277 100644
--- a/test/conv.jl
+++ b/test/conv.jl
@@ -305,7 +305,7 @@ ddims(x) = dropdims(x, dims=(ndims(x)-1, ndims(x)))
                 end
             end
 
-            # Test all in-place implementations/interfaces 
+            # Test all in-place implementations/interfaces
             convs = [NNlib.conv!, NNlib.conv_im2col!, NNlib.conv_direct!,]
             for conv! in convs
                 α, β = 2e0, -1e0
@@ -393,12 +393,12 @@ ddims(x) = dropdims(x, dims=(ndims(x)-1, ndims(x)))
             for beta in (-2.0, -1.0, 0.0, 0.5, 1.0, 2.0)
                 cache_dx, cache_dy, cache_w = ([0.17;;; 0.19;;; 0.23], [0.11;;; 0.13;;; 0.15], [1.0;;;])
                 dx_old = copy(cache_dx)
-                cdims = DenseConvDims(cache_dx, cache_w)            
+                cdims = DenseConvDims(cache_dx, cache_w)
                 NNlib.∇conv_data_im2col!(cache_dx, cache_dy, cache_w, cdims; alpha=1.0, beta)
                 @test isapprox(cache_dx, dx_old * beta + cache_dy, rtol = 1.0e-7)
             end
 
-            # Test all in-place implementations/interfaces 
+            # Test all in-place implementations/interfaces
             for (∇conv_filter!, ∇conv_data!) in (
                     (NNlib.∇conv_filter!,        NNlib.∇conv_data!),
                     (NNlib.∇conv_filter_im2col!, NNlib.∇conv_data_im2col!),
@@ -411,14 +411,14 @@ ddims(x) = dropdims(x, dims=(ndims(x)-1, ndims(x)))
                     # First, your basic convolution with no parameters
                     cdims = DenseConvDims(x, w)
                     dy = NNlib.conv(x, w, cdims)
-                    @test isapprox(ddims(∇conv_filter!(copy(w), x, dy, cdims; alpha=α, beta=β)), α*dw + β*w, rtol = 1.0e-7) 
+                    @test isapprox(ddims(∇conv_filter!(copy(w), x, dy, cdims; alpha=α, beta=β)), α*dw + β*w, rtol = 1.0e-7)
                     @test isapprox(ddims(∇conv_data!(copy(x), dy, w,   cdims; alpha=α, beta=β)), α*dx + β*x, rtol = 1.0e-7)
 
                     # Next, test convolution on views and alternate datatypes:
-                    @test isapprox(ddims(∇conv_filter!(copy(w), x, view(dy, repeat([:], ndims(dy))...), cdims; alpha=α, beta=β)), α*dw + β*w, rtol = 1.0e-7) 
+                    @test isapprox(ddims(∇conv_filter!(copy(w), x, view(dy, repeat([:], ndims(dy))...), cdims; alpha=α, beta=β)), α*dw + β*w, rtol = 1.0e-7)
                     @test isapprox(ddims(∇conv_data!(copy(x), view(dy, repeat([:], ndims(dy))...), w,   cdims; alpha=α, beta=β)), α*dx + β*x, rtol = 1.0e-7)
 
-                    @test isapprox(ddims(∇conv_filter!(Float32.(copy(w)), Float32.(x), Float32.(dy), cdims; alpha=Float32(α), beta=Float32(β))), α*dw + β*w, rtol = 1.0e-7) 
+                    @test isapprox(ddims(∇conv_filter!(Float32.(copy(w)), Float32.(x), Float32.(dy), cdims; alpha=Float32(α), beta=Float32(β))), α*dw + β*w, rtol = 1.0e-7)
                     @test isapprox(ddims(∇conv_data!(Float32.(copy(x)), Float32.(dy),  Float32.(w),  cdims; alpha=Float32(α), beta=Float32(β))), α*dx + β*x, rtol = 1.0e-7)
 
                     # Next, introduce stride:
@@ -432,19 +432,19 @@ ddims(x) = dropdims(x, dims=(ndims(x)-1, ndims(x)))
                     cdims = DenseConvDims(x, w; dilation=2)
                     dy = NNlib.conv(x, w, cdims)
                     flag_ = ∇conv_data! == NNlib.∇conv_data_direct! && rank == 3
-                    @test isapprox(ddims(∇conv_filter!(copy(w), x, dy, cdims; alpha=α, beta=β)), α*dw_dil + β*w, rtol = 1.0e-7) 
+                    @test isapprox(ddims(∇conv_filter!(copy(w), x, dy, cdims; alpha=α, beta=β)), α*dw_dil + β*w, rtol = 1.0e-7)
                     @test isapprox(ddims(∇conv_data!(copy(x), dy, w,   cdims; alpha=α, beta=β)), α*dx_dil + β*x, rtol = 1.0e-7) broken=flag_
 
                     # Next, introduce padding:
                     cdims = DenseConvDims(x, w; padding=1)
                     dy = NNlib.conv(x, w, cdims)
-                    @test isapprox(ddims(∇conv_filter!(copy(w), x, dy, cdims; alpha=α, beta=β)), α*dw_pad + β*w, rtol = 1.0e-7) 
+                    @test isapprox(ddims(∇conv_filter!(copy(w), x, dy, cdims; alpha=α, beta=β)), α*dw_pad + β*w, rtol = 1.0e-7)
                     @test isapprox(ddims(∇conv_data!(copy(x), dy, w,   cdims; alpha=α, beta=β)), α*dx_pad + β*x, rtol = 1.0e-7)
 
                     # Next, test crosscor/conv with a flipped kernel
                     cdims = DenseConvDims(x, w; flipkernel=true)
                     dy = NNlib.conv(x, w, cdims)
-                    @test isapprox(ddims(∇conv_filter!(copy(w), x, dy, cdims; alpha=α, beta=β)), α*dw_flip + β*w, rtol = 1.0e-7) 
+                    @test isapprox(ddims(∇conv_filter!(copy(w), x, dy, cdims; alpha=α, beta=β)), α*dw_flip + β*w, rtol = 1.0e-7)
                     @test isapprox(ddims(∇conv_data!(copy(x), dy, w,   cdims; alpha=α, beta=β)), α*dx_flip + β*x, rtol = 1.0e-7)
                 end
             end
@@ -856,22 +856,22 @@ end
     struct MyFloat <: Real
         set::Set{Float32}
     end
-    
+
     # Test that direct indexing fails when undefined.
     v = Array{MyFloat}(undef, 3)
     @test_throws UndefRefError v[1]
-    
+
     # Define minimal set of functions required for conv_direct!
     MyFloat(x::MyFloat) = x
     MyFloat(x::Real) = MyFloat(Set(Float32(x)))
-    
+
     Base.:+(x::MyFloat, y::MyFloat) = MyFloat(only(x.set) + only(y.set))
     Base.:*(x::MyFloat, y::MyFloat) = MyFloat(only(x.set) * only(y.set))
     Base.promote_rule(::Type{MyFloat}, ::Type{Float32})   = MyFloat
     Base.rand(::AbstractRNG, ::SamplerType{MyFloat}) = MyFloat(rand(Float32))
     Base.zero(::MyFloat) = MyFloat(zero(Float32))
     Base.zero(::Type{MyFloat}) = MyFloat(zero(Float32))
-    
+
     # Test conv_direct!
     x_size = (6, 7, 8, 5, 3)
     y_size = (5, 6, 7, 4, 3)
@@ -913,7 +913,7 @@ end
 @testset "EnzymeRules: conv! spatial_rank=$spatial_rank" for spatial_rank in (1, 2, 3)
   x = rand(rng, repeat([5], spatial_rank)..., 3, 2)
   w = rand(rng, repeat([3], spatial_rank)..., 3, 3)
-  
+
   cdims = DenseConvDims(x, w)
 
   curconv = conv
@@ -935,7 +935,7 @@ end
   x = rand(rng, repeat([5], spatial_rank)..., 3, 2)
   w = rand(rng, repeat([3], spatial_rank)..., 3, 3)
   y = conv(x, w, cdims)
-  
+
   cdims = DenseConvDims(x, w)
 
   curconv = ∇conv_data
@@ -957,7 +957,7 @@ end
   x = rand(rng, repeat([5], spatial_rank)..., 3, 2)
   w = rand(rng, repeat([3], spatial_rank)..., 3, 3)
   y = conv(x, w, cdims)
-  
+
   cdims = DenseConvDims(x, w)
 
   curconv = ∇conv_filter
@@ -978,7 +978,7 @@ end
 @testset "EnzymeRules: depthwiseconv! spatial_rank=$spatial_rank" for spatial_rank in (1, 2, 3)
   x = rand(rng, repeat([5], spatial_rank)..., 3, 2)
   w = rand(rng, repeat([3], spatial_rank)..., 3, 3)
-  
+
   cdims = DepthwiseConvDims(x, w)
 
   curconv = depthwiseconv
diff --git a/test/runtests.jl b/test/runtests.jl
index 86329999..b8080b6b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -81,6 +81,15 @@ end
 
             nnlib_testsuite(CPU)
 
+            if Threads.nthreads(:default) > 1
+                @test NNlib.should_use_spawn()
+                NNlib.@disallow_spawns begin
+                    @test NNlib.should_use_spawn() == false
+                end
+            else
+                @test NNlib.should_use_spawn() == false
+            end
+
             @testset "Activation Functions" begin
                 include("activations.jl")
                 include("bias_act.jl")


Reviews:
user: login: IanButterworth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T14:24:11Z
user: login: ToucheSir
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-12T19:30:28Z
user: login: IanButterworth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T19:34:47Z
user: login: ToucheSir
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-12T19:59:19Z
user: login: IanButterworth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T20:01:49Z
user: login: ToucheSir
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-17T03:41:47Z
user: login: IanButterworth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T15:21:05Z
user: login: IanButterworth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T15:25:59Z
user: login: ToucheSir
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-17T16:04:03Z
user: login: IanButterworth
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T16:58:40Z
user: login: ToucheSir
body: No, I just hadn't the time to make sure CI was all good and to do another read-through.

Can you bump the version, and then we'll merge + release?
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-18T22:02:50Z

Review comments:
diff_hunk: @@ -288,11 +304,19 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (wc, xc, yc) in zip(dw_cs, x_cs, dy_cs)
+            function do_work(wc, xc, yc)
                 x = @view in1[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
                 dy = @view in2[ntuple(i -> i == 4 ? yc : Colon(), 5)...]
-                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...]
-                Threads.@spawn $(Symbol("$(front_name)_$(backend)!"))(dw, x, dy, cdims2; kwargs...)
+                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...] # TODO: Is this supposed to use wc?
path: src/conv.jl
user: login: IanButterworth
body: I noticed that `yc` probably should be `wc` here?
created_at: 2025-03-12T14:24:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -288,11 +304,19 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (wc, xc, yc) in zip(dw_cs, x_cs, dy_cs)
+            function do_work(wc, xc, yc)
                 x = @view in1[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
                 dy = @view in2[ntuple(i -> i == 4 ? yc : Colon(), 5)...]
-                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...]
-                Threads.@spawn $(Symbol("$(front_name)_$(backend)!"))(dw, x, dy, cdims2; kwargs...)
+                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...] # TODO: Is this supposed to use wc?
path: src/conv.jl
user: login: ToucheSir
body: I agree. Just to be sure, can you test that `length(dy_cs) >= length(dw_cs)` holds? Maybe via a temporary test assert. I'm pretty sure it does, but the conv code is convoluted enough that it might not.
created_at: 2025-03-12T19:30:28Z
author_association: MEMBER
diff_hunk: @@ -288,11 +304,19 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (wc, xc, yc) in zip(dw_cs, x_cs, dy_cs)
+            function do_work(wc, xc, yc)
                 x = @view in1[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
                 dy = @view in2[ntuple(i -> i == 4 ? yc : Colon(), 5)...]
-                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...]
-                Threads.@spawn $(Symbol("$(front_name)_$(backend)!"))(dw, x, dy, cdims2; kwargs...)
+                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...] # TODO: Is this supposed to use wc?
path: src/conv.jl
user: login: IanButterworth
body: Added
created_at: 2025-03-12T19:34:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -288,11 +304,19 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (wc, xc, yc) in zip(dw_cs, x_cs, dy_cs)
+            function do_work(wc, xc, yc)
                 x = @view in1[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
                 dy = @view in2[ntuple(i -> i == 4 ? yc : Colon(), 5)...]
-                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...]
-                Threads.@spawn $(Symbol("$(front_name)_$(backend)!"))(dw, x, dy, cdims2; kwargs...)
+                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...] # TODO: Is this supposed to use wc?
path: src/conv.jl
user: login: ToucheSir
body: Sorry, I meant checking the length of the partitioned iterators before the loop. i.e. `length(dy_cs) >= length(dw_cs) || error("Inconsistent dimensions") # TODO: remove this check`.
created_at: 2025-03-12T19:59:18Z
author_association: MEMBER
diff_hunk: @@ -288,11 +304,19 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (wc, xc, yc) in zip(dw_cs, x_cs, dy_cs)
+            function do_work(wc, xc, yc)
                 x = @view in1[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
                 dy = @view in2[ntuple(i -> i == 4 ? yc : Colon(), 5)...]
-                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...]
-                Threads.@spawn $(Symbol("$(front_name)_$(backend)!"))(dw, x, dy, cdims2; kwargs...)
+                dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...] # TODO: Is this supposed to use wc?
path: src/conv.jl
user: login: IanButterworth
body: Ah yes. Fixed
created_at: 2025-03-12T20:01:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -13,11 +13,27 @@ using LinearAlgebra
 using LinearAlgebra.BLAS: @blasfunc, BlasInt
 using LinearAlgebra: AdjOrTransAbsMat, Adjoint, BlasFloat, Transpose
 using Random
+using ScopedValues
 using Statistics
 using Statistics: mean
 
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 
+const ALLOW_SPAWNS = ScopedValue(true)
+should_use_spawn() = Threads.nthreads(:default) > 1 && ALLOW_SPAWNS[]
+"""
+    @disallow_spawns ex
+
+Disallow NNlib to use `@spawn` on divisible workloads. i.e. within `conv` etc.
+"""
+macro disallow_spawns(ex)
path: src/NNlib.jl
user: login: ToucheSir
body: Would it make sense to have this as a function instead? At first I thought the macro would be better because it doesn't introduce new scope, but on second glance the callback to `with` is.

My worry would be that users run into closure boxing problems because they're not aware of this scope. At least with a function `disallow_spawns() do ... end`, that closure piece is explicit?
created_at: 2025-03-17T03:27:59Z
author_association: MEMBER
diff_hunk: @@ -13,11 +13,27 @@ using LinearAlgebra
 using LinearAlgebra.BLAS: @blasfunc, BlasInt
 using LinearAlgebra: AdjOrTransAbsMat, Adjoint, BlasFloat, Transpose
 using Random
+using ScopedValues
 using Statistics
 using Statistics: mean
 
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 
+const ALLOW_SPAWNS = ScopedValue(true)
path: src/NNlib.jl
user: login: ToucheSir
body: Another thought: would it make sense to have this be an integer param that controls the max degree of spawn parallelism NNlib can use? Setting that param <= 1 (or to `false`) would disable spawning as before.
created_at: 2025-03-17T03:30:40Z
author_association: MEMBER
diff_hunk: @@ -202,11 +202,21 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (xc, wc) in zip(x_cs, w_cs)
+            function do_work(xc, wc)
path: src/conv.jl
user: login: ToucheSir
body: ```suggestion
            function conv_group(xc, wc)
```
For ease of debugging, I think it would help to give each `do_work` function a unique name. Ideally a more informative one too.
created_at: 2025-03-17T03:32:26Z
author_association: MEMBER
diff_hunk: @@ -246,11 +256,21 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (xc, yc, wc) in zip(dx_cs, dy_cs, w_cs)
+            function do_work(xc, yc, wc)
path: src/conv.jl
user: login: ToucheSir
body: ```suggestion
            function ∇conv_data_group(xc, yc, wc)
```
Ditto here.
created_at: 2025-03-17T03:33:10Z
author_association: MEMBER
diff_hunk: @@ -288,11 +308,21 @@ for (front_name, backend, signature) in (
                                 C_in = channels_in(cdims) ÷ groupcount(cdims),
                                 C_out = channels_out(cdims) ÷ groupcount(cdims))
 
-            Threads.@sync for (wc, xc, yc) in zip(dw_cs, x_cs, dy_cs)
+            function do_work(wc, xc, yc)
path: src/conv.jl
user: login: ToucheSir
body: ```suggestion
            function ∇conv_filter_group(wc, xc, yc)
```
created_at: 2025-03-17T03:34:31Z
author_association: MEMBER
diff_hunk: @@ -104,22 +104,34 @@ for (gemm, elt) in gemm_datatype_mappings
 
                 old_threads = get_num_threads()
                 set_num_threads(1)
-                Threads.@sync for ks in Iterators.partition(1:size(C, 3), cld(size(C, 3), n_threads))
-                    Threads.@spawn for k in ks
+
+                parts = Iterators.partition(1:size(C, 3), cld(size(C, 3), n_threads))
+
+                function do_work(ks)
path: src/gemm.jl
user: login: ToucheSir
body: ```suggestion
                function gemm!_part(ks)
```
created_at: 2025-03-17T03:38:00Z
author_association: MEMBER
diff_hunk: @@ -47,20 +47,28 @@ function conv_im2col!(
 
     parts = Iterators.partition(axes(x, 5), ceil(Int, size(x, 5) / ntasks))
 
-    @sync for (task_n, part) in enumerate(parts)
-        Threads.@spawn begin
-            col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
-            for batch_idx in part
-                im2col!(col_slice, view(x, :, :, :, :, batch_idx), cdims)
-                GC.@preserve col_slice w y begin
-                    col_ptr = pointer(col_slice)
-                    w_ptr = pointer(w)
-                    y_ptr = pointer(y, (batch_idx - 1)*M*N + 1)
-                    gemm!(Val(false), Val(false), M, N, K, alpha, col_ptr, w_ptr, beta, y_ptr)
-                end
+    function do_work(task_n, part)
path: src/impl/conv_im2col.jl
user: login: ToucheSir
body: ```suggestion
    function conv_part(task_n, part)
```
created_at: 2025-03-17T03:39:14Z
author_association: MEMBER
diff_hunk: @@ -152,18 +160,25 @@ function ∇conv_data_im2col!(
 
     parts = Iterators.partition(axes(dx, 5), ceil(Int, size(dx, 5) / ntasks))
 
-    @sync for (task_n, part) in enumerate(parts)
-        Threads.@spawn begin
-            col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
-            for batch_idx in part
-                GC.@preserve col_slice w dy begin
-                    dy_ptr = pointer(dy, (batch_idx - 1)*M*K + 1)
-                    w_ptr = pointer(w)
-                    col_ptr = pointer(col_slice)
-                    gemm!(Val(false), Val(true), M, N, K, alpha, dy_ptr, w_ptr, T(0), col_ptr)
-                end
-                col2im!(view(dx, :, :, :, :, batch_idx), col_slice, cdims, beta)
+    function do_work(task_n, part)
path: src/impl/conv_im2col.jl
user: login: ToucheSir
body: ```suggestion
    function ∇conv_data_part(task_n, part)
```
created_at: 2025-03-17T03:39:43Z
author_association: MEMBER
diff_hunk: @@ -30,25 +30,32 @@ function depthwiseconv_im2col!(
 
     dcdims = DenseConvDims(cdims)
 
-    @sync for (task_n, part) in enumerate(parts)
-        Threads.@spawn begin
-            col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
-            for batch_idx in part
-                im2col!(col_slice, view(x, :, :, :, :, batch_idx), dcdims)
-
-                # We do a separate convolution for each channel in x, as we must
-                for c_in in 1:channels_in(cdims)
-                    # Walk each pointer forward as we process each input channel
-                    GC.@preserve col_slice w y begin
-                        col_ptr = pointer(col_slice, (c_in-1)*M*K+1)
-                        w_ptr = pointer(w, (c_in-1)*K*N+1)
-                        y_ptr = pointer(y, ((batch_idx - 1)*channels_in(cdims) + c_in - 1)*M*N + 1)
-                        gemm!(Val(false), Val(false), M, N, K, alpha, col_ptr, w_ptr, beta, y_ptr)
-                    end
+    function do_work(task_n, part)
path: src/impl/depthwiseconv_im2col.jl
user: login: ToucheSir
body: ```suggestion
    function depthwiseconv_part(task_n, part)
```
created_at: 2025-03-17T03:40:51Z
author_association: MEMBER
diff_hunk: @@ -117,22 +124,29 @@ function ∇depthwiseconv_data_im2col!(
 
     parts = Iterators.partition(axes(dx)[end], ceil(Int, size(dx, 5) / ntasks))
 
-    @sync for (task_n, part) in enumerate(parts)
-        Threads.@spawn begin
-            col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
-            for batch_idx in part
-                # We do a separate convolution for each channel in x, as we must
-                for cidx in 1:channels_in(cdims)
-                    GC.@preserve col_slice w dy begin
-                        # Walk each pointer forward as we process each input channel
-                        dy_ptr = pointer(dy, (batch_idx - 1)*M*K*channels_in(cdims)+(cidx - 1)*K*M + 1)
-                        w_ptr = pointer(w, (cidx - 1)*K*N + 1)
-                        col_ptr = pointer(col_slice, (cidx - 1)*M*N + 1)
-                        gemm!(Val(false), Val(true), M, N, K, alpha, dy_ptr, w_ptr, T(0), col_ptr)
-                    end
+    function do_work(task_n, part)
path: src/impl/depthwiseconv_im2col.jl
user: login: ToucheSir
body: ```suggestion
    function ∇depthwiseconv_data_part(task_n, part)
```
created_at: 2025-03-17T03:41:17Z
author_association: MEMBER
diff_hunk: @@ -13,11 +13,27 @@ using LinearAlgebra
 using LinearAlgebra.BLAS: @blasfunc, BlasInt
 using LinearAlgebra: AdjOrTransAbsMat, Adjoint, BlasFloat, Transpose
 using Random
+using ScopedValues
 using Statistics
 using Statistics: mean
 
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 
+const ALLOW_SPAWNS = ScopedValue(true)
path: src/NNlib.jl
user: login: IanButterworth
body: That could make sense. I'll comment that this is internal, so it can be changed when that design is figured out. 
created_at: 2025-03-17T15:21:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -13,11 +13,27 @@ using LinearAlgebra
 using LinearAlgebra.BLAS: @blasfunc, BlasInt
 using LinearAlgebra: AdjOrTransAbsMat, Adjoint, BlasFloat, Transpose
 using Random
+using ScopedValues
 using Statistics
 using Statistics: mean
 
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 
+const ALLOW_SPAWNS = ScopedValue(true)
+should_use_spawn() = Threads.nthreads(:default) > 1 && ALLOW_SPAWNS[]
+"""
+    @disallow_spawns ex
+
+Disallow NNlib to use `@spawn` on divisible workloads. i.e. within `conv` etc.
+"""
+macro disallow_spawns(ex)
path: src/NNlib.jl
user: login: IanButterworth
body: Good point. I moved to using `@with`
created_at: 2025-03-17T15:25:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -13,11 +13,27 @@ using LinearAlgebra
 using LinearAlgebra.BLAS: @blasfunc, BlasInt
 using LinearAlgebra: AdjOrTransAbsMat, Adjoint, BlasFloat, Transpose
 using Random
+using ScopedValues
 using Statistics
 using Statistics: mean
 
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 
+const ALLOW_SPAWNS = ScopedValue(true)
path: src/NNlib.jl
user: login: ToucheSir
body: Yes, I should've elaborated on this thought. The idea would be to have a way to control the degree of parallelism via a context manager/macro like `disallow_spawns`, but using a numeric param instead of an on/off flag.

But this design can be fleshed out later. Then there are more QoL things (not for this PR) we can consider around threading that would complement it. For example, deciding the number of tasks to spawn based on # of threads / # of BLAS threads.
created_at: 2025-03-17T16:04:03Z
author_association: MEMBER
diff_hunk: @@ -13,11 +13,27 @@ using LinearAlgebra
 using LinearAlgebra.BLAS: @blasfunc, BlasInt
 using LinearAlgebra: AdjOrTransAbsMat, Adjoint, BlasFloat, Transpose
 using Random
+using ScopedValues
 using Statistics
 using Statistics: mean
 
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 
+const ALLOW_SPAWNS = ScopedValue(true)
path: src/NNlib.jl
user: login: IanButterworth
body: Sounds good. Tbh the who julia threading story needs more work on the composability of threading like this. It's not defined how best to orchestrate it from the `Base.Threads` API. The ecosystem might be better, but adoption isn't strong.
created_at: 2025-03-17T16:58:40Z
author_association: CONTRIBUTOR

PR comments:
user: login: IanButterworth
created_at: 2025-03-12T21:24:00Z
body: I'm also thinking it would be good to be able to opt out of threading.
author_association: CONTRIBUTOR
user: login: ToucheSir
created_at: 2025-03-12T23:58:37Z
body: An opt-out would be good, but IMO new ones should be done with `ScopedValue`s instead of global toggles. Is https://github.com/vchuravy/ScopedValues.jl a good enough compat package for those?
author_association: MEMBER
user: login: IanButterworth
created_at: 2025-03-13T00:01:53Z
body: Sounds good. ~I think that can be a follow-on PR~ I did it here
author_association: CONTRIBUTOR
user: login: IanButterworth
created_at: 2025-03-13T15:27:30Z
body: Just trying something out. Will switch `ALLOW_THREADING[]` to a ScopedValue next.
author_association: CONTRIBUTOR
user: login: IanButterworth
created_at: 2025-03-14T02:39:40Z
body: Just noting that I tried simplifying this into a macro
```julia
macro maybe_thread(ex)
    ex.head == :for || throw(ArgumentError("Expected a for loop"))
    it, iter = ex.args[1].args
    loop_body = ex.args[2]
    quote
        if should_use_spawn() && length($(esc(iter))) > 1
            Threads.@sync for $(esc(it)) in $(esc(iter))
                Threads.@spawn $(esc(loop_body))
            end
        else
            $(esc(ex))
        end
    end
end
```
i.e. simply
```julia
@maybe_thread for (xc, wc) in zip(x_cs, w_cs)
    x = @view in1[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
    w = @view in2[ntuple(i -> i == 5 ? wc : Colon(), 5)...]
    y = @view out[ntuple(i -> i == 4 ? wc : Colon(), 5)...]
    $(Symbol("$(front_name)_$(backend)!"))(y, x, w, cdims2; kwargs...)
end
```
but hit some macro hygiene issue with `@sync`
```
ERROR: LoadError: syntax: invalid let syntax around task.jl:497
Stacktrace:
 [1] top-level scope
   @ none:1
 [2] eval(m::Module, e::Any)
   @ Core ./boot.jl:430
 [3] top-level scope
   @ ~/Documents/GitHub/NNlib.jl/src/conv.jl:183
 [4] include(mod::Module, _path::String)
   @ Base ./Base.jl:557
 [5] include(x::String)
   @ NNlib ~/Documents/GitHub/NNlib.jl/src/NNlib.jl:1
 [6] top-level scope
   @ ~/Documents/GitHub/NNlib.jl/src/NNlib.jl:80
```
author_association: CONTRIBUTOR
user: login: IanButterworth
created_at: 2025-03-17T19:13:40Z
body: Hard for me to tell if any of these failures are related?
author_association: CONTRIBUTOR
user: login: IanButterworth
created_at: 2025-03-18T18:35:13Z
body: bump @ToucheSir is there anything further to do here?
author_association: CONTRIBUTOR