Skip to content

Allow StructArray arguments in CUDAnative kernels. #87

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.

Sign up for GitHub

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

Closed
lcw wants to merge 3 commits into from
Closed

Allow StructArray arguments in CUDAnative kernels. #87

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
efb75eb
Avoid dynamic code in get_ith
lcw Sep 18, 2019
1bc3dbe
Make foreachfield a static function
lcw Sep 23, 2019
08b0a12
Use Adapt for converting CPU to GPU StructArrays
lcw Sep 26, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/StructArrays.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,10 @@ include("sort.jl")
include("groupjoin.jl")
include("lazy.jl")

# Use Adapt allows for automatic conversion of CPU to GPU StructArrays
import Adapt
Adapt.adapt_storage(to, s::StructArray) = replace_storage(x->Adapt.adapt(to, x), s)

function __init__()
Requires.@require Tables="bd369af6-aec1-5ad0-b16a-f7cc5008161c" include("tables.jl")
Requires.@require WeakRefStrings="ea10d353-3f73-51f8-a26c-33c1cb351aa5" begin
Expand Down
8 changes: 4 additions & 4 deletions src/structarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,11 +134,11 @@ Base.axes(s::StructArray) = axes(fieldarrays(s)[1])
Base.axes(s::StructArray{<:Any, <:Any, <:EmptyTup}) = (1:0,)

get_ith(cols::NamedTuple, I...) = get_ith(Tuple(cols), I...)
function get_ith(cols::NTuple{N, Any}, I...) where N
ntuple(N) do i
Copy link

@vchuravy vchuravy Sep 23, 2019

Choose a reason for hiding this comment

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

Might need to be

