more fixes

Author: SimonDanisch

src/accel/bvh.jl

@@ -251,7 +251,8 @@ end
             if !ln.is_interior && ln.n_primitives > Int32(0)
                 # Intersect ray with primitives in node.
                 for i in Int32(0):ln.n_primitives - Int32(1)
-                    tmp_primitive = primitives[ln.offset+i]
+                    offset = ln.offset % Int32
+                    tmp_primitive = primitives[offset+i]
                     tmp_hit, ray, tmp_interaction = intersect_p!(
                         tmp_primitive, ray,
                     )
@@ -293,13 +294,15 @@ end
 
     to_visit_offset, current_node_i = Int32(1), Int32(1)
     nodes_to_visit = zeros(MVector{64,Int32})
+    primitives = bvh.primitives
     @inbounds while true
         ln = bvh.nodes[current_node_i]
         if intersect_p(ln.bounds, ray, inv_dir, dir_is_neg)
             if !ln.is_interior && ln.n_primitives > Int32(0)
                 for i in Int32(0):ln.n_primitives-Int32(1)
+                    offset = ln.offset % Int32
                     intersect_p(
-                        bvh.primitives[ln.offset + i], ray,
+                        primitives[offset + i], ray,
                     ) && return true
                 end
                 to_visit_offset == 1 && break

src/integrators/sampler.jl

@@ -88,9 +88,13 @@ function li(
     l = RGBSpectrum(0f0)
     # Find closest ray intersection or return background radiance.
     hit, shape, si = intersect!(scene, ray)
+    lights = scene.lights
     if !hit
-        for light in scene.lights
-            l += le(light, ray)
+        Base.Cartesian.@nexprs 8 i -> begin
+            if i <= length(lights)
+                light = lights[i]
+                l += le(light, ray)
+            end
         end
         return l
     end
@@ -112,7 +116,6 @@ function li(
     # Compute emitted light if ray hit an area light source.
     l += le(si, wo)
     # Add contribution of each light source.
-    lights = scene.lights
     Base.Cartesian.@nexprs 8 i -> begin
         if i <= length(lights)
             light = lights[i]
@@ -129,8 +132,8 @@ function li(
     end
     if depth + 1 ≤ max_depth
         # Trace rays for specular reflection & refraction.
-        # l += specular_reflect(bsdf, sampler, max_depth, ray, si, scene, depth)
-        # l += specular_transmit(bsdf, sampler, max_depth, ray, si, scene, depth)
+        l += specular_reflect(bsdf, sampler, max_depth, ray, si, scene, depth)
+        l += specular_transmit(bsdf, sampler, max_depth, ray, si, scene, depth)
     end
     l
 end
@@ -152,7 +155,7 @@ end
     if !(pdf > 0f0 && !is_black(f) && abs(wi ⋅ ns) != 0f0)
         return RGBSpectrum(0f0)
     end
-    # Compute ray differential for specular reflection.
+    # # Compute ray differential for specular reflection.
     rd = RayDifferentials(spawn_ray(si, wi))
     if ray.has_differentials
         rx_origin = si.core.p + si.∂p∂x
@@ -215,7 +218,7 @@ end
         # intersected. Compute the relative IOR by first out by assuming
         # that the ray is entering the object.
         η = 1f0 / bsdf.η
-        if (ns ⋅ ns) < 0
+        if (ns ⋅ ns) < 0f0
             # If the ray isn't entering the object, then we need to invert
             # the relative IOR and negate the normal and its derivatives.
             η = 1f0 / η

src/materials/bsdf.jl

@@ -121,6 +121,8 @@ function (b::BSDF)(
     return output
 end
 
+u_int32(x) = Base.unsafe_trunc(Int32, x)
+
 """
 Compute incident ray direction for a given outgoing direction and
 a given mode of light scattering corresponding
@@ -132,16 +134,16 @@ function sample_f(
 
     # Choose which BxDF to sample.
     matching_components = num_components(b, type)
-    matching_components == 0 && return (
+    matching_components == Int32(0) && return (
         Vec3f(0f0), RGBSpectrum(0f0), 0f0, BSDF_NONE,
     )
     component = min(
-        max(1, ceil(Int64, u[1] * matching_components)),
+        max(Int32(1), u_int32(ceil(u[1] * matching_components))),
         matching_components,
     )
     # Get BxDF for chosen component.
     count = component
-    component -= 1
+    component -= Int32(1)
     bxdf = UberBxDF{RGBSpectrum}()
     bxdfs = b.bxdfs
     Base.Cartesian.@nexprs 8 i -> begin
@@ -177,24 +179,24 @@ function sample_f(
     )
     wi_world = local_to_world(b, wi)
     # Compute overall PDF with all matching BxDFs.
-    if !(bxdf.type & BSDF_SPECULAR != 0) && matching_components > 1
+    if !(bxdf.type & BSDF_SPECULAR != Int32(0)) && matching_components > Int32(1)
         Base.Cartesian.@nexprs 8 i -> begin
             if i <= bxdfs.last && bxdfs[i] != bxdf && bxdfs[i] & type
                 pdf += compute_pdf(bxdfs[i], wo, wi)
             end
         end
     end
-    matching_components > 1 && (pdf /= matching_components)
+    matching_components > Int32(1) && (pdf /= matching_components)
     # Compute value of BSDF for sampled direction.
-    if !(bxdf.type & BSDF_SPECULAR != 0)
-        reflect = ((wi_world ⋅ b.ng) * (wo_world ⋅ b.ng)) > 0
+    if !(bxdf.type & BSDF_SPECULAR != Int32(0))
+        reflect = ((wi_world ⋅ b.ng) * (wo_world ⋅ b.ng)) > 0f0
         f = RGBSpectrum(0f0)
         Base.Cartesian.@nexprs 8 i -> begin
             if i <= bxdfs.last
                 bxdf = bxdfs[i]
                 if  ((bxdf & type) && (
-                        (reflect && (bxdf.type & BSDF_REFLECTION != 0)) ||
-                        (!reflect && (bxdf.type & BSDF_TRANSMISSION != 0))
+                        (reflect && (bxdf.type & BSDF_REFLECTION != Int32(0))) ||
+                        (!reflect && (bxdf.type & BSDF_TRANSMISSION != Int32(0)))
                     ))
                     f += bxdf(wo, wi)
                 end
@@ -226,11 +228,11 @@ function compute_pdf(
 end
 
 @inline function num_components(b::BSDF, flags::UInt8)::Int64
-    num = 0
+    num = Int32(0)
     bxdfs = b.bxdfs
     Base.Cartesian.@nexprs 8 i -> begin
         if i <= bxdfs.last && (bxdfs[i] & flags)
-            num += 1
+            num += Int32(1)
         end
     end
     return num

src/sampler/sampler.jl

@@ -25,36 +25,7 @@ function Sampler(samples_per_pixel::Integer)
     )
 end
 
-using Random
-using RandomNumbers.Xorshifts
-
-
-const TRNG = Xoroshiro128Plus[]
-# Reset the per-thread random seeds to make results reproducible
-reseed!() = foreach(i-> Random.seed!(TRNG[i], i), 1:Threads.maxthreadid())
-
-function __init__()
-    # Instantiate 1 RNG (Random Number Generator) per thread, for performance.
-    # This can't be done during precompilation since the number of threads isn't known then.
-    resize!(TRNG, Threads.maxthreadid())
-    for i in 1:Threads.nthreads()
-        TRNG[i] = Xoroshiro128Plus(i)
-    end
-    nothing
-end
-
-"Per-thread rand()"
-@inline function trand()
-    @inbounds rng = TRNG[Threads.threadid()]
-    rand(rng)
-end
-
-@inline function trand(::Type{T}) where {T}
-    @inbounds rng = TRNG[Threads.threadid()]
-    rand(rng, T)
-end
-
-function get_camera_sample(sampler::AbstractSampler, p_raster::Point2f)
+@inline function get_camera_sample(sampler::AbstractSampler, p_raster::Point2f)
     p_film = p_raster .+ get_2d(sampler)
     time = get_1d(sampler)
     p_lens = get_2d(sampler)
@@ -161,11 +132,13 @@ struct UniformSampler <: AbstractSampler
     UniformSampler(samples_per_pixel::Integer) = new(1, samples_per_pixel)
 end
 
-function get_camera_sample(::UniformSampler, p_raster::Point2f)
-    @inbounds rng = TRNG[Threads.threadid()]
-    p_film = p_raster .+ rand(rng, Point2f)
-    p_lens = rand(rng, Point2f)
-    CameraSample(p_film, p_lens, rand(rng, Float32))
+@inline rand2f() = Point2f(0.5f0, 0.5f0)
+
+@inline function get_camera_sample(::UniformSampler, p_raster::Point2f)
+    p = rand2f()
+    p_film = Point2f(p_raster[1] + p[1], p_raster[2] + p[2])
+    p_lens = rand2f()
+    CameraSample(p_film, p_lens, 0f0)
 end
 
 @inline function has_next_sample(u::UniformSampler)::Bool
@@ -178,6 +151,6 @@ end
     u.current_sample = 1
 end
 @inline get_1d(u::UniformSampler)::Float32 = rand(Float32)
-@inline get_2d(u::UniformSampler)::Point2f = rand(Point2f)
+@inline get_2d(u::UniformSampler)::Point2f = rand2f()
 
 # include("stratified.jl")

test/cuda-mwe.jl

@@ -51,7 +51,6 @@ begin
 end
 
 @inline function get_camera_sample(p_raster::Point2)
-
     p_film = p_raster .+ rand(Point2f)
     p_lens = rand(Point2f)
     Trace.CameraSample(p_film, p_lens, rand(Float32))

test/gpu-sampler.jl

@@ -5,10 +5,10 @@ using KernelAbstractions
 import KernelAbstractions as KA
 using KernelAbstractions.Extras.LoopInfo: @unroll
 
-# using AMDGPU
-# ArrayType = ROCArray
-using CUDA
-ArrayType = CuArray
+using AMDGPU
+ArrayType = ROCArray
+# using CUDA
+# ArrayType = CuArray
 
 include("./../src/gpu-support.jl")
 
@@ -57,15 +57,10 @@ begin
         Trace.PointLight(Vec3f(0, 0, 2), Trace.RGBSpectrum(10.0f0)),
         Trace.PointLight(Vec3f(0, 3, 3), Trace.RGBSpectrum(25.0f0)),
     )
+    scene = Trace.Scene(lights, bvh)
     img = zeros(RGBf, res, res)
 end
 
-@inline function get_camera_sample(p_raster::Point2)
-    p_film = p_raster .+ rand(Point2f)
-    p_lens = rand(Point2f)
-    Trace.CameraSample(p_film, p_lens, rand(Float32))
-end
-
 # ray = Trace.Ray(o=Point3f(0.5, 0.5, 1.0), d=Vec3f(0.0, 0.0, -1.0))
 # l = Trace.RGBSpectrum(0.0f0)
 # open("test3.llvm", "w") do io
@@ -74,12 +69,30 @@ end
 
 @inline function trace_pixel(camera, scene, xy)
     pixel = Point2f(Tuple(xy))
-    camera_sample = get_camera_sample(pixel)
+    s = Trace.UniformSampler(8)
+    camera_sample = @inline Trace.get_camera_sample(s, pixel)
     ray, ω = Trace.generate_ray_differential(camera, camera_sample)
     if ω > 0.0f0
-        hit, shape, si = Trace.intersect!(scene, ray)
+        l = @inline Trace.li(s, 5, ray, scene, 1)
+    end
+    return l
+end
+
+@inline function trace_pixel(camera, scene, xy)
+    pixel = Point2f(reverse(Tuple(xy)))
+    s = Trace.UniformSampler(8)
+    camera_sample = @inline Trace.get_camera_sample(s, pixel)
+    ray, ω = Trace.generate_ray_differential(camera, camera_sample)
+    l = Trace.RGBSpectrum(0.0f0)
+    if ω > 0.0f0
+        hit, prim, si = Trace.intersect!(scene, ray)
         if hit
-            l = Trace.li(Trace.UniformSampler(8), 5, ray, scene, 1)
+            m = Trace.get_material(scene, prim)
+            bsdf = m(si, false, Trace.Radiance)
+            l = Trace.specular_reflect(
+                bsdf, s, 8, ray,
+                si, scene, 8,
+            )
         end
     end
     return l
@@ -91,6 +104,7 @@ end
         l = trace_pixel(camera, scene, xy)
         @inbounds img[xy] = RGBf(l.c...)
     end
+    nothing
 end
 
 function launch_trace_image!(img, camera, scene)
@@ -108,4 +122,44 @@ gpu_img = ArrayType(zeros(RGBf, res, res));
 # @btime launch_trace_image!(img, cam, bvh, lights);
 # @btime launch_trace_image!(gpu_img, cam, gpu_bvh, lights);
 launch_trace_image!(gpu_img, cam, gpu_scene);
+Array(gpu_img)
+
 launch_trace_image!(img, cam, scene)
+GLMakie.activate!(float=true)
+image(img)
+
+ray = Trace.RayDifferentials(Trace.Ray(o=Point3f(0.5, 0.5, 1.0), d=Vec3f(0.0, 0.0, -1.0)))
+open("le-wt.jl", "w") do io
+    code_warntype(io, Trace.li, typeof.((Trace.UniformSampler(8), 5, ray, scene, 1)))
+end
+
+
+function launch_trace_image_ir!(img, camera, scene)
+    backend = KA.get_backend(img)
+    kernel! = ka_trace_image!(backend)
+    open("test.ir", "w") do io
+        try
+            @device_code_llvm io begin
+                kernel!(img, camera, scene, ndrange=size(img), workgroupsize=(16, 16))
+            end
+        catch e
+            println(e)
+        end
+    end
+    KA.synchronize(backend)
+    return img
+end
+launch_trace_image_ir!(gpu_img, cam, gpu_scene);
+
+code_llvm(Trace.intersect!, (typeof(bvh), Trace.RayDifferentials))
+
+
+function trace_image!(img, camera, scene)
+    for xy in CartesianIndices(size(img))
+        @inbounds img[xy] = RGBf(trace_pixel(camera, scene, xy).c...)
+    end
+    return img
+end
+
+
+@time launch_trace_image!(img, cam, scene)

test/gpu-threading-benchmarks.jl

@@ -74,13 +74,11 @@ end
 
 @inline function trace_pixel(camera, scene, xy)
     pixel = Point2f(Tuple(xy))
-    camera_sample = get_camera_sample(pixel)
+    s = Trace.UniformSampler(8)
+    camera_sample = @inline Trace.get_camera_sample(s, pixel)
     ray, ω = Trace.generate_ray_differential(camera, camera_sample)
     if ω > 0.0f0
-        hit, shape, si = Trace.intersect!(scene, ray)
-        if hit
-            l = Trace.li(Trace.UniformSampler(8), 5, ray, scene, 1)
-        end
+        l = @inline Trace.li(s, 5, ray, scene, 1)
     end
     return l
 end
@@ -147,9 +145,9 @@ Array(gpu_img)
 # workgroupsize=(16,16)
 # 31.022 ms (35 allocations: 5.89 KiB)
 
-function trace_image!(img, camera, bvh)
+function trace_image!(img, camera, scene)
     for xy in CartesianIndices(size(img))
-        @inbounds img[xy] = trace_pixel(camera, bvh, xy)
+        @inbounds img[xy] = RGBf(trace_pixel(camera, scene, xy).c...)
     end
     return img
 end