Broadcast over CuArray error

Specific Domains GPU
1

Hi, I would like to define the following parametric function type

struct myParamFun
       fun::Function
       param
end
(pf::myParamFun)(x...)=pf.fun(x...,pf.param)

which works fine. However, when broadcasting over CuArray I get the error:

julia> pf=myParamFun((x,y,z,v)->x+y-z*v,2);

julia> xs=(rand(100) for _ in 1:3);

julia> using CUDA

julia> pf.(xs...);

julia> pf.(CuArray.(xs)...);
ERROR: GPU broadcast resulted in non-concrete element type Any.
This probably means that the function you are broadcasting contains an error or type instability.

However, the following works without any errors

julia> pf.fun.(CuArray.(xs)...,Ref(pf.param));

The @code_warntype output of the calls are as follows:

julia> @code_warntype pf.(CuArray.(xs)...);
MethodInstance for (::var"##dotfunction#1253#40")(::Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}})
  from (::var"##dotfunction#1253#40")(x1) in Main
Arguments
  #self#::Core.Const(var"##dotfunction#1253#40"())
  x1::Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}
Body::Any
1 ─ %1 = Core.tuple(Main.pf)::Tuple{Any}
│   %2 = Core._apply_iterate(Base.iterate, Base.broadcasted, %1, x1)::Any
│   %3 = Base.materialize(%2)::Any
└──      return %3


julia> @code_warntype pf.(xs...);
MethodInstance for (::var"##dotfunction#1254#41")(::Base.Generator{UnitRange{Int64}, var"#38#39"})
  from (::var"##dotfunction#1254#41")(x1) in Main
Arguments
  #self#::Core.Const(var"##dotfunction#1254#41"())
  x1::Base.Generator{UnitRange{Int64}, var"#38#39"}
Body::Any
1 ─ %1 = Core.tuple(Main.pf)::Tuple{Any}
│   %2 = Core._apply_iterate(Base.iterate, Base.broadcasted, %1, x1)::Any
│   %3 = Base.materialize(%2)::Any
└──      return %3

julia> @code_warntype pf.fun.(CuArray.(xs)...,Ref(pf.param));
MethodInstance for (::var"##dotfunction#1255#42")(::Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, ::Base.RefValue{Int64})
  from (::var"##dotfunction#1255#42")(x1, x2) in Main
Arguments
  #self#::Core.Const(var"##dotfunction#1255#42"())
  x1::Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}
  x2::Base.RefValue{Int64}
Body::Any
1 ─ %1 = Base.getproperty(Main.pf, :fun)::Any
│   %2 = Core.tuple(%1)::Tuple{Any}
│   %3 = Core.tuple(x2)::Tuple{Base.RefValue{Int64}}
│   %4 = Core._apply_iterate(Base.iterate, Base.broadcasted, %2, x1, %3)::Any
│   %5 = Base.materialize(%4)::Any
└──      return %5

So is there a way to implement this parametric function type correctly for CuArray broadcasting? Thank you in advance!

Update-- even if I define my parameteric function pf as a const, the broadcast still fails.

2

Maybe this?

struct myParamFun{F,T}
       fun::F
       param::T
end

You’ll need @code_warntype pf.(xs...) to not give Any.

3 Likes
3

Thanks, it seems to work, even though @code_warntype still gives Any

julia> struct myParamFun{F,P}
       fun::F
       param::P
       end

julia> (pf::myParamFun)(x...) =pf.fun(pf.param,x...)

julia> pf=myParamFun((v,x,y,z)->x+y-z*v,2);

julia> typeof(pf)
myParamFun{var"#7#8", Int64}

julia> xs=Tuple(rand(100) for _ in 1:3);

julia> @code_warntype pf.(xs...);
MethodInstance for (::var"##dotfunction#407#11")(::Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}})
  from (::var"##dotfunction#407#11")(x1) in Main
Arguments
  #self#::Core.Const(var"##dotfunction#407#11"())
  x1::Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}}
Body::Any
1 ─ %1 = Core.tuple(Main.pf)::Tuple{Any}
│   %2 = Core._apply_iterate(Base.iterate, Base.broadcasted, %1, x1)::Any
│   %3 = Base.materialize(%2)::Any
└──      return %3

julia> pf.(xs...)
100-element Vector{Float64}:
 -0.7525209350984445
...
julia> pf.(CuArray.(xs)...)
100-element CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}:
 -0.7525209350984445
 -1.2313547408936314
...
4

That’s just an artifact of it being an untyped, non-const global variable. If you capture it in a local scope with a (new) local identifier, then it’s type-stable:

julia> let pf=pf; @code_warntype pf.(xs...) end
MethodInstance for (::var"##dotfunction#293#8"{myParamFun{var"#3#4", Int64}})(::Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}})
  from (::var"##dotfunction#293#8")(x1) @ Main none:0
Arguments
  #self#::var"##dotfunction#293#8"{myParamFun{var"#3#4", Int64}}
  x1::Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}}
Body::Vector{Float64}
1 ─ %1 = Core.getfield(#self#, :pf)::myParamFun{var"#3#4", Int64}
│   %2 = Core.tuple(%1)::Tuple{myParamFun{var"#3#4", Int64}}
│   %3 = Core._apply_iterate(Base.iterate, Base.broadcasted, %2, x1)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, myParamFun{var"#3#4", Int64}, Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}}}
│   %4 = Base.materialize(%3)::Vector{Float64}
└──      return %4
1 Like