Julia Version 1.13.1
I wrote some code for a raytracer:
module EfficientRayTracer
using LinearAlgebra
using StaticArrays
using Images
using ImageView
using Infiltrator
using Profile
using InteractiveUtils
struct LightProperties
id::Float64 # diffuse
is::Float64 # specular
α::Float64
color::SVector{3, Float64}
r::Float64
end
struct Sphere
o::SVector{3, Float64} # location
r::Float64 # radius
p::LightProperties
Sphere(o, r, p) = new(SVector{3, Float64}(o), Float64(r), p)
end
struct Screen
xpix::Int
ypix::Int
xlims::SVector{2, Float64}
ylims::SVector{2, Float64}
function Screen(pix, xlims, ylims)
xlen = xlims[2] - xlims[1]
ylen = ylims[2] - ylims[1]
xpix = √(xlen / ylen * pix)
ypix = ylen / xlen * xpix
new(round(xpix), round(ypix), xlims, ylims)
end
end
struct LightSource
r::SVector{3, Float64} # location
# ambient, diffuse and specular intensities
ia::Float64
id::Float64
is::Float64
color::SVector{3, Float64}
end
function solve_quadratic(a, b, c)
D = b^2 - 4a*c
if D < 0
return nothing
end
sD = √D
x1 = (-b + sD) / 2a
x2 = (-b - sD) / 2a
return (x1, x2)
end
"""
Return minimum parameter value t for which intersection occurs or -Inf if not intersecting.
"""
function ray_intersect(r, d, sphere::Sphere)
a = sum(d .^ 2)
b = 2 * (r - sphere.o) ⋅ d
c = sum((r - sphere.o).^2) - sphere.r .^2
res = solve_quadratic(a, b, c)
if res !== nothing
t1, t2 = res
if t1 > 0 && t2 <= 0
return t1
elseif t2 > 0 && t1 <= 0
return t2
else
return min(t1, t2)
end
end
# no intersections
-Inf
end
@inline function outward_normal(r, sphere::Sphere)
return normalize(r .- sphere.o)
end
function closest_intersect(r, d, objects)
t_min = Inf
closest_object = nothing
for object in objects
t = ray_intersect(r, d, object)
if t > 0 && t < t_min
t_min = t
closest_object = object
end
end
return (t_min, closest_object)
end
function new_d_and_n(r, d, t_min, closest_object)
if t_min < Inf
r .+= t_min * d
n = outward_normal(r + t_min * d, closest_object)
d_refl = d - 2 * d ⋅n * n
return (d_refl, n)
end
end
"""
Show the defined scenery.
Optimized for minimum memory footprint and maximum speed by using SVectors.
"""
function show_scenery(zcam, screen::Screen, sources::Tuple{LightSource}, objects::Tuple)
max_refl = Inf
total_ambient = SA[0.0, 0, 0]
for source in sources
total_ambient += source.ia .* source.color
end
total_ambient::SVector{3, Float64} = min.(SA[1.0, 1, 1], total_ambient)
# swap x, y for correct image
# intensities = fill(total_ambient, screen.ypix, screen.xpix)
# for mysterious reasons, below code will cause the compiler
# to allocate in the hot loop when updating the intensities ...
# Also, adding a SINGLE @time in some places causes allocations
intensities = [total_ambient for i in 1:screen.ypix, j in 1:screen.xpix]
nrays = screen.xpix * screen.ypix
xrange = range(screen.xlims..., screen.xpix)
# swap y-axis for correct image
yrange = range(screen.ylims[2], screen.ylims[1], screen.ypix)
# CRITICAL: removed these unnecessary allocations
# rays_d = [normalize(SVector{3, Float64}(x, y, -zcam)) for x in xrange, y in yrange]
nmax_refl = 0
# hot loop is allocation-free
for idx in eachindex(intensities)
it = SVector{3, Float64}(0, 0, 0)
cart_idx = CartesianIndices(intensities)[idx]
d = normalize(SVector{3, Float64}(xrange[cart_idx[2]], yrange[cart_idx[1]], -zcam))
r = SVector{3, Float64}(0, 0, zcam)
refl_factor = 1
nrefl = 0
(t, obj) = closest_intersect(r, d, objects)
while obj !== nothing && nrefl < max_refl
nrefl += 1
r_new::SVector{3, Float64} = r + t * d
n_obj::SVector{3, Float64} = outward_normal(r_new, obj)
d_new = d - 2 * n_obj ⋅ d * n_obj
r_new = r_new + 1e-8 * d_new
for source in sources
d_src::SVector{3, Float64} = normalize(source.r - r_new)
if closest_intersect(r_new, d_src, objects)[2] === nothing
i_d = source.id * obj.p.id * d_src ⋅ n_obj
i_s = source.is * obj.p.is * max(d_src ⋅ d_new, 0) ^ obj.p.α
i = max(0, i_d) + max(0, i_s)
c = i * refl_factor * obj.p.color .* source.color
it += c
end
end
r = r_new
d = d_new
refl_factor *= obj.p.r
(t, obj) = closest_intersect(r, d, objects)
if nrefl == max_refl
nmax_refl += 1
end
end
intensities[idx] += it
end
# CRITICAL: reinterpret and not convert or list comprehension for 0 allocations
image = reinterpret(RGB{Float64}, intensities)
# imshow(image)
end
function update_value!(mat::Matrix{SVector{3, Float64}}, val::SVector{3, Float64}, idx::Int)
mat[idx] += val
end
function main()
screen = Screen(1000000, (-2.0, 2.0), (-1, 1.2))
light_props1 = LightProperties(.4, .6, 4, (0, 1, 0), .8)
light_props2 = LightProperties(.4, .6, 4, (1, 0, 0), .8)
light_props3 = LightProperties(.4, .6, 4, (1, 0, 1), .8)
light_props_plane = LightProperties(.4, .6, 4, (1, 1, 1), .5)
sphere1 = Sphere((-1, 0, 1), 1.0, light_props1)
sphere2 = Sphere((1, 0, .7), .3, light_props2)
sphere3 = Sphere((.8, .5, .2), .1, light_props3)
plane = Sphere((0, -10^4, 0), 10^4 - 1, light_props_plane)
source1 = LightSource((4.0, 1, 0.0), .1, 1, 1, (1, 1, 1))
source2 = LightSource((2, 2, -2), .1, 1, 1, (1, 1, 1))
# @code_warntype show_scenery(-2, screen, (source2,), (sphere1, sphere2, sphere3, plane))
show_scenery(-2, screen, (source2,), (sphere1, sphere2, sphere3, plane))
end
end
using .EfficientRayTracer
The version above has no allocations in the hot loop of the “show_scenery” function.
However, I noticed that by removing the type label of “total_ambient”, the compiler fails to see that “total_ambient” must be of type SVector{3, Float64}.
Curiously, this only happens whenever I use list comprehension allocation
intensities = [total_ambient for i in 1:screen.ypix, j in 1:screen.xpix]
If I instead use the “fill” allocation (which should be exactly identical), as
intensities = fill(total_ambient, screen.ypix, screen.xpix)
then the type of total_ambient is correctly inferred without the SVector annotation.
As a sidenode, the line causing the allocations when total_ambient is not correctly inferred is
intensities[idx] += it
which suggests Julia is converting it to an intermediate representation before adding SVectors directly (which is allocation-free in principle).
Let it be clear that type-annotating total_ambient solves the issue; however, I feel like the compiler’s type inference should have also been able to infer the SVector type without it (and the performance cost is massive). My main confusion here is that changing the list apprehension to a fill allocation solves the issue completely even without the type annotation. Please enlighten me in how the type inference works, and why it could be influenced indirectly like so.