Tutorial fails using Metal.jl

[ctessum]

Hi,

I am trying to run this tutorial on my laptop, which has an M1 processor. My understanding is that to do this, I should just change CUDA to Metal:

using DiffEqGPU, DifferentialEquations, StaticArrays, Metal

function lorenz2(u, p, t)
    σ = p[1]
    ρ = p[2]
    β = p[3]
    du1 = σ * (u[2] - u[1])
    du2 = u[1] * (ρ - u[3]) - u[2]
    du3 = u[1] * u[2] - β * u[3]
    return SVector{3}(du1, du2, du3)
end

u0 = @SVector [1.0f0; 0.0f0; 0.0f0]
tspan = (0.0f0, 10.0f0)
p = @SVector [10.0f0, 28.0f0, 8 / 3.0f0]
prob = ODEProblem{false}(lorenz2, u0, tspan, p)
prob_func = (prob, i, repeat) -> remake(prob, p = (@SVector rand(Float32, 3)) .* p)
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy = false)
sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(Metal.MetalBackend()),
    trajectories = 10_000,
    saveat = 1.0f0)

However, when I run the code above, the last line gives the error:

ERROR: InvalidIRError: compiling MethodInstance for DiffEqGPU.gpu_ode_asolve_kernel(::KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicCheck, Nothing, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}, KernelAbstractions.NDIteration.NDRange{1, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}}}, ::MtlDeviceVector{DiffEqGPU.ImmutableODEProblem{SVector{3, Float32}, Tuple{Float32, Float32}, false, SVector{3, Float32}, ODEFunction{false, SciMLBase.AutoSpecialize, typeof(lorenz2), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, 1}, ::GPUTsit5, ::MtlDeviceMatrix{SVector{3, Float32}, 1}, ::MtlDeviceMatrix{Float32, 1}, ::Float32, ::CallbackSet{Tuple{}, Tuple{}}, ::Nothing, ::Float32, ::Float32, ::StepRangeLen{Float32, Float64, Float64, Int64}, ::Val{false}) resulted in invalid LLVM IR
Reason: unsupported use of double value
Reason: unsupported use of double value
Reason: unsupported use of double value

These are the package versions:

(esml_demo) pkg> status DiffEqGPU
  [071ae1c0] DiffEqGPU v3.3.0
(esml_demo) pkg> status Metal
  [dde4c033] Metal v0.5.1
(esml_demo) pkg> status DifferentialEquations
  [0c46a032] DifferentialEquations v7.11.0

Is this the expected behavior?

[ctessum]

More information in case relevant:

Metal.versioninfo()

macOS 14.0.0, Darwin 23.0.0

Toolchain:
- Julia: 1.9.0
- LLVM: 14.0.6

Julia packages: 
- Metal.jl: 0.5.1
- Metal_LLVM_Tools_jll: 0.5.1+0

1 device:
- Apple M1 (2.406 MiB allocated)

[utkarsh530]

The Apple M1 does not support Float64 values yet, which is causing some issues with type ::StepRangeLen{Float32, Float64, Float64, Int64} (it turns out some Float64 happens with your CPU's precision). If you remove saveat=1.0f0, it should work.

I am trying to fix it using #317. Thanks for bringing it up!