diff implementation on :aarch64

Author: Moelf

Project.toml

@@ -12,6 +12,7 @@ Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LorentzVectorHEP = "f612022c-142a-473f-8cfd-a09cf3793c6c"
 LorentzVectors = "3f54b04b-17fc-5cd4-9758-90c048d965e3"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
+SIMD = "fdea26ae-647d-5447-a871-4b548cad5224"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 
 [weakdeps]
@@ -35,6 +36,7 @@ LorentzVectorHEP = "0.1.6"
 LorentzVectors = "0.4.3"
 Makie = "0.20, 0.21, 0.22"
 MuladdMacro = "0.2.4"
+SIMD = "3.7.1"
 StructArrays = "0.6.18, 0.7"
 Test = "1.9"
 julia = "1.9"

src/JetReconstruction.jl

@@ -18,6 +18,7 @@ module JetReconstruction
 using LorentzVectorHEP
 using MuladdMacro
 using StructArrays
+using SIMD
 
 # Import from LorentzVectorHEP methods for those 4-vector types
 pt2(p::LorentzVector) = LorentzVectorHEP.pt2(p)

src/Utils.jl

@@ -146,8 +146,49 @@ array. The use of `@turbo` macro gives a significant performance boost.
 - `dij_min`: The minimum value in the first `n` elements of the `dij` array.
 - `best`: The index of the minimum value in the `dij` array.
 """
-fast_findmin(dij, n) = begin
-    x = @fastmath foldl(min, @view(dij[begin:n]))
-    i = findfirst(==(x), dij)::Int
-    x, i
+function fast_findmin end
+
+if Sys.ARCH == :aarch64
+    function fast_findmin(dij, n)
+        x = @fastmath foldl(min, @view(dij[begin:n]))
+        i = findfirst(==(x), dij)::Int
+        x, i
+    end
+else
+    function fast_findmin(dij::DenseVector{T}, n) where {T}
+        laneIndices = SIMD.Vec{8, Int}((1, 2, 3, 4, 5, 6, 7, 8))
+        minvals = SIMD.Vec{8, T}(Inf)
+        min_indices = SIMD.Vec{8, Int}(0)
+
+        n_batches, remainder = divrem(n, 8)
+        lane = VecRange{8}(0)
+        i = 1
+        @inbounds @fastmath for _ in 1:n_batches
+            dijs = dij[lane + i]
+            predicate = dijs < minvals
+            minvals = vifelse(predicate, dijs, minvals)
+            min_indices = vifelse(predicate, laneIndices, min_indices)
+
+            i += 8
+            laneIndices += 8
+        end
+
+        min_value = SIMD.minimum(minvals)
+        min_index = @inbounds min_value == minvals[1] ? min_indices[1] :
+                              min_value == minvals[2] ? min_indices[2] :
+                              min_value == minvals[3] ? min_indices[3] :
+                              min_value == minvals[4] ? min_indices[4] :
+                              min_value == minvals[5] ? min_indices[5] :
+                              min_value == minvals[6] ? min_indices[6] :
+                              min_value == minvals[7] ? min_indices[7] : min_indices[8]
+
+        @inbounds @fastmath for _ in 1:remainder
+            xi = dij[i]
+            pred = dij[i] < min_value
+            min_value = ifelse(pred, xi, min_value)
+            min_index = ifelse(pred, i, min_index)
+            i += 1
+        end
+        return min_value, min_index
+    end
 end