Merge pull request #35 from li3zhen1/dev2

Perf: Add `BufferedKDTree`

Package.swift

@@ -45,10 +45,10 @@ let package = Package(
             // swiftSettings: [.unsafeFlags(["-Xfrontend", "-disable-availability-checking"])]
         ),
 
-        // .testTarget(
-        //     name: "NDTreeTests",
-        //     dependencies: ["ForceSimulation"]
-        // ),
+        .testTarget(
+            name: "KDTreeTests",
+            dependencies: ["ForceSimulation"]
+        ),
 
         .testTarget(
             name: "ForceSimulationTests",

README.md

@@ -210,9 +210,9 @@ See [Example](https://github.com/li3zhen1/Grape/tree/main/Examples/ForceDirected
 
 #### Simulation
 
-Grape uses simd to calculate position and velocity. Currently it takes **~0.010** seconds to iterate 120 times over the example graph(2D). (77 vertices, 254 edges, with manybody, center, collide and link forces. Release build on a M1 Max, tested with command `swift test -c release`)
+Grape uses simd to calculate position and velocity. Currently it takes **~0.005** seconds to iterate 120 times over the example graph(2D). (77 vertices, 254 edges, with manybody, center, collide and link forces. Release build on a M1 Max, tested with command `swift test -c release`)
 
-For 3D simulation, it takes **~0.014** seconds for the same graph and same configs.
+For 3D simulation, it takes **~0.008** seconds for the same graph and same configs.
 
 > [!IMPORTANT]
 > Due to heavy use of generics (some are not specialized in Debug mode), the performance in Debug build is ~100x slower than Release build. Grape might ship a version with pre-inlined generics to address this problem.

Sources/ForceSimulation/ForceProtocol.swift

@@ -12,6 +12,14 @@ public protocol ForceProtocol<Vector> {
     /// Bind to a kinetic system that describes the state of all nodes in your simulation.
     /// This has to be called before `apply` is called.
     @inlinable mutating func bindKinetics(_ kinetics: Kinetics<Vector>)
+
+}
+
+public extension ForceProtocol {
+    @inlinable
+    func dispose() {
+
+    }
 }
 
 public protocol Force2D: ForceProtocol where Vector == SIMD2<Double> {}
@@ -67,4 +75,4 @@ extension ForceField {
 
 public protocol ForceField2D: ForceField & Force2D { }
 
-public protocol ForceField3D: ForceField & Force3D { }
+public protocol ForceField3D: ForceField & Force3D { }

Sources/ForceSimulation/Forces/CenterForce.swift

@@ -6,7 +6,7 @@ extension Kinetics {
     /// See [Collide Force - D3](https://d3js.org/d3-force/collide).
     public struct CenterForce: ForceProtocol {
 
-        @usableFromInline  var kinetics: Kinetics! = nil
+        @usableFromInline var kinetics: Kinetics! = nil
 
         @inlinable
         public func apply() {

Sources/ForceSimulation/Forces/CollideForce.swift

@@ -68,10 +68,28 @@ extension Kinetics {
         public mutating func bindKinetics(_ kinetics: Kinetics) {
             self.kinetics = kinetics
             self.calculatedRadius = self.radius.calculateUnsafe(for: kinetics.validCount)
+            self.tree = .allocate(capacity: 1)
+            self.tree.initialize(
+                to:
+                    BufferedKDTree(
+                        rootBox: .init(
+                            p0: .init(repeating: 0),
+                            p1: .init(repeating: 1)
+                        ),
+                        nodeCapacity: kinetics.validCount,
+                        rootDelegate: .init(
+                            radiusBufferPointer: self.calculatedRadius.mutablePointer)
+                    )
+            )
         }
 
         @usableFromInline
         var calculatedRadius: UnsafeArray<Vector.Scalar>! = nil
+
+        @usableFromInline
+        internal var tree:
+            UnsafeMutablePointer<BufferedKDTree<Vector, MaxRadiusNDTreeDelegate<Vector>>>! = nil
+
         @inlinable
         public func apply() {
             assert(self.kinetics != nil, "Kinetics not bound to force")
@@ -82,14 +100,27 @@ extension Kinetics {
             let positionBufferPointer = kinetics.position.mutablePointer
             let velocityBufferPointer = kinetics.velocity.mutablePointer
 
+            let tree = self.tree!
+
             for _ in 0..<iterationsPerTick {
 
-                var tree = KDTree<Vector, MaxRadiusNDTreeDelegate<Vector>>(
-                    covering: kinetics.position,
-                    rootDelegate: MaxRadiusNDTreeDelegate<Vector>(
-                        radiusBufferPointer: calculatedRadius
-                    )
+                let coveringBox = KDBox<Vector>.cover(of: self.kinetics.position)
+
+                tree.pointee.reset(
+                    rootBox: coveringBox,
+                    rootDelegate: .init(radiusBufferPointer: calculatedRadius)
                 )
+                assert(tree.pointee.validCount == 1)
+
+                for p in kinetics.range {
+                    //                    #if DEBUG
+                    //                    let validCountBeforeAdd = tree.pointee.validCount
+                    //                    #endif
+                    tree.pointee.add(nodeIndex: p, at: positionBufferPointer[p])
+                    //                    #if DEBUG
+                    //                    assert(validCountBeforeAdd >= tree.pointee.validCount - 8)
+                    //                    #endif
+                }
 
                 for i in kinetics.range {
                     let iOriginalPosition = positionBufferPointer[i]
@@ -99,39 +130,46 @@ extension Kinetics {
                     let iPosition = iOriginalPosition + iOriginalVelocity
                     let random = kinetics.randomGenerator
 
-                    tree.visit { t in
+                    tree.pointee.visit { t in
 
                         let maxRadiusOfQuad = t.delegate.maxNodeRadius
                         let deltaR = maxRadiusOfQuad + iR
 
-                        if t.nodePosition != nil {
-                            for j in t.nodeIndices {
+                        if var jNode = t.nodeIndices {
+                            while true {
+                                let j = jNode.index
                                 //                            print("\(i)<=>\(j)")
                                 // is leaf, make sure every collision happens once.
-                                guard j > i else { continue }
+                                if j > i {
 
-                                let jR = calculatedRadius[j]
-                                let jOriginalPosition = positionBufferPointer[j]
-                                let jOriginalVelocity = velocityBufferPointer[j]
-                                var deltaPosition =
-                                    iPosition - (jOriginalPosition + jOriginalVelocity)
-                                let l = (deltaPosition).lengthSquared()
+                                    let jR = calculatedRadius[j]
+                                    let jOriginalPosition = positionBufferPointer[j]
+                                    let jOriginalVelocity = velocityBufferPointer[j]
+                                    var deltaPosition =
+                                        iPosition - (jOriginalPosition + jOriginalVelocity)
+                                    let l = (deltaPosition).lengthSquared()
 
-                                let deltaR = iR + jR
-                                if l < deltaR * deltaR {
+                                    let deltaR = iR + jR
+                                    if l < deltaR * deltaR {
 
-                                    var l = /*simd_length*/ (deltaPosition.jiggled(by: random))
-                                        .length()
-                                    l = (deltaR - l) / l * strength
+                                        var l = /*simd_length*/ (deltaPosition.jiggled(by: random))
+                                            .length()
+                                        l = (deltaR - l) / l * strength
 
-                                    let jR2 = jR * jR
+                                        let jR2 = jR * jR
 
-                                    let k = jR2 / (iR2 + jR2)
+                                        let k = jR2 / (iR2 + jR2)
 
-                                    deltaPosition *= l
+                                        deltaPosition *= l
 
-                                    velocityBufferPointer[i] += deltaPosition * k
-                                    velocityBufferPointer[j] -= deltaPosition * (1 - k)
+                                        velocityBufferPointer[i] += deltaPosition * k
+                                        velocityBufferPointer[j] -= deltaPosition * (1 - k)
+                                    }
+                                }
+                                if jNode.next == nil {
+                                    break
+                                } else {
+                                    jNode = jNode.next!.pointee
                                 }
                             }
                             return false
@@ -175,5 +213,10 @@ extension Kinetics {
             }
         }
 
+        @inlinable
+        public func dispose() {
+            self.tree.deinitialize(count: 1)
+        }
+
     }
 }

Sources/ForceSimulation/Forces/ManyBodyForce.swift

@@ -1,3 +1,5 @@
+import simd
+
 public struct MassCentroidKDTreeDelegate<Vector>: KDTreeDelegate
 where Vector: SimulatableVector {
 
@@ -82,8 +84,6 @@ extension Kinetics {
         public var mass: NodeMass
         @usableFromInline var precalculatedMass: UnsafeArray<Vector.Scalar>! = nil
 
-        // @usableFromInline var forces: [Vector] = []
-
         @inlinable
         public init(
             strength: Vector.Scalar,
@@ -94,11 +94,13 @@ extension Kinetics {
             self.mass = nodeMass
             self.theta = theta
             self.theta2 = theta * theta
+
         }
 
         @inlinable
         public func apply() {
-
+
+            // Avoid capturing self
             let alpha = self.kinetics.alpha
             let theta2 = self.theta2
             let distanceMin2 = self.distanceMin2
@@ -107,16 +109,18 @@ extension Kinetics {
             let precalculatedMass = self.precalculatedMass.mutablePointer
             let positionBufferPointer = kinetics.position.mutablePointer
             let random = kinetics.randomGenerator
+            let tree = self.tree!
 
-            var tree = KDTree(
-                covering: self.kinetics.position,
-                rootDelegate: MassCentroidKDTreeDelegate<Vector>(massProvider: precalculatedMass)
-            )
+            let coveringBox = KDBox<Vector>.cover(of: self.kinetics.position)
+            tree.pointee.reset(rootBox: coveringBox, rootDelegate: .init(massProvider: precalculatedMass))
+            for p in kinetics.range {
+                tree.pointee.add(nodeIndex: p, at: positionBufferPointer[p])
+            }
 
             for i in self.kinetics.range {
                 let pos = positionBufferPointer[i]
                 var f = Vector.zero
-                tree.visit { t in
+                tree.pointee.visit { t in
 
                     guard t.delegate.accumulatedCount > 0 else { return false }
                     let centroid =
@@ -148,13 +152,13 @@ extension Kinetics {
                         f += vec * k
                         return false
 
-                    } else if t.children != nil {
+                    } else if t.childrenBufferPointer != nil {
                         return true
                     }
 
                     if t.isFilledLeaf {
 
-                        if t.nodeIndices.contains(i) { return false }
+                        if t.nodeIndices!.contains(i) { return false }
 
                         let massAcc = t.delegate.accumulatedMass
 
@@ -176,8 +180,29 @@ extension Kinetics {
         public mutating func bindKinetics(_ kinetics: Kinetics) {
             self.kinetics = kinetics
             self.precalculatedMass = self.mass.calculateUnsafe(for: (kinetics.validCount))
-            // self.forces = .init(repeating: .zero, count: kinetics.validCount)
 
+            self.tree = .allocate(capacity: 1)
+            self.tree.initialize(
+                to:
+                    BufferedKDTree(
+                        rootBox: .init(
+                            p0: .init(repeating: 0),
+                            p1: .init(repeating: 1)
+                        ),
+                        nodeCapacity: kinetics.validCount,
+                        rootDelegate: MassCentroidKDTreeDelegate<Vector>(
+                            massProvider: precalculatedMass.mutablePointer)
+                    )
+            )
+        }
+
+        @usableFromInline
+        internal var tree:
+            UnsafeMutablePointer<BufferedKDTree<Vector, MassCentroidKDTreeDelegate<Vector>>>! = nil
+
+        @inlinable
+        public func dispose() {
+            self.tree.deinitialize(count: 1)
         }
     }
 }