ntuple(Val(N))` do i
   Base.@_inline_meta
  @inbounds res = getfield(cols, i)[I...]
 end

But the generated function works as well.

Copy link
Contributor Author

@lcw lcw Sep 23, 2019 •
edited
Loading

Choose a reason for hiding this comment

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

I tried

function get_ith(cols::NTuple{N, Any}, I...) where N
    ntuple(Val(N)) do i
        Base.@_inline_meta
        @inbounds res = getfield(cols, i)[I...]
    end
end

but still got a dynamic function

ERROR: LoadError: InvalidIRError: compiling kernel!(StructArray{SHermitianCompact{2,Float64,3},1,NamedTuple{(:lowertriangle,),Tuple{StructArray{SArray{Tuple{3},Float64,1,3},1,NamedTuple{(:data,),Tuple{StructArray{Tuple{Float64,Float64,Float64},1,Tuple{CuDeviceArray{Float64,1,CUDAnative.AS.Global},CuDeviceArray{Float64,1,CUDAnative.AS.Global},CuDeviceArray{Float64,1,CUDAnative.AS.Global}},Int64}}},Int64}}},Int64}, StructArray{SHermitianCompact{2,Float64,3},1,NamedTuple{(:lowertriangle,),Tuple{StructArray{SArray{Tuple{3},Float64,1,3},1,NamedTuple{(:data,),Tuple{StructArray{Tuple{Float64,Float64,Float64},1,Tuple{CuDeviceArray{Float64,1,CUDAnative.AS.Global},CuDeviceArray{Float64,1,CUDAnative.AS.Global},CuDeviceArray{Float64,1,CUDAnative.AS.Global}},Int64}}},Int64}}},Int64}) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to ntuple(f, ::Val{1}) in Base at ntuple.jl:41)
Stacktrace:
 [1] get_ith at /home/lucas/julia/dev/StructArrays/src/structarray.jl:138
 [2] get_ith at /home/lucas/julia/dev/StructArrays/src/structarray.jl:136
 [3] getindex at /home/lucas/julia/dev/StructArrays/src/structarray.jl:153
 [4] kernel! at /home/lucas/research/code/Heptapus.jl/examples/structarrays/try.jl:14
Reason: unsupported call to the Julia runtime (call to jl_f_getfield)
Stacktrace:
 [1] getindex at /home/lucas/julia/dev/StructArrays/src/structarray.jl:153
 [2] kernel! at /home/lucas/research/code/Heptapus.jl/examples/structarrays/try.jl:14
Reason: unsupported dynamic function invocation (call to foreachfield)
Stacktrace:
 [1] getindex at /home/lucas/julia/dev/StructArrays/src/structarray.jl:153
 [2] kernel! at /home/lucas/research/code/Heptapus.jl/examples/structarrays/try.jl:14
Reason: unsupported dynamic function invocation (call to foreachfield)
Stacktrace:
 [1] setindex! at /home/lucas/julia/dev/StructArrays/src/structarray.jl:168
 [2] kernel! at /home/lucas/research/code/Heptapus.jl/examples/structarrays/try.jl:14

Copy link
Collaborator

@piever piever Sep 26, 2019 •
edited
Loading

Choose a reason for hiding this comment

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

Generally, I would prefer a normal function over a generated function. What is the cause of "dynamic function invocation"? Is it due to the I... splatting? Never mind, splatting seems fine.

Copy link
Contributor Author

@lcw lcw Sep 26, 2019

Choose a reason for hiding this comment

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

I agree it would be nice to avoid generated functions. I am not sure why this has to be generated but there is definitely something wrong with what I did for foreachfield as the test suite doesn't pass. Maybe that is causing issues. If you get a chance, could you take a peak at what I did and see if you can spot something astray?

@inbounds res = getfield(cols, i)[I...]
return res
@generated function get_ith(cols::NTuple{N, Any}, I...) where N
args = ntuple(N) do i
:(@inbounds res = getfield(cols, $i)[I...])
end
:(($(args...),))
end

Base.@propagate_inbounds function Base.getindex(x::StructArray{T, <:Any, <:Any, CartesianIndex{N}}, I::Vararg{Int, N}) where {T, N}
Expand Down
44 changes: 37 additions & 7 deletions src/utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,20 +28,50 @@ else
const _getproperty = getproperty
end

function _foreachfield(names, xs)
function _sstuple(::Type{<:NTuple{N, Any}}) where {N}
ntuple(j->Symbol(j), N)
end

function _sstuple(::Type{NT}) where {NT<:NamedTuple}
_map_params(x->_sstuple(staticschema(x)), NT)
end

function _getcolproperties!(exprs, s, es=[])
if typeof(s) <: Symbol
push!(exprs, es)
return
end
for key in keys(s)
_getcolproperties!(exprs, getproperty(s,key), vcat(es, key))
end
end

@generated function foreachfield(::Type{T}, f, xs...) where {T<:Tup}
Copy link
Collaborator

@piever piever Sep 29, 2019

Choose a reason for hiding this comment

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

This introduces a lot of code complexity, while the original implementation is just a few lines of code: I'm not completely sure I understand conceptually what this refactor is trying to do. Is there an intuitive explanation as to way the original foreachfield is not suitable for CUDAnative while this implementation works?

# TODO get columnsproperties directly from T without converting to the
# tuple s.
s = _sstuple(T)
columnsproperties = []
_getcolproperties!(columnsproperties, s)

exprs = Expr[]
for field in names
sym = QuoteNode(field)
args = [Expr(:call, :_getproperty, :(getfield(xs, $j)), sym) for j in 1:length(xs)]
for col in columnsproperties
args = Expr[]
for prop in col
sym = QuoteNode(prop)
if length(args) == 0
args = [Expr(:call, :_getproperty, :(getfield(xs, $j)), sym) for j in 1:length(xs)]
else
for j in 1:length(xs)
args[j] = Expr(:call, :_getproperty, args[j], sym)
end
end
end
push!(exprs, Expr(:call, :f, args...))
end
push!(exprs, :(return nothing))
return Expr(:block, exprs...)
end

@generated foreachfield(::Type{<:NamedTuple{names}}, f, xs...) where {names} = _foreachfield(names, xs)
@generated foreachfield(::Type{<:NTuple{N, Any}}, f, xs...) where {N} = _foreachfield(Base.OneTo(N), xs)

foreachfield(f, x::T, xs...) where {T} = foreachfield(staticschema(T), f, x, xs...)

"""
Expand Down