Merge pull request #1990 from CliMA/ck/high_res

Limit z-threads in columnwise_partition
CliMA · Sep 16, 2024 · 4cda6fb · 4cda6fb
2 parents 4e7c44d + 2c7b404
commit 4cda6fb
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/ext/cuda/data_layouts_threadblock.jl b/ext/cuda/data_layouts_threadblock.jl
@@ -176,7 +176,11 @@ end
     n_max_threads::Integer,
 )
     (Nij, _, _, _, Nh) = DataLayouts.universal_size(us)
-    Nh_thread = min(Int(fld(n_max_threads, Nij * Nij)), Nh)
+    Nh_thread = min(
+        Int(fld(n_max_threads, Nij * Nij)),
+        maximum_allowable_threads()[3],
+        Nh,
+    )
     Nh_blocks = cld(Nh, Nh_thread)
     @assert prod((Nij, Nij, Nh_thread)) ≤ n_max_threads "threads,n_max_threads=($(prod((Nij, Nij, Nh_thread))),$n_max_threads)"
     return (; threads = (Nij, Nij, Nh_thread), blocks = (Nh_blocks,))