Fix TLAS build stability + keepalive + sync

- Move _build_flat_blas_arrays! from adapt_structure to sync! so data is
  ready before any kernel dispatch.

- Add KA.synchronize after AK.sortperm in BLAS/TLAS build to ensure sort
  temp buffers aren't freed while GPU is still using them.

- Replace merge_sort_by_key! with sortperm in TLAS topology (avoids 64-bit
  Int payload corruption on Lava).

- Add _REBUILD_KEEPALIVE and synchronize in multitypeset rebuild_static!.

- Remove underscore-prefixed function names in RaycoreLavaExt.

- Add instanced BVH test.

Author: SimonDanisch

ext/RaycoreLavaExt/RaycoreLavaExt.jl

@@ -23,9 +23,9 @@ using Lava: LavaBackend, LavaArray, LavaBLAS, LavaTLAS,
             lava_global_invocation_id_x,
             lava_rt_primitive_id, lava_rt_instance_custom_index,
             lava_rt_launch_id_x,
-            _lava_rt_ignore_intersection, _lava_rt_terminate_ray,
-            _lava_rt_payload_store_f32_at, _lava_rt_payload_load_f32_at,
-            _lava_rt_trace_ray,
+            lava_rt_ignore_intersection, lava_rt_terminate_ray,
+            lava_rt_payload_store_f32_at, lava_rt_payload_load_f32_at,
+            lava_rt_trace_ray,
             mat4_to_vk_transform
 
 const LavaHWTLAS = HWTLAS{LavaBackend}
@@ -139,14 +139,14 @@ Raycore.rt_primitive_id(::LavaHWAdapted) = lava_rt_primitive_id()
 Raycore.rt_instance_custom_index(::LavaHWAdapted) = lava_rt_instance_custom_index()
 Raycore.rt_launch_id_x(::LavaHWAdapted) = lava_rt_launch_id_x()
 Raycore.rt_global_invocation_id_x(::LavaHWAdapted) = lava_global_invocation_id_x()
-Raycore.rt_ignore_intersection(::LavaHWAdapted) = _lava_rt_ignore_intersection()
-Raycore.rt_terminate_ray(::LavaHWAdapted) = _lava_rt_terminate_ray()
-Raycore.rt_payload_store!(::LavaHWAdapted, val, slot) = _lava_rt_payload_store_f32_at(val, slot)
-Raycore.rt_payload_load(::LavaHWAdapted, slot) = _lava_rt_payload_load_f32_at(slot)
+Raycore.rt_ignore_intersection(::LavaHWAdapted) = lava_rt_ignore_intersection()
+Raycore.rt_terminate_ray(::LavaHWAdapted) = lava_rt_terminate_ray()
+Raycore.rt_payload_store!(::LavaHWAdapted, val, slot) = lava_rt_payload_store_f32_at(val, slot)
+Raycore.rt_payload_load(::LavaHWAdapted, slot) = lava_rt_payload_load_f32_at(slot)
 
 function Raycore.rt_trace_ray!(::LavaHWAdapted, flags, mask, sbt_offset, sbt_stride, miss_idx,
                                 ox, oy, oz, tmin, dx, dy, dz, tmax)
