Update ReactantCUDAExt.jl

EnzymeAD · Feb 1, 2025 · 8adc3de · 8adc3de
1 parent a320d27
commit 8adc3de
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl
@@ -263,6 +263,15 @@ function Adapt.adapt_structure(
     )
 end
 
+function threads_to_workgroupsize(threads, ndrange)
+    total = 1
+    return map(ndrange) do n
+        x = min(div(threads, total), n)
+        total *= x
+        return x
+    end
+end
+
 function (obj::KA.Kernel{ReactantBackend})(args...; ndrange=nothing, workgroupsize=nothing)
     backend = KA.backend(obj)