hlbvh

build_hlbvh.go

package hlbvh

import (
	"math/bits"
	"sync/atomic"
)

// Mark: - Kernels

func calculateMortonCodeKernel(dispatcher workDispatcher,
	primitiveBounds []bBox, sceneBound bBox, mortonCodes []int) {
	globalId := dispatcher.getGlobalId(0)
	if globalId < 0 {
		return
	}
	numPrimitiveBounds := len(primitiveBounds)

	if globalId < numPrimitiveBounds {
		bound := primitiveBounds[globalId]

		center := bound.pMax.add(bound.pMin).xyz().mulValue(0.5)
		sceneMin := sceneBound.pMin.xyz()
		sceneExtends := sceneBound.pMax.xyz().sub(sceneBound.pMax.xyz())

		mortonCodes[globalId] = int(calculateMortonCode(center.sub(sceneMin).div(sceneExtends)))
	}
}

func EmitHierarchyKernel(dispatcher workDispatcher,
	mortonCodes []int, bounds []bBox, indices []int, nodes []HlbvhNode, boundsSorted []bBox) {
	globalId := dispatcher.getGlobalId(0)
	if globalId < 0 {
		return
	}
	numPrims := len(indices)
	leafIdx := func(i int) int {
		return (numPrims - 1) + i
	}
	nodeIdx := func(i int) int { return i }

	if globalId < numPrims {
		nodes[leafIdx(globalId)].left = indices[globalId]
		nodes[leafIdx(globalId)].right = indices[globalId]
		boundsSorted[leafIdx(globalId)] = bounds[indices[globalId]]
	}

	if globalId < numPrims-1 {
		rng := findSpan(mortonCodes, globalId)
		split := findSplit(mortonCodes, rng)

		var c1Idx, c2Idx int

		if split == rng.x {
			c1Idx = leafIdx(split)
		} else {
			c1Idx = nodeIdx(split)
		}

		if split+1 == rng.y {
			c2Idx = leafIdx(split + 1)
		} else {
			c1Idx = nodeIdx(split + 1)
		}

		nodes[nodeIdx(globalId)].left = c1Idx
		nodes[nodeIdx(globalId)].right = c2Idx
		nodes[c1Idx].parent = nodeIdx(globalId)
		nodes[c2Idx].parent = nodeIdx(globalId)
	}
}

func RefitBoundsKernel(dispatcher workDispatcher,
	bounds []bBox, nodes []HlbvhNode, flags []int32) {
	globalId := dispatcher.getGlobalId(0)
	if globalId < 0 {
		return
	}
	numPrims := len(bounds)
	leafIdx := func(i int) int {
		return (numPrims - 1) + i
	}

	if globalId < numPrims {
		idx := leafIdx(globalId)
		for {
			idx := nodes[idx].parent
			if atomic.CompareAndSwapInt32(&flags[idx], 0, 1) {
				lc := nodes[idx].left
				rc := nodes[idx].right
				b := bBoxUnion(bounds[lc], bounds[rc])
				bounds[idx] = b
			} else {
				break
			}
			if idx == 0 {
				break
			}
		}
	}
}

// MARK: - Helpers

type float float32

type float4 struct {
	x, y, z, w float
}

func (f float4) add(r float4) float4 {
	return float4{
		x: f.x + r.x,
		y: f.y + r.y,
		z: f.z + r.z,
		w: f.w + r.w,
	}
}

func (f float4) xyz() float3 {
	return float3{
		x: f.x,
		y: f.y,
		z: f.z,
	}
}

type float3 struct {
	x, y, z float
}

func (f float3) mulValue(v float) float3 {
	return float3{
		x: f.x * v,
		y: f.y * v,
		z: f.z * v,
	}
}

func (f float3) div(v float3) float3 {
	return float3{
		x: f.x / v.x,
		y: f.y / v.y,
		z: f.z / v.z,
	}
}

func (f float3) sub(r float3) float3 {
	return float3{
		x: f.x + r.x,
		y: f.y + r.y,
		z: f.z + r.z,
	}
}

type bBox struct {
	pMin, pMax float4
}

func minFloat(x, y float) float {
	if x >= y {
		return y
	}
	return x
}

