Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Switch to KernelAbstractions.jl #559

Merged
merged 2 commits into from
Oct 17, 2024
Merged

Switch to KernelAbstractions.jl #559

merged 2 commits into from
Oct 17, 2024

Conversation

maleadt
Copy link
Member

@maleadt maleadt commented Sep 19, 2024

Continuation of #525. Separate PR because I reverted the launch_configuration removal that's on @leios' branch.

@maleadt maleadt changed the title Tb/kernelabstractions Switch to KernelAbstractions.jl Sep 19, 2024
@maleadt maleadt force-pushed the tb/kernelabstractions branch from 2a541d0 to 184b36f Compare September 20, 2024 10:46
@maleadt
Copy link
Member Author

maleadt commented Sep 20, 2024

One of the consequences here is that the backend function disappears, and is now part of KernelAbstractions.jl, so cc'ing people who were involved in #420 and JuliaGPU/Adapt.jl#50: @N5N3 @piever.

NeuralAttentionlib.jl also relies on GPUArraysCore.backend, so cc @chengchingwen.

I also see that Bennu.jl uses GPUArrays.launch_configuration directly, so cc @lcw.

Comment on lines -224 to -225
# WrappedArray from Adapt for Base wrappers.
backend(::Type{WA}) where WA<:WrappedArray = backend(unwrap_type(WA))
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@vchuravy Does KA.jl already support recursing into wrapped arrays for get_backend queries?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Uhm I don't think so, but we can add that.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not entirely sure we want to; it would pull the whole Union mess of Adapt's WrappedArray into KA.jl. On the other hand, there's not much of an alternative right now if we want the ability to launch kernels on wrapped arrays...

@lcw
Copy link
Contributor

lcw commented Sep 20, 2024

Thanks for the ping. I have archived Bennu.jl as I have switched to developing Raven.jl. In Raven.jl I do use GPUArraysCore.backend here but I also already use KernelAbstractions.jl. So, I think it will be a simple transition.

KernelAbstractions.isgpu(b::JLBackend) = false

function convert_to_cpu(obj::Kernel{JLBackend, W, N, F}) where {W, N, F}
return Kernel{typeof(KernelAbstractions.CPU(; static = obj.backend.static)), W, N, F}(KernelAbstractions.CPU(; static = obj.backend.static), obj.f)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is clever xD

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you explain? I didn't get what this was for, and it seems unused?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unless I did something wrong, it's used for kernel configuration:

function (obj::Kernel{JLBackend})(args...; ndrange=nothing, workgroupsize=nothing)
    device_args = jlconvert.(args)
    new_obj = convert_to_cpu(obj)
    new_obj(device_args...; ndrange, workgroupsize)
end

And will essentially transform anything that has a JLBackend into a CPU(static) (default false) for KA execution so we can use the GPU kernels on CPU Arrays. It's needed because JLBackend is within KernelAbstractions.GPU, but needs to actually call CPU kernels:

struct JLBackend <: KernelAbstractions.GPU
    static::Bool
    JLBackend(;static::Bool=false) = new(static)
end

@maleadt maleadt force-pushed the tb/kernelabstractions branch 4 times, most recently from 6eaa774 to 2f8f813 Compare October 17, 2024 13:08
@maleadt maleadt force-pushed the tb/kernelabstractions branch 2 times, most recently from d8fb6d3 to 7348bba Compare October 17, 2024 13:58
@maleadt maleadt force-pushed the tb/kernelabstractions branch from 7348bba to f418d7a Compare October 17, 2024 14:01
@maleadt
Copy link
Member Author

maleadt commented Oct 17, 2024

Let's do this!

@maleadt maleadt merged commit 2769d6b into master Oct 17, 2024
12 checks passed
@maleadt maleadt deleted the tb/kernelabstractions branch October 17, 2024 14:29
@leios
Copy link
Contributor

leios commented Oct 17, 2024

Woo! Great work!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants