fixed corner case of traversing a single-node tree with a multi-node …

…tree of different bounding volumes; added threading settings; kernels do not need the count_zeros memory index computation for every index, it is computed only once per level, on the CPU - can now be translated verbatim to GPU code

Project.toml

@@ -1,7 +1,7 @@
 name = "ImplicitBVH"
 uuid = "932a18dc-bb55-4cd5-bdd6-1368ec9cea29"
 authors = ["Andrei Leonard Nicusan"]
-version = "0.3.0"
+version = "0.4.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

README.md

@@ -82,6 +82,39 @@ bvh = BVH(bounding_spheres, BBox{Float32}, UInt32, 2)
 traversal = traverse(bvh, 3, traversal)
 ```
 
+Compute contacts between two different BVH trees (e.g. two different robotic parts):
+
+```julia
+using ImplicitBVH
+using ImplicitBVH: BBox, BSphere
+
+# Generate some simple bounding spheres (will be BVH leaves)
+bounding_spheres1 = [
+    BSphere{Float32}([0., 0., 0.], 0.5),
+    BSphere{Float32}([0., 0., 3.], 0.4),
+]
+
+bounding_spheres2 = [
+    BSphere{Float32}([0., 0., 1.], 0.6),
+    BSphere{Float32}([0., 0., 2.], 0.5),
+    BSphere{Float32}([0., 0., 4.], 0.6),
+]
+
+# Build BVHs using bounding boxes for nodes
+bvh1 = BVH(bounding_spheres1, BBox{Float32}, UInt32)
+bvh2 = BVH(bounding_spheres2, BBox{Float32}, UInt32)
+
+# Traverse BVH for contact detection
+traversal = traverse(
+    bvh1,
+    bvh2,
+    default_start_level(bvh1),
+    default_start_level(bvh2),
+    # previous_traversal_cache,
+    # num_threads=4,
+)
+```
+
 Check out the `benchmark` folder for an example traversing an STL model.
 
 

src/ImplicitBVH.jl

@@ -22,6 +22,6 @@ include("morton.jl")
 include("implicit_tree.jl")
 include("bounding_volumes.jl")
 include("build.jl")
-include("traverse.jl")
+include("traverse/traverse.jl")
 
 end     # module ImplicitBVH

src/bounding_volumes.jl

@@ -1,6 +1,8 @@
 """
     iscontact(a::BSphere, b::BSphere)
     iscontact(a::BBox, b::BBox)
+    iscontact(a::BSphere, b::BBox)
+    iscontact(a::BBox, b::BSphere)
 
 Check if two bounding volumes are touching or inter-penetrating.
 """
@@ -351,3 +353,20 @@ function iscontact(a::BBox, b::BBox)
     (a.up[2] >= b.lo[2] && a.lo[2] <= b.up[2]) &&
     (a.up[3] >= b.lo[3] && a.lo[3] <= b.up[3])
 end
+
+
+# Contact detection between heterogeneous BVs - only needed when one BVH has exactly one leaf
+function iscontact(a::BSphere, b::BBox)
+    # This is an edge case, used for broad-phase collision detection, so we simply take the
+    # sphere's bounding box, as a full sphere-box contact detection is computationally heavy
+    ab = BBox(
+        (a.x[1] - a.r, a.x[2] - a.r, a.x[3] - a.r),
+        (a.x[1] + a.r, a.x[2] + a.r, a.x[3] + a.r),
+    )
+    iscontact(ab, b)
+end
+
+
+function iscontact(a::BBox, b::BSphere)
+    iscontact(b, a)
+end

src/build.jl

@@ -28,7 +28,8 @@ Tree Level          Nodes & Leaves               Build Up    Traverse Down
         bounding_volumes::AbstractVector{L},
         node_type::Type{N}=L,
         morton_type::Type{U}=UInt,
-        built_level::Integer=1,
+        built_level::Integer=1;
+        num_threads=Threads.nthreads(),
     ) where {L, N, U <: MortonUnsigned}
 
 # Fields
@@ -134,7 +135,8 @@ function BVH(
     bounding_volumes::AbstractVector{L},
     node_type::Type{N}=L,
     morton_type::Type{U}=UInt,
-    built_level=1,
+    built_level=1;
+    num_threads=Threads.nthreads(),
 ) where {L, N, U <: MortonUnsigned}
 
     # Ensure correctness
@@ -165,7 +167,7 @@ function BVH(
 
     # Compute morton codes for the bounding volumes
     mortons = similar(bounding_volumes, morton_type)
-    @inbounds morton_encode!(mortons, bounding_volumes)
+    @inbounds morton_encode!(mortons, bounding_volumes, num_threads=num_threads)
 
     # Compute indices that sort codes along the Z-curve - closer objects have closer Morton codes
     # TODO: check parallel SyncSort or ThreadsX.QuickSort
@@ -176,22 +178,22 @@ function BVH(
 
     # Aggregate bounding volumes up to root
     if tree.real_nodes >= 2
-        aggregate_oibvh!(bvh_nodes, bounding_volumes, tree, order, built_ilevel)
+        aggregate_oibvh!(bvh_nodes, bounding_volumes, tree, order, built_ilevel, num_threads)
     end
 
     BVH(built_ilevel, tree, bvh_nodes, bounding_volumes, order)
 end
 
 
 # Build ImplicitBVH nodes above the leaf-level from the bottom up, inplace
-function aggregate_oibvh!(bvh_nodes, bvh_leaves, tree, order, built_level=1)
+function aggregate_oibvh!(bvh_nodes, bvh_leaves, tree, order, built_level, num_threads)
 
     # Special case: aggregate level above leaves - might have different node types
-    aggregate_last_level!(bvh_nodes, bvh_leaves, tree, order)
+    aggregate_last_level!(bvh_nodes, bvh_leaves, tree, order, num_threads)
 
     level = tree.levels - 2
     while level >= built_level
-        aggregate_level!(bvh_nodes, level, tree)
+        aggregate_level!(bvh_nodes, level, tree, num_threads)
         level -= 1
     end
 
@@ -243,7 +245,7 @@ end
 end
 
 
-@inline function aggregate_last_level!(bvh_nodes, bvh_leaves, tree, order)
+@inline function aggregate_last_level!(bvh_nodes, bvh_leaves, tree, order, num_threads)
     # Memory index of first node on this level (i.e. first above leaf-level)
     level = tree.levels - 1
     start_pos = memory_index(tree, pow2(level - 1))
@@ -256,7 +258,7 @@ end
 
     # Split computation into contiguous ranges of minimum 100 elements each; if only single thread
     # is needed, inline call
-    tp = TaskPartitioner(num_nodes, Threads.nthreads(), 100)
+    tp = TaskPartitioner(num_nodes, num_threads, 100)
     if tp.num_tasks == 1
         @inbounds aggregate_last_level_range!(
             bvh_nodes, bvh_leaves, order,
@@ -303,7 +305,7 @@ end
 end
 
 
-@inline function aggregate_level!(bvh_nodes, level, tree)
+@inline function aggregate_level!(bvh_nodes, level, tree, num_threads)
     # Memory index of first node on this level
     start_pos = memory_index(tree, pow2(level - 1))
 
@@ -316,7 +318,7 @@ end
 
     # Split computation into contiguous ranges of minimum 100 elements each; if only single thread
     # is needed, inline call
-    tp = TaskPartitioner(num_nodes, Threads.nthreads(), 100)
+    tp = TaskPartitioner(num_nodes, num_threads, 100)
     if tp.num_tasks == 1
         @inbounds aggregate_level_range!(
             bvh_nodes, start_pos,

src/morton.jl

@@ -168,7 +168,8 @@ function morton_encode!(
     mortons::AbstractVector{U},
     bounding_volumes::AbstractVector,
     mins,
-    maxs,
+    maxs;
+    num_threads=Threads.nthreads(),
 ) where {U <: MortonUnsigned}
 
     # Bounds checking
@@ -180,7 +181,7 @@ function morton_encode!(
     length(bounding_volumes) == 0 && return nothing
 
     # Encode bounding volumes' centres across multiple threads using contiguous ranges
-    tp = TaskPartitioner(length(bounding_volumes), Threads.nthreads(), 1000)
+    tp = TaskPartitioner(length(bounding_volumes), num_threads, 1000)
     if tp.num_tasks == 1
         morton_encode_range!(
             mortons, bounding_volumes,
@@ -205,11 +206,15 @@ function morton_encode!(
 end
 
 
-function morton_encode!(mortons::AbstractVector{U}, bounding_volumes) where {U <: MortonUnsigned}
+function morton_encode!(
+    mortons::AbstractVector{U},
+    bounding_volumes;
+    num_threads=Threads.nthreads(),
+) where {U <: MortonUnsigned}
+
     # Compute exclusive bounds [xmin, ymin, zmin], [xmax, ymax, zmax].
-    # TODO: see uint_encode from Base.Sort, scaling might be better
     mins, maxs = bounding_volumes_extrema(bounding_volumes)
-    morton_encode!(mortons, bounding_volumes, mins, maxs)
+    morton_encode!(mortons, bounding_volumes, mins, maxs, num_threads=num_threads)
     nothing
 end
 
@@ -220,9 +225,14 @@ end
 Encode the centers of some `bounding_volumes` as Morton codes of type `U <: `
 [`MortonUnsigned`](@ref). See [`morton_encode!`](@ref) for full details. 
 """
-function morton_encode(bounding_volumes, ::Type{U}=UInt) where {U <: MortonUnsigned}
+function morton_encode(
+    bounding_volumes,
+    ::Type{U}=UInt;
+    num_threads=Threads.nthreads(),
+) where {U <: MortonUnsigned}
+
     # Pre-allocate vector of morton codes
     mortons = similar(bounding_volumes, U)
-    morton_encode!(mortons, bounding_volumes)
+    morton_encode!(mortons, bounding_volumes, num_threads=num_threads)
     mortons
 end

src/traverse/traverse.jl

@@ -0,0 +1,82 @@
+"""
+    $(TYPEDEF)
+
+Alias for a tuple of two indices representing e.g. a contacting pair.
+"""
+const IndexPair = Tuple{Int, Int}
+
+
+"""
+    $(TYPEDEF)
+
+Collected BVH traversal `contacts` vector, some stats, plus the two buffers `cache1` and `cache2`
+which can be reused for future traversals to minimise memory allocations.
+
+# Fields
+- `start_level1::Int`: the level at which the single/pair-tree traversal started for the first BVH.
+- `start_level2::Int`: the level at which the pair-tree traversal started for the second BVH.
+- `num_checks::Int`: the total number of contact checks done.
+- `num_contacts::Int`: the number of contacts found.
+- `contacts::view(cache1, 1:num_contacts)`: the contacting pairs found, as a view into `cache1`.
+- `cache1::C1{IndexPair} <: AbstractVector`: first BVH traversal buffer.
+- `cache2::C2{IndexPair} <: AbstractVector`: second BVH traversal buffer.
+"""
+struct BVHTraversal{C1 <: AbstractVector, C2 <: AbstractVector}
+    # Stats
+    start_level1::Int
+    start_level2::Int
+    num_checks::Int
+
+    # Data
+    num_contacts::Int
+    cache1::C1
+    cache2::C2
+end
+
+
+# Constructor in the case of single-tree traversal (e.g. traverse(bvh)), when we only have a
+# single start_level
+function BVHTraversal(
+    start_level::Int,
+    num_checks::Int,
+    num_contacts::Int,
+    cache1::AbstractVector,
+    cache2::AbstractVector,
+)
+    BVHTraversal(start_level, start_level, num_checks, num_contacts, cache1, cache2)
+end
+
+
+# Custom pretty-printing
+function Base.show(io::IO, t::BVHTraversal{C1, C2}) where {C1, C2}
+    print(
+        io,
+        """
+        BVHTraversal
+          start_level1: $(typeof(t.start_level1)) $(t.start_level1)
+          start_level2: $(typeof(t.start_level2)) $(t.start_level2)
+          num_checks:   $(typeof(t.num_checks)) $(t.num_checks)
+          num_contacts: $(typeof(t.num_contacts)) $(t.num_contacts)
+          contacts:     $(Base.typename(typeof(t.contacts)).wrapper){IndexPair}($(size(t.contacts)))
+          cache1:       $C1($(size(t.cache1)))
+          cache2:       $C2($(size(t.cache2)))
+        """
+    )
+end
+
+
+function Base.getproperty(bt::BVHTraversal, sym::Symbol)
+   if sym === :contacts
+       return @view bt.cache1[1:bt.num_contacts]
+   else
+       return getfield(bt, sym)
+   end
+end
+
+Base.propertynames(::BVHTraversal) = (:start_level1, :start_level2, :num_checks, :contacts,
+                                      :num_contacts, :cache1, :cache2)
+
+
+# Single BVH and BVH-BVH traversal
+include("traverse_single.jl")
+include("traverse_pair.jl")