fix more problematic gpu code

Author: SimonDanisch

docs/code/TraceMakie.jl

@@ -192,6 +192,7 @@ begin
     mesh!(scene, catmesh, color=load(Makie.assetpath("diffusemap.png")))
     center!(scene)
     # 1.024328 seconds (16.94 M allocations: 5.108 GiB, 46.19% gc time, 81 lock conflicts)
+    # 0.913530 seconds (16.93 M allocations: 5.108 GiB, 42.52% gc time, 57 lock conflicts)
     @time render_scene(scene)
 end
 
@@ -206,5 +207,6 @@ begin
     surface!(scene, xs, ys, zs)
     center!(scene)
     #  1.598740s
+    # 1.179450 seconds (17.30 M allocations: 5.126 GiB, 36.48% gc time, 94 lock conflicts)
     @time render_scene(scene)
 end

docs/code/basic-scene.jl

@@ -75,7 +75,8 @@ begin
     img = reverse(film.framebuffer, dims=1)
 end
 # 6.296157 seconds (17.64 k allocations: 19.796 MiB, 0.13% gc time, 45 lock conflicts)
-
+# After more GPU optimizations
+# 4.169616 seconds (17.37 k allocations: 19.777 MiB, 0.14% gc time, 20 lock conflicts)
 
 camera_sample = Trace.get_camera_sample(integrator.sampler, Point2f(512))
 ray, ω = Trace.generate_ray_differential(integrator.camera, camera_sample)

src/accel/bvh.jl

@@ -267,7 +267,7 @@ end
             else
                 if dir_is_neg[ln.split_axis] == Int32(2)
                     nodes_to_visit[to_visit_offset] = current_node_i + Int32(1)
-                    current_node_i = Int32(ln.offset)
+                    current_node_i = ln.offset % Int32
                 else
                     nodes_to_visit[to_visit_offset] = ln.offset % Int32
                     current_node_i += Int32(1)
@@ -293,7 +293,7 @@ end
 
     to_visit_offset, current_node_i = Int32(1), Int32(1)
     nodes_to_visit = zeros(MVector{64,Int32})
-    while true
+    @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)

src/integrators/sampler.jl

@@ -73,7 +73,7 @@ function (i::SamplerIntegrator)(scene::Scene, film)
 end
 
 function get_material(bvh::BVHAccel, shape::Triangle)
-    if shape.material_idx == 0
+    @inbounds if shape.material_idx == 0
         return bvh.materials[1]
     else
         return bvh.materials[shape.material_idx]

src/textures/basic.jl

@@ -26,12 +26,13 @@ function Base.convert(::Type{Texture{ElType,N,T}}, ::NoTexture) where {ElType,N,
 end
 
 function (c::Texture{T})(si::SurfaceInteraction)::T where {T<:TextureType}
-    if c.isconst
+    @inbounds if c.isconst
         return c.const_value
     else
-        uv = Vec2f(1.0-si.uv[2], si.uv[1])
-        idx = round.(Int, 1 .+ ((size(c.data) .- 1) .* uv))
-        idx = clamp.(idx, 1, size(c.data))
+        uv = Vec2f(1f0 - si.uv[2], si.uv[1])
+        s = unsafe_trunc.(Int32, size(c.data))
+        idx = map(x -> unsafe_trunc(Int32, x), Int32(1) .+ ((s .- Int32(1)) .* uv))
+        idx = clamp.(idx, Int32(1), s)
         return c.data[idx...]
     end
 end

test/gpu-threading-benchmarks.jl

@@ -100,7 +100,7 @@ end
     if ω > 0.0f0
         hit, shape, si = Trace.intersect!(bvh, ray)
         if hit
-            l = Trace.RGBSpectrum(si.core.n...)#simple_shading(bvh, shape, ray, si, l, 1, 8, lights)
+            l = simple_shading(bvh, shape, ray, si, l, 1, 8, lights)
         end
     end
     return RGBf(l.c...)
@@ -118,23 +118,37 @@ import KernelAbstractions as KA
     end
 end
 
+function launch_trace_image_ir!(img, camera, bvh, lights)
+    backend = KA.get_backend(img)
+    kernel! = ka_trace_image!(backend)
+    open("test2.ir", "w") do io
+        @device_code_llvm io begin
+            kernel!(img, camera, bvh, lights, ndrange = size(img), workgroupsize = (16, 16))
+        end
+    end
+    AMDGPU.synchronize(; stop_hostcalls=false)
+    return img
+end
 function launch_trace_image!(img, camera, bvh, lights)
     backend = KA.get_backend(img)
     kernel! = ka_trace_image!(backend)
-    kernel!(img, camera, bvh, lights, ndrange = size(img), workgroupsize = (16, 16))
+    kernel!(img, camera, bvh, lights, ndrange=size(img), workgroupsize=(16, 16))
     KA.synchronize(backend)
     return img
 end
-
-using CUDA
-ArrayType = CuArray
+using AMDGPU
+ArrayType = ROCArray
+# using CUDA
+# ArrayType = CuArray
 preserve = []
 gpu_bvh = to_gpu(ArrayType, bvh; preserve=preserve);
 gpu_img = ArrayType(zeros(RGBf, res, res));
 # launch_trace_image!(img, cam, bvh, lights);
 # @btime launch_trace_image!(img, cam, bvh, lights);
 # @btime launch_trace_image!(gpu_img, cam, gpu_bvh, lights);
-@btime launch_trace_image!(gpu_img, cam, gpu_bvh, lights);
+launch_trace_image!(gpu_img, cam, gpu_bvh, lights);
+@btime launch_trace_image!(img, cam, bvh, lights)
+# 76.420 ms (234 allocations: 86.05 KiB)
 Array(gpu_img)
 
 function cu_trace_image!(img, camera, bvh, lights)