Merge branch 'sd/more-improvements' into la/pow_issue

Author: SimonDanisch

docs/code/basic-scene.jl

@@ -12,31 +12,32 @@ function tmesh(prim, material)
     return Trace.GeometricPrimitive(m, material)
 end
 
-material_red = Trace.MatteMaterial(
-    Trace.ConstantTexture(Trace.RGBSpectrum(0.796f0, 0.235f0, 0.2f0)),
-    Trace.ConstantTexture(0.0f0),
-)
-material_blue = Trace.MatteMaterial(
-    Trace.ConstantTexture(Trace.RGBSpectrum(0.251f0, 0.388f0, 0.847f0)),
-    Trace.ConstantTexture(0.0f0),
-)
-material_white = Trace.MatteMaterial(
-    Trace.ConstantTexture(Trace.RGBSpectrum(1.0f0)),
-    Trace.ConstantTexture(0.0f0),
-)
-mirror = Trace.MirrorMaterial(Trace.ConstantTexture(Trace.RGBSpectrum(1.0f0)))
-glass = Trace.GlassMaterial(
-    Trace.ConstantTexture(Trace.RGBSpectrum(1.0f0)),
-    Trace.ConstantTexture(Trace.RGBSpectrum(1.0f0)),
-    Trace.ConstantTexture(0.0f0),
-    Trace.ConstantTexture(0.0f0),
-    Trace.ConstantTexture(1.5f0),
-    true,
-)
-
 LowSphere(radius, contact=Point3f(0)) = Sphere(contact .+ Point3f(0, 0, radius), radius)
 
 begin
+
+    material_red = Trace.MatteMaterial(
+        Trace.ConstantTexture(Trace.RGBSpectrum(0.796f0, 0.235f0, 0.2f0)),
+        Trace.ConstantTexture(0.0f0),
+    )
+    material_blue = Trace.MatteMaterial(
+        Trace.ConstantTexture(Trace.RGBSpectrum(0.251f0, 0.388f0, 0.847f0)),
+        Trace.ConstantTexture(0.0f0),
+    )
+    material_white = Trace.MatteMaterial(
+        Trace.ConstantTexture(Trace.RGBSpectrum(1.0f0)),
+        Trace.ConstantTexture(0.0f0),
+    )
+    mirror = Trace.MirrorMaterial(Trace.ConstantTexture(Trace.RGBSpectrum(1.0f0)))
+    glass = Trace.GlassMaterial(
+        Trace.ConstantTexture(Trace.RGBSpectrum(1.0f0)),
+        Trace.ConstantTexture(Trace.RGBSpectrum(1.0f0)),
+        Trace.ConstantTexture(0.0f0),
+        Trace.ConstantTexture(0.0f0),
+        Trace.ConstantTexture(1.5f0),
+        true,
+    )
+
     s1 = tmesh(LowSphere(0.5f0), material_white)
     s2 = tmesh(LowSphere(0.3f0, Point3f(0.5, 0.5, 0)), material_blue)
     s3 = tmesh(LowSphere(0.3f0, Point3f(-0.5, 0.5, 0)), mirror)
@@ -49,12 +50,12 @@ begin
 
     bvh = Trace.BVHAccel([s1, s2, s3, s4, ground, back, l, r], 1);
 
-    lights = [
+    lights = (
         # Trace.PointLight(Vec3f(0, -1, 2), Trace.RGBSpectrum(22.0f0)),
         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);
+    )
+    scene = Trace.Scene([lights...], bvh);
     resolution = Point2f(1024)
     f = Trace.LanczosSincFilter(Point2f(1.0f0), 3.0f0)
     film = Trace.Film(resolution,
@@ -77,6 +78,8 @@ 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)
+# After first shading running on GPU
+# 3.835527 seconds (17.36 k allocations: 19.779 MiB, 0.16% gc time, 41 lock conflicts)
 
 camera_sample = Trace.get_camera_sample(integrator.sampler, Point2f(512))
 ray, ω = Trace.generate_ray_differential(integrator.camera, camera_sample)

src/integrators/sampler.jl

@@ -129,8 +129,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

test/gpu-sampler.jl

@@ -109,3 +109,4 @@ gpu_img = ArrayType(zeros(RGBf, res, res));
 # @btime launch_trace_image!(gpu_img, cam, gpu_bvh, lights);
 launch_trace_image!(gpu_img, cam, gpu_scene);
 launch_trace_image!(img, cam, scene, lights)
+Array(gpu_img)

test/gpu-threading-benchmarks.jl

@@ -1,6 +1,15 @@
 using GeometryBasics, LinearAlgebra, Trace, BenchmarkTools
 using ImageShow
 using Makie
+using KernelAbstractions
+import KernelAbstractions as KA
+using KernelAbstractions.Extras.LoopInfo: @unroll
+using AMDGPU
+
+ArrayType = ROCArray
+# using CUDA
+# ArrayType = CuArray
+
 include("./../src/gpu-support.jl")
 
 LowSphere(radius, contact=Point3f(0)) = Sphere(contact .+ Point3f(0, 0, radius), radius)
@@ -18,6 +27,7 @@ material_red = Trace.MatteMaterial(
     Trace.ConstantTexture(0.0f0),
 )
 
+
 begin
     s1 = tmesh(LowSphere(0.5f0), material_red)
     s2 = tmesh(LowSphere(0.3f0, Point3f(0.5, 0.5, 0)), material_red)
