Fix the launch heuristic. (#40)

JuliaGPU · Jul 10, 2022 · ca5cb76 · ca5cb76
1 parent ef72706
commit ca5cb76
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/src/gpuarrays.jl b/src/gpuarrays.jl
@@ -14,11 +14,8 @@ struct mtlKernelContext <: AbstractKernelContext end
     kernel = @metal launch=false f(mtlKernelContext(), args...)
 
     # The pipeline state automatically computes occupancy stats
-    threads_needed = cld(elements, elements_per_thread)
-
-    # Limit the threadgroup size
-    threads = min(threads_needed, kernel.pipeline_state.maxTotalThreadsPerThreadgroup)
-    blocks  = cld(threads_needed, threads)
+    threads = min(elements, kernel.pipeline_state.maxTotalThreadsPerThreadgroup)
+    blocks  = cld(elements, threads)
 
     return (; threads, blocks)
 end