func maxFloat(x, y float) float {
	if x >= y {
		return x
	}
	return y
}

func minFloat4(a, b float4) float4 {
	return float4{
		x: minFloat(a.x, b.x),
		y: minFloat(a.y, b.y),
		z: minFloat(a.z, b.z),
		w: minFloat(a.w, b.w),
	}
}

func maxFloat4(a, b float4) float4 {
	return float4{
		x: maxFloat(a.x, b.x),
		y: maxFloat(a.y, b.y),
		z: maxFloat(a.z, b.z),
		w: maxFloat(a.w, b.w),
	}
}

func expandBits(v uint) uint {
	v = (v * 0x00010001) & 0xFF0000FF
	v = (v * 0x00000101) & 0x0F00F00F
	v = (v * 0x00000011) & 0xC30C30C3
	v = (v * 0x00000005) & 0x49249249
	return v
}

func calculateMortonCode(p float3) uint {
	x := minFloat(maxFloat(p.x*1024, 0), 1023)
	y := minFloat(maxFloat(p.y*1024, 0), 1023)
	z := minFloat(maxFloat(p.z*1024, 0), 1023)
	xx := expandBits(uint(x))
	yy := expandBits(uint(y))
	zz := expandBits(uint(z))
	return xx*4 + yy*2 + zz
}

func bBoxUnion(a, b bBox) bBox {
	return bBox{
		pMin: minFloat4(a.pMin, b.pMin),
		pMax: maxFloat4(a.pMax, b.pMax),
	}
}

type workDispatcher interface {
	getGlobalId(dim int) int
}

func minInt(x, y int) int {
	if x >= y {
		return y
	}
	return x
}

func maxInt(x, y int) int {
	if x >= y {
		return x
	}
	return y
}

func clz(v int) int {
	return bits.LeadingZeros(uint(v))
}

func delta(mortonCodes []int, i1, i2 int) int {
	numPrims := len(mortonCodes)

	left := minInt(i1, i2)
	right := maxInt(i1, i2)

	if left < 0 || right >= numPrims {
		return -1
	}

	leftCode := mortonCodes[left]
	rightCode := mortonCodes[right]

	if leftCode != rightCode {
		return clz(leftCode ^ rightCode)
	}
	return 32 + clz(leftCode^rightCode)
}

type int2 struct {
	x, y int
}

func signInt(v int) int {
	switch {
	case v > 0:
		return 1
	case v < 0:
		return -1
	default:
		return 0
	}
}

func findSpan(mortonCodes []int, idx int) int2 {
	d := signInt(delta(mortonCodes, idx, idx+1) - delta(mortonCodes, idx, idx-1))
	deltaMin := delta(mortonCodes, idx, idx-d)
	lmax := 2
	for delta(mortonCodes, idx, idx+lmax*d) > deltaMin {
		lmax *= 2
	}
	l := 0
	t := lmax
	for {
		t /= 2
		if delta(mortonCodes, idx, idx+(l+t)*d) > deltaMin {
			l = l + t
		}
		if !(t > 1) {
			break
		}
	}

	span := int2{
		x: minInt(idx, idx+l*d),
		y: maxInt(idx, idx+l*d),
	}

	return span
}

func findSplit(mortonCodes []int, span int2) int {
	left := span.x
	right := span.y

	numIdentical := delta(mortonCodes, left, right)
	for {
		newSplit := (left + right) / 2
		if delta(mortonCodes, left, newSplit) > numIdentical {
			left = newSplit
		} else {
			right = newSplit
		}
		if !(right > left+1) {
			break
		}
	}

	return left
}

type HlbvhNode struct {
	parent, left, right, next int
}

type simpleDispatcher struct {
	max     int
	current int
}

func newSimpleDispatcher(max int) *simpleDispatcher {
	return &simpleDispatcher{
		max:     max,
		current: max - 1,
	}
}

func (s *simpleDispatcher) init(v ...int) {
	if len(v) > 0 {
		s.max = v[0]
	}
	s.current = s.max - 1
}

func (s *simpleDispatcher) getGlobalId(dim int) int {
	if s.current < 0 {
		return -1
	}
	ret := s.current
	s.current--
	return ret
}