You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
functionrun_julia(height, width)
y =range(-1.0f0, 0.0f0; length = height) # need Float32 because Jax defaults to it
x =range(-1.5f0, 0.0f0; length = width)
c = x'.+ y*im
fractal =fill(Int32(20), height, width)
# this checks if indicies are compatible between `c` and `fractal`@inboundsfor idx ineachindex(c, fractal)
_c = c[idx]
z = _c
m =true
Base.Cartesian.@nexprs20 i ->begin
z = z^2+ _c
az4 =abs2(z) >4f0
fractal[idx] =ifelse(m&az4, Int32(i), fractal[idx]) # 32-bit Int, same reason as above
m &= (!az4)
endendreturn fractal
endusing KernelAbstractions
@kernelfunctionjulia_kernel!(c, fractal)
I =@index(Global)
_c = c[I]
z = _c
@inboundsfor i =1:20
z = z^2+ _c
ifabs2(z) >4f0
fractal[I] =Int32(i)
breakendendendfunctionrun_julia_cpu_jaxstype(height, width)
y =range(-1.0f0, 0.0f0; length = height)
x =range(-1.5f0, 0.0f0; length = width)
c = x'.+ y*im
fractal =fill(Int32(20), height, width)
kernel! =julia_kernel!(CPU(), length(c)÷Threads.nthreads()) # we're using 1-thread here
event =kernel!(c, fractal; ndrange=length(c))
wait(event) # not copying back, need to block herereturn fractal
end
we have
julia> Threads.nthreads()
1
julia>@benchmarkrun_julia(2000,3000)
BenchmarkTools.Trial:100 samples with 1 evaluation.
Range (min … max):49.380 ms …52.606 ms ┊ GC (min … max):0.37%…1.04%
Time (median):49.789 ms ┊ GC (median):0.98%
Time (mean ± σ):49.982 ms ±583.818 μs ┊ GC (mean ± σ):0.98%±0.07%
▅█▇▄ ▃
▃▁▅▇█████▃▇██▅▇▅▁▆▆▃▆▃▁▆▃▅▁▃▁▁▁▁▁▃▁▁▁▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▃▁▁▆ ▃
49.4 ms Histogram: frequency by time 52 ms <
Memory estimate:68.66 MiB, allocs estimate:4.
julia>@benchmarkrun_julia_cpu_jaxstype(2000,3000)
BenchmarkTools.Trial:29 samples with 1 evaluation.
Range (min … max):176.674 ms …179.932 ms ┊ GC (min … max):0.27%…0.08%
Time (median):177.228 ms ┊ GC (median):0.26%
Time (mean ± σ):177.528 ms ±780.443 μs ┊ GC (mean ± σ):0.26%±0.03%
▃▃ ▃▃ ▃ █
██▇██▇▁█▇▇▇▁▇▁▁▁▁▇▁▁▇▇▁▁▇▇█▇▁▇▇▁▁▁▁▁▁▁▁▁▁▁▁▁▇▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▇ ▁
177 ms Histogram: frequency by time 180 ms <
Memory estimate:68.67 MiB, allocs estimate:28.
The text was updated successfully, but these errors were encountered:
consider this example from a Jax discussion:
Source code
we have
The text was updated successfully, but these errors were encountered: