KA without cuda backend (#670)

* KA without cuda backend * fix * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
EnzymeAD · Feb 1, 2025 · c2786dd · c2786dd
1 parent e9471bd
commit c2786dd
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl
@@ -272,7 +272,7 @@ function threads_to_workgroupsize(threads, ndrange)
     end
 end
 
-function (obj::KA.Kernel{ReactantBackend})(args...; ndrange=nothing, workgroupsize=nothing)
+function ka_with_reactant(ndrange, workgroupsize, obj, args...)
     backend = KA.backend(obj)
 
     ndrange, workgroupsize, iterspace, dynamic = KA.launch_config(
@@ -325,6 +325,14 @@ function (obj::KA.Kernel{ReactantBackend})(args...; ndrange=nothing, workgroupsi
     return nothing
 end
 
+Reactant.@reactant_overlay @noinline function (obj::KA.Kernel{ReactantBackend})(
+    args...; ndrange=nothing, workgroupsize=nothing
+)
+    return Reactant.call_with_reactant(
+        ka_with_reactant, ndrange, workgroupsize, obj, args...
+    )
+end
+
 Adapt.adapt_storage(to::KA.ConstAdaptor, a::CuTracedArray) = Base.Experimental.Const(a)
 
 function recudaconvert(arg)

diff --git a/ext/ReactantKernelAbstractionsExt.jl b/ext/ReactantKernelAbstractionsExt.jl
@@ -83,4 +83,12 @@ function KA.priority!(::ReactantBackend, prio::Symbol)
     return nothing
 end
 
+function tokw(ndrange, workgroupsize, obj, args...)
+    @inline obj(args...; ndrange, workgroupsize)
+end
+
+function (obj::KA.Kernel{ReactantBackend})(args...; ndrange=nothing, workgroupsize=nothing)
+    @jit tokw(ndrange, workgroupsize, obj, args...)
+end
+
 end