@@ -51,88 +61,42 @@ 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))
 end
 
-using KernelAbstractions.Extras.LoopInfo: @unroll
-
-function simple_shading(bvh, shape, ray, si, l, depth, max_depth, lights)
-    core = si.core
-    n = si.shading.n
-    wo = core.wo
-    # Compute scattering functions for surface interaction.
-    si = Trace.compute_differentials(si, ray)
-    mat = Trace.get_material(bvh, shape)
-    if mat.type === Trace.NO_MATERIAL
-        return l
-    end
-    bsdf = mat(si, false, Trace.Radiance)
-    # Compute emitted light if ray hit an area light source.
-    l += Trace.le(si, wo)
-    # Add contribution of each light source.
-    @unroll for light in lights
-        sampled_li, wi, pdf, vt = Trace.sample_li(
-            light, core, rand(Point2f),
-        )
-        (Trace.is_black(sampled_li) || pdf ≈ 0.0f0) && continue
-        f = bsdf(wo, wi)
-        if !Trace.is_black(f) && !Trace.intersect_p(bvh, Trace.spawn_ray(vt.p0, vt.p1))
-            l += f * sampled_li * abs(wi ⋅ n) / pdf
-        end
-    end
-    # if depth + 1 <= max_depth
-    #     # Trace rays for specular reflection & refraction.
-    #     l += specular_reflect(bsdf, i, ray, si, scene, depth)
-    #     l += specular_transmit(bsdf, i, ray, si, scene, depth)
-    # end
-    return l
-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
+#     code_llvm(io, simple_shading, typeof.((bvh, bvh.primitives[1], Trace.RayDifferentials(ray), Trace.SurfaceInteraction(), l, 1, 1, lights)))
+# end
 
-@inline function trace_pixel(camera, bvh, xy, lights)
+@inline function trace_pixel(camera, scene, xy)
     pixel = Point2f(Tuple(xy))
     camera_sample = get_camera_sample(pixel)
     ray, ω = Trace.generate_ray_differential(camera, camera_sample)
-    l = Trace.RGBSpectrum(0.0f0)
     if ω > 0.0f0
-        hit, shape, si = Trace.intersect!(bvh, ray)
+        hit, shape, si = Trace.intersect!(scene, ray)
         if hit
-            l = simple_shading(bvh, shape, ray, si, l, 1, 8, lights)
+            l = Trace.li(Trace.UniformSampler(8), 5, ray, scene, 1)
         end
     end
-    return RGBf(l.c...)
+    return l
 end
 
-using KernelAbstractions
-import KernelAbstractions as KA
-
-
-@kernel function ka_trace_image!(img, camera, bvh, lights)
-    idx = @index(Global, Linear)
-    if checkbounds(Bool, img, idx)
-        xy = Tuple(divrem(idx, size(img, 1)))
-        @inbounds img[idx] = trace_pixel(camera, bvh, xy, lights)
+@kernel function ka_trace_image!(img, camera, scene)
+    xy = @index(Global, Cartesian)
+    if checkbounds(Bool, img, xy)
+        l = trace_pixel(camera, scene, xy)
+        @inbounds img[xy] = RGBf(l.c...)
     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)
+function launch_trace_image!(img, camera, scene)
     backend = KA.get_backend(img)
     kernel! = ka_trace_image!(backend)
-    kernel!(img, camera, bvh, lights, ndrange=size(img), workgroupsize=(16, 16))
+    kernel!(img, camera, scene, lights, ndrange=size(img), workgroupsize=(16, 16))
     KA.synchronize(backend)
     return img
 end
@@ -141,17 +105,19 @@ end
 # using CUDA
 # ArrayType = CuArray
 
-using Metal
-ArrayType = MtlArray
+# using Metal
+# ArrayType = MtlArray
+
 preserve = []
-gpu_bvh = to_gpu(ArrayType, bvh; preserve=preserve);
+gpu_scene = to_gpu(ArrayType, scene; 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);
-launch_trace_image!(gpu_img, cam, gpu_bvh, lights);
-@btime launch_trace_image!(img, cam, bvh, lights)
+launch_trace_image!(gpu_img, cam, gpu_scene);
+launch_trace_image!(img, cam, scene, lights)
 # 76.420 ms (234 allocations: 86.05 KiB)
+# 75.973 ms (234 allocations: 86.05 KiB)
 Array(gpu_img)
 
 function cu_trace_image!(img, camera, bvh, lights)
@@ -255,3 +221,24 @@ t = Trace.Triangle(m, 1)
 r = Trace.Ray(o=Point3f(ray_origin), d=ray_direction)
 Trace.intersect_p(t, r)
 Trace.intersect_triangle(r.o, r.d, t.vertices...)
+
+# 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
+
+ray = Trace.RayDifferentials(Trace.Ray(o=Point3f(0.5, 0.5, 1.0), d=Vec3f(0.0, 0.0, -1.0)))
+open("li.llvm", "w") do io
+    code_llvm(io, Trace.li, typeof.((Trace.UniformSampler(8), 5, ray, scene, 1)))
+end
+
+open("li-wt.jl", "w") do io
+    code_warntype(io, Trace.li, typeof.((Trace.UniformSampler(8), 5, ray, scene, 1)))
+end