-    _lava_rt_trace_ray(flags, mask, sbt_offset, sbt_stride, miss_idx,
+    lava_rt_trace_ray(flags, mask, sbt_offset, sbt_stride, miss_idx,
                        ox, oy, oz, tmin, dx, dy, dz, tmax)
 end
 

src/instanced-bvh.jl

@@ -830,6 +830,10 @@ function _rebuild_bvh!(tlas::TLAS)
     tlas.nodes = nodes
     tlas.root_aabb = root_aabb
 
+    # Build flat BLAS arrays during sync (not during adapt_structure).
+    # This ensures the data is ready before any kernel dispatch.
+    _build_flat_blas_arrays!(tlas)
+
     tlas.dirty = false
     return
 end
@@ -933,9 +937,7 @@ The TLAS must stay alive while the StaticTLAS is in use.
 """
 function Adapt.adapt_structure(to, tlas::TLAS)
     sync!(tlas)
-
-    # Build flat BLAS arrays for traversal (avoids pointer-in-buffer on Metal)
-    _build_flat_blas_arrays!(tlas)
+    # Flat BLAS arrays are built during sync! -- no rebuild here.
 
     if tlas._flat_blas_nodes === nothing
         # Empty scene — need correct types for StaticTLAS type parameters
@@ -1259,6 +1261,7 @@ function build_blas(
     # Sort primitives by Morton codes
     # AcceleratedKernels only supports GPU backends, use Julia's sortperm for CPU
     perm = AK.sortperm(morton_codes)
+    KA.synchronize(backend)  # Ensure sort temp buffers aren't freed while GPU is still using them
     morton_codes = morton_codes[perm]
     primitives = primitives[perm]
 
@@ -1387,17 +1390,17 @@ function _build_tlas_topology(blas_array, instances, backend)
     calc_kernel!(morton_codes, instances, blas_array, scene_min, scene_extent, ndrange=n)
     KA.synchronize(backend)
 
-    # Sort indices by Morton codes
-    # AcceleratedKernels only supports GPU backends, use Julia's sortperm for CPU
+    # Sort instances by Morton codes.
+    # On Lava, merge_sort_by_key! with a 64-bit Int payload can corrupt the
+    # permutation vector, which later sends TLAS leaf creation out of bounds.
+    # Use sortperm like the BLAS path so the permutation type matches the backend.
     if backend isa KA.CPU
         sorted_indices = sortperm(morton_codes)
         morton_codes .= morton_codes[sorted_indices]
     else
-        sorted_indices = KA.allocate(backend, Int, n)
-        iota_k! = iota_kernel!(backend)
-        iota_k!(sorted_indices, ndrange=n)
-        KA.synchronize(backend)
-        AK.merge_sort_by_key!(morton_codes, sorted_indices)
+        sorted_indices = AK.sortperm(morton_codes)
+        KA.synchronize(backend)  # Ensure sort temp buffers aren't freed while GPU is still using them
+        morton_codes = morton_codes[sorted_indices]
     end
 
     # Allocate nodes and initialize with empty values

src/multitypeset.jl

@@ -265,7 +265,17 @@ to_tuple(mts::MultiTypeSet) = to_tuple(get_static(mts))
 # Internal: Rebuild the static tuple - converts CPU vectors to GPU
 # ============================================================================
 
+const _REBUILD_KEEPALIVE = Any[]
+
 function rebuild_static!(dhv::MultiTypeSet)
+    # Keep old static arrays alive until the NEXT rebuild completes.
+    # Without this, GC can free old buffers during adapt() of new data,
+    # causing use-after-free (DEVICE_LOST on large scenes like killeroo).
+    old_static = dhv.static
+    if old_static !== nothing
+        push!(_REBUILD_KEEPALIVE, old_static)
+    end
+
     # Convert CPU data vectors to GPU
     data_tuple = if isempty(dhv.data_order)
         ()
@@ -279,6 +289,10 @@ function rebuild_static!(dhv::MultiTypeSet)
         Tuple(Adapt.adapt(dhv.backend, dhv.texture_isbits[T]) for T in dhv.texture_order)
     end
     dhv.static = StaticMultiTypeSet(data_tuple, tex_tuple)
+
+    # Material/light rebuilds enqueue GPU uploads. Make them visible before any
+    # subsequent BLAS/TLAS kernels or later rebuilds can reuse/finalize backing storage.
+    KA.synchronize(dhv.backend)
 end
 
 # ============================================================================

test/test_instanced_bvh.jl

@@ -783,6 +783,25 @@ else
         @test cl_tlas.nodes isa CLArray
     end
 
+    @testset "TLAS sync with many instances" begin
+        mesh = make_triangle_mesh()
+        transforms = [Mat4f(1, 0, 0, 0,
+                            0, 1, 0, 0,
+                            0, 0, 1, 0,
+                            Float32(mod(i - 1, 9)) * 1.5f0,
+                            Float32((i - 1) ÷ 9) * 1.25f0,
+                            0,
+                            1) for i in 1:81]
+
+        tlas = Raycore.TLAS(cl_backend)
+        push!(tlas, mesh, transforms)
+        sync!(tlas)
+
+        @test length(tlas.instances) == 81
+        @test length(tlas.nodes) == 161
+        @test Raycore.world_bound(tlas) isa Bounds3
+    end
+
     @testset "closest_hit_kernel! - basic intersection" begin
         mesh = make_triangle_mesh()
         tlas, _ = TLAS([mesh]; backend=cl_backend)