I have 2 question about Metal.jl and Flux.jl

[math4mad]

1 from flux document demo

using Metal,Flux

    function cu(par)
        MtlArray(par)
    end

  W = cu(rand(2, 5)) 
  b = cu(rand(2))

  predict(x) = W*x .+ b
  loss(x, y) = sum((predict(x) .- y).^2)

 x, y = cu(rand(5)), cu(rand(2)) 
 loss(x, y) 

ERROR: LoadError: InvalidIRError: compiling kernel #63#64(Metal.mtlKernelContext, MtlDeviceVector{Float64, 1}, MtlDeviceMatrix{Float64, 1}, MtlDeviceVector{Float64, 1}) resulted in invalid LLVM IR
Reason: unsupported use of double floating-point value
Reason: unsupported use of double floating-point value

Just replace function's name
MtlArray work fine! But loss function seems has problem , I don't know how to solve

2 there are built-in gpu function like CUDA.jl?

  julia> using Flux, CUDA

julia> m = Dense(10, 5) |> gpu
Dense(10 => 5)      # 55 parameters

julia> x = rand(10) |> gpu

it can be define model's layer one by one , but fmap not working.

like this

using Metal,Flux

    function cu(par)
        MtlArray(par)
    end
d = Dense(10 => 5, σ)
d = fmap(cu, d)
d.weight 
d(cu(rand(10))) 

m = Chain(Dense(10 => 5, σ), Dense(5 => 2), softmax)
m = fmap(cu, m)
d(cu(rand(10)))

[math4mad]

must be Float32 ?

 W = cu(rand(Float32,2, 5)) # a 2×5 CuArray
 b = cu(rand(Float32,2))

it works

[maleadt]

Metal GPUs do not support Float32, indeed.

[max-Hawkins]

Float64^

[maleadt]

Oh, yeah, thanks for the correction!

[maleadt]

The gpu function is defined by Flux, so Metal support would need to be added there. There is currently no cu-like function in Metal, but that should be easy to define using Adapt.jl.

[maleadt]

#114

[rgobbel]

Verified that this works, on an M2 MacBook Air, MacOS 13.1.1, Julia 1.9.0-rc2:

using Metal,Flux

W = mtl(rand(2, 5)) 
b = mtl(rand(2))

predict(x) = W*x .+ b
loss(x, y) = sum((predict(x) .- y).^2)

x, y = mtl(rand(5)), mtl(rand(2)) 
loss(x, y)

Output:

julia> using Metal,Flux

julia> W = mtl(rand(2, 5))
2×5 MtlMatrix{Float32}:
 0.437989   0.329312  0.0235297  0.0711579  0.237615
 0.0544768  0.26934   0.265755   0.822853   0.081311

julia> b = mtl(rand(2))
2-element MtlVector{Float32}:
 0.37443423
 0.7773641

julia> predict(x) = W*x .+ b
predict (generic function with 1 method)

julia> loss(x, y) = sum((predict(x) .- y).^2)
loss (generic function with 1 method)

julia> x, y = mtl(rand(5)), mtl(rand(2))
(Float32[0.79346555, 0.32377717, 0.78265744, 0.10197092, 0.9622683], Float32[0.30823693, 0.08668886])

julia> loss(x, y)
2.019197f0