From 63c40a75cb6e655f985289a1fbb67bc53be9fb69 Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Sat, 1 Feb 2025 08:19:25 -0600 Subject: [PATCH 1/3] KA without cuda backend --- ext/ReactantCUDAExt.jl | 2 +- ext/ReactantKernelAbstractionsExt.jl | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl index 6256fa7bbd..bf7a4599e6 100644 --- 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) +Reactant.@reactant_overlay @noinline function (obj::KA.Kernel{ReactantBackend})(args...; ndrange=nothing, workgroupsize=nothing) backend = KA.backend(obj) ndrange, workgroupsize, iterspace, dynamic = KA.launch_config( diff --git a/ext/ReactantKernelAbstractionsExt.jl b/ext/ReactantKernelAbstractionsExt.jl index e225cbfae2..c9070cb1b2 100644 --- a/ext/ReactantKernelAbstractionsExt.jl +++ b/ext/ReactantKernelAbstractionsExt.jl @@ -83,4 +83,9 @@ function KA.priority!(::ReactantBackend, prio::Symbol) return nothing end +function (obj::KA.Kernel{ReactantBackend})(args...; ndrange=nothing, workgroupsize=nothing) + obj2 = KA.Kernel{KA.CPU, KA.workgroupsize(obj), KA.ndrange(obj), typeof(obj.f)}(KA.CPU(), obj.f) + obj2(args...; ndrange, workgroupsize) +end + end From f6d6c1e62f46543c50b530ccc7a82610dd4c6ff3 Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Sat, 1 Feb 2025 08:35:40 -0600 Subject: [PATCH 2/3] fix --- ext/ReactantCUDAExt.jl | 6 +++++- ext/ReactantKernelAbstractionsExt.jl | 7 +++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl index bf7a4599e6..5490cbb2f8 100644 --- a/ext/ReactantCUDAExt.jl +++ b/ext/ReactantCUDAExt.jl @@ -272,7 +272,7 @@ function threads_to_workgroupsize(threads, ndrange) end end -Reactant.@reactant_overlay @noinline 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,10 @@ Reactant.@reactant_overlay @noinline function (obj::KA.Kernel{ReactantBackend})( return nothing end +Reactant.@reactant_overlay @noinline function (obj::KA.Kernel{ReactantBackend})(args...; ndrange=nothing, workgroupsize=nothing) + 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 index c9070cb1b2..b3e4bc499c 100644 --- a/ext/ReactantKernelAbstractionsExt.jl +++ b/ext/ReactantKernelAbstractionsExt.jl @@ -83,9 +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) - obj2 = KA.Kernel{KA.CPU, KA.workgroupsize(obj), KA.ndrange(obj), typeof(obj.f)}(KA.CPU(), obj.f) - obj2(args...; ndrange, workgroupsize) + @jit tokw(ndrange, workgroupsize, obj, args...) end end From 3b30e05d1b3f8f5c684c3df50c8a551460eeb6dd Mon Sep 17 00:00:00 2001 From: William Moses Date: Sat, 1 Feb 2025 08:57:05 -0600 Subject: [PATCH 3/3] Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- ext/ReactantCUDAExt.jl | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl index 5490cbb2f8..391be0bd9e 100644 --- a/ext/ReactantCUDAExt.jl +++ b/ext/ReactantCUDAExt.jl @@ -325,8 +325,12 @@ function ka_with_reactant(ndrange, workgroupsize, obj, args...) return nothing end -Reactant.@reactant_overlay @noinline function (obj::KA.Kernel{ReactantBackend})(args...; ndrange=nothing, workgroupsize=nothing) - Reactant.call_with_reactant(ka_with_reactant, ndrange, workgroupsize, obj, args...) +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)