simpler way to calculate maximum grain size

JuliaGPU · Mar 9, 2023 · 4af969b · 4af969b
1 parent c214237
commit 4af969b
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/src/mapreduce.jl b/src/mapreduce.jl
@@ -211,9 +211,7 @@ function GPUArrays.mapreducedim!(f::F, op::OP, R::WrappedMtlArray{T},
     #
     # also, make sure the grain size is not too high so as to starve threads of work.
     other_groups = length(Rother)
-    while grain > 1 && length(Rreduce) <= reduce_threads * prevpow(2, grain)
-        grain >>= 1
-    end
+    grain = min(grain, prevpow(2, cld(length(Rreduce), reduce_threads)))
     reduce_groups = cld(length(Rreduce), reduce_threads * grain)
 
     # determine the launch configuration