fix nthreads on 1.9 (#496)

FluxML · Jun 14, 2023 · e5cff84 · e5cff84 · CarloLucibello · Jun 14, 2023
1 parent fbef0d9
commit e5cff84
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/dim_helpers/ConvDims.jl b/src/dim_helpers/ConvDims.jl
@@ -73,7 +73,7 @@ function im2col_dims(c::ConvDims)
         # Size of single dotproduct within convolution
         prod(kernel_size(c))*channels_in(c),
         # One workspace per thread
-        Threads.nthreads(),
+        VERSION > v"1.9.0-0" ? Threads.maxthreadid() : Threads.nthreads(),
     )
 end
 

diff --git a/src/gemm.jl b/src/gemm.jl
@@ -88,7 +88,9 @@ for (gemm, elt) in gemm_datatype_mappings
             strB = size(B, 3) == 1 ? 0 : Base.stride(B, 3)
             strC = Base.stride(C, 3)
 
-            n_threads = min(Threads.nthreads(), 1 + max(length(A), length(B)) ÷ 8000)
+            n_threads = min(
+                VERSION > v"1.9.0-0" ? Threads.maxthreadid() : Threads.nthreads(),
+                1 + max(length(A), length(B)) ÷ 8000)
             # In some tests, size (20,20,20) is worth splitting between two threads,
             # as is size (32,32,8).